Narrow type in when case bodies

2 files changed, 48 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 8aafe2eb..14992468 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1147,9 +1147,17 @@ public class TypeChecker {
             );
         }
 
+        var variableName = untyped.expression() instanceof UntypedReferenceNode expressionReference
+            ? Optional.of(expressionReference.name())
+            : Optional.<String>empty();
+
         var typedCasesResult = context.mapWithContext(
             untyped.cases(),
-            TypeChecker::typeCheckWhenCase
+            (untypedWhenCase, whenCaseContext) -> typeCheckWhenCase(
+                untypedWhenCase,
+                variableName,
+                whenCaseContext
+            )
         );
         var typedCases = typedCasesResult.value();
         context = typedCasesResult.context();
@@ -1191,6 +1199,7 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedWhenCaseNode> typeCheckWhenCase(
         UntypedWhenCaseNode untyped,
+        Optional<String> variableName,
         TypeCheckerContext context
     ) {
         var typedCaseType = typeCheckTypeLevelExpression(
@@ -1199,10 +1208,24 @@ public class TypeChecker {
             context
         );
 
-        var typedBodyResult = typeCheckBlockNested(untyped.body(), context);
+        var declaredVariables = DeclaredVariables.forFunctionStatements(
+            untyped.body().statements(),
+            P.vector()
+        );
+        var bodyEnvironment = EnvironmentBuilder.EMPTY.addUndefined(declaredVariables);
+
+        context = context.enterBlock(bodyEnvironment.build());
+
+        if (variableName.isPresent()) {
+            context = context.updateType(variableName.get(), typedCaseType.value(), untyped.source());
+        }
+
+        var typedBodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.names(), context);
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
+        context = context.exit();
+
         return context.toResult(new TypedWhenCaseNode(
             typedCaseType,
             typedBody,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 7d1f77cc..2fc19b80 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -240,6 +240,29 @@ public class TypeCheckerContext {
         return Optional.ofNullable(environment().get(name));
     }
 
+    public TypeCheckerContext updateType(String name, Type type, Source source) {
+        var variable = this.environment().get(name);
+
+        if (variable == null) {
+            throw new UnknownVariableError(name, source);
+        }
+
+        return this.updateLexicalContext(
+            lexicalContext -> new LexicalContext(
+                lexicalContext.namespaceName,
+                lexicalContext.environment.plus(name, new Variable(
+                    Optional.of(type),
+                    variable.variableType()
+                )),
+                lexicalContext.effect,
+                lexicalContext.returnType,
+                lexicalContext.boundVariables,
+                lexicalContext.freeVariables,
+                lexicalContext.state
+            )
+        );
+    }
+
     public TypeCheckerContext addLexicalState(String name, Object state) {
         return this.updateLexicalContext(
             lexicalContext -> new LexicalContext(