Handle shadowing of variables in function

5 files changed, 174 insertions, 29 deletions

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 b450158b..cdaacf27 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -6,7 +6,6 @@ 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(
@@ -57,6 +56,7 @@ public record TypedFunctionDefinitionNode(
             Optional.empty(),
             new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
             FreeVariables.builder(),
+            Empty.vector(),
             Empty.vector()
         );
     }
@@ -68,6 +68,7 @@ public record TypedFunctionDefinitionNode(
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
         FreeVariables.Builder freeVariables,
+        PVector<String> boundVariableNames,
         PVector<TypedFunctionStatementNode> body
     ) {
         public TypedFunctionDefinitionNode build() {
@@ -77,7 +78,7 @@ public record TypedFunctionDefinitionNode(
                 effect,
                 returnType,
                 freeVariables.build(),
-                new TypedBlockNode(body, List.of(), Sources.NULL),
+                new TypedBlockNode(body, boundVariableNames, Sources.NULL),
                 Sources.NULL
             );
         }
@@ -90,6 +91,7 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 this.returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body
             );
         }
@@ -102,8 +104,9 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 this.returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body
-            );
+            ).addBoundVariable(name);
         }
 
         public TypedFunctionDefinitionNode.Builder addNamedParam(String name, TypedTypeLevelExpressionNode<Type> type) {
@@ -114,8 +117,9 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 this.returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body
-            );
+            ).addBoundVariable(name);
         }
 
         public TypedFunctionDefinitionNode.Builder withEffect(TypedTypeLevelExpressionNode<Effect> effect) {
@@ -126,6 +130,7 @@ public record TypedFunctionDefinitionNode(
                 Optional.of(effect),
                 this.returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body
             );
         }
@@ -138,6 +143,7 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body
             );
         }
@@ -150,6 +156,20 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 this.returnType,
                 this.freeVariables.plus(freeVariable, Sources.NULL),
+                this.boundVariableNames,
+                this.body
+            );
+        }
+
+        public Builder addBoundVariable(String name) {
+            return new TypedFunctionDefinitionNode.Builder(
+                this.name,
+                this.positionalParams,
+                this.namedParams,
+                this.effect,
+                this.returnType,
+                this.freeVariables,
+                this.boundVariableNames.plus(name),
                 this.body
             );
         }
