Handle free variables in inner scopes

3 files changed, 71 insertions, 6 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 99357563..c05a2475 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -154,10 +154,21 @@ public class TypeCheckerContext {
     }
 
     private TypeCheckerContext withVariable(String name) {
-        // TODO: walk up stack
-        return boundVariables().contains(name)
-            ? this
-            : updateLexicalContext(lexicalContext -> lexicalContext.withFreeVariable(name));
+        var newLexicalContextStack = this.lexicalContextStack;
+
+        for (var index = 0; index < this.lexicalContextStack.size(); index++) {
+            var lexicalContext = this.lexicalContextStack.get(index);
+            if (lexicalContext.boundVariables.contains(name)) {
+                break;
+            } else{
+                newLexicalContextStack = newLexicalContextStack.with(
+                    index,
+                    lexicalContext.withFreeVariable(name)
+                );
+            }
+        }
+
+        return new TypeCheckerContext(namespaceTypeInfo, newLexicalContextStack);
     }
 
     public Optional<Variable> compileTimeLookup(String name) {
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