Store state of function definition type checking in context

3 files changed, 61 insertions, 25 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
index a1b121b0..4e1e46ec 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
@@ -8,6 +8,6 @@ public interface TypeCheckSteps<T> {
     TypeCheckerResult<T> define(TypeCheckerContext context);
 
     default TypeCheckerResult<T> typeCheckAllPhases(TypeCheckerContext context) {
-        return define(context);
+        return define(declareFunctions(context));
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index fadd1c7c..a8ef4fd1 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -396,6 +396,13 @@ public class TypeChecker {
         };
     }
 
+    private record TypeCheckFunctionDefinitionState(
+        FunctionType functionType,
+        TypedParamsNode typedParams,
+        Optional<TypedTypeLevelExpressionNode<Effect>> typedEffect,
+        TypedTypeLevelExpressionNode<Type> typedReturnType
+    ) {}
+
     private static TypeCheckerContext declareFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
@@ -423,23 +430,23 @@ public class TypeChecker {
             typedReturnType.value()
         );
 
-        return context.declare(untyped.name(), functionType, untyped.source());
+        return context
+            .declare(untyped.name(), functionType, untyped.source())
+            .addLexicalState(untyped.name(), new TypeCheckFunctionDefinitionState(
+                functionType,
+                typedParams,
+                typedEffect,
+                typedReturnType
+            ));
     }
 
     private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedParams = typeCheckParams(untyped.params(), context);
-
-        Optional<TypedTypeLevelExpressionNode<Effect>> typedEffect;
-        if (untyped.effect().isPresent()) {
-            typedEffect = Optional.of(typeCheckEffect(untyped.effect().get(), context));
-        } else {
-            typedEffect = Optional.empty();
-        }
-
-        var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
+        var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
+        var functionType = state.functionType;
+        var typedParams = state.typedParams;
 
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
         // TODO: check for duplicate param names
@@ -458,8 +465,8 @@ public class TypeChecker {
 
         context = context.enterFunction(
             bodyEnvironment.build(),
-            typedEffect.map(e -> e.value()).orElse(Effects.NONE),
-            typedReturnType.value()
+            functionType.effect(),
+            functionType.returnType()
         );
 
         var bodyResult = typeCheckBlock(untyped.body(), context);
@@ -473,18 +480,18 @@ public class TypeChecker {
 
         switch (typedBody.behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                if (!isSubType(blockReturn.type(), typedReturnType.value(), context)) {
+                if (!isSubType(blockReturn.type(), functionType.returnType(), context)) {
                     throw new UnexpectedTypeError(
-                        TypeSets.type(typedReturnType.value()),
+                        TypeSets.type(functionType.returnType()),
                         blockReturn.type(),
                         untyped.source()
                     );
                 }
             }
             case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                if (!isSubType(Types.UNIT, typedReturnType.value(), context)) {
+                if (!isSubType(Types.UNIT, functionType.returnType(), context)) {
                     throw new UnexpectedTypeError(
-                        TypeSets.type(typedReturnType.value()),
+                        TypeSets.type(functionType.returnType()),
                         Types.UNIT,
                         untyped.source()
                     );
@@ -497,8 +504,8 @@ public class TypeChecker {
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
-            typedEffect,
-            typedReturnType,
+            state.typedEffect,
+            state.typedReturnType,
             freeVariables,
             typedBody,
             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 348007cf..5cb2d6c6 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -25,7 +25,8 @@ public class TypeCheckerContext {
             Effects.NONE,
             Optional.empty(),
             Set.of(),
-            FreeVariables.builder()
+            FreeVariables.builder(),
+            Empty.map()
         ))
     );
 
@@ -48,19 +49,22 @@ public class TypeCheckerContext {
         private final Optional<Type> returnType;
         private final Set<String> boundVariables;
         private final FreeVariables.Builder freeVariables;
+        private final PMap<String, Object> state;
 
         public LexicalContext(
             PMap<String, Variable> environment,
             Effect effect,
             Optional<Type> returnType,
             Set<String> boundVariables,
-            FreeVariables.Builder freeVariables
+            FreeVariables.Builder freeVariables,
+            PMap<String, Object> state
         ) {
             this.environment = environment;
             this.effect = effect;
             this.returnType = returnType;
             this.boundVariables = boundVariables;
             this.freeVariables = freeVariables;
+            this.state = state;
         }
 
         private LexicalContext withFreeVariable(String name) {
@@ -69,7 +73,8 @@ public class TypeCheckerContext {
                 this.effect,
                 this.returnType,
                 this.boundVariables,
-                this.freeVariables.plus(name)
+                this.freeVariables.plus(name),
+                this.state
             );
         }
     }
@@ -125,7 +130,8 @@ public class TypeCheckerContext {
                 newEffect,
                 newReturnType,
                 newEnvironment.keySet(),
-                FreeVariables.builder()
+                FreeVariables.builder(),
+                Empty.map()
             ))
         );
     }
@@ -175,6 +181,28 @@ public class TypeCheckerContext {
         return Optional.ofNullable(environment().get(name));
     }
 
+    public TypeCheckerContext addLexicalState(String name, Object state) {
+        return this.updateLexicalContext(
+            lexicalContext -> new LexicalContext(
+                lexicalContext.environment,
+                lexicalContext.effect,
+                lexicalContext.returnType,
+                lexicalContext.boundVariables,
+                lexicalContext.freeVariables,
+                lexicalContext.state.plus(name, state)
+            )
+        );
+    }
+
+    public <T> T getLexicalState(Class<T> clazz, String name) {
+        var state = this.currentLexicalContext().state.get(name);
+        if (state == null) {
+            // TODO: better error
+            throw new RuntimeException("Missing state");
+        }
+        return clazz.cast(state);
+    }
+
     public TypeCheckerContext declare(
         String name,
         Type type,
@@ -210,7 +238,8 @@ public class TypeCheckerContext {
                         lexicalContext.effect,
                         lexicalContext.returnType,
                         lexicalContext.boundVariables,
-                        lexicalContext.freeVariables
+                        lexicalContext.freeVariables,
+                        lexicalContext.state
                     )
                 );
             }