Throw error if function reaches end with inferred non-unit type

5 files changed, 178 insertions, 49 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
index 8f617ca8..415d4502 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
 
 import java.util.List;
 
@@ -16,4 +17,8 @@ public record TypedFunctionBodyNode(
     public TypedFunctionStatementBehaviour behaviour() {
         return block.behaviour();
     }
+
+    public Type type() {
+        return block.type();
+    }
 }
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 d394b971..ad22b7ea 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionSignatureNode.java
@@ -73,5 +73,14 @@ public record TypedFunctionSignatureNode(
                 Optional.of(returnType)
             );
         }
+
+        public Builder withoutReturnType() {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                Optional.empty()
+            );
+        }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
new file mode 100644
index 00000000..bb96cfa7
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
@@ -0,0 +1,21 @@
+package org.shedlang.compiler.typechecker;
+
+import org.shedlang.compiler.types.Type;
+
+public sealed interface ReturnTypeState {
+    ReturnTypeState NOT_IN_FUNCTION = new NotInFunction();
+    ReturnTypeState IMPLICIT_UNKNOWN = new ImplicitUnknown();
+
+    static ReturnTypeState explicit(Type type) {
+        return new Explicit(type);
+    }
+
+    static ReturnTypeState implicitLowerBound(Type type) {
+        return new ImplicitLowerBound(type);
+    }
+
+    record NotInFunction() implements ReturnTypeState {}
+    record Explicit(Type type) implements ReturnTypeState {}
+    record ImplicitUnknown() implements ReturnTypeState {}
+    record ImplicitLowerBound(Type type) implements ReturnTypeState {}
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 02f4702b..a68cc23c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -486,47 +486,79 @@ public class TypeChecker {
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
-        var returnType = functionSignatureToReturnType(typedSignature)
-            .orElse(Types.UNIT);
+        var explicitReturnType = functionSignatureToReturnType(typedSignature);
 
         context = context.enterFunction(
             bodyEnvironment.build(),
             functionSignatureToEffect(typedSignature),
-            returnType
+            explicitReturnType
         );
 
         var bodyResult = typeCheckBlock(untyped, bodyEnvironment.names(), context);
         var typedBody = bodyResult.value();
         context = bodyResult.context();
 
-        // TODO: get free variables on exit?
-        var freeVariables = context.freeVariables();
-
-        context = context.exit();
-
-        switch (typedBody.behaviour()) {
-            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                if (!isSubType(blockReturn.type(), returnType, context)) {
-                    throw new UnexpectedTypeError(
-                        TypeSets.type(returnType),
-                        blockReturn.type(),
-                        untyped.source()
-                    );
+        if (explicitReturnType.isPresent()) {
+            switch (typedBody.behaviour()) {
+                case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                    if (!isSubType(blockReturn.type(), explicitReturnType.get(), context)) {
+                        throw new UnexpectedTypeError(
+                            TypeSets.type(explicitReturnType.get()),
+                            blockReturn.type(),
+                            untyped.source()
+                        );
+                    }
                 }
-            }
-            case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                if (!isSubType(Types.UNIT, returnType, context)) {
-                    throw new UnexpectedTypeError(
-                        TypeSets.type(returnType),
-                        Types.UNIT,
-                        untyped.source()
-                    );
+                case TypedFunctionStatementBehaviour.Continue continue_ -> {
+                    if (!isSubType(Types.UNIT, explicitReturnType.get(), context)) {
+                        throw new UnexpectedTypeError(
+                            TypeSets.type(explicitReturnType.get()),
+                            Types.UNIT,
+                            untyped.source()
+                        );
+                    }
+                }
+                case TypedFunctionStatementBehaviour.Terminate terminate -> {
                 }
             }
-            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+        } else {
+            switch (typedBody.behaviour()) {
+                case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                    throw new UnsupportedOperationException("TODO");
+                }
+                case TypedFunctionStatementBehaviour.Continue continue_ -> {
+                    switch (context.returnTypeState()) {
+                        case ReturnTypeState.Explicit explicit -> {
+                            throw new UnsupportedOperationException("TODO");
+                        }
+                        case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
+                            // No explicit returns.
+                        }
+                        case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
+                            checkIsSubType(
+                                Types.UNIT,
+                                implicitLowerBound.type(),
+                                typedBody.source(),
+                                context
+                            );
+                        }
+                        case ReturnTypeState.NotInFunction notInFunction -> {
+                            // TODO: better error
+                            throw new IllegalStateException();
+                        }
+                    }
+                }
+                case TypedFunctionStatementBehaviour.Terminate terminate -> {
+                    throw new UnsupportedOperationException("TODO");
+                }
             }
         }
 
