Use function type parsing for native function declarations

5 files changed, 43 insertions, 105 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java
index f94ef1fc..f5dbe28e 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNativeFunctionDeclarationNode.java
@@ -2,35 +2,28 @@ package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
-import org.shedlang.compiler.types.*;
-
-import java.util.List;
-import java.util.Optional;
+import org.shedlang.compiler.types.FunctionType;
 
 public record TypedNativeFunctionDeclarationNode(
     String name,
-    TypedFunctionSignatureNode signature,
-    FunctionType type,
+    TypedFunctionTypeNode signature,
     Source source
 ) implements TypedNamespaceStatementNode {
     public static Builder builder() {
         return new Builder(
             "f",
-            TypedFunctionSignatureNode.builder(),
-            new FunctionType(Optional.empty(), List.of(), List.of(), Effects.NONE, Types.UNIT)
+            new TypedFunctionTypeNode(FunctionType.builder().build(), Sources.NULL)
         );
     }
 
     public record Builder(
         String name,
-        TypedFunctionSignatureNode.Builder signature,
-        FunctionType type
+        TypedFunctionTypeNode signature
     ) {
         public TypedNativeFunctionDeclarationNode build() {
             return new TypedNativeFunctionDeclarationNode(
                 name,
-                signature.build(),
-                type,
+                signature,
                 Sources.NULL
             );
         }
@@ -38,61 +31,16 @@ public record TypedNativeFunctionDeclarationNode(
         public TypedNativeFunctionDeclarationNode.Builder withName(String name) {
             return new TypedNativeFunctionDeclarationNode.Builder(
                 name,
-                this.signature,
-                this.type
+                this.signature
             );
         }
 
-        public TypedNativeFunctionDeclarationNode.Builder addPositionalParam(
-            String name,
-            TypedTypeLevelExpressionNode<Type> type
-        ) {
-            return addPositionalParam(name, Optional.of(type));
-        }
-
-        public TypedNativeFunctionDeclarationNode.Builder addPositionalParam(
-            String name,
-            Optional<TypedTypeLevelExpressionNode<Type>> type
+        public TypedNativeFunctionDeclarationNode.Builder withSignature(
+            FunctionType signature
         ) {
             return new TypedNativeFunctionDeclarationNode.Builder(
                 this.name,
-                this.signature.addPositionalParam(name, type),
-                this.type
-            );
-        }
-
-        public TypedNativeFunctionDeclarationNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
-            return new TypedNativeFunctionDeclarationNode.Builder(
-                this.name,
-                this.signature.addNamedParam(name, type),
-                this.type
-            );
-        }
-
-        public TypedNativeFunctionDeclarationNode.Builder addNamedParam(
-            String name,
-            Optional<TypedTypeLevelExpressionNode<Type>> type
-        ) {
-            return new TypedNativeFunctionDeclarationNode.Builder(
-                this.name,
-                this.signature.addNamedParam(name, type),
-                this.type
-            );
-        }
-
-        public TypedNativeFunctionDeclarationNode.Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
-            return new TypedNativeFunctionDeclarationNode.Builder(
-                this.name,
-                this.signature.withEffect(effect),
-                this.type
-            );
-        }
-
-        public TypedNativeFunctionDeclarationNode.Builder withReturnType(TypedTypeLevelExpressionNode<Type> returnType) {
-            return new TypedNativeFunctionDeclarationNode.Builder(
-                this.name,
-                this.signature.withReturnType(returnType),
-                this.type
+                new TypedFunctionTypeNode(signature, Sources.NULL)
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNativeFunctionDeclarationNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNativeFunctionDeclarationNode.java
index 541d7703..6be99af4 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNativeFunctionDeclarationNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNativeFunctionDeclarationNode.java
@@ -4,27 +4,24 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.Visibility;
 
-import java.util.Optional;
-
 public record UntypedNativeFunctionDeclarationNode(
     String name,
     Visibility visibility,
-    // TODO: signature should probably be UntypedFunctionTypeNode (e.g. no names for positional args)
-    UntypedFunctionSignatureNode signature,
+    UntypedFunctionTypeNode signature,
     Source source
 ) implements UntypedNamespaceStatementNode, UntypedVariableDeclarationNode {
     public static Builder builder() {
         return new Builder(
             "f",
             Visibility.PRIVATE,
-            UntypedFunctionSignatureNode.builder()
+            UntypedFunctionTypeNode.builder()
         );
     }
 
     public record Builder(
         String name,
         Visibility visibility,
-        UntypedFunctionSignatureNode.Builder signature
+        UntypedFunctionTypeNode.Builder signature
     ) {
         public UntypedNativeFunctionDeclarationNode build() {
             return new UntypedNativeFunctionDeclarationNode(
@@ -43,18 +40,13 @@ public record UntypedNativeFunctionDeclarationNode(
             );
         }
 
-        public Builder addPositionalParam(String name, UntypedTypeLevelExpressionNode type) {
-            return addPositionalParam(name, Optional.of(type));
-        }
-
         public Builder addPositionalParam(
-            String name,
-            Optional<UntypedTypeLevelExpressionNode> type
+            UntypedTypeLevelExpressionNode type
         ) {
             return new Builder(
                 this.name,
                 this.visibility,
-                this.signature.addPositionalParam(name, type)
+                this.signature.addPositionalParam(type)
             );
         }
 
@@ -74,10 +66,6 @@ public record UntypedNativeFunctionDeclarationNode(
             );
         }
 
-        public Builder withReturnTypeUnit() {
-            return withReturnType(new UntypedTypeLevelReferenceNode("Unit", Sources.NULL));
-        }
-
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
                 this.name,
@@ -86,14 +74,6 @@ public record UntypedNativeFunctionDeclarationNode(
             );
         }
 
-        public Builder withoutReturnType() {
-            return new Builder(
-                this.name,
-                this.visibility,
-                this.signature.withoutReturnType()
-            );
-        }
-
         public Builder withVisibility(Visibility visibility) {
             return new Builder(
                 this.name,
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index b6391e2f..003d909c 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -573,6 +573,23 @@ public class Parser {
         var start = tokens.startPosition();
 
         tokens.skip(TokenType.KEYWORD_FN);
+        var signature = parseFunctionTypeSignature(tokens);
+
+        var end = tokens.endPosition();
+        var source = new FileFragmentSource(start.to(end));
+
+        return new UntypedFunctionTypeNode(
+            signature.params(),
+            signature.effect(),
+            signature.returnType(),
+            source
+        );
+    }
+
+    private static UntypedFunctionTypeNode parseFunctionTypeSignature(
+        TokenIterator<TokenType> tokens
+    ) {
+        var start = tokens.startPosition();
 
         tokens.skip(TokenType.SYMBOL_PAREN_OPEN);
         var params = parseParamTypes(tokens);
@@ -851,7 +868,7 @@ public class Parser {
         tokens.skip(TokenType.KEYWORD_FUN);
         var name = tokens.nextValue(TokenType.IDENTIFIER);
 
-        var signature = parseFunctionSignature(tokens);
+        var signature = parseFunctionTypeSignature(tokens);
 
         tokens.skip(TokenType.SYMBOL_SEMICOLON);
 
diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index 11827861..356323cc 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -14,6 +14,7 @@ import java.text.MessageFormat;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 public class ShidCodeGen {
@@ -660,14 +661,15 @@ public class ShidCodeGen {
         TypedNativeFunctionDeclarationNode node,
         Context context
     ) {
+        var signature = node.signature();
         context = context.addInstruction(new ShidDeclareNativeFunction(
             context.namespaceName(),
             node.name(),
-            node.signature().params().positional().stream()
-                .map(param -> new ShidParam(param.name(), param.source()))
+            IntStream.range(0, signature.value().positionalParams().size())
+                .mapToObj(paramIndex -> new ShidParam("arg" + paramIndex, signature.source()))
                 .toList(),
-            node.signature().params().named().stream()
-                .map(param -> new ShidParam(param.name(), param.source()))
+            signature.value().namedParams().stream()
+                .map(param -> new ShidParam(param.name(), signature.source()))
                 .toList(),
             node.source()
         ));
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index a9d98ff2..3e8e5ddf 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1625,33 +1625,24 @@ public class TypeChecker {
         return new TypeCheckSteps<TypedNativeFunctionDeclarationNode>() {
             @Override
             public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
-                var declareResult = declareFunction(
-                    untyped.name(),
-                    untyped.signature(),
-                    untyped.source(),
-                    Optional.empty(),
-                    context
-                );
-                var state = declareResult.value();
-                context = declareResult.context();
+                var typedNode = typeCheckFunctionType(untyped.signature(), context);
 
                 context = context.declareVariableType(
                     untyped.name(),
-                    state.type(),
+                    typedNode.value(),
                     untyped.source()
                 );
 
-                return context.addLexicalState(untyped.name(), state);
+                return context.addLexicalState(untyped.name(), typedNode);
             }
 
             @Override
             public TypeCheckerResult<TypedNativeFunctionDeclarationNode> finish(TypeCheckerContext context) {
-                var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
+                var typedNode = context.getLexicalState(TypedFunctionTypeNode.class, untyped.name());
 
                 return context.toResult(new TypedNativeFunctionDeclarationNode(
                     untyped.name(),
-                    state.typedSignatureNode,
-                    state.type,
+                    typedNode,
                     untyped.source()
                 ));
             }