Extract typeCheckFunctionReturnValue()

1 files changed, 45 insertions, 58 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index a24ca306..0ceca376 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -498,67 +498,37 @@ public class TypeChecker {
         var typedBody = bodyResult.value();
         context = bodyResult.context();
 
+        switch (typedBody.behaviour()) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                context = typeCheckFunctionReturnValue(
+                    blockReturn.type(),
+                    // TODO: source should be last statement?
+                    typedBody.source(),
+                    context
+                );
+            }
+            case TypedFunctionStatementBehaviour.Continue continue_ -> {
+                if (context.returnTypeState() instanceof ReturnTypeState.ImplicitLowerBound implicitLowerBound) {
+                    checkIsSubType(Types.UNIT, implicitLowerBound.type(), untyped.source(), context);
+                } else {
+                    context = typeCheckFunctionReturnValue(Types.UNIT, untyped.source(), context);
+                }
+            }
+            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+            }
+        }
+
         Type returnType;
 
         switch (context.returnTypeState()) {
             case ReturnTypeState.Explicit explicit -> {
                 returnType = explicit.type();
-
-                switch (typedBody.behaviour()) {
-                    case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                        if (!isSubType(blockReturn.type(), explicit.type(), context)) {
-                            throw new UnexpectedTypeError(
-                                TypeSets.type(explicit.type()),
-                                blockReturn.type(),
-                                untyped.source()
-                            );
-                        }
-                    }
-                    case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                        if (!isSubType(Types.UNIT, explicit.type(), context)) {
-                            throw new UnexpectedTypeError(
-                                TypeSets.type(explicit.type()),
-                                Types.UNIT,
-                                untyped.source()
-                            );
-                        }
-                    }
-                    case TypedFunctionStatementBehaviour.Terminate terminate -> {
-                    }
-                }
             }
             case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
-                switch (typedBody.behaviour()) {
-                    case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                        throw new UnsupportedOperationException("TODO");
-                    }
-                    case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                        returnType = implicitLowerBound.type();
-
-                        checkIsSubType(
-                            Types.UNIT,
-                            implicitLowerBound.type(),
-                            typedBody.source(),
-                            context
-                        );
-                    }
-                    case TypedFunctionStatementBehaviour.Terminate terminate -> {
-                        returnType = implicitLowerBound.type();
-                    }
-                }
+                returnType = implicitLowerBound.type();
             }
             case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
-                switch (typedBody.behaviour()) {
-                    case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                        throw new UnsupportedOperationException("TODO");
-                    }
-                    case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                        returnType = Types.UNIT;
-                    }
-                    case TypedFunctionStatementBehaviour.Terminate terminate -> {
-                        throw new UnsupportedOperationException("TODO");
-                    }
-                }
+                throw new UnsupportedOperationException("TODO");
             }
             case ReturnTypeState.NotInFunction notInFunction -> {
                 // TODO: better error
@@ -690,29 +660,46 @@ public class TypeChecker {
         var expression = expressionResult.value();
         context = expressionResult.context();
 
+        context = typeCheckFunctionReturnValue(expression.type(), expression.source(), context);
+
+        var typed = new TypedFunctionReturnNode(expression, untyped.source());
+
+        return context.toResult(typed);
+    }
+
+    private static TypeCheckerContext typeCheckFunctionReturnValue(
+        Type returnValueType,
+        Source returnValueSource,
+        TypeCheckerContext context
+    ) {
         switch (context.returnTypeState()) {
             case ReturnTypeState.Explicit explicit -> {
-                checkExpressionType(expression, explicit.type(), context);
+                checkIsSubType(returnValueType, explicit.type(), returnValueSource, context);
             }
             case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
                 context = context.updateReturnTypeState(
-                    ReturnTypeState.implicitLowerBound(expression.type())
+                    ReturnTypeState.implicitLowerBound(returnValueType)
                 );
             }
             case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
                 throw new UnsupportedOperationException("TODO");
+//                var supertype = Types.commonSupertype(List.of(
+//                    implicitLowerBound.type(),
+//                    returnValueType
+//                ));
+//                context = context.updateReturnTypeState(
+//                    ReturnTypeState.implicitLowerBound(supertype)
+//                );
             }
             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());
+                throw new CannotReturnOutsideOfFunctionError(returnValueSource);
             }
         };
 
-        var typed = new TypedFunctionReturnNode(expression, untyped.source());
-
-        return context.toResult(typed);
+        return context;
     }
 
     static TypedFunctionSignatureNode typeCheckFunctionSignature(