Store free variables used in function

5 files changed, 165 insertions, 59 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedArgsNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedArgsNode.java
index 473fa197..6f03295d 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedArgsNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedArgsNode.java
@@ -5,7 +5,7 @@ import org.shedlang.compiler.sources.Source;
 import java.util.List;
 
 public record TypedArgsNode(
-    List<TypedExpressionNode> positional,
+    List<? extends TypedExpressionNode> positional,
     List<TypedNamedArgNode> named,
     Source source
 ) implements TypedNode {
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index 4bf840f9..79d13152 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.pcollections.Empty;
+import org.pcollections.POrderedSet;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
@@ -10,12 +11,14 @@ import org.shedlang.compiler.types.Types;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 
 public record TypedFunctionDefinitionNode(
     String name,
     TypedParamsNode params,
     Optional<TypedTypeLevelExpressionNode<Effect>> effect,
     TypedTypeLevelExpressionNode<Type> returnType,
+    Set<String> freeVariables,
     List<TypedFunctionStatementNode> body,
     Source source
 ) implements TypedNamespaceStatementNode {
@@ -26,6 +29,7 @@ public record TypedFunctionDefinitionNode(
             Empty.vector(),
             Optional.empty(),
             new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
+            Empty.orderedSet(),
             Empty.vector()
         );
     }
@@ -36,6 +40,7 @@ public record TypedFunctionDefinitionNode(
         PVector<TypedParamNode> namedParams,
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
+        POrderedSet<String> freeVariables,
         PVector<TypedFunctionStatementNode> body
     ) {
         public TypedFunctionDefinitionNode build() {
@@ -44,6 +49,7 @@ public record TypedFunctionDefinitionNode(
                 new TypedParamsNode(positionalParams, namedParams, Sources.NULL),
                 effect,
                 returnType,
+                freeVariables,
                 body,
                 Sources.NULL
             );
@@ -56,6 +62,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 this.effect,
                 this.returnType,
+                this.freeVariables,
                 this.body
             );
         }
@@ -67,6 +74,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 this.effect,
                 this.returnType,
+                this.freeVariables,
                 this.body
             );
         }
@@ -78,6 +86,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams.plus(new TypedParamNode(name, type, Sources.NULL)),
                 this.effect,
                 this.returnType,
+                this.freeVariables,
                 this.body
             );
         }
@@ -89,6 +98,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 Optional.of(effect),
                 this.returnType,
+                this.freeVariables,
                 this.body
             );
         }
@@ -100,6 +110,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 this.effect,
                 returnType,
+                this.freeVariables,
                 this.body
             );
         }