+        // TODO: get free variables on exit?
+        var freeVariables = context.freeVariables();
+
+        context = context.exit();
+
         var typed = new TypedFunctionBodyNode(
             freeVariables,
             typedBody,
@@ -645,14 +677,25 @@ public class TypeChecker {
         var expression = expressionResult.value();
         context = expressionResult.context();
 
-        var expectedReturnType = context.returnType().orElseThrow(() -> {
-            // This should be syntactically impossible, but since we only have
-            // one context type, we check for this possibility and throw a
-            // reasonable error.
-            return new CannotReturnOutsideOfFunctionError(untyped.source());
-        });
-
-        checkExpressionType(expression, expectedReturnType, context);
+        switch (context.returnTypeState()) {
+            case ReturnTypeState.Explicit explicit -> {
+                checkExpressionType(expression, explicit.type(), context);
+            }
+            case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
+                context = context.updateReturnTypeState(
+                    ReturnTypeState.implicitLowerBound(expression.type())
+                );
+            }
+            case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
+                throw new UnsupportedOperationException("TODO");
+            }
+            case ReturnTypeState.NotInFunction notInFunction -> {
+                // This should be syntactically impossible, but since we only have
+                // one context type, we check for this possibility and throw a
+                // reasonable error.
+                throw new CannotReturnOutsideOfFunctionError(untyped.source());
+            }
+        };
 
         var typed = new TypedFunctionReturnNode(expression, untyped.source());
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 030586b1..832c5137 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -22,7 +22,7 @@ public class TypeCheckerContext {
             Optional.empty(),
             Empty.map(),
             Effects.NONE,
-            Optional.empty(),
+            Optional.of(ReturnTypeState.NOT_IN_FUNCTION),
             Set.of(),
             FreeVariables.builder(),
             Empty.map()
@@ -58,7 +58,7 @@ public class TypeCheckerContext {
         private final Optional<NamespaceName> namespaceName;
         private final PMap<String, Variable> environment;
         private final Effect effect;
-        private final Optional<Type> returnType;
+        private final Optional<ReturnTypeState> returnTypeState;
         private final Set<String> boundVariables;
         private final FreeVariables.Builder freeVariables;
         private final PMap<String, Object> state;
@@ -67,7 +67,7 @@ public class TypeCheckerContext {
             Optional<NamespaceName> namespaceName,
             PMap<String, Variable> environment,
             Effect effect,
-            Optional<Type> returnType,
+            Optional<ReturnTypeState> returnTypeState,
             Set<String> boundVariables,
             FreeVariables.Builder freeVariables,
             PMap<String, Object> state
@@ -75,12 +75,24 @@ public class TypeCheckerContext {
             this.namespaceName = namespaceName;
             this.environment = environment;
             this.effect = effect;
-            this.returnType = returnType;
+            this.returnTypeState = returnTypeState;
             this.boundVariables = boundVariables;
             this.freeVariables = freeVariables;
             this.state = state;
         }
 
+        public StackElement updateReturnTypeState(ReturnTypeState newReturnTypeState) {
+            return new LexicalContext(
+                this.namespaceName,
+                this.environment,
+                this.effect,
+                Optional.of(newReturnTypeState),
+                this.boundVariables,
+                this.freeVariables,
+                this.state
+            );
+        }
+
         @Override
         public boolean isBoundInElement(String name) {
             return this.boundVariables.contains(name);
@@ -92,7 +104,7 @@ public class TypeCheckerContext {
                 this.namespaceName,
                 this.environment,
                 this.effect,
-                this.returnType,
+                this.returnTypeState,
                 this.boundVariables,
                 this.freeVariables.plus(name, referenceSource),
                 this.state
@@ -145,11 +157,11 @@ public class TypeCheckerContext {
             .map(x -> Map.entry(x.getKey(), Variable.defined(x.getValue(), VariableType.VAL)))
             .collect(PCollectors.toMap());
 
-        return enter(Optional.empty(), environment, Effects.NONE, Optional.empty());
+        return enter(Optional.empty(), environment, Effects.NONE, Optional.of(ReturnTypeState.NOT_IN_FUNCTION));
     }
 
     public TypeCheckerContext enterNamespace(NamespaceName namespaceName, PMap<String, Variable> environment) {
-        return enter(Optional.of(namespaceName), environment, Effects.NONE, Optional.empty());
+        return enter(Optional.of(namespaceName), environment, Effects.NONE, Optional.of(ReturnTypeState.NOT_IN_FUNCTION));
     }
 
     public TypeCheckerContext enterFunction(
@@ -157,18 +169,35 @@ public class TypeCheckerContext {
         Effect effect,
         Type returnType
     ) {
-        return enter(this.namespaceName(), environment, effect, Optional.of(returnType));
+        return enterFunction(
+            environment,
+            effect,
+            Optional.of(returnType)
+        );
+    }
+
+    public TypeCheckerContext enterFunction(
+        PMap<String, Variable> environment,
+        Effect effect,
+        Optional<Type> returnType
+    ) {
+        return enter(
+            this.namespaceName(),
+            environment,
+            effect,
+            Optional.of(returnType.map(ReturnTypeState::explicit).orElse(ReturnTypeState.IMPLICIT_UNKNOWN))
+        );
     }
 
     public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
-        return enter(this.namespaceName(), environment, this.effect(), this.returnType());
+        return enter(this.namespaceName(), environment, this.effect(), Optional.empty());
     }
 
     private TypeCheckerContext enter(
         Optional<NamespaceName> newNamespaceName,
         PMap<String, Variable> newEnvironment,
         Effect newEffect,
-        Optional<Type> newReturnType
+        Optional<ReturnTypeState> newReturnTypeState
     ) {
         return new TypeCheckerContext(
             this.typeInfos,
@@ -176,7 +205,7 @@ public class TypeCheckerContext {
                 newNamespaceName,
                 this.environment().plusAll(newEnvironment),
                 newEffect,
-                newReturnType,
+                newReturnTypeState,
                 newEnvironment.keySet(),
                 FreeVariables.builder(),
                 Empty.map()
@@ -255,7 +284,7 @@ public class TypeCheckerContext {
                     variable.variableType()
                 )),
                 lexicalContext.effect,
-                lexicalContext.returnType,
+                lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
                 lexicalContext.state
@@ -269,7 +298,7 @@ public class TypeCheckerContext {
                 lexicalContext.namespaceName,
                 lexicalContext.environment,
                 lexicalContext.effect,
-                lexicalContext.returnType,
+                lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
                 lexicalContext.state.plus(name, state)
@@ -313,7 +342,7 @@ public class TypeCheckerContext {
                     variable.variableType()
                 )),
                 lexicalContext.effect,
-                lexicalContext.returnType,
+                lexicalContext.returnTypeState,
                 lexicalContext.boundVariables,
                 lexicalContext.freeVariables,
                 lexicalContext.state
@@ -337,8 +366,30 @@ public class TypeCheckerContext {
         return currentLexicalContext().effect;
     }
 
-    public Optional<Type> returnType() {
-        return currentLexicalContext().returnType;
+    public ReturnTypeState returnTypeState() {
+        for (var element : this.stack) {
+            if (element instanceof LexicalContext lexicalContext && lexicalContext.returnTypeState.isPresent()) {
+                return lexicalContext.returnTypeState.get();
+            }
+        }
+
+        // TODO: better error
+        throw new RuntimeException("missing lexical context");
+    }
+
+    public TypeCheckerContext updateReturnTypeState(ReturnTypeState newReturnTypeState) {
+        var newStack = this.stack;
+        for (var i = 0; i < this.stack.size(); i++) {
+            if (this.stack.get(i) instanceof LexicalContext lexicalContext && lexicalContext.returnTypeState.isPresent()) {
+                newStack = newStack.with(i, lexicalContext.updateReturnTypeState(newReturnTypeState));
+                break;
+            }
+        }
+
+        return new TypeCheckerContext(
+            this.typeInfos,
+            newStack
+        );
     }
 
     public FreeVariables freeVariables() {