Store stack of lexical contexts

3 files changed, 139 insertions, 92 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 8758ad6a..c4571a90 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -318,13 +318,13 @@ public class TypeChecker {
         checkExpressionType(typedCondition, Types.BOOL, context);
 
         // TODO: get environment and pass to enterBlock()
-        var bodyContext = context.enterBlock(P.map());
+        context = context.enterBlock(P.map());
 
-        var typedBodyResult = typeCheckBlock(conditionalBranch.body(), bodyContext);
+        var typedBodyResult = typeCheckBlock(conditionalBranch.body(), context);
         var typedBody = typedBodyResult.value();
-        bodyContext = typedBodyResult.context();
+        context = typedBodyResult.context();
 
-        context = context.exit(bodyContext);
+        context = context.exit();
 
         return context.toResult(new TypedConditionalBranchNode(
             typedCondition,
@@ -410,15 +410,20 @@ public class TypeChecker {
             bodyEnvironment = declareVariables(statement, bodyEnvironment);
         }
 
-        var bodyContext = context.enterFunction(
+        context = context.enterFunction(
             bodyEnvironment.build(),
             typedEffect.map(e -> e.value()).orElse(Effects.NONE),
             typedReturnType.value()
         );
 
-        var bodyResult = typeCheckBlock(untyped.body(), bodyContext);
+        var bodyResult = typeCheckBlock(untyped.body(), context);
         var typedBody = bodyResult.value();
-        bodyContext = bodyResult.context();
+        context = bodyResult.context();
+
+        // TODO: get free variables on exit?
+        var freeVariables = context.freeVariables();
+
+        context = context.exit();
 
         switch (typedBody.behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
@@ -448,7 +453,7 @@ public class TypeChecker {
             typedParams,
             typedEffect,
             typedReturnType,
-            bodyContext.freeVariables(),
+            freeVariables,
             typedBody,
             untyped.source()
         );
@@ -509,13 +514,13 @@ public class TypeChecker {
         context = typedConditionalBranchesResult.context();
 
         // TODO: get environment and pass to enterBlock()
-        var bodyContext = context.enterBlock(P.map());
+        context = context.enterBlock(P.map());
 
-        var typedElseBodyResult = typeCheckBlock(untyped.elseBody(), bodyContext);
+        var typedElseBodyResult = typeCheckBlock(untyped.elseBody(), context);
         var typedElseBody = typedElseBodyResult.value();
-        bodyContext = typedElseBodyResult.context();
+        context = typedElseBodyResult.context();
 
-        context = context.exit(bodyContext);
+        context = context.exit();
 
         var branchTypes = Stream.concat(
             typedConditionalBranches.stream()
@@ -553,14 +558,14 @@ public class TypeChecker {
             namespaceEnvironment = declareVariables(statement, namespaceEnvironment);
         }
 
-        var namespaceContext = context.enterNamespace(namespaceEnvironment.build());
+        context = context.enterNamespace(namespaceEnvironment.build());
 
         var typedBody = new ArrayList<TypedNamespaceStatementNode>();
 
         for (var untypedStatement : untyped.body()) {
-            var typedStatementResult = typeCheckNamespaceStatement(untypedStatement, namespaceContext);
+            var typedStatementResult = typeCheckNamespaceStatement(untypedStatement, context);
             var typedStatement = typedStatementResult.value();
-            namespaceContext = typedStatementResult.context();
+            context = typedStatementResult.context();
 
             typedBody.add(typedStatement);
         }
@@ -583,7 +588,7 @@ public class TypeChecker {
         );
 
         return context
-            .exit(namespaceContext)
+            .exit()
             .addNamespaceFields(typed.namespaceName(), fieldTypes)
             .toResult(typed);
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 2a5423db..99357563 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -2,35 +2,33 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
+import org.pcollections.PStack;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.CannotDefineVariableDeclaredInOuterScopeError;
 import org.shedlang.compiler.typechecker.errors.UnknownVariableError;
 import org.shedlang.compiler.typechecker.errors.VariableIsAlreadyDefinedError;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.*;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 
 public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
-        Effects.NONE,
-        Optional.empty(),
-        Set.of(),
-        FreeVariables.builder(),
-        Empty.map()
+        P.stack(new LexicalContext(
+            Empty.map(),
+            Effects.NONE,
+            Optional.empty(),
+            Set.of(),
+            FreeVariables.builder()
+        ))
     );
 
-    private final PMap<String, Variable> environment;
-    private final Effect effect;
-    private final Optional<Type> returnType;
-    private final Set<String> boundVariables;
-    private final FreeVariables.Builder freeVariables;
-    private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
-
     public sealed interface Variable {
         record Undefined() implements Variable {}
         record Defined(Type type) implements Variable {}
@@ -44,20 +42,51 @@ public class TypeCheckerContext {
         return new Variable.Undefined();
     }
 
+    public static class LexicalContext {
+        private final PMap<String, Variable> environment;
+        private final Effect effect;
+        private final Optional<Type> returnType;
+        private final Set<String> boundVariables;
+        private final FreeVariables.Builder freeVariables;
+
+        public LexicalContext(
+            PMap<String, Variable> environment,
+            Effect effect,
+            Optional<Type> returnType,
+            Set<String> boundVariables,
+            FreeVariables.Builder freeVariables
+        ) {
+            this.environment = environment;
+            this.effect = effect;
+            this.returnType = returnType;
+            this.boundVariables = boundVariables;
+            this.freeVariables = freeVariables;
+        }
+
+        private LexicalContext withFreeVariable(String name) {
+            return new LexicalContext(
+                this.environment,
+                this.effect,
+                this.returnType,
+                this.boundVariables,
+                this.freeVariables.plus(name)
+            );
+        }
+    }
+
+    private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
+    private final PStack<LexicalContext> lexicalContextStack;
+
     public TypeCheckerContext(
-        PMap<String, Variable> environment,
-        Effect effect,
-        Optional<Type> returnType,
-        Set<String> boundVariables,
-        FreeVariables.Builder freeVariables,
-        PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo
+        PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo,
+        PStack<LexicalContext> lexicalContextStack
     ) {
-        this.environment = environment;
-        this.effect = effect;
-        this.returnType = returnType;
-        this.boundVariables = boundVariables;
-        this.freeVariables = freeVariables;
         this.namespaceTypeInfo = namespaceTypeInfo;
+        this.lexicalContextStack = lexicalContextStack;
+    }
+
+    private LexicalContext currentLexicalContext() {
+        return this.lexicalContextStack.get(0);
     }
 
     public TypeCheckerContext enterNamespaceWithDefinedVariables(PMap<String, Type> variables) {
@@ -69,7 +98,7 @@ public class TypeCheckerContext {
     }
 
     public TypeCheckerContext enterNamespace(PMap<String, Variable> environment) {
-        return enter(environment, this.effect, this.returnType);
+        return enter(environment, Effects.NONE, Optional.empty());
     }
 
     public TypeCheckerContext enterFunction(
@@ -81,50 +110,58 @@ public class TypeCheckerContext {
     }
 
     public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
-        return enter(environment, this.effect, this.returnType);
+        return enter(environment, this.effect(), this.returnType());
     }
 
     private TypeCheckerContext enter(
-        PMap<String, Variable> environment,
-        Effect effect,
-        Optional<Type> returnType
+        PMap<String, Variable> newEnvironment,
+        Effect newEffect,
+        Optional<Type> newReturnType
     ) {
         return new TypeCheckerContext(
-            this.environment.plusAll(environment),
-            effect,
-            returnType,
-            environment.keySet(),
-            FreeVariables.builder(),
-            this.namespaceTypeInfo
+            this.namespaceTypeInfo,
+            this.lexicalContextStack.plus(new LexicalContext(
+                this.environment().plusAll(newEnvironment),
+                newEffect,
+                newReturnType,
+                newEnvironment.keySet(),
+                FreeVariables.builder()
+            ))
         );
     }
 
-    public TypeCheckerContext exit(TypeCheckerContext childContext) {
-        // TODO: we need to track free variables, but it's easy to discard the
-        // child context without calling exit. What's the right way of handling
-        // this?  Does using a callback work when we have type phases?
-        var context = this;
-
-        for (var freeVariable : childContext.freeVariables.build()) {
-            context = context.withVariable(freeVariable.name());
-        }
+    public TypeCheckerContext exit() {
+        return new TypeCheckerContext(
+            this.namespaceTypeInfo,
+            this.lexicalContextStack.minus(0)
+        );
+    }
 
-        return context;
+    private TypeCheckerContext updateLexicalContext(
+        Function<LexicalContext, LexicalContext> update
+    ) {
+        return new TypeCheckerContext(
+            this.namespaceTypeInfo,
+            P.stackUpdateTop(this.lexicalContextStack, update)
+        );
     }
 
     public TypeCheckerResult<Optional<Variable>> runtimeLookup(String name) {
-        var variable = Optional.ofNullable(environment.get(name));
+        var variable = Optional.ofNullable(environment().get(name));
 
         var context = withVariable(name);
         return context.toResult(variable);
     }
 
     private TypeCheckerContext withVariable(String name) {
-        return boundVariables.contains(name) ? this : withFreeVariable(name);
+        // TODO: walk up stack
+        return boundVariables().contains(name)
+            ? this
+            : updateLexicalContext(lexicalContext -> lexicalContext.withFreeVariable(name));
     }
 
     public Optional<Variable> compileTimeLookup(String name) {
-        return Optional.ofNullable(environment.get(name));
+        return Optional.ofNullable(environment().get(name));
     }
 
     public TypeCheckerContext define(
@@ -132,7 +169,7 @@ public class TypeCheckerContext {
         Type type,
         Source source
     ) {
-        var variable = this.environment.get(name);
+        var variable = this.environment().get(name);
 
         if (variable == null) {
             throw new UnknownVariableError(name, source);
@@ -143,43 +180,41 @@ public class TypeCheckerContext {
                 throw new VariableIsAlreadyDefinedError(name, source);
             }
             case Variable.Undefined undefined -> {
-                if (!this.boundVariables.contains(name)) {
+                if (!this.boundVariables().contains(name)) {
                     throw new CannotDefineVariableDeclaredInOuterScopeError(name, source);
                 }
 
-                return new TypeCheckerContext(
-                    this.environment.plus(name, new Variable.Defined(type)),
-                    this.effect,
-                    this.returnType,
-                    this.boundVariables,
-                    this.freeVariables,
-                    this.namespaceTypeInfo
+                return this.updateLexicalContext(
+                    lexicalContext -> new LexicalContext(
+                        lexicalContext.environment.plus(name, new Variable.Defined(type)),
+                        lexicalContext.effect,
+                        lexicalContext.returnType,
+                        lexicalContext.boundVariables,
+                        lexicalContext.freeVariables
+                    )
                 );
             }
         }
     }
 
+    private PMap<String, Variable> environment() {
+        return currentLexicalContext().environment;
+    }
+
+    private Set<String> boundVariables() {
+        return currentLexicalContext().boundVariables;
+    }
+
     public Effect effect() {
-        return effect;
+        return currentLexicalContext().effect;
     }
 
     public Optional<Type> returnType() {
-        return returnType;
-    }
-    
-    private TypeCheckerContext withFreeVariable(String name) {
-        return new TypeCheckerContext(
-            environment,
-            this.effect,
-            this.returnType,
-            boundVariables,
-            freeVariables.plus(name),
-            this.namespaceTypeInfo
-        );
+        return currentLexicalContext().returnType;
     }
 
     public FreeVariables freeVariables() {
-        return freeVariables.build();
+        return currentLexicalContext().freeVariables.build();
     }
 
     public TypeCheckerContext addNamespaceFields(
@@ -189,12 +224,8 @@ public class TypeCheckerContext {
         var namespaceTypeInfo = new NamespaceTypeInfo(fields);
 
         return new TypeCheckerContext(
-            this.environment,
-            this.effect,
-            this.returnType,
-            this.boundVariables,
-            this.freeVariables,
-            this.namespaceTypeInfo.plus(namespaceName, namespaceTypeInfo)
+            this.namespaceTypeInfo.plus(namespaceName, namespaceTypeInfo),
+            this.lexicalContextStack
         );
     }
 
diff --git a/src/main/java/org/shedlang/compiler/util/P.java b/src/main/java/org/shedlang/compiler/util/P.java
index 5985a767..fc822614 100644
--- a/src/main/java/org/shedlang/compiler/util/P.java
+++ b/src/main/java/org/shedlang/compiler/util/P.java
@@ -4,6 +4,7 @@ import org.pcollections.*;
 
 import java.util.Arrays;
 import java.util.Map;
+import java.util.function.Function;
 
 public class P {
     private P() {
@@ -33,6 +34,16 @@ public class P {
     }
 
     @SafeVarargs
+    public static <T> PStack<T> stack(T... values) {
+        return ConsPStack.from(Arrays.asList(values));
+    }
+
+    public static <T> PStack<T> stackUpdateTop(PStack<T> stack, Function<T, T> update) {
+        var newElement = update.apply(stack.get(0));
+        return stack.with(0, newElement);
+    }
+
+    @SafeVarargs
     public static <T> PVector<T> vector(T... values) {
         return TreePVector.from(Arrays.asList(values));
     }