Introduce UntypedFunctionSignatureNode

5 files changed, 144 insertions, 117 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
index c81a4409..d9abdcdb 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -6,14 +6,10 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.Visibility;
 
-import java.util.Optional;
-
 public record UntypedFunctionDefinitionNode(
     String name,
     Visibility visibility,
-    UntypedParamsNode params,
-    Optional<UntypedTypeLevelExpressionNode> effect,
-    UntypedTypeLevelExpressionNode returnType,
+    UntypedFunctionSignatureNode signature,
     UntypedBlockNode body,
     Source source
 ) implements UntypedNamespaceStatementNode, UntypedVariableDeclarationNode {
@@ -21,10 +17,7 @@ public record UntypedFunctionDefinitionNode(
         return new Builder(
             "f",
             Visibility.PRIVATE,
-            Empty.vector(),
-            Empty.vector(),
-            Optional.empty(),
-            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL),
+            UntypedFunctionSignatureNode.builder(),
             Empty.vector()
         );
     }
@@ -32,19 +25,14 @@ public record UntypedFunctionDefinitionNode(
     public record Builder(
         String name,
         Visibility visibility,
-        PVector<UntypedParamNode> positionalParams,
-        PVector<UntypedParamNode> namedParams,
-        Optional<UntypedTypeLevelExpressionNode> effect,
-        UntypedTypeLevelExpressionNode returnType,
+        UntypedFunctionSignatureNode.Builder signature,
         PVector<UntypedFunctionStatementNode> body
     ) {
         public UntypedFunctionDefinitionNode build() {
             return new UntypedFunctionDefinitionNode(
                 name,
                 visibility,
-                new UntypedParamsNode(positionalParams, namedParams, Sources.NULL),
-                effect,
-                returnType,
+                signature.build(),
                 new UntypedBlockNode(body, Sources.NULL),
                 Sources.NULL
             );
@@ -54,10 +42,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 name,
                 this.visibility,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.body
             );
         }
@@ -66,10 +51,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 this.visibility,
-                this.positionalParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature.addPositionalParam(name, type),
                 this.body
             );
         }
@@ -78,10 +60,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 this.visibility,
-                this.positionalParams,
-                this.namedParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                this.effect,
-                this.returnType,
+                this.signature.addNamedParam(name, type),
                 this.body
             );
         }
@@ -90,10 +69,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 this.visibility,
-                this.positionalParams,
-                this.namedParams,
-                Optional.of(effect),
-                this.returnType,
+                this.signature.withEffect(effect),
                 this.body
             );
         }
@@ -102,10 +78,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 this.visibility,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                returnType,
+                this.signature.withReturnType(returnType),
                 this.body
             );
         }
@@ -114,10 +87,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 this.visibility,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature.addStatement(statement),
                 this.body.plus(statement)
             );
         }
