Handle free variables in inner scopes

2 files changed, 56 insertions, 2 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerContextTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerContextTests.java
index 13f2486b..7c6e709b 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerContextTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerContextTests.java
@@ -12,8 +12,7 @@ import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.zwobble.precisely.AssertThat.assertThat;
-import static org.zwobble.precisely.Matchers.equalTo;
-import static org.zwobble.precisely.Matchers.isOptionalOf;
+import static org.zwobble.precisely.Matchers.*;
 
 public class TypeCheckerContextTests {
     @Test
@@ -75,4 +74,26 @@ public class TypeCheckerContextTests {
 
         assertThat(error.variableName(), equalTo("x"));
     }
+
+    @Test
+    public void runtimeLookupUpdatesFreeVariablesUntilLexicalContextInWhichVariableIsDefined() {
+        var context = TypeCheckerContext.EMPTY
+            .enterBlock(P.map())
+            .enterBlock(P.map(
+                Map.entry("x", TypeCheckerContext.undefined())
+            ))
+            .enterBlock(P.map())
+            .enterBlock(P.map());
+
+        var result = context.runtimeLookup("x").context();
+
+        assertThat(result.freeVariables(), containsExactly(
+            has("name", x -> x.name(), equalTo("x"))
+        ));
+        assertThat(result.exit().freeVariables(), containsExactly(
+            has("name", x -> x.name(), equalTo("x"))
+        ));
+        assertThat(result.exit().exit().freeVariables(), containsExactly());
+        assertThat(result.exit().exit().exit().freeVariables(), containsExactly());
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index ee88b457..794dd2c3 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -315,6 +315,39 @@ public class TypeCheckerFunctionDefinitionTests {
         );
     }
 
+    @Test
+    public void freeVariablesIncludeVariablesReferencedInBlockInFunctionButDefinedOutsideOfFunction() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withReturnType(Untyped.typeLevelReference("Unit"))
+            .addStatement(Untyped.expressionStatement(Untyped.if_(
+                List.of(
+                    Untyped.conditionalBranch(
+                        Untyped.boolLiteral(true),
+                        Untyped.block(List.of())
+                    )
+                ),
+                Untyped.block(List.of(
+                    Untyped.expressionStatement(Untyped.reference("x"))
+                ))
+            )))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespaceWithDefinedVariables(P.map(
+                Map.entry("Unit", Types.metaType(Types.UNIT)),
+                Map.entry("x", Types.STRING)
+            ));
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        assertThat(
+            typed.value(),
+            isTypedFunctionDefinitionNode()
+                .withFreeVariables(containsExactly(
+                    has("name", x -> x.name(), equalTo("x"))
+                ))
+        );
+    }
+
     private static TypeCheckerResult<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context