Add explicit locals

7 files changed, 104 insertions, 34 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 d8ca3fc2..7a1aba01 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/NameAllocator.java
@@ -21,12 +21,7 @@ public record NameAllocator<TPrefix>(PSet<String> existingNames) {
     }
 
     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 fullName = allocateFixedName(existingNames::contains, prefix, name);
         var newAllocator = new NameAllocator<TPrefix>(existingNames.plus(fullName));
         return new Result<>(fullName, newAllocator);
     }
@@ -52,6 +47,16 @@ public record NameAllocator<TPrefix>(PSet<String> existingNames) {
         }
     }
 
+    public static <TPrefix> String allocateFixedName(Predicate<String> isUsed, TPrefix prefix, String name) {
+        var fullName = prefix + "__" + name;
+
+        if (isUsed.test(fullName)) {
+            throw new RuntimeException("name is already allocated: " + fullName);
+        }
+
+        return fullName;
+    }
+
     public record Result<TPrefix>(String name, NameAllocator<TPrefix> allocator) {
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 781cece5..d31e92e5 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -1,7 +1,7 @@
 package org.shedlang.compiler.patio;
 
-import org.shedlang.compiler.patio.ast.instructions.PatioIntGet;
 import org.shedlang.compiler.patio.ast.PatioLocalRef;
+import org.shedlang.compiler.patio.ast.instructions.PatioIntGet;
 import org.shedlang.compiler.patio.ast.instructions.PatioReturn;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.stackir.instructions.*;
@@ -37,7 +37,8 @@ public class PatioCodeGen {
                 for (var argName : reversed(call.namedArgNames())) {
                     var argResult = context.allocateLocal(
                         PatioCodeGenLocalSpace.CALL_ARG_NAMED,
-                        argName
+                        argName,
+                        PatioCodeGenTypes.VALUE_TYPE
                     );
                     var arg = argResult.value();
                     context = argResult.context();
@@ -54,7 +55,8 @@ public class PatioCodeGen {
                 for (var i = 0; i < call.positionalArgCount(); i++) {
                     var argResult = context.allocateLocal(
                         PatioCodeGenLocalSpace.CALL_ARG_POSITIONAL,
-                        Integer.toString(call.positionalArgCount() - i - 1)
+                        Integer.toString(call.positionalArgCount() - i - 1),
+                        PatioCodeGenTypes.VALUE_TYPE
                     );
                     var arg = argResult.value();
                     context = argResult.context();
@@ -68,7 +70,11 @@ public class PatioCodeGen {
                 }
                 Collections.reverse(positionalArgs);
 
-                var receiverResult = context.allocateLocal(PatioCodeGenLocalSpace.RECEIVER);
+                var receiverResult = context.allocateLocal(
+                    PatioCodeGenLocalSpace.RECEIVER,
+                    // TODO: should be pointer type
+                    PatioCodeGenTypes.VALUE_TYPE
+                );
                 var receiver = receiverResult.value();
                 context = receiverResult.context();
 
@@ -77,7 +83,10 @@ public class PatioCodeGen {
                     stackIrInstruction.source()
                 );
 
-                var targetResult = context.allocateLocal(PatioCodeGenLocalSpace.TEMPORARY);
+                var targetResult = context.allocateLocal(
+                    PatioCodeGenLocalSpace.TEMPORARY,
+                    PatioCodeGenTypes.VALUE_TYPE
+                );
                 var target = targetResult.value();
                 context = targetResult.context();
 
@@ -123,13 +132,13 @@ public class PatioCodeGen {
             }
 
             case PushBoolInstruction pushBool ->
-                pushInt(pushBool.value() ? 1 : 0, pushBool.source(), context);
+                pushValue(pushBool.value() ? 1 : 0, pushBool.source(), context);
 
             case PushIntInstruction pushInt ->
-                pushInt(pushInt.value().intValueExact(), pushInt.source(), context);
+                pushValue(pushInt.value().intValueExact(), pushInt.source(), context);
 
             case PushUnitInstruction pushUnit ->
-                pushInt(0, pushUnit.source(), context);
+                pushValue(0, pushUnit.source(), context);
 
             case ReturnInstruction return_ -> {
                 var valueResult = context.temporaryPop();
@@ -144,8 +153,8 @@ public class PatioCodeGen {
         };
     }
 
-    private static PatioCodeGenContext pushInt(int value, Source source, PatioCodeGenContext context) {
-        var localResult = context.allocateLocal(PatioCodeGenLocalSpace.TEMPORARY);
+    private static PatioCodeGenContext pushValue(int value, Source source, PatioCodeGenContext context) {
+        var localResult = context.allocateLocal(PatioCodeGenLocalSpace.TEMPORARY, PatioCodeGenTypes.VALUE_TYPE);
         var local = localResult.value();
         context = localResult.context();
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 6308c59a..451c4d5e 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -4,6 +4,7 @@ import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.patio.ast.instructions.PatioCallPointer;
 import org.shedlang.compiler.patio.ast.PatioLocalRef;
 import org.shedlang.compiler.patio.ast.instructions.PatioMemoryLoad;
+import org.shedlang.compiler.patio.ast.types.PatioTypes;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.stackir.instructions.DefineFunctionInstruction;
 
@@ -50,7 +51,11 @@ class PatioCodeGenClosures {
         Source source,
         PatioCodeGenContext context
     ) {
-        var functionPointerResult = context.allocateLocal(PatioCodeGenLocalSpace.CALL_FUNCTION_POINTER);
+        var functionPointerResult = context.allocateLocal(
+            PatioCodeGenLocalSpace.CALL_FUNCTION_POINTER,
+            // TODO: should be function pointer
+            PatioTypes.INT
+        );
         var functionPointer = functionPointerResult.value();
         context = functionPointerResult.context();
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index c4d14e2c..8139a022 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -7,9 +7,11 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.NameAllocator;
 import org.shedlang.compiler.patio.ast.PatioCompilationUnitElement;
 import org.shedlang.compiler.patio.ast.PatioFunctionDefinition;
+import org.shedlang.compiler.patio.ast.PatioLocal;
 import org.shedlang.compiler.patio.ast.PatioLocalRef;
 import org.shedlang.compiler.patio.ast.instructions.PatioInstruction;
 import org.shedlang.compiler.patio.ast.instructions.PatioLocalCopy;
+import org.shedlang.compiler.patio.ast.types.PatioType;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.util.P;
 
@@ -56,13 +58,13 @@ public record PatioCodeGenContext(
     public record FunctionContext(
         PStack<PatioLocalRef> temporaryStack,
         PVector<PatioInstruction> instructions,
-        NameAllocator<PatioCodeGenLocalSpace> localNames,
+        OrderedPMap<String, PatioLocal> localsByName,
         PMap<String, PatioLocalRef> stackIrNameToLocal
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
             P.stack(),
             P.vector(),
-            NameAllocator.empty(),
+            OrderedPMap.empty(),
             P.map()
         );
 
@@ -70,7 +72,7 @@ public record PatioCodeGenContext(
             return new FunctionContext(
                 this.temporaryStack.minus(0),
                 this.instructions,
-                this.localNames,
+                this.localsByName,
                 this.stackIrNameToLocal
             ).toResult(this.temporaryStack.get(0));
         }
@@ -79,7 +81,7 @@ public record PatioCodeGenContext(
             return new FunctionContext(
                 this.temporaryStack.plus(value),
                 this.instructions,
-                this.localNames,
+                this.localsByName,
                 this.stackIrNameToLocal
             );
         }
@@ -88,33 +90,49 @@ public record PatioCodeGenContext(
             return new FunctionContext(
                 this.temporaryStack,
                 this.instructions.plus(instruction),
-                this.localNames,
+                this.localsByName,
                 this.stackIrNameToLocal
             );
         }
 
-        public Result<PatioLocalRef> allocateLocal(PatioCodeGenLocalSpace space, Optional<String> name) {
-            var result = this.localNames.allocateName(space, name);
+        public Result<PatioLocalRef> allocateLocal(
+            PatioCodeGenLocalSpace space,
+            Optional<String> name,
+            PatioType type
+        ) {
+            var fullName = NameAllocator.allocateName(
+                this.localsByName::containsKey,
+                space,
+                name
+            );
 
             return new FunctionContext(
                 this.temporaryStack,
                 this.instructions,
-                result.allocator(),
+                this.localsByName.plus(fullName, new PatioLocal(fullName, type)),
                 this.stackIrNameToLocal
-            ).toResult(new PatioLocalRef(result.name()));
+            ).toResult(new PatioLocalRef(fullName));
         }
 
         public Result<PatioLocalRef> stackIrLocalNameToLocal(String name) {
             var context = this;
 
             if (!this.stackIrNameToLocal.containsKey(name)) {
-                var result = this.localNames.allocateFixedName(PatioCodeGenLocalSpace.LOCAL, name);
+                // TODO: the mapping is redundant given we fix the mapping
+                var fullName = NameAllocator.allocateFixedName(
+                    this.localsByName::containsKey,
+                    PatioCodeGenLocalSpace.LOCAL,
+                    name
+                );
 
                 context = new FunctionContext(
                     this.temporaryStack,
                     this.instructions,
-                    result.allocator(),
-                    this.stackIrNameToLocal.plus(name, new PatioLocalRef(result.name()))
+                    this.localsByName.plus(
+                        fullName,
+                        new PatioLocal(fullName, PatioCodeGenTypes.VALUE_TYPE)
+                    ),
+                    this.stackIrNameToLocal.plus(name, new PatioLocalRef(fullName))
                 );
             }
 
@@ -168,14 +186,21 @@ public record PatioCodeGenContext(
         return withFunctionContext(this.functionContext.addInstruction(instruction));
     }
 
-    public PatioCodeGenResult<PatioLocalRef> allocateLocal(PatioCodeGenLocalSpace space, String name) {
-        var result = this.functionContext.allocateLocal(space, Optional.of(name));
+    public PatioCodeGenResult<PatioLocalRef> allocateLocal(
+        PatioCodeGenLocalSpace space,
+        String name,
+        PatioType type
+    ) {
+        var result = this.functionContext.allocateLocal(space, Optional.of(name), type);
 
         return withFunctionContextResult(result);
     }
 
-    public PatioCodeGenResult<PatioLocalRef> allocateLocal(PatioCodeGenLocalSpace space) {
-        var result = this.functionContext.allocateLocal(space, Optional.empty());
+    public PatioCodeGenResult<PatioLocalRef> allocateLocal(
+        PatioCodeGenLocalSpace space,
+        PatioType type
+    ) {
+        var result = this.functionContext.allocateLocal(space, Optional.empty(), type);
 
         return withFunctionContextResult(result);
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTypes.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTypes.java
new file mode 100644
index 00000000..d399fa08
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTypes.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.patio;
+
+import org.shedlang.compiler.patio.ast.types.PatioType;
+import org.shedlang.compiler.patio.ast.types.PatioTypes;
+
+public class PatioCodeGenTypes {
+    private PatioCodeGenTypes() {}
+
+    public static final PatioType VALUE_TYPE = PatioTypes.INT;
+}
diff --git a/src/main/java/org/shedlang/compiler/patio/ast/PatioLocal.java b/src/main/java/org/shedlang/compiler/patio/ast/PatioLocal.java
new file mode 100644
index 00000000..1cb7433b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/patio/ast/PatioLocal.java
@@ -0,0 +1,6 @@
+package org.shedlang.compiler.patio.ast;
+
+import org.shedlang.compiler.patio.ast.types.PatioType;
+
+public record PatioLocal(String name, PatioType type) {
+}
diff --git a/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java b/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
index 3b4dfde5..f5a90775 100644
--- a/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
+++ b/src/main/java/org/shedlang/compiler/patio/ast/PatioSerializer.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.patio.ast;
 
 import org.shedlang.compiler.patio.ast.instructions.*;
 import org.shedlang.compiler.patio.ast.types.PatioIntType;
+import org.shedlang.compiler.patio.ast.types.PatioType;
 import org.shedlang.compiler.treetext.Tree;
 
 import java.util.List;
@@ -30,6 +31,11 @@ public class PatioSerializer {
         serializeTree(tree, builder, 0);
     }
 
+    public static void serializeLocal(PatioLocal local, StringBuilder builder) {
+        var tree = localToTree(local);
+        serializeTree(tree, builder, 0);
+    }
+
     public static void serializeType(PatioIntType type, StringBuilder builder) {
         var tree = typeToTree(type);
         serializeTree(tree, builder, 0);
@@ -120,7 +126,11 @@ public class PatioSerializer {
         };
     }
 
-    private static Tree typeToTree(PatioIntType type) {
+    private static Tree localToTree(PatioLocal local) {
+        return treeList(treeString(local.name()), typeToTree(local.type()));
+    }
+
+    private static Tree typeToTree(PatioType type) {
         return switch (type) {
             case PatioIntType intType ->
                 treeIdentifier("int");