Tidy up return type logic

1 files changed, 54 insertions, 65 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index b444a3a5..a24ca306 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -499,82 +499,71 @@ public class TypeChecker {
         context = bodyResult.context();
 
         Type returnType;
-        if (explicitReturnType.isPresent()) {
-            returnType = explicitReturnType.get();
-
-            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()
-                        );
+
+        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 TypedFunctionStatementBehaviour.Continue continue_ -> {
-                    if (!isSubType(Types.UNIT, explicitReturnType.get(), context)) {
-                        throw new UnexpectedTypeError(
-                            TypeSets.type(explicitReturnType.get()),
+            }
+            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,
-                            untyped.source()
+                            implicitLowerBound.type(),
+                            typedBody.source(),
+                            context
                         );
                     }
-                }
-                case TypedFunctionStatementBehaviour.Terminate terminate -> {
+                    case TypedFunctionStatementBehaviour.Terminate terminate -> {
+                        returnType = implicitLowerBound.type();
+                    }
                 }
             }
-        } else {
-            switch (typedBody.behaviour()) {
-                case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
-                    throw new UnsupportedOperationException("TODO");
-                }
-                case TypedFunctionStatementBehaviour.Continue continue_ -> {
-                    switch (context.returnTypeState()) {
-                        case ReturnTypeState.Explicit explicit -> {
-                            // TODO: better error
-                            throw new IllegalStateException();
-                        }
-                        case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
-                            // No explicit returns.
-                            returnType = Types.UNIT;
-                        }
-                        case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
-                            returnType = implicitLowerBound.type();
-
-                            checkIsSubType(
-                                Types.UNIT,
-                                implicitLowerBound.type(),
-                                typedBody.source(),
-                                context
-                            );
-                        }
-                        case ReturnTypeState.NotInFunction notInFunction -> {
-                            // TODO: better error
-                            throw new IllegalStateException();
-                        }
+            case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
+                switch (typedBody.behaviour()) {
+                    case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                        throw new UnsupportedOperationException("TODO");
                     }
-                }
-                case TypedFunctionStatementBehaviour.Terminate terminate -> {
-                    switch (context.returnTypeState()) {
-                        case ReturnTypeState.Explicit explicit -> {
-                            // TODO: better error
-                            throw new IllegalStateException();
-                        }
-                        case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
-                            // TODO: better error
-                            throw new IllegalStateException();
-                        }
-                        case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
-                            returnType = implicitLowerBound.type();
-                        }
-                        case ReturnTypeState.NotInFunction notInFunction -> {
-                            // TODO: better error
-                            throw new IllegalStateException();
-                        }
+                    case TypedFunctionStatementBehaviour.Continue continue_ -> {
+                        returnType = Types.UNIT;
+                    }
+                    case TypedFunctionStatementBehaviour.Terminate terminate -> {
+                        throw new UnsupportedOperationException("TODO");
                     }
                 }
             }
+            case ReturnTypeState.NotInFunction notInFunction -> {
+                // TODO: better error
+                throw new IllegalStateException();
+            }
         }
 
         // TODO: get free variables on exit?