Remove type from typed function signature node

4 files changed, 78 insertions, 51 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index c98a7b1e..889b4fee 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -4,13 +4,15 @@ import org.pcollections.Empty;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
-import org.shedlang.compiler.types.Effect;
-import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.*;
+
+import java.util.List;
 
 public record TypedFunctionDefinitionNode(
     String name,
     TypedFunctionSignatureNode signature,
     TypedFunctionBodyNode body,
+    FunctionType type,
     Source source
 ) implements TypedNamespaceStatementNode {
     public static Builder builder() {
@@ -19,7 +21,8 @@ public record TypedFunctionDefinitionNode(
             TypedFunctionSignatureNode.builder(),
             FreeVariables.builder(),
             Empty.vector(),
-            Empty.vector()
+            Empty.vector(),
+            new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT)
         );
     }
 
@@ -28,7 +31,8 @@ public record TypedFunctionDefinitionNode(
         TypedFunctionSignatureNode.Builder signature,
         FreeVariables.Builder freeVariables,
         PVector<String> boundVariableNames,
-        PVector<TypedFunctionStatementNode> body
+        PVector<TypedFunctionStatementNode> body,
+        FunctionType type
     ) {
         public TypedFunctionDefinitionNode build() {
             return new TypedFunctionDefinitionNode(
@@ -39,6 +43,7 @@ public record TypedFunctionDefinitionNode(
                     new TypedBlockNode(body, boundVariableNames, Sources.NULL),
                     Sources.NULL
                 ),
+                type,
                 Sources.NULL
             );
         }
@@ -49,7 +54,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature,
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -59,7 +65,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature.addPositionalParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             ).addBoundVariable(name);
         }
 
@@ -69,7 +76,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature.addNamedParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             ).addBoundVariable(name);
         }
 
@@ -79,7 +87,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature.withEffect(effect),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -89,7 +98,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature.withReturnType(returnType),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -99,7 +109,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature,
                 this.freeVariables.plus(freeVariable, Sources.NULL),
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -109,7 +120,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature,
                 this.freeVariables,
                 this.boundVariableNames.plus(name),
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -119,7 +131,8 @@ public record TypedFunctionDefinitionNode(
                 this.signature,
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body.plus(statement)
+                this.body.plus(statement),
+                this.type
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
index 1bc46fdb..401f4009 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
@@ -4,25 +4,23 @@ import org.pcollections.Empty;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
-import org.shedlang.compiler.types.Effect;
-import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.types.*;
+
+import java.util.List;
 
 public record TypedFunctionExpressionNode(
     TypedFunctionSignatureNode signature,
     TypedFunctionBodyNode body,
+    FunctionType type,
     Source source
 ) implements TypedExpressionNode {
-    @Override
-    public Type type() {
-        return signature.type();
-    }
-
     public static Builder builder() {
         return new Builder(
             TypedFunctionSignatureNode.builder(),
             FreeVariables.builder(),
             Empty.vector(),
-            Empty.vector()
+            Empty.vector(),
+            new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT)
         );
     }
 
@@ -30,7 +28,8 @@ public record TypedFunctionExpressionNode(
         TypedFunctionSignatureNode.Builder signature,
         FreeVariables.Builder freeVariables,
         PVector<String> boundVariableNames,
-        PVector<TypedFunctionStatementNode> body
+        PVector<TypedFunctionStatementNode> body,
+        FunctionType type
     ) {
         public TypedFunctionExpressionNode build() {
             return new TypedFunctionExpressionNode(
@@ -40,6 +39,7 @@ public record TypedFunctionExpressionNode(
                     new TypedBlockNode(body, boundVariableNames, Sources.NULL),
                     Sources.NULL
                 ),
+                type,
                 Sources.NULL
             );
         }
@@ -49,7 +49,8 @@ public record TypedFunctionExpressionNode(
                 this.signature.addPositionalParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             ).addBoundVariable(name);
         }
 
@@ -58,7 +59,8 @@ public record TypedFunctionExpressionNode(
                 this.signature.addNamedParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             ).addBoundVariable(name);
         }
 
