Store bound variable names in typed block node

8 files changed, 120 insertions, 12 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
index 5c0c8435..9753f974 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
@@ -7,8 +7,14 @@ import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
 
+/**
+ * @param boundVariableNames The names of variables that are bound in the
+ *                           block's scope, not including any children scopes
+ *                           such as the bodies of <code>if</code> expressions.
+ */
 public record TypedBlockNode(
     List<? extends TypedFunctionStatementNode> statements,
+    List<String> boundVariableNames,
     Source source
 ) implements TypedNode {
     public TypedFunctionStatementBehaviour behaviour() {
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index 2b14e345..b450158b 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -6,6 +6,7 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.*;
 
+import java.util.List;
 import java.util.Optional;
 
 public record TypedFunctionDefinitionNode(
@@ -76,7 +77,7 @@ public record TypedFunctionDefinitionNode(
                 effect,
                 returnType,
                 freeVariables.build(),
-                new TypedBlockNode(body, Sources.NULL),
+                new TypedBlockNode(body, List.of(), Sources.NULL),
                 Sources.NULL
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
index 79c0a25e..53853305 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
@@ -6,6 +6,8 @@ import org.shedlang.compiler.ast.untyped.UntypedVariableDeclarationNode;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.DuplicateVariableName;
 
+import java.util.List;
+
 public class EnvironmentBuilder {
     public static final EnvironmentBuilder EMPTY = new EnvironmentBuilder(Empty.map());
 
@@ -42,4 +44,8 @@ public class EnvironmentBuilder {
     public PMap<String, TypeCheckerContext.Variable> build() {
         return environment;
     }
+
+    public List<String> names() {
+        return environment.keySet().stream().toList();
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 4ec4d5dc..bd02d88b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -10,10 +10,7 @@ import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -186,6 +183,7 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedBlockNode> typeCheckBlock(
         UntypedBlockNode untyped,
+        List<String> declaredVariableNames,
         TypeCheckerContext context
     ) {
         var typedStatementsResult = context.mapWithContext(
@@ -195,7 +193,11 @@ public class TypeChecker {
         var typedStatements = typedStatementsResult.value();
         context = typedStatementsResult.context();
 
-        var typed = new TypedBlockNode(typedStatements, untyped.source());
+        var typed = new TypedBlockNode(
+            typedStatements,
+            declaredVariableNames,
+            untyped.source()
+        );
 
         return context.toResult(typed);
     }
@@ -212,7 +214,7 @@ public class TypeChecker {
 
         context = context.enterBlock(bodyEnvironment.build());
 
-        var typedBodyResult = typeCheckBlock(untyped, context);
+        var typedBodyResult = typeCheckBlock(untyped, bodyEnvironment.names(), context);
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
@@ -450,7 +452,7 @@ public class TypeChecker {
             functionType.returnType()
         );
 
-        var bodyResult = typeCheckBlock(untyped.body(), context);
+        var bodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.names(), context);
         var typedBody = bodyResult.value();
         context = bodyResult.context();
 
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
index 160950b6..7ce49957 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
@@ -14,7 +14,7 @@ public class Typed {
     }
 
     public static TypedBlockNode block(List<TypedFunctionStatementNode> statements) {
-        return new TypedBlockNode(statements, Sources.NULL);
+        return new TypedBlockNode(statements, List.of(), Sources.NULL);
     }
 
     public static TypedBlockReturnNode blockReturn(TypedExpressionNode expression) {
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index 8ce4b168..3a46fa54 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -35,6 +35,12 @@ public class TypedNodeMatchers {
         return has("statements", x -> x.statements(), statements);
     }
 
+    public static Matcher<TypedBlockNode> isTypedBlockNodeWithBoundVariableNames(
+        Matcher<Iterable<String>> boundVariableNames
+    ) {
+        return has("boundVariableNames", x -> x.boundVariableNames(), boundVariableNames);
+    }
+
     public static Matcher<TypedFunctionStatementNode> isTypedBlockReturnNode(
         Matcher<TypedExpressionNode> expression
     ) {
@@ -56,7 +62,7 @@ public class TypedNodeMatchers {
     }
 
     public static Matcher<TypedConditionalBranchNode> isTypedConditionalBranchNode(
-        Matcher<TypedExpressionNode> condition,
+        Matcher<? super TypedExpressionNode> condition,
         Matcher<TypedBlockNode> body
     ) {
         return allOf(
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index ab213695..6e77563f 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -348,6 +348,51 @@ public class TypeCheckerFunctionDefinitionTests {
         );
     }
 
+    @Test
+    public void boundVariablesAreSet() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .addPositionalParam("argPositional", Untyped.typeLevelReference("Unit"))
+            .addNamedParam("argNamed", Untyped.typeLevelReference("Unit"))
+            .withReturnType(Untyped.typeLevelReference("Unit"))
+            .addStatement(Untyped.val(
+                Untyped.targetVariable("local1"),
+                Untyped.intLiteral(1)
+            ))
+            .addStatement(Untyped.expressionStatement(Untyped.reference("y")))
+            .addStatement(Untyped.expressionStatement(Untyped.if_(
+                List.of(
+                    Untyped.conditionalBranch(
+                        Untyped.boolLiteral(true),
+                        Untyped.block(List.of(
+                            Untyped.val(
+                                Untyped.targetVariable("local2"),
+                                Untyped.intLiteral(2)
+                            )
+                        ))
+                    )
+                ),
+                Untyped.block(List.of())
+            )))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespaceWithDefinedVariables(P.map(
+                Map.entry("Unit", Types.metaType(Types.UNIT)),
+                Map.entry("y", Types.STRING)
+            ));
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        assertThat(
+            typed.value(),
+            isTypedFunctionDefinitionNode()
+                .withBody(isTypedBlockNodeWithBoundVariableNames(containsExactly(
+                    equalTo("argPositional"),
+                    equalTo("argNamed"),
+                    equalTo("local1")
+                )))
+        );
+    }
+
     private static TypeCheckerResult<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(
         UntypedFunctionDefinitionNode untyped,
         TypeCheckerContext context
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIfTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIfTests.java
index 0aaf3d5e..a52fdfac 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIfTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerIfTests.java
@@ -13,8 +13,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 TypeCheckerIfTests {
     @Test
@@ -149,4 +148,47 @@ public class TypeCheckerIfTests {
             ))
         ));
     }
+
+    @Test
+    public void boundVariablesAreSet() {
+        var untyped = Untyped.if_(
+            List.of(
+                Untyped.conditionalBranch(
+                    Untyped.boolLiteral(true),
+                    Untyped.block(List.of(
+                        Untyped.val(
+                            Untyped.targetVariable("local1"),
+                            Untyped.intLiteral(1)
+                        )
+                    ))
+                )
+            ),
+            Untyped.block(List.of(
+                Untyped.val(
+                    Untyped.targetVariable("local2"),
+                    Untyped.intLiteral(2)
+                )
+            ))
+        );
+        var context = TypeCheckerContext.EMPTY;
+
+        var result = TypeChecker.typeCheckExpression(untyped, context);
+
+        assertThat(
+            result.value(),
+            isTypedIfNode(
+                isSequence(
+                    isTypedConditionalBranchNode(
+                        anything(),
+                        isTypedBlockNodeWithBoundVariableNames(containsExactly(
+                            equalTo("local1")
+                        ))
+                    )
+                ),
+                isTypedBlockNodeWithBoundVariableNames(containsExactly(
+                    equalTo("local2")
+                ))
+            )
+        );
+    }
 }