Check function definition has return type

5 files changed, 55 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);
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index b44b0f96..352e46b6 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.typechecker;
 import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.untyped.Untyped;
 import org.shedlang.compiler.ast.untyped.UntypedFunctionDefinitionNode;
+import org.shedlang.compiler.typechecker.errors.FunctionSignatureMissingReturnTypeError;
 import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.FunctionType;
 import org.shedlang.compiler.types.Types;
@@ -11,6 +12,7 @@ import org.shedlang.compiler.util.P;
 import java.util.List;
 import java.util.Map;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
 import static org.shedlang.compiler.typechecker.TypeCheckerTesting.typeCheckNamespaceStatement;
 import static org.zwobble.precisely.AssertThat.assertThat;
@@ -40,4 +42,18 @@ public class TypeCheckerFunctionDefinitionTests {
                 .withType(new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT))
         );
     }
+
+    @Test
+    public void whenSignatureIsMissingReturnTypeThenErrorIsThrown() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withoutReturnType()
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(Map.entry("Unit", Types.metaType(Types.UNIT))));
+
+        assertThrows(
+            FunctionSignatureMissingReturnTypeError.class,
+            () -> typeCheckNamespaceStatement(untyped, context)
+        );
+    }
 }