Parse function signatures without return type

4 files changed, 22 insertions, 7 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 d9abdcdb..2fd648a0 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -74,6 +74,10 @@ public record UntypedFunctionDefinitionNode(
             );
         }
 
+        public Builder withReturnTypeUnit() {
+            return withReturnType(new UntypedTypeLevelReferenceNode("Unit", Sources.NULL));
+        }
+
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
                 this.name,
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
index 179b4f1e..cd45f465 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
@@ -10,7 +10,7 @@ import java.util.Optional;
 public record UntypedFunctionSignatureNode(
     UntypedParamsNode params,
     Optional<UntypedTypeLevelExpressionNode> effect,
-    UntypedTypeLevelExpressionNode returnType,
+    Optional<UntypedTypeLevelExpressionNode> returnType,
     Source source
 ) implements UntypedNode {
     public static Builder builder() {
@@ -18,7 +18,7 @@ public record UntypedFunctionSignatureNode(
             Empty.vector(),
             Empty.vector(),
             Optional.empty(),
-            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL),
+            Optional.empty(),
             Empty.vector()
         );
     }
@@ -27,7 +27,7 @@ public record UntypedFunctionSignatureNode(
         PVector<UntypedParamNode> positionalParams,
         PVector<UntypedParamNode> namedParams,
         Optional<UntypedTypeLevelExpressionNode> effect,
-        UntypedTypeLevelExpressionNode returnType,
+        Optional<UntypedTypeLevelExpressionNode> returnType,
         PVector<UntypedFunctionStatementNode> body
     ) {
         public UntypedFunctionSignatureNode build() {
@@ -74,7 +74,7 @@ public record UntypedFunctionSignatureNode(
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
-                returnType,
+                Optional.of(returnType),
                 this.body
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index ac7f5bb1..6cb1b134 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -352,8 +352,7 @@ public class Parser {
 
         var effect = parseFunctionTypeEffect(tokens);
 
-        tokens.skip(TokenType.SYMBOL_THIN_ARROW);
-        var returnType = parseTopLevelTypeLevelExpression(tokens);
+        var returnType = parseFunctionSignatureReturnType(tokens);
 
         var end = tokens.position();
         var source = start.to(end);
@@ -366,6 +365,15 @@ public class Parser {
         );
     }
 
+    private static Optional<UntypedTypeLevelExpressionNode> parseFunctionSignatureReturnType(TokenIterator<TokenType> tokens) {
+        if (tokens.trySkip(TokenType.SYMBOL_THIN_ARROW)) {
+            var expression = parseTopLevelTypeLevelExpression(tokens);
+            return Optional.of(expression);
+        } else {
+            return Optional.empty();
+        }
+    }
+
     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 fac340ed..5ad434eb 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -640,9 +640,12 @@ public class TypeChecker {
         var typedEffect = untyped.effect()
             .map(untypedEffect -> typeCheckEffect(untypedEffect, context));
 
+        // TODO: handle optional return type
+        var untypedReturnTypeNode = untyped.returnType().orElseThrow();
+
         var typedReturnType = typeCheckTypeLevelExpression(
             TypeSets.META_TYPE,
-            untyped.returnType(),
+            untypedReturnTypeNode,
             context
         );