Introduce TypedFunctionSignatureNode

4 files changed, 117 insertions, 110 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 cdaacf27..e03018c6 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -4,57 +4,20 @@ 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.*;
-
-import java.util.Optional;
+import org.shedlang.compiler.types.Effect;
+import org.shedlang.compiler.types.Type;
 
 public record TypedFunctionDefinitionNode(
     String name,
-    TypedParamsNode params,
-    Optional<TypedTypeLevelExpressionNode<Effect>> effect,
-    TypedTypeLevelExpressionNode<Type> returnType,
+    TypedFunctionSignatureNode signature,
     FreeVariables freeVariables,
     TypedBlockNode body,
-    FunctionType type,
     Source source
 ) implements TypedNamespaceStatementNode {
-    public TypedFunctionDefinitionNode(
-        String name,
-        TypedParamsNode params,
-        Optional<TypedTypeLevelExpressionNode<Effect>> effect,
-        TypedTypeLevelExpressionNode<Type> returnType,
-        FreeVariables freeVariables,
-        TypedBlockNode body,
-        Source source
-    ) {
-        this(
-            name,
-            params,
-            effect,
-            returnType,
-            freeVariables,
-            body,
-            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()
-            ),
-            source
-        );
-    }
-
     public static Builder builder() {
         return new Builder(
             "f",
-            Empty.vector(),
-            Empty.vector(),
-            Optional.empty(),
-            new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
+            TypedFunctionSignatureNode.builder(),
             FreeVariables.builder(),
             Empty.vector(),
             Empty.vector()
@@ -63,10 +26,7 @@ public record TypedFunctionDefinitionNode(
 
     public record Builder(
         String name,
-        PVector<TypedParamNode> positionalParams,
-        PVector<TypedParamNode> namedParams,
-        Optional<TypedTypeLevelExpressionNode<Effect>> effect,
-        TypedTypeLevelExpressionNode<Type> returnType,
+        TypedFunctionSignatureNode.Builder signature,
         FreeVariables.Builder freeVariables,
         PVector<String> boundVariableNames,
         PVector<TypedFunctionStatementNode> body
@@ -74,9 +34,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode build() {
             return new TypedFunctionDefinitionNode(
                 name,
-                new TypedParamsNode(positionalParams, namedParams, Sources.NULL),
-                effect,
-                returnType,
+                signature.build(),
                 freeVariables.build(),
                 new TypedBlockNode(body, boundVariableNames, Sources.NULL),
                 Sources.NULL
@@ -86,10 +44,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder withName(String name) {
             return new TypedFunctionDefinitionNode.Builder(
                 name,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body
@@ -99,10 +54,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams.plus(new TypedParamNode(name, type, Sources.NULL)),
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature.addPositionalParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body
@@ -112,10 +64,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams.plus(new TypedParamNode(name, type, Sources.NULL)),
-                this.effect,
-                this.returnType,
+                this.signature.addNamedParam(name, type),
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body
@@ -125,10 +74,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams,
-                Optional.of(effect),
-                this.returnType,
+                this.signature.withEffect(effect),
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body
@@ -138,10 +84,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                returnType,
+                this.signature.withReturnType(returnType),
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body
@@ -151,10 +94,7 @@ public record TypedFunctionDefinitionNode(
         public Builder addFreeVariable(String freeVariable) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.freeVariables.plus(freeVariable, Sources.NULL),
                 this.boundVariableNames,
                 this.body
@@ -164,10 +104,7 @@ public record TypedFunctionDefinitionNode(
         public Builder addBoundVariable(String name) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.freeVariables,
                 this.boundVariableNames.plus(name),
                 this.body
@@ -177,10 +114,7 @@ public record TypedFunctionDefinitionNode(
         public TypedFunctionDefinitionNode.Builder addStatement(TypedFunctionStatementNode statement) {
             return new TypedFunctionDefinitionNode.Builder(
                 this.name,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.freeVariables,
                 this.boundVariableNames,
                 this.body.plus(statement)
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
new file mode 100644
index 00000000..ff7cfb9a
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
@@ -0,0 +1,91 @@
+package org.shedlang.compiler.ast.typed;
+
+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.*;
+
+import java.util.Optional;
+
+public record TypedFunctionSignatureNode(
+    TypedParamsNode params,
+    Optional<TypedTypeLevelExpressionNode<Effect>> effect,
+    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(),
+            Empty.vector(),
+            Optional.empty(),
+            new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL)
+        );
+    }
+
+    public record Builder(
+        PVector<TypedParamNode> positionalParams,
+        PVector<TypedParamNode> namedParams,
+        Optional<TypedTypeLevelExpressionNode<Effect>> effect,
+        TypedTypeLevelExpressionNode<Type> returnType
+    ) {
+        public TypedFunctionSignatureNode build() {
+            return new TypedFunctionSignatureNode(
+                new TypedParamsNode(positionalParams, namedParams, Sources.NULL),
+                effect,
+                returnType,
+                Sources.NULL
+            );
+        }
+
+        public Builder addPositionalParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new Builder(
+                this.positionalParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.namedParams,
+                this.effect,
+                this.returnType
+            );
+        }
+
+        public Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams.plus(new TypedParamNode(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType
+            );
+        }
+
+        public Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                Optional.of(effect),
+                this.returnType
+            );
+        }
+
+        public Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                returnType
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 130b5fdf..33d84149 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -189,14 +189,14 @@ public class StackCompiler {
 
         var c = context;
 
-        var positionalParams = node.params().positional().stream()
+        var positionalParams = node.signature().params().positional().stream()
             .map(param -> new DefineFunctionInstruction.Param(
                 c.shedVariableToLocal(param.name()),
                 param.source()
             ))
             .toList();
 
-        var namedParams = node.params().named().stream()
+        var namedParams = node.signature().params().named().stream()
             .map(param -> new DefineFunctionInstruction.Param(
                 c.shedVariableToLocal(param.name()),
                 param.source()
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 9cd0d474..3dd5e138 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -495,13 +495,12 @@ public class TypeChecker {
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedSignature = context.getLexicalState(TypedFunctionSignature.class, untyped.name());
+        var typedSignature = context.getLexicalState(TypedFunctionSignatureNode.class, untyped.name());
         var functionType = typedSignature.type();
-        var typedParams = typedSignature.params;
 
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
 
-        for (var param : typedParams.all()) {
+        for (var param : typedSignature.params().all()) {
             bodyEnvironment = bodyEnvironment.add(
                 param.name(),
                 Variable.defined(param.type().value(), VariableType.PARAM),
@@ -556,9 +555,7 @@ public class TypeChecker {
 
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
-            typedParams,
-            typedSignature.effect,
-            typedSignature.returnType,
+            typedSignature,
             freeVariables,
             typedBody,
             untyped.source()
@@ -590,27 +587,7 @@ public class TypeChecker {
         return context.toResult(typed);
     }
 
-    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()
-            );
-        }
-    }
-
-    static TypedFunctionSignature typeCheckFunctionSignature(
+    static TypedFunctionSignatureNode typeCheckFunctionSignature(
         UntypedFunctionSignatureNode untyped,
         TypeCheckerContext context
     ) {
@@ -625,7 +602,12 @@ public class TypeChecker {
             context
         );
 
-        return new TypedFunctionSignature(typedParams, typedEffect, typedReturnType);
+        return new TypedFunctionSignatureNode(
+            typedParams,
+            typedEffect,
+            typedReturnType,
+            untyped.source()
+        );
     }
 
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(