Store free variables used in function

17 files changed, 210 insertions, 78 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;
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedArgsNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedArgsNodeMatcher.java
index 05037526..0d70a40d 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedArgsNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedArgsNodeMatcher.java
@@ -26,7 +26,7 @@ public class TypedArgsNodeMatcher extends MatcherBuilder<TypedArgsNode, TypedArg
         return new TypedArgsNodeMatcher(submatchers);
     }
 
-    public TypedArgsNodeMatcher withPositional(Matcher<? super Iterable<TypedExpressionNode>> positional) {
+    public TypedArgsNodeMatcher withPositional(Matcher<? super Iterable<? extends TypedExpressionNode>> positional) {
         return with(has("positional", x -> x.positional(), positional));
     }
 
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
index 04a06f64..b1c6e59a 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNodeMatcher.java
@@ -29,7 +29,7 @@ public class TypedFunctionDefinitionNodeMatcher extends MatcherBuilder<TypedName
         return new TypedFunctionDefinitionNodeMatcher(submatchers);
     }
 
-    public TypedFunctionDefinitionNodeMatcher withBody(Matcher<Iterable<TypedFunctionStatementNode>> body) {
+    public TypedFunctionDefinitionNodeMatcher withBody(Matcher<Iterable<? extends TypedFunctionStatementNode>> body) {
         return with(has("body", x -> x.body(), body));
     }
 
@@ -37,6 +37,10 @@ public class TypedFunctionDefinitionNodeMatcher extends MatcherBuilder<TypedName
         return with(has("effect", x -> x.effect(), effect));
     }
 
+    public TypedFunctionDefinitionNodeMatcher withFreeVariables(Matcher<Iterable<String>> freeVariables) {
+        return with(has("freeVariables", x -> x.freeVariables(), freeVariables));
+    }
+
     public TypedFunctionDefinitionNodeMatcher withName(String name) {
         return with(has("name", x -> x.name(), equalTo(name)));
     }
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index ebf1ef42..dad777f6 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -62,7 +62,7 @@ public class TypedNodeMatchers {
     }
 
     public static Matcher<TypedNamespaceNode> isTypedNamespaceNode(
-        Matcher<Iterable<TypedNamespaceStatementNode>> body
+        Matcher<Iterable<? extends TypedNamespaceStatementNode>> body
     ) {
         return has("body", x -> x.body(), body);
     }
