Support fetching of variable declarations without environment builder

5 files changed, 37 insertions, 17 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
index a8a6fa71..ab29001d 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -14,7 +14,7 @@ public record UntypedFunctionDefinitionNode(
     UntypedTypeLevelExpressionNode returnType,
     UntypedBlockNode body,
     Source source
-) implements UntypedNamespaceStatementNode {
+) implements UntypedNamespaceStatementNode, UntypedVariableDeclarationNode {
     public static Builder builder() {
         return new Builder(
             "f",
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
index 514b794a..a70965bf 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
@@ -3,5 +3,5 @@ package org.shedlang.compiler.ast.untyped;
 import org.shedlang.compiler.sources.Source;
 
 public sealed interface UntypedTargetNode extends UntypedNode {
-    record Variable(String name, Source source) implements UntypedTargetNode {}
+    record Variable(String name, Source source) implements UntypedTargetNode, UntypedVariableDeclarationNode {}
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVariableDeclarationNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVariableDeclarationNode.java
new file mode 100644
index 00000000..b805b701
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVariableDeclarationNode.java
@@ -0,0 +1,5 @@
+package org.shedlang.compiler.ast.untyped;
+
+public interface UntypedVariableDeclarationNode extends UntypedNode {
+    String name();
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
index a78f45ec..79c0a25e 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
+import org.shedlang.compiler.ast.untyped.UntypedVariableDeclarationNode;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.DuplicateVariableName;
 
@@ -14,6 +15,16 @@ public class EnvironmentBuilder {
         this.environment = environment;
     }
 
+    public EnvironmentBuilder addUndefined(Iterable<UntypedVariableDeclarationNode> nodes) {
+        var builder = this;
+
+        for (var node : nodes) {
+            builder = builder.addUndefined(node.name(), node.source());
+        }
+
+        return builder;
+    }
+
     public EnvironmentBuilder addUndefined(String name, Source source) {
         if (environment.containsKey(name)) {
             throw new DuplicateVariableName(name, source);
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 6600f9fa..73bd8eb7 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.typechecker;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.ast.untyped.*;
 import org.shedlang.compiler.sources.Source;
@@ -68,45 +69,45 @@ public class TypeChecker {
         });
     }
 
-    static EnvironmentBuilder declareVariables(
+    static PVector<UntypedVariableDeclarationNode> declaredVariables(
         UntypedFunctionStatementNode statement,
-        EnvironmentBuilder environment
+        PVector<UntypedVariableDeclarationNode> variables
     ) {
         return switch (statement) {
             case UntypedBlockReturnNode blockReturn ->
-                environment;
+                variables;
 
             case UntypedExpressionStatementNode expressionStatement ->
-                environment;
+                variables;
 
             case UntypedFunctionReturnNode functionReturn ->
-                environment;
+                variables;
 
             case UntypedValNode val ->
-                declareVariables(val.target(), environment);
+                declaredVariables(val.target(), variables);
         };
     }
 
-    static EnvironmentBuilder declareVariables(
+    static PVector<UntypedVariableDeclarationNode> declaredVariables(
         UntypedNamespaceStatementNode statement,
-        EnvironmentBuilder environment
+        PVector<UntypedVariableDeclarationNode> variables
     ) {
         return switch (statement) {
             case UntypedFunctionDefinitionNode functionDefinition ->
-                environment.addUndefined(functionDefinition.name(), functionDefinition.source());
+                variables.plus(functionDefinition);
 
             case UntypedValNode val ->
-                declareVariables(val.target(), environment);
+                declaredVariables(val.target(), variables);
         };
     }
 
-    private static EnvironmentBuilder declareVariables(
+    private static PVector<UntypedVariableDeclarationNode> declaredVariables(
         UntypedTargetNode target,
-        EnvironmentBuilder environment
+        PVector<UntypedVariableDeclarationNode> variables
     ) {
         return switch (target) {
             case UntypedTargetNode.Variable variable ->
-                environment.addUndefined(variable.name(), variable.source());
+                variables.plus(variable);
         };
     }
 
@@ -460,7 +461,9 @@ public class TypeChecker {
         }
 
         for (var statement : untyped.body().statements()) {
-            bodyEnvironment = declareVariables(statement, bodyEnvironment);
+            // TODO: pass interface to collect variables?
+            var declaredVariables = declaredVariables(statement, P.vector());
+            bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
         }
 
         context = context.enterFunction(
@@ -609,7 +612,8 @@ public class TypeChecker {
         var namespaceEnvironment = EnvironmentBuilder.EMPTY;
 
         for (var statement : untyped.body()) {
-            namespaceEnvironment = declareVariables(statement, namespaceEnvironment);
+            var declaredVariables = declaredVariables(statement, P.vector());
+            namespaceEnvironment = namespaceEnvironment.addUndefined(declaredVariables);
         }
 
         context = context.enterNamespace(namespaceEnvironment.build());