@@ -126,10 +96,7 @@ public record UntypedFunctionDefinitionNode(
             return new Builder(
                 this.name,
                 visibility,
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.body
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
index bc931a5a..1ffddf69 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
@@ -5,37 +5,25 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 
-import java.util.Optional;
-
 public record UntypedFunctionExpressionNode(
-    UntypedParamsNode params,
-    Optional<UntypedTypeLevelExpressionNode> effect,
-    UntypedTypeLevelExpressionNode returnType,
+    UntypedFunctionSignatureNode signature,
     UntypedBlockNode body,
     Source source
 ) implements UntypedExpressionNode {
     public static Builder builder() {
         return new Builder(
-            Empty.vector(),
-            Empty.vector(),
-            Optional.empty(),
-            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL),
+            UntypedFunctionSignatureNode.builder(),
             Empty.vector()
         );
     }
 
     public record Builder(
-        PVector<UntypedParamNode> positionalParams,
-        PVector<UntypedParamNode> namedParams,
-        Optional<UntypedTypeLevelExpressionNode> effect,
-        UntypedTypeLevelExpressionNode returnType,
+        UntypedFunctionSignatureNode.Builder signature,
         PVector<UntypedFunctionStatementNode> body
     ) {
         public UntypedFunctionExpressionNode build() {
             return new UntypedFunctionExpressionNode(
-                new UntypedParamsNode(positionalParams, namedParams, Sources.NULL),
-                effect,
-                returnType,
+                signature.build(),
                 new UntypedBlockNode(body, Sources.NULL),
                 Sources.NULL
             );
@@ -43,50 +31,35 @@ public record UntypedFunctionExpressionNode(
 
         public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
             return new Builder(
-                this.positionalParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature.addPositionalParam(name, type),
                 this.body
             );
         }
 
         public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
             return new Builder(
-                this.positionalParams,
-                this.namedParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
-                this.effect,
-                this.returnType,
+                this.signature.addNamedParam(name, type),
                 this.body
             );
         }
 
         public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
             return new Builder(
-                this.positionalParams,
-                this.namedParams,
-                Optional.of(effect),
-                this.returnType,
+                this.signature.withEffect(effect),
                 this.body
             );
         }
 
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                returnType,
+                this.signature.withReturnType(returnType),
                 this.body
             );
         }
 
         public Builder addStatement(UntypedFunctionStatementNode statement) {
             return new Builder(
-                this.positionalParams,
-                this.namedParams,
-                this.effect,
-                this.returnType,
+                this.signature,
                 this.body.plus(statement)
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
new file mode 100644
index 00000000..179b4f1e
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
@@ -0,0 +1,92 @@
+package org.shedlang.compiler.ast.untyped;
+
+import org.pcollections.Empty;
+import org.pcollections.PVector;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
+
+import java.util.Optional;
+
+public record UntypedFunctionSignatureNode(
+    UntypedParamsNode params,
+    Optional<UntypedTypeLevelExpressionNode> effect,
+    UntypedTypeLevelExpressionNode returnType,
+    Source source
+) implements UntypedNode {
+    public static Builder builder() {
+        return new Builder(
+            Empty.vector(),
+            Empty.vector(),
+            Optional.empty(),
+            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL),
+            Empty.vector()
+        );
+    }
+
+    public record Builder(
+        PVector<UntypedParamNode> positionalParams,
+        PVector<UntypedParamNode> namedParams,
+        Optional<UntypedTypeLevelExpressionNode> effect,
+        UntypedTypeLevelExpressionNode returnType,
+        PVector<UntypedFunctionStatementNode> body
+    ) {
+        public UntypedFunctionSignatureNode build() {
+            return new UntypedFunctionSignatureNode(
+                new UntypedParamsNode(positionalParams, namedParams, Sources.NULL),
+                effect,
+                returnType,
+                Sources.NULL
+            );
+        }
+
+        public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
+            return new Builder(
+                this.positionalParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
+        public Builder addNamedParam(String name, UntypedTypeLevelExpressionNode type) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams.plus(new UntypedParamNode(name, type, Sources.NULL)),
+                this.effect,
+                this.returnType,
+                this.body
+            );
+        }
+
+        public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                Optional.of(effect),
+                this.returnType,
+                this.body
+            );
+        }
+
+        public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                returnType,
+                this.body
+            );
+        }
+
+        public Builder addStatement(UntypedFunctionStatementNode statement) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.body.plus(statement)
+            );
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index fefd6d62..ac7f5bb1 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -289,14 +289,7 @@ public class Parser {
         tokens.skip(TokenType.KEYWORD_FUN);
         var name = tokens.nextValue(TokenType.IDENTIFIER);
 
-        tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
-        var params = parseParams(tokens);
-        tokens.skip(TokenType.SYMBOL_PAREN_CLOSE);
-
-        var effect = parseFunctionTypeEffect(tokens);
-
-        tokens.skip(TokenType.SYMBOL_THIN_ARROW);
-        var returnType = parseTopLevelTypeLevelExpression(tokens);
+        var signature = parseFunctionSignature(tokens);
 
         var body = parseBlock(tokens);
 
@@ -306,9 +299,7 @@ public class Parser {
         return new UntypedFunctionDefinitionNode(
             name,
             visibility,
-            params,
-            effect,
-            returnType,
+            signature,
             body,
             source
         );
@@ -321,14 +312,7 @@ public class Parser {
 
         tokens.skip(TokenType.KEYWORD_FN);
 
-        tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
-        var params = parseParams(tokens);
-        tokens.skip(TokenType.SYMBOL_PAREN_CLOSE);
-
-        var effect = parseFunctionTypeEffect(tokens);
-
-        tokens.skip(TokenType.SYMBOL_THIN_ARROW);
-        var returnType = parseTopLevelTypeLevelExpression(tokens);
+        var signature = parseFunctionSignature(tokens);
 
         var body = parseBlock(tokens);
 
@@ -336,9 +320,7 @@ public class Parser {
         var source = start.to(end);
 
         return new UntypedFunctionExpressionNode(
-            params,
-            effect,
-            returnType,
+            signature,
             body,
             source
         );
@@ -359,6 +341,31 @@ public class Parser {
         return new UntypedFunctionReturnNode(expression, source);
     }
 
+    static UntypedFunctionSignatureNode parseFunctionSignature(
+        TokenIterator<TokenType> tokens
+    ) {
+        var start = tokens.position();
+
+        tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
+        var params = parseParams(tokens);
+        tokens.skip(TokenType.SYMBOL_PAREN_CLOSE);
+
+        var effect = parseFunctionTypeEffect(tokens);
+
+        tokens.skip(TokenType.SYMBOL_THIN_ARROW);
+        var returnType = parseTopLevelTypeLevelExpression(tokens);
+
+        var end = tokens.position();
+        var source = start.to(end);
+
+        return new UntypedFunctionSignatureNode(
+            params,
+            effect,
+            returnType,
+            source
+        );
+    }
+
     static UntypedFunctionStatementNode parseFunctionStatement(TokenIterator<TokenType> tokens) {
         return switch (tokens.peek().tokenType()) {
             case KEYWORD_ASSIGN, KEYWORD_VAL, KEYWORD_VAR, SYMBOL_UNDERSCORE ->
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 0c3fa89d..9cd0d474 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -484,11 +484,7 @@ public class TypeChecker {
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedSignature = typeCheckFunctionSignature(new UntypedFunctionSignature(
-            untyped.params(),
-            untyped.effect(),
-            untyped.returnType()
-        ), context);
+        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
         return context
             .define(untyped.name(), typedSignature.type(), untyped.source())
@@ -594,14 +590,6 @@ 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,
@@ -622,8 +610,8 @@ public class TypeChecker {
         }
     }
 
-    private static TypedFunctionSignature typeCheckFunctionSignature(
-        UntypedFunctionSignature untyped,
+    static TypedFunctionSignature typeCheckFunctionSignature(
+        UntypedFunctionSignatureNode untyped,
         TypeCheckerContext context
     ) {
         var typedParams = typeCheckParams(untyped.params(), context);