Support capture of free variables

3 files changed, 92 insertions, 18 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
index 218fb696..37f044cc 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
@@ -45,6 +45,11 @@ public class FreeVariables implements Iterable<FreeVariable> {
         return freeVariables.size();
     }
 
+    @Override
+    public String toString() {
+        return freeVariables.toString();
+    }
+
     public static Builder builder() {
         return new Builder(OrderedPSet.empty());
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index a8ef4fd1..6600f9fa 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -625,10 +625,16 @@ public class TypeChecker {
         }
 
         for (var typeCheckSteps : typeCheckStepSets) {
+            context = context.enterFreeVariableCapture();
+
             var typedStatementResult = typeCheckSteps.define(context);
             var typedStatement = typedStatementResult.value();
             context = typedStatementResult.context();
 
+            var freeVariablesResult = context.exitFreeVariableCapture();
+            var freeVariables = freeVariablesResult.value();
+            context = freeVariablesResult.context();
+
             typedBody.add(typedStatement);
         }
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 5cb2d6c6..f64292d4 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -43,7 +43,32 @@ public class TypeCheckerContext {
         return new Variable.Undefined();
     }
 
-    public static class LexicalContext {
+    private sealed interface StackElement {
+        boolean isBoundInElement(String name);
+        StackElement withFreeVariable(String name);
+    }
+
+    private static final class FreeVariableCapture implements StackElement {
+        public static final FreeVariableCapture EMPTY = new FreeVariableCapture(FreeVariables.builder());
+
+        private final FreeVariables.Builder freeVariables;
+
+        private FreeVariableCapture(FreeVariables.Builder freeVariables) {
+            this.freeVariables = freeVariables;
+        }
+
+        @Override
+        public boolean isBoundInElement(String name) {
+            return false;
+        }
+
+        @Override
+        public StackElement withFreeVariable(String name) {
+            return new FreeVariableCapture(this.freeVariables.plus(name));
+        }
+    }
+
+    private static final class LexicalContext implements StackElement {
         private final PMap<String, Variable> environment;
         private final Effect effect;
         private final Optional<Type> returnType;
@@ -67,7 +92,13 @@ public class TypeCheckerContext {
             this.state = state;
         }
 
-        private LexicalContext withFreeVariable(String name) {
+        @Override
+        public boolean isBoundInElement(String name) {
+            return this.boundVariables.contains(name);
+        }
+
+        @Override
+        public LexicalContext withFreeVariable(String name) {
             return new LexicalContext(
                 this.environment,
                 this.effect,
@@ -80,18 +111,42 @@ public class TypeCheckerContext {
     }
 
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
-    private final PStack<LexicalContext> lexicalContextStack;
+    private final PStack<StackElement> stack;
 
     public TypeCheckerContext(
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo,
-        PStack<LexicalContext> lexicalContextStack
+        PStack<StackElement> stack
     ) {
         this.namespaceTypeInfo = namespaceTypeInfo;
-        this.lexicalContextStack = lexicalContextStack;
+        this.stack = stack;
     }
 
     private LexicalContext currentLexicalContext() {
-        return this.lexicalContextStack.get(0);
+        for (var element : this.stack) {
+            if (element instanceof LexicalContext lexicalContext) {
+                return lexicalContext;
+            }
+        }
+
+        // TODO: better error
+        throw new RuntimeException("missing lexical context");
+    }
+
+    private TypeCheckerContext updateLexicalContext(
+        Function<LexicalContext, LexicalContext> update
+    ) {
+        var newStack = this.stack;
+
+        for (var i = 0; i < this.stack.size(); i++) {
+            if (this.stack.get(i) instanceof LexicalContext lexicalContext) {
+                newStack = newStack.with(i, update.apply(lexicalContext));
+            }
+        }
+
+        return new TypeCheckerContext(
+            this.namespaceTypeInfo,
+            newStack
+        );
     }
 
     public TypeCheckerContext enterNamespaceWithDefinedVariables(PMap<String, Type> variables) {
@@ -125,7 +180,7 @@ public class TypeCheckerContext {
     ) {
         return new TypeCheckerContext(
             this.namespaceTypeInfo,
-            this.lexicalContextStack.plus(new LexicalContext(
+            this.stack.plus(new LexicalContext(
                 this.environment().plusAll(newEnvironment),
                 newEffect,
                 newReturnType,
@@ -139,19 +194,27 @@ public class TypeCheckerContext {
     public TypeCheckerContext exit() {
         return new TypeCheckerContext(
             this.namespaceTypeInfo,
-            this.lexicalContextStack.minus(0)
+            this.stack.minus(0)
         );
     }
 
-    private TypeCheckerContext updateLexicalContext(
-        Function<LexicalContext, LexicalContext> update
-    ) {
+    public TypeCheckerContext enterFreeVariableCapture() {
+        // TODO: do this properly (should not affect environment, possibly other bits of lexical context)
         return new TypeCheckerContext(
             this.namespaceTypeInfo,
-            P.stackUpdateTop(this.lexicalContextStack, update)
+            this.stack.plus(FreeVariableCapture.EMPTY)
         );
     }
 
+    public TypeCheckerResult<FreeVariables> exitFreeVariableCapture() {
+        if (this.stack.get(0) instanceof FreeVariableCapture freeVariableCapture) {
+            return this.exit().toResult(freeVariableCapture.freeVariables.build());
+        } else {
+            // TODO: better error
+            throw new RuntimeException("top of stack is not free variable capture");
+        }
+    }
+
     public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name) {
         var variable = Optional.ofNullable(environment().get(name));
 
@@ -160,16 +223,16 @@ public class TypeCheckerContext {
     }
 
     private TypeCheckerContext withVariable(String name) {
-        var newLexicalContextStack = this.lexicalContextStack;
+        var newLexicalContextStack = this.stack;
 
-        for (var index = 0; index < this.lexicalContextStack.size(); index++) {
-            var lexicalContext = this.lexicalContextStack.get(index);
-            if (lexicalContext.boundVariables.contains(name)) {
+        for (var index = 0; index < this.stack.size(); index++) {
+            var stackElement = this.stack.get(index);
+            if (stackElement.isBoundInElement(name)) {
                 break;
             } else{
                 newLexicalContextStack = newLexicalContextStack.with(
                     index,
-                    lexicalContext.withFreeVariable(name)
+                    stackElement.withFreeVariable(name)
                 );
             }
         }
@@ -274,7 +337,7 @@ public class TypeCheckerContext {
 
         return new TypeCheckerContext(
             this.namespaceTypeInfo.plus(namespaceName, namespaceTypeInfo),
-            this.lexicalContextStack
+            this.stack
         );
     }