Require variables to be declared before definition

4 files changed, 54 insertions, 17 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
index 62322251..a232031b 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
@@ -1,5 +1,7 @@
 package org.shedlang.compiler.ast.untyped;
 
+import java.util.List;
+
 public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedBlockReturnNode, UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedValNode {
     UntypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
index 4ac46d2a..a78f45ec 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
@@ -14,6 +14,13 @@ public class EnvironmentBuilder {
         this.environment = environment;
     }
 
+    public EnvironmentBuilder addUndefined(String name, Source source) {
+        if (environment.containsKey(name)) {
+            throw new DuplicateVariableName(name, source);
+        }
+        return new EnvironmentBuilder(environment.plus(name, new TypeCheckerContext.Variable.Undefined()));
+    }
+
     public EnvironmentBuilder add(String name, TypeCheckerContext.Variable variable, 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 5563d9bb..391d8e0c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -68,26 +68,45 @@ public class TypeChecker {
         });
     }
 
-    private static List<UntypedTargetNode.Variable> targetVariables(UntypedFunctionStatementNode statement) {
+    static EnvironmentBuilder declareVariables(
+        UntypedFunctionStatementNode statement,
+        EnvironmentBuilder environment
+    ) {
         return switch (statement) {
             case UntypedBlockReturnNode blockReturn ->
-                List.of();
+                environment;
 
             case UntypedExpressionStatementNode expressionStatement ->
-                List.of();
+                environment;
 
             case UntypedFunctionReturnNode functionReturn ->
-                List.of();
+                environment;
+
+            case UntypedValNode val ->
+                declareVariables(val.target(), environment);
+        };
+    }
+
+    static EnvironmentBuilder declareVariables(
+        UntypedNamespaceStatementNode statement,
+        EnvironmentBuilder environment
+    ) {
+        return switch (statement) {
+            case UntypedFunctionDefinitionNode functionDefinition ->
+                environment.addUndefined(functionDefinition.name(), functionDefinition.source());
 
             case UntypedValNode val ->
-                targetVariables(val.target());
+                declareVariables(val.target(), environment);
         };
     }
 
-    private static List<UntypedTargetNode.Variable> targetVariables(UntypedTargetNode target) {
+    private static EnvironmentBuilder declareVariables(
+        UntypedTargetNode target,
+        EnvironmentBuilder environment
+    ) {
         return switch (target) {
             case UntypedTargetNode.Variable variable ->
-                List.of(variable);
+                environment.addUndefined(variable.name(), variable.source());
         };
     }
 
@@ -386,13 +405,7 @@ public class TypeChecker {
         }
 
         for (var statement : untyped.body().statements()) {
-            for (var targetVariable : targetVariables(statement)) {
-                bodyEnvironment = bodyEnvironment.add(
-                    targetVariable.name(),
-                    TypeCheckerContext.undefined(),
-                    targetVariable.source()
-                );
-            }
+            bodyEnvironment = declareVariables(statement, bodyEnvironment);
         }
 
         var bodyContext = context.enterFunction(
@@ -527,12 +540,20 @@ public class TypeChecker {
         UntypedNamespaceNode untyped,
         TypeCheckerContext context
     ) {
+        var namespaceEnvironment = EnvironmentBuilder.EMPTY;
+
+        for (var statement : untyped.body()) {
+            namespaceEnvironment = declareVariables(statement, namespaceEnvironment);
+        }
+
+        var namespaceContext = context.enterNamespace(namespaceEnvironment.build());
+
         var typedBody = new ArrayList<TypedNamespaceStatementNode>();
 
         for (var untypedStatement : untyped.body()) {
-            var typedStatementResult = typeCheckNamespaceStatement(untypedStatement, context);
+            var typedStatementResult = typeCheckNamespaceStatement(untypedStatement, namespaceContext);
             var typedStatement = typedStatementResult.value();
-            context = typedStatementResult.context();
+            namespaceContext = typedStatementResult.context();
 
             typedBody.add(typedStatement);
         }
@@ -554,7 +575,10 @@ public class TypeChecker {
             untyped.source()
         );
 
-        return context.addNamespaceFields(typed.namespaceName(), fieldTypes).toResult(typed);
+        return context
+            .exit(namespaceContext)
+            .addNamespaceFields(typed.namespaceName(), fieldTypes)
+            .toResult(typed);
     }
 
     static TypeCheckerResult<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 6b24dd6b..8e9aa206 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -4,6 +4,7 @@ import org.pcollections.Empty;
 import org.pcollections.PMap;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.typechecker.errors.UnknownVariableError;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.PCollectors;
 import org.shedlang.compiler.util.POrderedMap;
@@ -129,6 +130,9 @@ public class TypeCheckerContext {
         Type type,
         Source source
     ) {
+        if (!this.environment.containsKey(name)) {
+            throw new UnknownVariableError(name, source);
+        }
         // TODO: check if already defined
         // TODO: check is not defined
         // TODO: check is defined in this scope