Extract NameAllocator

2 files changed, 47 insertions, 23 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
index f6f4c9c6..3ebf7b98 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
@@ -18,7 +18,7 @@ public record LlvmCompilerContext(
     PMap<String, LlvmOperand> locals,
     PVector<LlvmInstruction> instructions,
     POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName,
-    PSet<String> localVariableNames
+    NameAllocator<LlvmCompilerLocalVariableSpace> localVariableNames
 ) {
     public static LlvmCompilerContext EMPTY = fromTypeInfo(TypeInfo.EMPTY);
 
@@ -29,7 +29,7 @@ public record LlvmCompilerContext(
             Empty.map(),
             Empty.vector(),
             POrderedMap.empty(),
-            Empty.set()
+            NameAllocator.empty()
         );
     }
 
@@ -134,27 +134,15 @@ public record LlvmCompilerContext(
         LlvmCompilerLocalVariableSpace space,
         Optional<String> name
     ) {
-        var prefix = space.toString();
-        if (name.isPresent()) {
-            prefix += "__" + name.get();
-        }
-
-        var index = 0;
-        while (true) {
-            var fullName = prefix + "__" + index;
-            if (localVariableNames.contains(fullName)) {
-                index++;
-            } else {
-                return new LlvmCompilerContext(
-                    typeInfo,
-                    stack,
-                    locals,
-                    instructions,
-                    topLevelEntitiesByName,
-                    localVariableNames.plus(fullName)
-                ).toResult(new LlvmVariableLocal(fullName));
-            }
-        }
+        var result = localVariableNames.generateName(space, name);
+        return new LlvmCompilerContext(
+            typeInfo,
+            stack,
+            locals,
+            instructions,
+            topLevelEntitiesByName,
+            result.allocator()
+        ).toResult(new LlvmVariableLocal(result.name()));
     }
 
     public <T> LlvmCompilerResult<T> toResult(T value) {
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java b/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
new file mode 100644
index 00000000..d8e37d48
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
@@ -0,0 +1,36 @@
+package org.shedlang.compiler.backends.llvm;
+
+import org.pcollections.Empty;
+import org.pcollections.PSet;
+
+import java.util.Optional;
+
+public record NameAllocator<TPrefix>(PSet<String> localVariableNames) {
+    public static <TPrefix> NameAllocator<TPrefix> empty() {
+        return new NameAllocator<>(Empty.set());
+    }
+
+    public Result<TPrefix> generateName(
+        TPrefix prefix,
+        Optional<String> name
+    ) {
+        var fullPrefix = prefix.toString();
+        if (name.isPresent()) {
+            fullPrefix += "__" + name.get();
+        }
+
+        var index = 0;
+        while (true) {
+            var fullName = fullPrefix + "__" + index;
+            if (localVariableNames.contains(fullName)) {
+                index++;
+            } else {
+                var newAllocator = new NameAllocator<TPrefix>(localVariableNames.plus(fullName));
+                return new Result<>(fullName, newAllocator);
+            }
+        }
+    }
+
+    public record Result<TPrefix>(String name, NameAllocator<TPrefix> allocator) {
+    }
+}