Set up environment properly in nested blocks

2 files changed, 42 insertions, 17 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
index 0d7a6a62..99bd6ca2 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
@@ -3,7 +3,20 @@ package org.shedlang.compiler.typechecker;
 import org.pcollections.PVector;
 import org.shedlang.compiler.ast.untyped.*;
 
+import java.util.List;
+
 class DeclaredVariables {
+    static PVector<UntypedVariableDeclarationNode> forFunctionStatements(
+        List<UntypedFunctionStatementNode> statements,
+        PVector<UntypedVariableDeclarationNode> variables
+    ) {
+        for (var statement : statements) {
+            variables = declaredVariables(statement, variables);
+        }
+
+        return variables;
+    }
+
     static PVector<UntypedVariableDeclarationNode> declaredVariables(
         UntypedFunctionStatementNode statement,
         PVector<UntypedVariableDeclarationNode> variables
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 94bcd8d6..6c460431 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -199,6 +199,27 @@ public class TypeChecker {
         return context.toResult(typed);
     }
 
+    private static TypeCheckerResult<TypedBlockNode> typeCheckBlockNested(
+        UntypedBlockNode untyped,
+        TypeCheckerContext context
+    ) {
+        var declaredVariables = DeclaredVariables.forFunctionStatements(
+            untyped.statements(),
+            P.vector()
+        );
+        var bodyEnvironment = EnvironmentBuilder.EMPTY.addUndefined(declaredVariables);
+
+        context = context.enterBlock(bodyEnvironment.build());
+
+        var typedBodyResult = typeCheckBlock(untyped, context);
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        context = context.exit();
+
+        return context.toResult(typedBody);
+    }
+
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
         UntypedBlockReturnNode untyped,
         TypeCheckerContext context
@@ -278,15 +299,10 @@ public class TypeChecker {
 
         checkExpressionType(typedCondition, Types.BOOL, context);
 
-        // TODO: get environment and pass to enterBlock()
-        context = context.enterBlock(P.map());
-
-        var typedBodyResult = typeCheckBlock(conditionalBranch.body(), context);
+        var typedBodyResult = typeCheckBlockNested(conditionalBranch.body(), context);
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
-        context = context.exit();
-
         return context.toResult(new TypedConditionalBranchNode(
             typedCondition,
             typedBody,
@@ -420,11 +436,12 @@ public class TypeChecker {
             );
         }
 
-        for (var statement : untyped.body().statements()) {
-            // TODO: pass interface to collect variables?
-            var declaredVariables = declaredVariables(statement, P.vector());
-            bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
-        }
+        // TODO: pass interface to collect variables?
+        var declaredVariables = DeclaredVariables.forFunctionStatements(
+            untyped.body().statements(),
+            P.vector()
+        );
+        bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
 
         context = context.enterFunction(
             bodyEnvironment.build(),
@@ -530,15 +547,10 @@ public class TypeChecker {
         var typedConditionalBranches = typedConditionalBranchesResult.value();
         context = typedConditionalBranchesResult.context();
 
-        // TODO: get environment and pass to enterBlock()
-        context = context.enterBlock(P.map());
-
-        var typedElseBodyResult = typeCheckBlock(untyped.elseBody(), context);
+        var typedElseBodyResult = typeCheckBlockNested(untyped.elseBody(), context);
         var typedElseBody = typedElseBodyResult.value();
         context = typedElseBodyResult.context();
 
-        context = context.exit();
-
         var branchTypes = Stream.concat(
             typedConditionalBranches.stream()
                 .map(conditionalBranch -> conditionalBranch.body().type()),