Get type info for when expression in single expression

1 files changed, 10 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 89a2f926..73076a74 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -2052,12 +2052,16 @@ public class TypeChecker {
         var typedExpression = typedExpressionResult.value();
         context = typedExpressionResult.context();
 
-        if (!(typedExpression.type() instanceof SumType expressionType)) {
-            throw new WhenExpressionIsNotInstanceOfSumTypeError(
-                typedExpression.type(),
-                typedExpression.source()
-            );
-        }
+        var expressionTypeInfo = switch (typedExpression.type()) {
+            case SumType expressionType ->
+                context.sumTypeInfo(expressionType, untyped.source());
+
+            default ->
+                throw new WhenExpressionIsNotInstanceOfSumTypeError(
+                    typedExpression.type(),
+                    typedExpression.source()
+                );
+        };
 
         var variableName = untyped.expression() instanceof UntypedReferenceNode expressionReference
             ? Optional.of(expressionReference.name())
@@ -2074,7 +2078,6 @@ public class TypeChecker {
         var typedCases = typedCasesResult.value();
         context = typedCasesResult.context();
 
-        var expressionTypeInfo = context.sumTypeInfo(expressionType, untyped.source());
         var variantTypes = expressionTypeInfo.variants();
         var handledVariantTypes = new HashSet<Type>();