Narrow type in when case bodies

1 files changed, 57 insertions, 2 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
index f593a710..cc731458 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhenTests.java
@@ -15,8 +15,7 @@ import java.util.Map;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
 import static org.zwobble.precisely.AssertThat.assertThat;
-import static org.zwobble.precisely.Matchers.equalTo;
-import static org.zwobble.precisely.Matchers.isSequence;
+import static org.zwobble.precisely.Matchers.*;
 
 public class TypeCheckerWhenTests {
     private static final NamespaceName NAMESPACE_NAME = NamespaceName.of("Example", "Project");
@@ -253,4 +252,60 @@ public class TypeCheckerWhenTests {
 
         assertThat(result.value().type(), equalTo(Types.TOP));
     }
+
+    @Test
+    public void whenExpressionIsReferenceThenTypeIsNarrowedInCaseBodies() {
+        var untyped = Untyped.when(
+            Untyped.reference("x"),
+            List.of(
+                Untyped.whenCase(
+                    Untyped.typeLevelReference(VARIANT_TYPE_1.name()),
+                    Untyped.block(List.of(
+                        Untyped.blockReturn(Untyped.fieldAccess(
+                            Untyped.reference("x"),
+                            "a"
+                        ))
+                    ))
+                ),
+                Untyped.whenCase(
+                    Untyped.typeLevelReference(VARIANT_TYPE_2.name()),
+                    Untyped.block(List.of(
+                        Untyped.blockReturn(Untyped.fieldAccess(
+                            Untyped.reference("x"),
+                            "b"
+                        ))
+                    ))
+                )
+            )
+        );
+        var context = TypeCheckerContext.EMPTY
+            .addSumTypeInfo(SUM_TYPE, new SumTypeInfo(P.orderedSet(VARIANT_TYPE_1, VARIANT_TYPE_2)))
+            .addStructTypeInfo(VARIANT_TYPE_1, StructTypeInfo.fromFields(P.orderedMap(
+                Map.entry("a", Types.STRING)
+            )))
+            .addStructTypeInfo(VARIANT_TYPE_2, StructTypeInfo.fromFields(P.orderedMap(
+                Map.entry("b", Types.INT)
+            )))
+            .enterGlobal(P.map(
+                Map.entry("x", SUM_TYPE),
+                Map.entry(VARIANT_TYPE_1.name(), Types.metaType(VARIANT_TYPE_1)),
+                Map.entry(VARIANT_TYPE_2.name(), Types.metaType(VARIANT_TYPE_2))
+            ));
+
+        var result = TypeChecker.typeCheckExpression(untyped, context);
+
+        assertThat(result.value(), isTypedWhenNode(
+            isTypedReferenceNode("x", SUM_TYPE),
+            isSequence(
+                isTypedWhenCaseNode(
+                    isTypedTypeLevelReferenceNode("Variant1", VARIANT_TYPE_1),
+                    has("type", x -> x.type(), equalTo(Types.STRING))
+                ),
+                isTypedWhenCaseNode(
+                    isTypedTypeLevelReferenceNode("Variant2", VARIANT_TYPE_2),
+                    has("type", x -> x.type(), equalTo(Types.INT))
+                )
+            )
+        ));
+    }
 }