Extract POrderedMap

2 files changed, 56 insertions, 16 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 44f14d73..b24025b6 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
@@ -1,24 +1,28 @@
 package org.shedlang.compiler.backends.llvm;
 
-import org.pcollections.*;
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+import org.pcollections.PStack;
+import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmGlobalDefinition;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmTopLevelEntity;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmInstruction;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
+import org.shedlang.compiler.util.POrderedMap;
+
+import java.util.Collection;
 
 public record LlvmCompilerContext(
     PStack<LlvmOperand> stack,
     PMap<String, LlvmOperand> locals,
     PVector<LlvmInstruction> instructions,
-    PVector<LlvmTopLevelEntity> topLevelEntities,
-    PSet<String> topLevelEntityNames
+    POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName
 ) {
     public static LlvmCompilerContext EMPTY = new LlvmCompilerContext(
         Empty.stack(),
         Empty.map(),
         Empty.vector(),
-        Empty.vector(),
-        Empty.set()
+        POrderedMap.empty()
     );
 
     public LlvmCompilerContext push(LlvmOperand operand) {
@@ -26,8 +30,7 @@ public record LlvmCompilerContext(
             stack.plus(operand),
             locals,
             instructions,
-            topLevelEntities,
-            topLevelEntityNames
+            topLevelEntitiesByName
         );
     }
 
@@ -36,8 +39,7 @@ public record LlvmCompilerContext(
             stack.minus(0),
             locals,
             instructions,
-            topLevelEntities,
-            topLevelEntityNames
+            topLevelEntitiesByName
         ).toResult(stack.get(0));
     }
 
@@ -46,8 +48,7 @@ public record LlvmCompilerContext(
             stack,
             locals.plus(name, value),
             instructions,
-            topLevelEntities,
-            topLevelEntityNames
+            topLevelEntitiesByName
         );
     }
 
@@ -60,8 +61,7 @@ public record LlvmCompilerContext(
             stack,
             locals,
             instructions.plus(instruction),
-            topLevelEntities,
-            topLevelEntityNames
+            topLevelEntitiesByName
         );
     }
 
@@ -81,15 +81,14 @@ public record LlvmCompilerContext(
             stack,
             locals,
             instructions,
-            topLevelEntities.plus(globalDefinition),
-            topLevelEntityNames.plus(globalDefinition.name())
+            topLevelEntitiesByName.plus(globalDefinition.name(), globalDefinition)
         ).toResult(globalDefinition);
     }
 
     private String generateTopLevelEntityName(String prefix) {
         var name = prefix;
         var nextSuffix = 1;
-        while (topLevelEntityNames.contains(name)) {
+        while (topLevelEntitiesByName().containsKey(name)) {
             name = prefix + "_" + nextSuffix;
             nextSuffix++;
         }
@@ -99,4 +98,8 @@ public record LlvmCompilerContext(
     public <T> LlvmCompilerResult<T> toResult(T value) {
         return new LlvmCompilerResult<>(value, this);
     }
+
+    public Collection<LlvmTopLevelEntity> topLevelEntities() {
+        return topLevelEntitiesByName.values();
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/util/POrderedMap.java b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
new file mode 100644
index 00000000..400a4413
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
@@ -0,0 +1,37 @@
+package org.shedlang.compiler.util;
+
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+import org.pcollections.PVector;
+
+import java.util.Collection;
+
+public class POrderedMap<K, V> {
+    public static <K, V> POrderedMap<K, V> empty() {
+        return new POrderedMap<>(Empty.map(), Empty.vector());
+    }
+
+    private final PMap<K, V> map;
+    private final PVector<K> keys;
+
+    private POrderedMap(
+        PMap<K, V> map,
+        PVector<K> keys
+    ) {
+
+        this.map = map;
+        this.keys = keys;
+    }
+
+    public POrderedMap<K, V> plus(K key, V value) {
+        return new POrderedMap<>(map.plus(key, value), keys.plus(key));
+    }
+
+    public boolean containsKey(K key) {
+        return map.containsKey(key);
+    }
+
+    public Collection<V> values() {
+        return map.values();
+    }
+}