@@ -162,6 +182,7 @@ public record TypedFunctionDefinitionNode(
                 this.effect,
                 this.returnType,
                 this.freeVariables,
+                this.boundVariableNames,
                 this.body.plus(statement)
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
index 359d91b1..628cffd2 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceNode.java
@@ -8,6 +8,7 @@ import java.util.List;
 public record TypedNamespaceNode(
     NamespaceName namespaceName,
     List<TypedNamespaceStatementNode> body,
+    List<String> boundVariableNames,
     List<String> fieldNames,
     Source source
 ) implements TypedNode {
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 7d1ed4b0..167983bb 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -1,7 +1,6 @@
 package org.shedlang.compiler.stackir;
 
-import org.pcollections.PStack;
-import org.pcollections.PVector;
+import org.pcollections.*;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
@@ -37,7 +36,7 @@ public class StackCompiler {
         return context;
     }
 
-    private static Context compileBlock(TypedBlockNode node, Context context) {
+    static Context compileBlock(TypedBlockNode node, Context context) {
         for (var statement : node.statements()) {
             context = compileFunctionStatement(statement, context);
         }
@@ -125,16 +124,24 @@ public class StackCompiler {
         TypedFunctionDefinitionNode node,
         Context context
     ) {
+        context = context.enterBlock(node.body().boundVariableNames());
+
+        var c = context;
+
         var positionalParams = node.params().positional().stream()
-            .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+            .map(param -> new DefineFunctionInstruction.Param(
+                c.shedVariableToLocal(param.name()),
+                param.source()
+            ))
             .toList();
 
         var namedParams = node.params().named().stream()
-            .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
+            .map(param -> new DefineFunctionInstruction.Param(
+                c.shedVariableToLocal(param.name()),
+                param.source()
+            ))
             .toList();
 
-        context = context.enterBlock();
-
         context = compileBlock(node.body(), context);
 
         switch (node.body().behaviour()) {
@@ -152,16 +159,28 @@ public class StackCompiler {
 
         context = context.exitBlock();
 
+        var freeVariables = FreeVariables.builder();
+        for (var shedFreeVariable : node.freeVariables()) {
+            var freeVariableLocalName = context.shedVariableToLocal(shedFreeVariable.name());
+            for (var freeVariableSource : shedFreeVariable.referenceSources()) {
+                freeVariables = freeVariables.plus(
+                    freeVariableLocalName,
+                    freeVariableSource
+                );
+            }
+        }
+
         var defineFunctionInstruction = new DefineFunctionInstruction(
             node.name(),
             positionalParams,
             namedParams,
-            node.freeVariables(),
+            freeVariables.build(),
             body,
             node.source()
         );
 
-        var storeInstruction = new LocalStoreInstruction(node.name(), node.source());
+        var localName = context.shedVariableToLocal(node.name());
+        var storeInstruction = new LocalStoreInstruction(localName, node.source());
 
         return context
             .addInstruction(defineFunctionInstruction)
@@ -210,17 +229,23 @@ public class StackCompiler {
         Context context
     ) {
         if (conditionalBranches.isEmpty()) {
-            return compileBlock(elseBody, context);
+            context = context.enterBlock(elseBody.boundVariableNames());
+            context = compileBlock(elseBody, context);
+            var elseInstructions = context.instructions;
+            context = context.exitBlock();
+            for (var instruction : elseInstructions) {
+                context = context.addInstruction(instruction);
+            }
         } else {
             var conditionalBranch = conditionalBranches.get(0);
             context = compileExpression(conditionalBranch.condition(), context);
 
-            context = context.enterBlock();
+            context = context.enterBlock(conditionalBranch.body().boundVariableNames());
             context = compileBlock(conditionalBranch.body(), context);
             var ifTrueInstructions = context.instructions;
             context = context.exitBlock();
 
-            context = context.enterBlock();
+            context = context.enterBlock(List.of());
             context = compileIf(
                 conditionalBranches.subList(1, conditionalBranches.size()),
                 elseBody,
@@ -246,18 +271,27 @@ public class StackCompiler {
     public static StackIrNamespace compileNamespace(TypedNamespaceNode node) {
         var context = Context.EMPTY;
 
+        context = context.addShedVariable(ShedBuiltins.PRINT_NAME);
+        var printLocalName = context.shedVariableToLocal(ShedBuiltins.PRINT_NAME);
+
         context = context
             .addInstruction(new NamespaceImportInstruction(ShedBuiltins.NAMESPACE_NAME, node.source()))
             .addInstruction(new FieldLoadInstruction(ShedBuiltins.NAMESPACE_NAME, ShedBuiltins.PRINT_NAME, node.source()))
-            .addInstruction(new LocalStoreInstruction(ShedBuiltins.PRINT_NAME, node.source()));
+            .addInstruction(new LocalStoreInstruction(printLocalName, node.source()));
+
+        for (var boundVariableName : node.boundVariableNames()) {
+            context = context.addShedVariable(boundVariableName);
+        }
 
         for (var statement : node.body()) {
             context = compileNamespaceStatement(statement, context);
         }
 
         for (var fieldName : node.fieldNames()) {
+            var localName = context.shedVariableToLocal(fieldName);
+
             context = context
-                .addInstruction(new LocalLoadInstruction(fieldName, node.source()))
+                .addInstruction(new LocalLoadInstruction(localName, node.source()))
                 .addInstruction(new NamespaceFieldStoreInstruction(
                     node.namespaceName(),
                     fieldName,
@@ -279,7 +313,9 @@ public class StackCompiler {
     }
 
     private static Context compileReference(TypedReferenceNode node, Context context) {
-        return context.addInstruction(new LocalLoadInstruction(node.name(), node.source()));
+        var localName = context.shedVariableToLocal(node.name());
+
+        return context.addInstruction(new LocalLoadInstruction(localName, node.source()));
     }
 
     private static Context compileStringLiteral(
@@ -292,8 +328,16 @@ public class StackCompiler {
 
     private static Context compileTarget(TypedTargetNode target, Context context) {
         return switch (target) {
-            case TypedTargetNode.Variable variable ->
-                context.addInstruction(new LocalStoreInstruction(variable.name(), variable.source()));
+            case TypedTargetNode.Variable variable -> {
+                context = context.addInstruction(
+                    new LocalStoreInstruction(
+                        context.shedVariableToLocal(variable.name()),
+                        variable.source()
+                    )
+                );
+
+                yield context;
+            }
         };
     }
 
@@ -306,38 +350,101 @@ public class StackCompiler {
     public static class Context {
         public static final Context EMPTY = new Context(
             P.stack(),
-            P.vector()
+            P.vector(),
+            Empty.set(),
+            P.stack(),
+            P.map()
         );
 
         private final PStack<PVector<Instruction>> instructionsStack;
         private final PVector<Instruction> instructions;
-
-        private Context(PStack<PVector<Instruction>> instructionsStack, PVector<Instruction> instructions) {
+        private final PSet<String> localNames;
+        private final PStack<PMap<String, String>> shedVariableNameToLocalNameStack;
+        private final PMap<String, String> shedVariableNameToLocalName;
+
+        private Context(
+            PStack<PVector<Instruction>> instructionsStack,
+            PVector<Instruction> instructions,
+            PSet<String> localNames,
+            PStack<PMap<String, String>> shedVariableNameToLocalNameStack, PMap<String, String> shedVariableNameToLocalName
+        ) {
             this.instructionsStack = instructionsStack;
             this.instructions = instructions;
+            this.localNames = localNames;
+            this.shedVariableNameToLocalNameStack = shedVariableNameToLocalNameStack;
+            this.shedVariableNameToLocalName = shedVariableNameToLocalName;
         }
 
-        public Context enterBlock() {
-            return new Context(
+        public Context enterBlock(List<String> shedVariableNames) {
+            var context = new Context(
                 this.instructionsStack.plus(this.instructions),
-                P.vector()
+                P.vector(),
+                this.localNames,
+                this.shedVariableNameToLocalNameStack.plus(this.shedVariableNameToLocalName),
+                this.shedVariableNameToLocalName
             );
+
+            for (var shedVariableName : shedVariableNames) {
+                context = context.addShedVariable(shedVariableName);
+            }
+
+            return context;
         }
 
         public Context exitBlock() {
             return new Context(
                 this.instructionsStack.minus(0),
-                this.instructionsStack.get(0)
+                this.instructionsStack.get(0),
+                this.localNames,
+                this.shedVariableNameToLocalNameStack.minus(0),
+                this.shedVariableNameToLocalNameStack.get(0)
             );
         }
 
         private Context addInstruction(Instruction instruction) {
             return new Context(
                 this.instructionsStack,
-                this.instructions.plus(instruction)
+                this.instructions.plus(instruction),
+                this.localNames,
+                this.shedVariableNameToLocalNameStack,
+                this.shedVariableNameToLocalName
             );
         }
 
+        public String shedVariableToLocal(String variableName) {
+            var localName = this.shedVariableNameToLocalName.get(variableName);
+
+            if (localName == null) {
+                throw new RuntimeException("No local name for Shed variable: " + variableName);
+            }
+
+            return localName;
+        }
+
+        public Context addShedVariable(String variableName) {
+            var localName = generateLocalNameForVariable(variableName);
+
+            return new Context(
+                this.instructionsStack,
+                this.instructions,
+                this.localNames.plus(localName),
+                this.shedVariableNameToLocalNameStack,
+                this.shedVariableNameToLocalName.plus(variableName, localName)
+            );
+        }
+
+        private String generateLocalNameForVariable(String variableName) {
+            var index = 0;
+            while (true) {
+                var localName = variableName + "__" + index;
+                if (this.localNames.contains(localName)) {
+                    index++;
+                } else {
+                    return localName;
+                }
+            }
+        }
+
         public List<Instruction> instructions() {
             return this.instructions;
         }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
index af7c0828..5627d462 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
@@ -36,6 +36,17 @@ class DeclaredVariables {
         };
     }
 
+    public static PVector<UntypedVariableDeclarationNode> forNamespaceStatements(
+        List<UntypedNamespaceStatementNode> statements,
+        PVector<UntypedVariableDeclarationNode> variables
+    ) {
+        for (var statement : statements) {
+            variables = DeclaredVariables.forNamespaceStatement(statement, variables);
+        }
+
+        return variables;
+    }
+
     static PVector<UntypedVariableDeclarationNode> forNamespaceStatement(
         UntypedNamespaceStatementNode 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 bd02d88b..6d96a98f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -648,6 +648,8 @@ public class TypeChecker {
 
         dependencyChecker.check();
 
+        var declaredVariables = DeclaredVariables.forNamespaceStatements(untyped.body(), P.vector());
+
         var fieldTypes = typedBody.stream()
             .flatMap(bodyStatement ->
                 bodyStatement instanceof TypedFunctionDefinitionNode functionDefinition
@@ -661,6 +663,9 @@ public class TypeChecker {
         var typed = new TypedNamespaceNode(
             untyped.namespaceName(),
             typedBody,
+            declaredVariables.stream()
+                .map(boundVariable -> boundVariable.name())
+                .toList(),
             fieldNames,
             untyped.source()
         );