Allocate separate space for stack IR locals in Patio locals

4 files changed, 60 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java b/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
index 90228a02..d8ca3fc2 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
@@ -20,6 +20,17 @@ public record NameAllocator<TPrefix>(PSet<String> existingNames) {
         return new Result<>(fullName, newAllocator);
     }
 
+    public Result<TPrefix> allocateFixedName(TPrefix prefix, String name) {
+        var fullName = prefix + "__" + name;
+
+        if (existingNames.contains(fullName)) {
+            throw new RuntimeException("name is already allocated: " + fullName);
+        }
+
+        var newAllocator = new NameAllocator<TPrefix>(existingNames.plus(fullName));
+        return new Result<>(fullName, newAllocator);
+    }
+
     public static <TPrefix> String allocateName(
         Predicate<String> isUsed,
         TPrefix prefix,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 1419402a..d4a31c3f 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -72,19 +72,27 @@ public class PatioCodeGen {
             case DiscardInstruction discard ->
                 context.addInstruction(new PatioDiscard(discard.source()));
 
-            case LocalLoadInstruction localLoad ->
-                context.addInstruction(new PatioLocalLoad(
-                    // TODO: allocate local name
-                    new PatioLocal(localLoad.name()),
+            case LocalLoadInstruction localLoad -> {
+                var localResult = context.stackIrLocalNameToLocal(localLoad.name());
+                var local = localResult.value();
+                context = localResult.context();
+
+                yield context.addInstruction(new PatioLocalLoad(
+                    local,
                     localLoad.source()
                 ));
+            }
 
-            case LocalStoreInstruction localStore ->
-                // TODO: allocate local name
-                context.addInstruction(new PatioLocalStore(
-                    new PatioLocal(localStore.name()),
+            case LocalStoreInstruction localStore -> {
+                var localResult = context.stackIrLocalNameToLocal(localStore.name());
+                var local = localResult.value();
+                context = localResult.context();
+
+                yield context.addInstruction(new PatioLocalStore(
+                    local,
                     localStore.source()
                 ));
+            }
 
             case PushBoolInstruction pushBool ->
                 context.addInstruction(new PatioPushInt(pushBool.value() ? 1 : 0, pushBool.source()));
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index dbb1b686..c10eced7 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.patio;
 
+import org.pcollections.PMap;
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.NameAllocator;
 import org.shedlang.compiler.patio.instructions.PatioInstruction;
@@ -13,17 +14,20 @@ public record PatioCodeGenContext(
 ) {
     public record FunctionContext(
         PVector<PatioInstruction> instructions,
-        NameAllocator<PatioCodeGenLocalSpace> localNames
+        NameAllocator<PatioCodeGenLocalSpace> localNames,
+        PMap<String, PatioLocal> stackIrNameToLocal
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
             P.vector(),
-            NameAllocator.empty()
+            NameAllocator.empty(),
+            P.map()
         );
 
         public FunctionContext addInstruction(PatioInstruction instruction) {
             return new FunctionContext(
                 this.instructions.plus(instruction),
-                this.localNames
+                this.localNames,
+                this.stackIrNameToLocal
             );
         }
 
@@ -32,10 +36,27 @@ public record PatioCodeGenContext(
 
             return new FunctionContext(
                 this.instructions,
-                result.allocator()
+                result.allocator(),
+                this.stackIrNameToLocal
             ).toResult(new PatioLocal(result.name()));
         }
 
+        public Result<PatioLocal> stackIrLocalNameToLocal(String name) {
+            var context = this;
+
+            if (!this.stackIrNameToLocal.containsKey(name)) {
+                var result = this.localNames.allocateFixedName(PatioCodeGenLocalSpace.LOCAL, name);
+
+                context = new FunctionContext(
+                    this.instructions,
+                    result.allocator(),
+                    this.stackIrNameToLocal.plus(name, new PatioLocal(result.name()))
+                );
+            }
+
+            return context.toResult(context.stackIrNameToLocal.get(name));
+        }
+
         private <T> Result<T> toResult(T value) {
             return new Result<>(value, this);
         }
@@ -65,6 +86,12 @@ public record PatioCodeGenContext(
         return new PatioCodeGenContext(result.context).toResult(result.value);
     }
 
+    public PatioCodeGenResult<PatioLocal> stackIrLocalNameToLocal(String name) {
+        var result = this.functionContext.stackIrLocalNameToLocal(name);
+
+        return new PatioCodeGenContext(result.context).toResult(result.value);
+    }
+
     private <T> PatioCodeGenResult<T> toResult(T value) {
         return new PatioCodeGenResult<>(value, this);
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
index 04a56095..942ae6f3 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
@@ -4,7 +4,8 @@ public enum PatioCodeGenLocalSpace {
     CALL_ARG_POSITIONAL("callArgPositional"),
     CALL_ARG_NAMED("callArgNamed"),
     CALL_FUNCTION_POINTER("callFunctionPointer"),
-    RECEIVER("receiver");;
+    LOCAL("local"),
+    RECEIVER("receiver");
 
     private final String name;