Create FreeVariable type with index and name

6 files changed, 57 insertions, 29 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java
new file mode 100644
index 00000000..ff56adf5
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariable.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.ast.typed;
+
+public record FreeVariable(int index, String name) {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
index 6f5cbd44..218fb696 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
@@ -6,35 +6,62 @@ import java.util.Arrays;
 import java.util.Iterator;
 import java.util.stream.Stream;
 
-public class FreeVariables implements Iterable<String> {
+public class FreeVariables implements Iterable<FreeVariable> {
     public static FreeVariables empty() {
         return new FreeVariables(OrderedPSet.empty());
     }
 
-    public static FreeVariables of(String... freeVariables) {
-        return new FreeVariables(OrderedPSet.from(Arrays.asList(freeVariables)));
+    public static FreeVariables of(String... freeVariableNames) {
+        return of(Arrays.asList(freeVariableNames));
     }
 
-    private final OrderedPSet<String> freeVariables;
-
-    private FreeVariables(OrderedPSet<String> freeVariables) {
-        this.freeVariables = freeVariables;
+    public static FreeVariables of(Iterable<String> freeVariableNames) {
+        var freeVariables = OrderedPSet.<FreeVariable>empty();
+        var freeVariableIndex = 0;
+        for (var freeVariableName : freeVariableNames) {
+            var freeVariable = new FreeVariable(freeVariableIndex, freeVariableName);
+            freeVariables = freeVariables.plus(freeVariable);
+            freeVariableIndex++;
+        }
+        return new FreeVariables(freeVariables);
     }
 
-    public FreeVariables plus(String name) {
-        return new FreeVariables(this.freeVariables.plus(name));
+    private final OrderedPSet<FreeVariable> freeVariables;
+
+    private FreeVariables(OrderedPSet<FreeVariable> freeVariables) {
+        this.freeVariables = freeVariables;
     }
 
     @Override
-    public Iterator<String> iterator() {
+    public Iterator<FreeVariable> iterator() {
         return freeVariables.iterator();
     }
 
-    public Stream<String> stream() {
+    public Stream<FreeVariable> stream() {
         return freeVariables.stream();
     }
 
     public int size() {
         return freeVariables.size();
     }
+
+    public static Builder builder() {
+        return new Builder(OrderedPSet.empty());
+    }
+
+    public static class Builder {
+        private final OrderedPSet<String> freeVariables;
+
+        private Builder(OrderedPSet<String> freeVariables) {
+            this.freeVariables = freeVariables;
+        }
+
+        public Builder plus(String name) {
+            return new Builder(this.freeVariables.plus(name));
+        }
+
+        public FreeVariables build() {
+            return FreeVariables.of(this.freeVariables);
+        }
+    }
 }
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 739c1ec3..417b9ad6 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -55,7 +55,7 @@ public record TypedFunctionDefinitionNode(
             Empty.vector(),
             Optional.empty(),
             new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
-            FreeVariables.empty(),
+            FreeVariables.builder(),
             Empty.vector()
         );
     }
@@ -66,7 +66,7 @@ public record TypedFunctionDefinitionNode(
         PVector<TypedParamNode> namedParams,
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
-        FreeVariables freeVariables,
+        FreeVariables.Builder freeVariables,
         PVector<TypedFunctionStatementNode> body
     ) {
         public TypedFunctionDefinitionNode build() {
@@ -75,7 +75,7 @@ public record TypedFunctionDefinitionNode(
                 new TypedParamsNode(positionalParams, namedParams, Sources.NULL),
                 effect,
                 returnType,
-                freeVariables,
+                freeVariables.build(),
                 body,
                 Sources.NULL
             );
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
index 45b7e920..8e17ea1b 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerClosures.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.backends.llvm;
 
+import org.shedlang.compiler.ast.typed.FreeVariable;
 import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDefinition;
@@ -114,16 +115,13 @@ public class LlvmCompilerClosures {
             context
         );
 
-        var freeVariableIndex = 0;
         for (var freeVariable : freeVariables) {
             context = storeClosureFreeVariablePointer(
                 closurePointer,
-                freeVariableIndex,
                 freeVariable,
                 closureTypes,
                 context
             );
-            freeVariableIndex++;
         }
 
         var closureValueResult = castClosurePointerToValue(closurePointer, context);
@@ -299,14 +297,13 @@ public class LlvmCompilerClosures {
 
     private static LlvmCompilerContext storeClosureFreeVariablePointer(
         LlvmOperand closurePointer,
-        int freeVariableIndex,
-        String freeVariable,
+        FreeVariable freeVariable,
         Types closureTypes,
         LlvmCompilerContext context
     ) {
         var closureEnvironmentPointerResult = getClosureFreeVariablePointer(
             closurePointer,
-            freeVariableIndex,
+            freeVariable,
             closureTypes,
             context
         );
@@ -315,14 +312,14 @@ public class LlvmCompilerClosures {
 
         return context.addInstruction(new LlvmStore(
             LlvmCompilerTypes.VALUE_TYPE,
-            context.local(freeVariable),
+            context.local(freeVariable.name()),
             closureEnvironmentPointer
         ));
     }
 
     private static LlvmCompilerResult<LlvmVariableLocal> getClosureFreeVariablePointer(
         LlvmOperand closurePointer,
-        int freeVariableIndex,
+        FreeVariable freeVariable,
         Types closureTypes,
         LlvmCompilerContext context
     ) {
@@ -339,7 +336,7 @@ public class LlvmCompilerClosures {
             List.of(
                 LlvmTypedOperand.arrayIndex(0),
                 LlvmTypedOperand.structureIndex(1),
-                LlvmTypedOperand.arrayIndex(freeVariableIndex)
+                LlvmTypedOperand.arrayIndex(freeVariable.index())
             )
         ));
 
diff --git a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
index a9e6350f..12592b42 100644
--- a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
+++ b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
@@ -60,7 +60,7 @@ public class InstructionSerializer {
                     )),
                     treeNamed("free-variables", treeList(
                         defineFunctionInstruction.freeVariables().stream()
-                            .map(freeVariable -> treeString(freeVariable))
+                            .map(freeVariable -> treeString(freeVariable.name()))
                             .toList()
                     )),
                     treeNamed("body", treeListIndented(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 8e63189b..69080884 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -18,19 +18,19 @@ public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
         Set.of(),
-        FreeVariables.empty(),
+        FreeVariables.builder(),
         Empty.map()
     );
 
     private final PMap<String, Type> environment;
     private final Set<String> boundVariables;
-    private final FreeVariables freeVariables;
+    private final FreeVariables.Builder freeVariables;
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
 
     public TypeCheckerContext(
         PMap<String, Type> environment,
         Set<String> boundVariables,
-        FreeVariables freeVariables,
+        FreeVariables.Builder freeVariables,
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo
     ) {
         this.environment = environment;
@@ -43,7 +43,7 @@ public class TypeCheckerContext {
         return new TypeCheckerContext(
             this.environment.plusAll(environment),
             environment.keySet(),
-            FreeVariables.empty(),
+            FreeVariables.builder(),
             this.namespaceTypeInfo
         );
     }
@@ -69,7 +69,7 @@ public class TypeCheckerContext {
     }
 
     public FreeVariables freeVariables() {
-        return freeVariables;
+        return freeVariables.build();
     }
 
     public TypeCheckerContext addNamespaceFields(