Allow function expression to have implicit unit return type

2 files changed, 46 insertions, 22 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
index 1ffddf69..c624aab7 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionExpressionNode.java
@@ -57,6 +57,13 @@ public record UntypedFunctionExpressionNode(
             );
         }
 
+        public Builder withoutReturnType() {
+            return new Builder(
+                this.signature.withoutReturnType(),
+                this.body
+            );
+        }
+
         public Builder addStatement(UntypedFunctionStatementNode statement) {
             return new Builder(
                 this.signature,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 39dcbd66..02f4702b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -486,7 +486,8 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
-        var returnType = functionSignatureToReturnType(typedSignature);
+        var returnType = functionSignatureToReturnType(typedSignature)
+            .orElse(Types.UNIT);
 
         context = context.enterFunction(
             bodyEnvironment.build(),
@@ -551,26 +552,40 @@ public class TypeChecker {
         };
     }
 
+    private record TypeCheckFunctionDefinitionState(
+        TypedFunctionSignatureNode signature,
+        FunctionType type
+    ) {
+    }
+
     private static TypeCheckerContext declareFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
-        var type = functionSignatureToFunctionType(typedSignature);
+        var type = new FunctionType(
+            functionSignatureToPositionalParams(typedSignature),
+            functionSignatureToNamedParams(typedSignature),
+            functionSignatureToEffect(typedSignature),
+            functionSignatureToReturnType(typedSignature)
+                .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(untyped.signature().source()))
+        );
+
+        var state = new TypeCheckFunctionDefinitionState(typedSignature, type);
 
         return context
             .declareVariableType(untyped.name(), type, untyped.source())
-            .addLexicalState(untyped.name(), typedSignature);
+            .addLexicalState(untyped.name(), state);
     }
 
     private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedSignature = context.getLexicalState(TypedFunctionSignatureNode.class, untyped.name());
-
-        var type = functionSignatureToFunctionType(typedSignature);
+        var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
+        var typedSignature = state.signature;
+        var type = state.type;
 
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
@@ -597,7 +612,12 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
-        var type = functionSignatureToFunctionType(typedSignature);
+        var type = new FunctionType(
+            functionSignatureToPositionalParams(typedSignature),
+            functionSignatureToNamedParams(typedSignature),
+            functionSignatureToEffect(typedSignature),
+            functionSignatureToReturnType(typedSignature).orElse(Types.UNIT)
+        );
 
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
@@ -663,27 +683,24 @@ public class TypeChecker {
         );
     }
 
-    private static FunctionType functionSignatureToFunctionType(TypedFunctionSignatureNode signature) {
-        return new FunctionType(
-            signature.params().positional().stream()
-                .map(param -> param.type().value())
-                .toList(),
-            signature.params().named().stream()
-                .map(param -> new NamedParam(param.name(), param.type().value()))
-                .toList(),
-            functionSignatureToEffect(signature),
-            functionSignatureToReturnType(signature)
-        );
+    private static List<Type> functionSignatureToPositionalParams(TypedFunctionSignatureNode signature) {
+        return signature.params().positional().stream()
+            .map(param -> param.type().value())
+            .toList();
+    }
+
+    private static List<NamedParam> functionSignatureToNamedParams(TypedFunctionSignatureNode signature) {
+        return signature.params().named().stream()
+            .map(param -> new NamedParam(param.name(), param.type().value()))
+            .toList();
     }
 
     private static Effect functionSignatureToEffect(TypedFunctionSignatureNode signature) {
         return signature.effect().map(e -> e.value()).orElse(Effects.NONE);
     }
 
-    private static Type functionSignatureToReturnType(TypedFunctionSignatureNode signature) {
-        return signature.returnType()
-            .orElseThrow(() -> new FunctionSignatureMissingReturnTypeError(signature.source()))
-            .value();
+    private static Optional<Type> functionSignatureToReturnType(TypedFunctionSignatureNode signature) {
+        return signature.returnType().map(node -> node.value());
     }
 
     public static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckFunctionStatement(