Introduce FreeVariables collection type

5 files changed, 57 insertions, 20 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
new file mode 100644
index 00000000..6f5cbd44
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/FreeVariables.java
@@ -0,0 +1,40 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.pcollections.OrderedPSet;
+
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.stream.Stream;
+
+public class FreeVariables implements Iterable<String> {
+    public static FreeVariables empty() {
+        return new FreeVariables(OrderedPSet.empty());
+    }
+
+    public static FreeVariables of(String... freeVariables) {
+        return new FreeVariables(OrderedPSet.from(Arrays.asList(freeVariables)));
+    }
+
+    private final OrderedPSet<String> freeVariables;
+
+    private FreeVariables(OrderedPSet<String> freeVariables) {
+        this.freeVariables = freeVariables;
+    }
+
+    public FreeVariables plus(String name) {
+        return new FreeVariables(this.freeVariables.plus(name));
+    }
+
+    @Override
+    public Iterator<String> iterator() {
+        return freeVariables.iterator();
+    }
+
+    public Stream<String> stream() {
+        return freeVariables.stream();
+    }
+
+    public int size() {
+        return freeVariables.size();
+    }
+}
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 524e11a3..739c1ec3 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -1,7 +1,6 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.pcollections.Empty;
-import org.pcollections.OrderedPSet;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
@@ -9,14 +8,13 @@ import org.shedlang.compiler.types.*;
 
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 
 public record TypedFunctionDefinitionNode(
     String name,
     TypedParamsNode params,
     Optional<TypedTypeLevelExpressionNode<Effect>> effect,
     TypedTypeLevelExpressionNode<Type> returnType,
-    Set<String> freeVariables,
+    FreeVariables freeVariables,
     List<TypedFunctionStatementNode> body,
     FunctionType type,
     Source source
@@ -26,7 +24,7 @@ public record TypedFunctionDefinitionNode(
         TypedParamsNode params,
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
-        Set<String> freeVariables,
+        FreeVariables freeVariables,
         List<TypedFunctionStatementNode> body,
         Source source
     ) {
@@ -57,7 +55,7 @@ public record TypedFunctionDefinitionNode(
             Empty.vector(),
             Optional.empty(),
             new TypedTypeLevelReferenceNode<>("Unit", Types.metaType(Types.UNIT), Sources.NULL),
-            OrderedPSet.empty(),
+            FreeVariables.empty(),
             Empty.vector()
         );
     }
@@ -68,7 +66,7 @@ public record TypedFunctionDefinitionNode(
         PVector<TypedParamNode> namedParams,
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
-        OrderedPSet<String> freeVariables,
+        FreeVariables freeVariables,
         PVector<TypedFunctionStatementNode> body
     ) {
         public TypedFunctionDefinitionNode build() {
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 eece1c39..9db59c56 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.FreeVariables;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDefinition;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmParam;
@@ -13,13 +14,12 @@ import org.shedlang.compiler.backends.llvm.ir.ast.types.*;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 
 public class LlvmCompilerClosures {
     private LlvmCompilerClosures() {
     }
 
-    public record Types(Optional<Set<String>> freeVariables) {
+    public record Types(Optional<FreeVariables> freeVariables) {
         public LlvmTypePointer closurePointer() {
             return LlvmTypes.pointer(closureStruct());
         }
@@ -53,17 +53,17 @@ public class LlvmCompilerClosures {
 
     public static final Types CALLABLE_TYPES = new Types(Optional.empty());
 
-    public static Types types(Set<String> freeVariables) {
+    public static Types types(FreeVariables freeVariables) {
         return new Types(Optional.of(freeVariables));
     }
 
-    public static Types types(Optional<Set<String>> freeVariables) {
+    public static Types types(Optional<FreeVariables> freeVariables) {
         return new Types(freeVariables);
     }
 
     public static LlvmCompilerContext create(
         String name,
-        Set<String> freeVariables,
+        FreeVariables freeVariables,
         List<LlvmBasicBlock> body,
         LlvmCompilerContext context
     ) {
@@ -129,7 +129,6 @@ public class LlvmCompilerClosures {
         ));
 
         var freeVariableIndex = 0;
-        // TODO: change type of free variables to be list?
         for (var freeVariable : freeVariables) {
             var closureEnvironmentPointerResult = context.generateVariableLocal(
                 LlvmCompilerLocalVariableSpace.CLOSURE_ENVIRONMENT_POINTER
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
index 1c4a6ab5..07815898 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/DefineFunctionInstruction.java
@@ -1,9 +1,9 @@
 package org.shedlang.compiler.stackir.instructions;
 
+import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.sources.Source;
 
 import java.util.List;
-import java.util.Set;
 
 /**
  * Push a function value onto the stack with the given definition.
@@ -12,7 +12,7 @@ public record DefineFunctionInstruction(
     String name,
     List<Param> positionalParams,
     List<Param> namedParams,
-    Set<String> freeVariables,
+    FreeVariables freeVariables,
     List<Instruction> body,
     Source source
 ) implements Instruction {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index a3f63a2d..8e63189b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -1,8 +1,8 @@
 package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
-import org.pcollections.OrderedPSet;
 import org.pcollections.PMap;
+import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.NamespaceTypeInfo;
 import org.shedlang.compiler.types.Type;
@@ -18,19 +18,19 @@ public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
         Set.of(),
-        OrderedPSet.empty(),
+        FreeVariables.empty(),
         Empty.map()
     );
 
     private final PMap<String, Type> environment;
     private final Set<String> boundVariables;
-    private final OrderedPSet<String> freeVariables;
+    private final FreeVariables freeVariables;
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
 
     public TypeCheckerContext(
         PMap<String, Type> environment,
         Set<String> boundVariables,
-        OrderedPSet<String> freeVariables,
+        FreeVariables freeVariables,
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo
     ) {
         this.environment = environment;
@@ -43,7 +43,7 @@ public class TypeCheckerContext {
         return new TypeCheckerContext(
             this.environment.plusAll(environment),
             environment.keySet(),
-            OrderedPSet.empty(),
+            FreeVariables.empty(),
             this.namespaceTypeInfo
         );
     }
@@ -68,7 +68,7 @@ public class TypeCheckerContext {
         );
     }
 
-    public Set<String> freeVariables() {
+    public FreeVariables freeVariables() {
         return freeVariables;
     }