Set ReturnTypeState in PartialFunctionSignature

3 files changed, 27 insertions, 16 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
index c0501856..4edfe582 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/PartialFunctionSignature.java
@@ -10,13 +10,12 @@ import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
-import java.util.Optional;
 
 public record PartialFunctionSignature(
     List<Param> positionalParams,
     List<Param> namedParams,
     Effect effect,
-    Optional<Type> returnType
+    ReturnTypeState returnTypeState
 ) {
     public record Param(String name, Type type, Source source) {}
 
@@ -29,7 +28,7 @@ public record PartialFunctionSignature(
             this.positionalParams,
             this.namedParams,
             effect,
-            this.returnType
+            this.returnTypeState
         );
     }
 
@@ -38,7 +37,7 @@ public record PartialFunctionSignature(
             P.vector(),
             P.vector(),
             Effects.NONE,
-            Optional.empty()
+            ReturnTypeState.IMPLICIT_UNKNOWN
         );
     }
 
@@ -46,14 +45,14 @@ public record PartialFunctionSignature(
         PVector<Param> positionalParams,
         PVector<Param> namedParams,
         Effect effect,
-        Optional<Type> returnType
+        ReturnTypeState returnTypeState
     ) {
         public PartialFunctionSignature build() {
             return new PartialFunctionSignature(
                 positionalParams,
                 namedParams,
                 effect,
-                returnType
+                returnTypeState
             );
         }
 
@@ -62,7 +61,7 @@ public record PartialFunctionSignature(
                 this.positionalParams.plus(new Param(name, type, Sources.NULL)),
                 this.namedParams,
                 this.effect,
-                this.returnType
+                this.returnTypeState
             );
         }
 
@@ -71,7 +70,7 @@ public record PartialFunctionSignature(
                 this.positionalParams,
                 this.namedParams.plus(new Param(name, type, Sources.NULL)),
                 this.effect,
-                this.returnType
+                this.returnTypeState
             );
         }
 
@@ -80,7 +79,7 @@ public record PartialFunctionSignature(
                 this.positionalParams,
                 this.namedParams,
                 effect,
-                this.returnType
+                this.returnTypeState
             );
         }
 
@@ -89,7 +88,7 @@ public record PartialFunctionSignature(
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
-                Optional.of(returnType)
+                ReturnTypeState.explicit(returnType)
             );
         }
 
@@ -98,7 +97,7 @@ public record PartialFunctionSignature(
                 this.positionalParams,
                 this.namedParams,
                 this.effect,
-                Optional.empty()
+                ReturnTypeState.IMPLICIT_UNKNOWN
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 5e1fc6e2..4cf49b47 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -569,12 +569,10 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
-        var explicitReturnType = signature.returnType();
-
         context = context.enterFunction(
             bodyEnvironment.build(),
             signature.effect(),
-            explicitReturnType
+            signature.returnTypeState()
         );
 
         var bodyResult = typeCheckBlock(untyped, bodyEnvironment.names(), context);
@@ -713,7 +711,7 @@ public class TypeChecker {
             positionalParams,
             namedParams,
             effect,
-            Optional.of(returnType)
+            ReturnTypeState.explicit(returnType)
         );
 
         var type = new FunctionType(
@@ -928,6 +926,8 @@ public class TypeChecker {
             namedParams,
             effect,
             typedSignature.returnType().map(node -> node.value())
+                .map(ReturnTypeState::explicit)
+                .orElse(ReturnTypeState.IMPLICIT_UNKNOWN)
         );
     }
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 3f3147f0..3fb95910 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -180,11 +180,23 @@ public class TypeCheckerContext {
         Effect effect,
         Optional<Type> returnType
     ) {
+        return enterFunction(
+            environment,
+            effect,
+            returnType.map(ReturnTypeState::explicit).orElse(ReturnTypeState.IMPLICIT_UNKNOWN)
+        );
+    }
+
+    public TypeCheckerContext enterFunction(
+        PMap<String, Variable> environment,
+        Effect effect,
+        ReturnTypeState returnTypeState
+    ) {
         return enter(
             this.namespaceName(),
             environment,
             effect,
-            Optional.of(returnType.map(ReturnTypeState::explicit).orElse(ReturnTypeState.IMPLICIT_UNKNOWN))
+            Optional.of(returnTypeState)
         );
     }