Check function definition has return type

4 files changed, 39 insertions, 11 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 2fd648a0..769861a0 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -87,6 +87,15 @@ public record UntypedFunctionDefinitionNode(
             );
         }
 
+        public Builder withoutReturnType() {
+            return new Builder(
+                this.name,
+                this.visibility,
+                this.signature.withoutReturnType(),
+                this.body
+            );
+        }
+
         public Builder addStatement(UntypedFunctionStatementNode statement) {
             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 cd45f465..dea78312 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionSignatureNode.java
@@ -79,6 +79,16 @@ public record UntypedFunctionSignatureNode(
             );
         }
 
+        public Builder withoutReturnType() {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                Optional.empty(),
+                this.body
+            );
+        }
+
         public Builder addStatement(UntypedFunctionStatementNode statement) {
             return new Builder(
                 this.positionalParams,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 10236fee..39dcbd66 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -648,19 +648,17 @@ 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,
-            untypedReturnTypeNode,
-            context
-        );
+        var typedReturnType = untyped.returnType()
+            .map(untypedReturnType -> typeCheckTypeLevelExpression(
+                TypeSets.META_TYPE,
+                untypedReturnType,
+                context
+            ));
 
         return new TypedFunctionSignatureNode(
             typedParams,
             typedEffect,
-            Optional.of(typedReturnType),
+            typedReturnType,
             untyped.source()
         );
     }
@@ -683,8 +681,9 @@ public class TypeChecker {
     }
 
     private static Type functionSignatureToReturnType(TypedFunctionSignatureNode signature) {
-        // TODO: raise proper error
-        return signature.returnType().orElseThrow().value();
+        return signature.returnType()
+            .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(signature.source()))
+            .value();
     }
 
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/FunctionSignatureMissingReturnTypeError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/FunctionSignatureMissingReturnTypeError.java
new file mode 100644
index 00000000..0b31be99
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/FunctionSignatureMissingReturnTypeError.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class FunctionSignatureMissingReturnTypeError extends SourceError {
+    public FunctionSignatureMissingReturnTypeError(Source source) {
+        super("Function signature is missing return type", source);
+    }
+}