@@ -78,8 +78,8 @@ public class TypedNodeMatchers {
     }
 
     public static Matcher<TypedParamsNode> isTypedParamsNode(
-        Matcher<Iterable<TypedParamNode>> positional,
-        Matcher<Iterable<TypedParamNode>> named
+        Matcher<Iterable<? extends TypedParamNode>> positional,
+        Matcher<Iterable<? extends TypedParamNode>> named
     ) {
         return allOf(
             has("positional", x -> x.positional(), positional),
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNodeMatcher.java
index c88111df..8dae4a31 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNodeMatcher.java
@@ -29,7 +29,7 @@ public class UntypedFunctionDefinitionNodeMatcher extends MatcherBuilder<Untyped
         return new UntypedFunctionDefinitionNodeMatcher(submatchers);
     }
 
-    public UntypedFunctionDefinitionNodeMatcher withBody(Matcher<Iterable<UntypedFunctionStatementNode>> body) {
+    public UntypedFunctionDefinitionNodeMatcher withBody(Matcher<Iterable<? extends UntypedFunctionStatementNode>> body) {
         return with(has("body", x -> x.body(), body));
     }
 
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
index 6c06e52a..95d30cc6 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/UntypedNodeMatchers.java
@@ -40,8 +40,8 @@ public class UntypedNodeMatchers {
     }
 
     public static Matcher<UntypedParamsNode> isUntypedParamsNode(
-        Matcher<Iterable<UntypedParamNode>> positional,
-        Matcher<Iterable<UntypedParamNode>> named
+        Matcher<Iterable<? extends UntypedParamNode>> positional,
+        Matcher<Iterable<? extends UntypedParamNode>> named
     ) {
         return allOf(
             has("positional", x -> x.positional(), positional),
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBoolLiteralTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBoolLiteralTests.java
index 44258695..c39718de 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBoolLiteralTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBoolLiteralTests.java
@@ -15,7 +15,7 @@ public class TypeCheckerBoolLiteralTests {
 
         var typed = TypeChecker.typeCheckExpression(untyped, TypeCheckerContext.EMPTY);
 
-        assertThat(typed, isTypedBoolLiteralNode(true));
-        assertThat(typed.type(), equalTo(Types.BOOL));
+        assertThat(typed.value(), isTypedBoolLiteralNode(true));
+        assertThat(typed.value().type(), equalTo(Types.BOOL));
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerCallTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerCallTests.java
index 19c08b14..24ded26f 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerCallTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerCallTests.java
@@ -33,7 +33,7 @@ public class TypeCheckerCallTests {
         var typed = TypeChecker.typeCheckExpression(untyped, context);
 
         assertThat(
-            typed,
+            typed.value(),
             isTypedCallNode()
                 .withReceiver(isTypedReferenceNode("f", functionType))
                 .withArgs(isTypedArgsNode().withPositional(isSequence()).withNamed(isSequence()))
@@ -79,7 +79,7 @@ public class TypeCheckerCallTests {
         var typed = TypeChecker.typeCheckExpression(untyped, context);
 
         assertThat(
-            typed,
+            typed.value(),
             isTypedCallNode()
                 .withArgs(isTypedArgsNode().withPositional(isSequence(
                     isTypedStringLiteralNode("hello"),
@@ -108,7 +108,7 @@ public class TypeCheckerCallTests {
         var typed = TypeChecker.typeCheckExpression(untyped, context);
 
         assertThat(
-            typed,
+            typed.value(),
             isTypedCallNode()
                 .withArgs(isTypedArgsNode().withNamed(isSequence(
                     isTypedNamedArgNode("x", isTypedStringLiteralNode("hello")),
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
index a475ed6c..b7ffca1b 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerExpressionStatementTests.java
@@ -14,6 +14,6 @@ public class TypeCheckerExpressionStatementTests {
 
         var typed = TypeChecker.typeCheckFunctionStatement(untyped, TypeCheckerContext.EMPTY);
 
-        assertThat(typed, isTypedExpressionStatementNode(isTypedStringLiteralNode("hello")));
+        assertThat(typed.value(), isTypedExpressionStatementNode(isTypedStringLiteralNode("hello")));
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index 1671bed4..7ce8c7db 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -103,4 +103,26 @@ public class TypeCheckerFunctionDefinitionTests {
                 ))
         );
     }
+
+    @Test
+    public void freeVariablesIncludeVariablesReferencedInFunctionButDefinedOutsideOfFunction() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withReturnType(Untyped.typeLevelReference("Unit"))
+            .addStatement(Untyped.expressionStatement(Untyped.reference("y")))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enter(P.map(
+                Map.entry("Unit", Types.metaType(Types.UNIT)),
+                Map.entry("x", Types.STRING),
+                Map.entry("y", Types.STRING)
+            ));
+
+        var typed = TypeChecker.typeCheckNamespaceStatement(untyped, context);
+
+        assertThat(
+            typed,
+            isTypedFunctionDefinitionNode()
+                .withFreeVariables(containsExactly(equalTo("y")))
+        );
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIntLiteralTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIntLiteralTests.java
index ca47b336..ce007495 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIntLiteralTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIntLiteralTests.java
@@ -15,7 +15,7 @@ public class TypeCheckerIntLiteralTests {
 
         var typed = TypeChecker.typeCheckExpression(untyped, TypeCheckerContext.EMPTY);
 
-        assertThat(typed, isTypedIntLiteralNode(42));
-        assertThat(typed.type(), equalTo(Types.INT));
+        assertThat(typed.value(), isTypedIntLiteralNode(42));
+        assertThat(typed.value().type(), equalTo(Types.INT));
     }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerReferenceTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerReferenceTests.java
index 659514f2..9d067cc1 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerReferenceTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerReferenceTests.java
@@ -22,7 +22,7 @@ public class TypeCheckerReferenceTests {
 
         var typed = TypeChecker.typeCheckExpression(untyped, context);
 
-        assertThat(typed, isTypedReferenceNode("x", Types.STRING));
+        assertThat(typed.value(), isTypedReferenceNode("x", Types.STRING));
     }
 
     @Test
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStringLiteralTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStringLiteralTests.java
index 352fbef9..7ede9a67 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStringLiteralTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStringLiteralTests.java
@@ -15,7 +15,7 @@ public class TypeCheckerStringLiteralTests {
 
         var typed = TypeChecker.typeCheckExpression(untyped, TypeCheckerContext.EMPTY);
 
-        assertThat(typed, isTypedStringLiteralNode("hello"));
-        assertThat(typed.type(), equalTo(Types.STRING));
+        assertThat(typed.value(), isTypedStringLiteralNode("hello"));
+        assertThat(typed.value().type(), equalTo(Types.STRING));
     }
 }