@@ -67,7 +69,8 @@ public record TypedFunctionExpressionNode(
                 this.signature.withEffect(effect),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -76,7 +79,8 @@ public record TypedFunctionExpressionNode(
                 this.signature.withReturnType(returnType),
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -85,7 +89,8 @@ public record TypedFunctionExpressionNode(
                 this.signature,
                 this.freeVariables.plus(freeVariable, Sources.NULL),
                 this.boundVariableNames,
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -94,7 +99,8 @@ public record TypedFunctionExpressionNode(
                 this.signature,
                 this.freeVariables,
                 this.boundVariableNames.plus(name),
-                this.body
+                this.body,
+                this.type
             );
         }
 
@@ -103,7 +109,8 @@ public record TypedFunctionExpressionNode(
                 this.signature,
                 this.freeVariables,
                 this.boundVariableNames,
-                this.body.plus(statement)
+                this.body.plus(statement),
+                this.type
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
index ff7cfb9a..7c83dd1c 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
@@ -14,20 +14,6 @@ public record TypedFunctionSignatureNode(
     TypedTypeLevelExpressionNode<Type> returnType,
     Source source
 ) implements TypedNode {
-    public FunctionType type() {
-        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()
-        );
-    }
-
-
     public static Builder builder() {
         return new Builder(
             Empty.vector(),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 5ad434eb..530b1d7c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -469,8 +469,6 @@ public class TypeChecker {
         TypedFunctionSignatureNode typedSignature,
         TypeCheckerContext context
     ) {
-        var functionType = typedSignature.type();
-
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
 
         for (var param : typedSignature.params().all()) {
@@ -488,10 +486,12 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
+        var returnType = typedSignature.returnType().value();
+
         context = context.enterFunction(
             bodyEnvironment.build(),
-            functionType.effect(),
-            functionType.returnType()
+            typedSignature.effect().map(e -> e.value()).orElse(Effects.NONE),
+            returnType
         );
 
         var bodyResult = typeCheckBlock(untyped, bodyEnvironment.names(), context);
@@ -505,18 +505,18 @@ public class TypeChecker {
 
         switch (typedBody.behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                if (!isSubType(blockReturn.type(), functionType.returnType(), context)) {
+                if (!isSubType(blockReturn.type(), returnType, context)) {
                     throw new UnexpectedTypeError(
-                        TypeSets.type(functionType.returnType()),
+                        TypeSets.type(returnType),
                         blockReturn.type(),
                         untyped.source()
                     );
                 }
             }
             case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                if (!isSubType(Types.UNIT, functionType.returnType(), context)) {
+                if (!isSubType(Types.UNIT, returnType, context)) {
                     throw new UnexpectedTypeError(
-                        TypeSets.type(functionType.returnType()),
+                        TypeSets.type(returnType),
                         Types.UNIT,
                         untyped.source()
                     );
@@ -557,8 +557,10 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
+        var type = functionSignatureToFunctionType(typedSignature);
+
         return context
-            .declareVariableType(untyped.name(), typedSignature.type(), untyped.source())
+            .declareVariableType(untyped.name(), type, untyped.source())
             .addLexicalState(untyped.name(), typedSignature);
     }
 
@@ -568,6 +570,8 @@ public class TypeChecker {
     ) {
         var typedSignature = context.getLexicalState(TypedFunctionSignatureNode.class, untyped.name());
 
+        var type = functionSignatureToFunctionType(typedSignature);
+
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
             typedSignature,
@@ -580,6 +584,7 @@ public class TypeChecker {
             untyped.name(),
             typedSignature,
             typedBody,
+            type,
             untyped.source()
         );
 
@@ -592,6 +597,8 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
+        var type = functionSignatureToFunctionType(typedSignature);
+
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
             typedSignature,
@@ -603,6 +610,7 @@ public class TypeChecker {
         var typed = new TypedFunctionExpressionNode(
             typedSignature,
             typedBody,
+            type,
             untyped.source()
         );
 
@@ -657,6 +665,19 @@ public class TypeChecker {
         );
     }
 
+    private static FunctionType functionSignatureToFunctionType(TypedFunctionSignatureNode typedSignature) {
+        return new FunctionType(
+            typedSignature.params().positional().stream()
+                .map(param -> param.type().value())
+                .toList(),
+            typedSignature.params().named().stream()
+                .map(param -> new NamedParam(param.name(), param.type().value()))
+                .toList(),
+            typedSignature.effect().map(e -> e.value()).orElse(Effects.NONE),
+            typedSignature.returnType().value()
+        );
+    }
+
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
         TypeCheckerContext context