Make return type optional in typed function signature node

2 files changed, 21 insertions, 12 deletions

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 7c83dd1c..d394b971 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
@@ -11,7 +11,7 @@ import java.util.Optional;
 public record TypedFunctionSignatureNode(
     TypedParamsNode params,
     Optional<TypedTypeLevelExpressionNode<Effect>> effect,
-    TypedTypeLevelExpressionNode<Type> returnType,
+    Optional<TypedTypeLevelExpressionNode<Type>> returnType,
     Source source
 ) implements TypedNode {
     public static Builder builder() {
@@ -19,7 +19,7 @@ public record TypedFunctionSignatureNode(
             Empty.vector(),
             Empty.vector(),
             Optional.empty(),
-            new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL)
+            Optional.empty()
         );
     }
 
@@ -27,7 +27,7 @@ public record TypedFunctionSignatureNode(
         PVector<TypedParamNode> positionalParams,
         PVector<TypedParamNode> namedParams,
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
-        TypedTypeLevelExpressionNode<Type> returnType
+        Optional<TypedTypeLevelExpressionNode<Type>> returnType
     ) {
         public TypedFunctionSignatureNode build() {
             return new TypedFunctionSignatureNode(
@@ -70,7 +70,7 @@ public record TypedFunctionSignatureNode(
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
-                returnType
+                Optional.of(returnType)
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 530b1d7c..10236fee 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -486,11 +486,11 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
-        var returnType = typedSignature.returnType().value();
+        var returnType = functionSignatureToReturnType(typedSignature);
 
         context = context.enterFunction(
             bodyEnvironment.build(),
-            typedSignature.effect().map(e -> e.value()).orElse(Effects.NONE),
+            functionSignatureToEffect(typedSignature),
             returnType
         );
 
@@ -660,24 +660,33 @@ public class TypeChecker {
         return new TypedFunctionSignatureNode(
             typedParams,
             typedEffect,
-            typedReturnType,
+            Optional.of(typedReturnType),
             untyped.source()
         );
     }
 
-    private static FunctionType functionSignatureToFunctionType(TypedFunctionSignatureNode typedSignature) {
+    private static FunctionType functionSignatureToFunctionType(TypedFunctionSignatureNode signature) {
         return new FunctionType(
-            typedSignature.params().positional().stream()
+            signature.params().positional().stream()
                 .map(param -> param.type().value())
                 .toList(),
-            typedSignature.params().named().stream()
+            signature.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()
+            functionSignatureToEffect(signature),
+            functionSignatureToReturnType(signature)
         );
     }
 
+    private static Effect functionSignatureToEffect(TypedFunctionSignatureNode signature) {
+        return signature.effect().map(e -> e.value()).orElse(Effects.NONE);
+    }
+
+    private static Type functionSignatureToReturnType(TypedFunctionSignatureNode signature) {
+        // TODO: raise proper error
+        return signature.returnType().orElseThrow().value();
+    }
+
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
         UntypedFunctionStatementNode untyped,
         TypeCheckerContext context