Extract function signature types

1 files changed, 58 insertions, 37 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 0dd93a38..f3cc142c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -480,53 +480,28 @@ 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
     ) {
-        var typedParams = typeCheckParams(untyped.params(), context);
-
-        var typedEffect = untyped.effect()
-            .map(untypedEffect -> typeCheckEffect(untypedEffect, context));
-
-        var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
-
-        // TODO: remove duplication with TypedFunctionDefinitionNode constructor.
-        var functionType = new FunctionType(
-            typedParams.positional().stream()
-                .map(param -> param.type().value())
-                .toList(),
-            typedParams.named().stream()
-                .map(param -> new NamedParam(param.name(), param.type().value()))
-                .toList(),
-            typedEffect.map(e -> e.value()).orElse(Effects.NONE),
-            typedReturnType.value()
-        );
+        var typedSignature = typeCheckFunctionSignature(new UntypedFunctionSignature(
+            untyped.params(),
+            untyped.effect(),
+            untyped.returnType()
+        ), context);
 
         return context
-            .define(untyped.name(), functionType, untyped.source())
-            .addLexicalState(untyped.name(), new TypeCheckFunctionDefinitionState(
-                functionType,
-                typedParams,
-                typedEffect,
-                typedReturnType
-            ));
+            .define(untyped.name(), typedSignature.type(), untyped.source())
+            .addLexicalState(untyped.name(), typedSignature);
     }
 
     private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
-        var functionType = state.functionType;
-        var typedParams = state.typedParams;
+        var typedSignature = context.getLexicalState(TypedFunctionSignature.class, untyped.name());
+        var functionType = typedSignature.type();
+        var typedParams = typedSignature.params;
 
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
         // TODO: check for duplicate param names
@@ -587,8 +562,8 @@ public class TypeChecker {
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
-            state.typedEffect,
-            state.typedReturnType,
+            typedSignature.effect,
+            typedSignature.returnType,
             freeVariables,
             typedBody,
             untyped.source()
@@ -620,6 +595,52 @@ public class TypeChecker {
         return context.toResult(typed);
     }
 
+    private record UntypedFunctionSignature(
+        UntypedParamsNode params,
+        Optional<UntypedTypeLevelExpressionNode> effect,
+        UntypedTypeLevelExpressionNode returnType
+    ) {
+
+    }
+
+    private record TypedFunctionSignature(
+        TypedParamsNode params,
+        Optional<TypedTypeLevelExpressionNode<Effect>> effect,
+        TypedTypeLevelExpressionNode<Type> returnType
+    ) {
+        private FunctionType type() {
+            // TODO: remove duplication with TypedFunctionDefinitionNode constructor.
+            return new FunctionType(
+                params.positional().stream()
+                    .map(param -> param.type().value())
+                    .toList(),
+                params.named().stream()
+                    .map(param -> new NamedParam(param.name(), param.type().value()))
+                    .toList(),
+                effect.map(e -> e.value()).orElse(Effects.NONE),
+                returnType.value()
+            );
+        }
+    }
+
+    private static TypedFunctionSignature typeCheckFunctionSignature(
+        UntypedFunctionSignature untyped,
+        TypeCheckerContext context
+    ) {
+        var typedParams = typeCheckParams(untyped.params(), context);
+
+        var typedEffect = untyped.effect()
+            .map(untypedEffect -> typeCheckEffect(untypedEffect, context));
+
+        var typedReturnType = typeCheckTypeLevelExpression(
+            TypeSets.META_TYPE,
+            untyped.returnType(),
+            context
+        );
+
+        return new TypedFunctionSignature(typedParams, typedEffect, typedReturnType);
+    }
+
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
         TypeCheckerContext context