Check for reference before definition

4 files changed, 94 insertions, 18 deletions

diff --git a/src/main/java/org/shedlang/compiler/ShedCompiler.java b/src/main/java/org/shedlang/compiler/ShedCompiler.java
index 6f72328f..e91f6655 100644
--- a/src/main/java/org/shedlang/compiler/ShedCompiler.java
+++ b/src/main/java/org/shedlang/compiler/ShedCompiler.java
@@ -39,7 +39,7 @@ public class ShedCompiler {
             .addNamespaceFields(ShedBuiltins.NAMESPACE_NAME, P.orderedMap(
                 Map.entry(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
             ))
-            .enterBlock(P.map(
+            .enterBlockWithDefinedVariables(P.map(
                 Map.entry("Io", Types.effectType(Effects.IO)),
                 Map.entry("Unit", Types.metaType(Types.UNIT)),
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 1dd4b118..feafa5f4 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -10,6 +10,7 @@ import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.PCollectors;
 
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
@@ -43,7 +44,13 @@ public class TypeChecker {
             throw new UnknownVariableError(name, source);
         }
 
-        return type.get();
+        return switch (type.get()) {
+            case TypeCheckerContext.Variable.Defined defined ->
+                defined.type();
+
+            case TypeCheckerContext.Variable.Undefined undefined ->
+                throw new ReferenceBeforeDefinitionError(name, source);
+        };
     }
 
     private static TypeCheckerResult<Type> runtimeLookup(String name, Source source, TypeCheckerContext context) {
@@ -52,10 +59,33 @@ public class TypeChecker {
                 throw new UnknownVariableError(name, source);
             }
 
-            return type.get();
+            return switch (type.get()) {
+                case TypeCheckerContext.Variable.Defined defined ->
+                    defined.type();
+
+                case TypeCheckerContext.Variable.Undefined undefined ->
+                    throw new ReferenceBeforeDefinitionError(name, source);
+            };
         });
     }
 
+    private static List<String> definedVariableNames(UntypedFunctionStatementNode statement) {
+        return switch (statement) {
+            case UntypedExpressionStatementNode expressionStatement ->
+                List.of();
+
+            case UntypedValNode val ->
+                definedVariableNames(val.target());
+        };
+    }
+
+    private static List<String> definedVariableNames(UntypedTargetNode target) {
+        return switch (target) {
+            case UntypedTargetNode.Variable variable ->
+                List.of(variable.name());
+        };
+    }
+
     private static TypeCheckerResult<TypedArgsNode> typeCheckArgs(
         UntypedArgsNode args,
         TypeCheckerContext context,
@@ -230,8 +260,18 @@ public class TypeChecker {
         var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
 
         var bodyEnvironment = typedParams.all()
-            .map(param -> Map.entry(param.name(), param.type().value()))
+            .map(param -> Map.entry(
+                param.name(),
+                TypeCheckerContext.defined(param.type().value())
+            ))
             .collect(PCollectors.toMap());
+
+        for (var statement : untyped.body()) {
+            for (var definedVariableName : definedVariableNames(statement)) {
+                bodyEnvironment = bodyEnvironment.plus(definedVariableName, TypeCheckerContext.undefined());
+            }
+        }
+
         var bodyContext = context.enterFunction(bodyEnvironment, typedEffect.map(e -> e.value()).orElse(Effects.NONE));
 
         var bodyResult = bodyContext.mapWithContext(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 2cf99f9f..f05cb7e1 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -4,12 +4,10 @@ import org.pcollections.Empty;
 import org.pcollections.PMap;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.PCollectors;
 import org.shedlang.compiler.util.POrderedMap;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
 import java.util.function.BiFunction;
 
 public class TypeCheckerContext {
@@ -21,14 +19,27 @@ public class TypeCheckerContext {
         Empty.map()
     );
 
-    private final PMap<String, Type> environment;
+    private final PMap<String, Variable> environment;
     private final Effect effect;
     private final Set<String> boundVariables;
     private final FreeVariables.Builder freeVariables;
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
 
+    public sealed interface Variable {
+        record Undefined() implements Variable {}
+        record Defined(Type type) implements Variable {}
+    }
+
+    public static Variable defined(Type type) {
+        return new Variable.Defined(type);
+    }
+
+    public static Variable undefined() {
+        return new Variable.Undefined();
+    }
+
     public TypeCheckerContext(
-        PMap<String, Type> environment,
+        PMap<String, Variable> environment,
         Effect effect,
         Set<String> boundVariables,
         FreeVariables.Builder freeVariables,
@@ -41,15 +52,23 @@ public class TypeCheckerContext {
         this.namespaceTypeInfo = namespaceTypeInfo;
     }
 
-    public TypeCheckerContext enterBlock(PMap<String, Type> environment) {
+    public TypeCheckerContext enterBlockWithDefinedVariables(PMap<String, Type> variables) {
+        var environment = variables.entrySet().stream()
+            .map(x -> Map.entry(x.getKey(), TypeCheckerContext.defined(x.getValue())))
+            .collect(PCollectors.toMap());
+
+        return enter(environment, this.effect);
+    }
+
+    public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
         return enter(environment, this.effect);
     }
 
-    public TypeCheckerContext enterFunction(PMap<String, Type> environment, Effect effect) {
+    public TypeCheckerContext enterFunction(PMap<String, Variable> environment, Effect effect) {
         return enter(environment, effect);
     }
 
-    private TypeCheckerContext enter(PMap<String, Type> environment, Effect effect) {
+    private TypeCheckerContext enter(PMap<String, Variable> environment, Effect effect) {
         return new TypeCheckerContext(
             this.environment.plusAll(environment),
             effect,
@@ -60,14 +79,14 @@ public class TypeCheckerContext {
 
     }
 
-    public TypeCheckerResult<Optional<Type>> runtimeLookup(String name) {
-        var type = Optional.ofNullable(environment.get(name));
+    public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name) {
+        var variable = Optional.ofNullable(environment.get(name));
 
         var context = boundVariables.contains(name) ? this : withFreeVariable(name);
-        return context.toResult(type);
+        return context.toResult(variable);
     }
 
-    public Optional<Type> compileTimeLookup(String name) {
+    public Optional<Variable> compileTimeLookup(String name) {
         return Optional.ofNullable(environment.get(name));
     }
 
@@ -77,7 +96,7 @@ public class TypeCheckerContext {
     ) {
         // TODO: check if already defined
         return new TypeCheckerContext(
-            this.environment.plus(name, type),
+            this.environment.plus(name, new Variable.Defined(type)),
             this.effect,
             this.boundVariables,
             this.freeVariables,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/ReferenceBeforeDefinitionError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/ReferenceBeforeDefinitionError.java
new file mode 100644
index 00000000..edf52543
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/ReferenceBeforeDefinitionError.java
@@ -0,0 +1,17 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class ReferenceBeforeDefinitionError extends SourceError {
+    private final String variableName;
+
+    public ReferenceBeforeDefinitionError(String variableName, Source source) {
+        super("Cannot reference variable before definition: " + variableName, source);
+        this.variableName = variableName;
+    }
+
+    public String variableName() {
+        return variableName;
+    }
+}