@@ -111,6 +122,7 @@ public record TypedFunctionDefinitionNode(
                 this.namedParams,
                 this.effect,
                 this.returnType,
+                this.freeVariables,
                 this.body.plus(statement)
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index ff56de3b..b6bfb399 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -29,8 +29,8 @@ public class TypeChecker {
         return subType.equals(superType);
     }
 
-    private static Type lookupVariable(String name, Source source, TypeCheckerContext context) {
-        var type = context.environment(name);
+    private static Type compileTimeLookup(String name, Source source, TypeCheckerContext context) {
+        var type = context.compileTimeLookup(name);
 
         if (type.isEmpty()) {
             throw new UnknownVariableError(name, source);
@@ -39,7 +39,17 @@ public class TypeChecker {
         return type.get();
     }
 
-    private static TypedArgsNode typeCheckArgs(
+    private static TypeCheckerResult<Type> runtimeLookup(String name, Source source, TypeCheckerContext context) {
+        return context.runtimeLookup(name).map(type -> {
+            if (type.isEmpty()) {
+                throw new UnknownVariableError(name, source);
+            }
+
+            return type.get();
+        });
+    }
+
+    private static TypeCheckerResult<TypedArgsNode> typeCheckArgs(
         UntypedArgsNode args,
         TypeCheckerContext context,
         FunctionType receiverType
@@ -52,17 +62,20 @@ public class TypeChecker {
             );
         }
 
-        var typedPositionalArgs = args.positional().stream()
-            .map(arg -> typeCheckExpression(arg, context))
-            .toList();
-
-        var typedNamedArgs = args.named().stream()
-            .map(arg -> new TypedNamedArgNode(
-                arg.name(),
-                typeCheckExpression(arg.expression(), context),
-                arg.source()
-            ))
-            .toList();
+        var positionalResult = context.mapWithContext(args.positional(), TypeChecker::typeCheckExpression);
+        var typedPositionalArgs = positionalResult.value();
+        context = positionalResult.context();
+
+        var namedResult = context.mapWithContext(args.named(), (arg, currentContext) -> {
+            return typeCheckExpression(arg.expression(), currentContext)
+                .map(typedExpression -> new TypedNamedArgNode(
+                    arg.name(),
+                    typedExpression,
+                    arg.source()
+                ));
+        });
+        var typedNamedArgs = namedResult.value();
+        context = namedResult.context();
 
         for (var i = 0; i < typedPositionalArgs.size(); i++) {
             var typedArg = typedPositionalArgs.get(i);
@@ -93,23 +106,26 @@ public class TypeChecker {
             }
         }
 
-        return new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source());
+        return context.toResult(new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source()));
     }
 
-    public static TypedExpressionNode typeCheckBoolLiteral(
-        UntypedBoolLiteralNode untyped
+    public static TypeCheckerResult<TypedExpressionNode> typeCheckBoolLiteral(
+        UntypedBoolLiteralNode untyped,
+        TypeCheckerContext context
     ) {
-        return new TypedBoolLiteralNode(
+        return context.toResult(new TypedBoolLiteralNode(
             untyped.value(),
             untyped.source()
-        );
+        ));
     }
 
-    private static TypedExpressionNode typeCheckCall(
+    private static TypeCheckerResult<TypedExpressionNode> typeCheckCall(
         UntypedCallNode untyped,
         TypeCheckerContext context
     ) {
-        var typedReceiver = typeCheckExpression(untyped.receiver(), context);
+        var receiverResult = typeCheckExpression(untyped.receiver(), context);
+        var typedReceiver = receiverResult.value();
+        context = receiverResult.context();
 
         if (!(typedReceiver.type() instanceof FunctionType receiverType)) {
             throw new UnexpectedTypeError(
@@ -119,13 +135,17 @@ public class TypeChecker {
             );
         }
 
-        return new TypedCallNode(
+        var argsResult = typeCheckArgs(untyped.args(), context, receiverType);
+        var typedArgs = argsResult.value();
+        context = argsResult.context();
+
+        return context.toResult(new TypedCallNode(
             typedReceiver,
-            typeCheckArgs(untyped.args(), context, receiverType),
+            typedArgs,
             false,
             receiverType.returnType(),
             untyped.source()
-        );
+        ));
     }
 
     private static TypedTypeLevelExpressionNode<Effect> typeCheckEffect(
@@ -135,36 +155,38 @@ public class TypeChecker {
         return typeCheckTypeLevelExpression(TypeSets.EFFECT, untyped, context);
     }
 
-    public static TypedExpressionNode typeCheckExpression(
+    public static TypeCheckerResult<? extends TypedExpressionNode> typeCheckExpression(
         UntypedExpressionNode untyped,
         TypeCheckerContext context
     ) {
         return switch (untyped) {
             case UntypedBoolLiteralNode untypedBoolLiteral ->
-                typeCheckBoolLiteral(untypedBoolLiteral);
+                typeCheckBoolLiteral(untypedBoolLiteral, context);
 
             case UntypedCallNode untypedCall ->
                 typeCheckCall(untypedCall, context);
 
             case UntypedIntLiteralNode untypedIntLiteral ->
-                typeCheckIntLiteral(untypedIntLiteral);
+                typeCheckIntLiteral(untypedIntLiteral, context);
 
             case UntypedReferenceNode untypedReference ->
                 typeCheckReference(untypedReference, context);
 
             case UntypedStringLiteralNode untypedStringLiteral ->
-                typeCheckStringLiteral(untypedStringLiteral);
+                typeCheckStringLiteral(untypedStringLiteral, context);
         };
     }
 
-    private static TypedFunctionStatementNode typeCheckExpressionStatement(
+    private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckExpressionStatement(
         UntypedExpressionStatementNode untyped,
         TypeCheckerContext context
     ) {
-        return new TypedExpressionStatementNode(
-            typeCheckExpression(untyped.expression(), context),
-            untyped.source()
-        );
+        return typeCheckExpression(untyped.expression(), context).map(typedExpression -> {
+            return new TypedExpressionStatementNode(
+                typedExpression,
+                untyped.source()
+            );
+        });
     }
 
     private static TypedNamespaceStatementNode typeCheckFunctionDefinition(
@@ -182,21 +204,25 @@ public class TypeChecker {
 
         var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
 
-        var typedBody = untyped.body().stream()
-            .map(statement -> typeCheckFunctionStatement(statement, context))
-            .toList();
+        var bodyResult = context.mapWithContext(
+            untyped.body(),
+            TypeChecker::typeCheckFunctionStatement
+        );
+        var typedBody = bodyResult.value();
+        context = bodyResult.context();
 
         return new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
             typedEffect,
             typedReturnType,
+            context.freeVariables(),
             typedBody,
             untyped.source()
         );
     }
 
-    public static TypedFunctionStatementNode typeCheckFunctionStatement(
+    public static TypeCheckerResult<TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
         TypeCheckerContext context
     ) {
@@ -206,13 +232,14 @@ public class TypeChecker {
         };
     }
 
-    private static TypedExpressionNode typeCheckIntLiteral(
-        UntypedIntLiteralNode untyped
+    private static TypeCheckerResult<TypedExpressionNode> typeCheckIntLiteral(
+        UntypedIntLiteralNode untyped,
+        TypeCheckerContext context
     ) {
-        return new TypedIntLiteralNode(
+        return context.toResult(new TypedIntLiteralNode(
             untyped.value(),
             untyped.source()
-        );
+        ));
     }
 
     static TypedNamespaceNode typeCheckNamespace(
@@ -263,26 +290,27 @@ public class TypeChecker {
         return new TypedParamsNode(typedPositional, typedNamed, untyped.source());
     }
 
-    private static TypedReferenceNode typeCheckReference(
+    private static TypeCheckerResult<TypedReferenceNode> typeCheckReference(
         UntypedReferenceNode untyped,
         TypeCheckerContext context
     ) {
-        var type = lookupVariable(untyped.name(), untyped.source(), context);
-
-        return new TypedReferenceNode(
-            untyped.name(),
-            type,
-            untyped.source()
-        );
+        return runtimeLookup(untyped.name(), untyped.source(), context).map(type -> {
+            return new TypedReferenceNode(
+                untyped.name(),
+                type,
+                untyped.source()
+            );
+        });
     }
 
-    private static TypedStringLiteralNode typeCheckStringLiteral(
-        UntypedStringLiteralNode untyped
+    private static TypeCheckerResult<TypedStringLiteralNode> typeCheckStringLiteral(
+        UntypedStringLiteralNode untyped,
+        TypeCheckerContext context
     ) {
-        return new TypedStringLiteralNode(
+        return context.toResult(new TypedStringLiteralNode(
             untyped.value(),
             untyped.source()
-        );
+        ));
     }
 
     static <T extends TypeLevelValue> TypedTypeLevelExpressionNode<T> typeCheckTypeLevelExpression(
@@ -301,7 +329,7 @@ public class TypeChecker {
         UntypedTypeLevelReferenceNode untyped,
         TypeCheckerContext context
     ) {
-        var type = lookupVariable(untyped.name(), untyped.source(), context);
+        var type = compileTimeLookup(untyped.name(), untyped.source(), context);
 
         if (!(type instanceof TypeLevelValueType typeLevelValueType)) {
             throw new UnexpectedTypeError(valueType, type, untyped.source());
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index c130b0d3..91d1daac 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -2,26 +2,67 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
+import org.pcollections.POrderedSet;
 import org.shedlang.compiler.types.Type;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
+import java.util.Set;
+import java.util.function.BiFunction;
 
 public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
-        Empty.map()
+        Empty.map(),
+        Empty.orderedSet()
     );
 
     private final PMap<String, Type> environment;
+    private final POrderedSet<String> freeVariables;
 
-    public TypeCheckerContext(PMap<String, Type> environment) {
+    public TypeCheckerContext(PMap<String, Type> environment, POrderedSet<String> freeVariables) {
         this.environment = environment;
+        this.freeVariables = freeVariables;
     }
 
     public TypeCheckerContext enter(PMap<String, Type> environment) {
-        return new TypeCheckerContext(this.environment.plusAll(environment));
+        return new TypeCheckerContext(this.environment.plusAll(environment), Empty.orderedSet());
     }
 
-    public Optional<Type> environment(String name) {
+    public TypeCheckerResult<Optional<Type>> runtimeLookup(String name) {
+        var type = Optional.ofNullable(environment.get(name));
+        return withFreeVariable(name).toResult(type);
+    }
+
+    public Optional<Type> compileTimeLookup(String name) {
         return Optional.ofNullable(environment.get(name));
     }
+    
+    private TypeCheckerContext withFreeVariable(String name) {
+        return new TypeCheckerContext(environment, freeVariables.plus(name));
+    }
+
+    public <T> TypeCheckerResult<T> toResult(T value) {
+        return new TypeCheckerResult<>(value, this);
+    }
+
+    public <T, U> TypeCheckerResult<List<U>> mapWithContext(
+        Iterable<T> iterable,
+        BiFunction<T, TypeCheckerContext, TypeCheckerResult<U>> func
+    ) {
+        var mappedValues = new ArrayList<U>();
+        var context = this;
+
+        for (var value : iterable) {
+            var result = func.apply(value, context);
+            mappedValues.add(result.value());
+            context = result.context();
+        }
+
+        return new TypeCheckerResult<>(mappedValues, context);
+    }
+
+    public Set<String> freeVariables() {
+        return freeVariables;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java
new file mode 100644
index 00000000..8c4fa2b1
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java
@@ -0,0 +1,25 @@
+package org.shedlang.compiler.typechecker;
+
+import java.util.function.Function;
+
+public class TypeCheckerResult<T> {
+    private final T value;
+    private final TypeCheckerContext context;
+
+    public TypeCheckerResult(T value, TypeCheckerContext context) {
+        this.value = value;
+        this.context = context;
+    }
+
+    public <U> TypeCheckerResult<U> map(Function<T, U> func) {
+        return new TypeCheckerResult<>(func.apply(value), context);
+    }
+
+    public T value() {
+        return value;
+    }
+
+    public TypeCheckerContext context() {
+        return context;
+    }
+}