Compile namespace-field-store to LLVM IR

9 files changed, 193 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
index 322c9a4f..8b933bad 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompiler.java
@@ -1,6 +1,10 @@
 package org.shedlang.compiler.backends.llvm;
 
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmGetElementPtr;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmStore;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperandInt;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmTypedOperand;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.NamespaceName;
 
@@ -34,8 +38,29 @@ public class LlvmCompiler {
             }
 
             case NamespaceFieldStoreInstruction namespaceFieldStore -> {
-                throw new UnsupportedOperationException("TODO");
+                var popResult = context.pop();
+                var fieldValue = popResult.value();
+                context = popResult.context();
+
+                var namespaceFields = context.typeInfo().namespaceFields(
+                    namespaceFieldStore.namespaceName()
+                );
+                var fieldPointerName = "fieldPointer__" + namespaceFieldStore.fieldName();
+                yield context.addInstruction(new LlvmGetElementPtr(
+                    new LlvmVariableLocal(fieldPointerName),
+                    LlvmCompilerNamespaces.typeNamespacePointer(namespaceFields),
+                    LlvmCompilerNamespaces.operandNamespaceValue(namespaceFieldStore.namespaceName()),
+                    List.of(
+                        LlvmTypedOperand.arrayIndex(0),
+                        LlvmTypedOperand.structureIndex(namespaceFields.keys().indexOf(namespaceFieldStore.fieldName()))
+                    )
+                )).addInstruction(new LlvmStore(
+                    LlvmCompilerTypes.VALUE_TYPE,
+                    fieldValue,
+                    new LlvmVariableLocal(fieldPointerName)
+                ));
             }
+
             case PushBoolInstruction pushBool ->
                 context.push(new LlvmOperandInt(pushBool.value() ? 1 : 0));
 
@@ -49,6 +74,7 @@ public class LlvmCompiler {
 
                 yield context.push(LlvmCompilerStrings.operandAsShedObject(globalDefinition));
             }
+
             case PushUnitInstruction pushUnit ->
                 context.push(new LlvmOperandInt(0));
         };
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 b24025b6..144575c0 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerContext.java
@@ -8,25 +8,33 @@ 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.typechecker.TypeInfo;
 import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.Collection;
 
 public record LlvmCompilerContext(
+    TypeInfo typeInfo,
     PStack<LlvmOperand> stack,
     PMap<String, LlvmOperand> locals,
     PVector<LlvmInstruction> instructions,
     POrderedMap<String, LlvmTopLevelEntity> topLevelEntitiesByName
 ) {
-    public static LlvmCompilerContext EMPTY = new LlvmCompilerContext(
-        Empty.stack(),
-        Empty.map(),
-        Empty.vector(),
-        POrderedMap.empty()
-    );
+    public static LlvmCompilerContext EMPTY = fromTypeInfo(TypeInfo.EMPTY);
+
+    public static LlvmCompilerContext fromTypeInfo(TypeInfo typeInfo) {
+        return new LlvmCompilerContext(
+            typeInfo,
+            Empty.stack(),
+            Empty.map(),
+            Empty.vector(),
+            POrderedMap.empty()
+        );
+    }
 
     public LlvmCompilerContext push(LlvmOperand operand) {
         return new LlvmCompilerContext(
+            typeInfo,
             stack.plus(operand),
             locals,
             instructions,
@@ -36,6 +44,7 @@ public record LlvmCompilerContext(
 
     public LlvmCompilerResult<LlvmOperand> pop() {
         return new LlvmCompilerContext(
+            typeInfo,
             stack.minus(0),
             locals,
             instructions,
@@ -45,6 +54,7 @@ public record LlvmCompilerContext(
 
     public LlvmCompilerContext addLocal(String name, LlvmOperand value) {
         return new LlvmCompilerContext(
+            typeInfo,
             stack,
             locals.plus(name, value),
             instructions,
@@ -58,6 +68,7 @@ public record LlvmCompilerContext(
 
     public LlvmCompilerContext addInstruction(LlvmInstruction instruction) {
         return new LlvmCompilerContext(
+            typeInfo,
             stack,
             locals,
             instructions.plus(instruction),
@@ -78,6 +89,7 @@ public record LlvmCompilerContext(
         );
 
         return new LlvmCompilerContext(
+            typeInfo,
             stack,
             locals,
             instructions,
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaces.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaces.java
new file mode 100644
index 00000000..85caf95e
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaces.java
@@ -0,0 +1,34 @@
+package org.shedlang.compiler.backends.llvm;
+
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypePointer;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypeStructure;
+import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.util.POrderedMap;
+
+public class LlvmCompilerNamespaces {
+    private LlvmCompilerNamespaces() {
+    }
+
+    public static LlvmTypePointer typeNamespacePointer(POrderedMap<String, Type> namespaceFields) {
+        return new LlvmTypePointer(typeNamespaceStruct(namespaceFields));
+    }
+
+    public static LlvmType typeNamespaceStruct(POrderedMap<String, Type> namespaceFields) {
+        return new LlvmTypeStructure(
+            namespaceFields.values().stream()
+                .map(x -> LlvmCompilerTypes.VALUE_TYPE)
+                .toList()
+        );
+    }
+
+    public static LlvmVariableGlobal operandNamespaceValue(NamespaceName namespaceName) {
+        return new LlvmVariableGlobal(globalNameNamespaceValue(namespaceName));
+    }
+
+    private static String globalNameNamespaceValue(NamespaceName namespaceName) {
+        return "shed__namespace__" + String.join("_", namespaceName.parts());
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/operands/LlvmTypedOperand.java b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/operands/LlvmTypedOperand.java
index 91be449f..85e39cad 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/operands/LlvmTypedOperand.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/operands/LlvmTypedOperand.java
@@ -1,9 +1,17 @@
 package org.shedlang.compiler.backends.llvm.ir.ast.operands;
 
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 
 public record LlvmTypedOperand(
     LlvmType type,
     LlvmOperand operand
 ) {
+    public static LlvmTypedOperand arrayIndex(long index) {
+        return new LlvmTypedOperand(LlvmTypes.INT_64, new LlvmOperandInt(index));
+    }
+
+    public static LlvmTypedOperand structureIndex(long index) {
+        return new LlvmTypedOperand(LlvmTypes.INT_32, new LlvmOperandInt(index));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 3429be63..93207cd3 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -113,4 +113,8 @@ public class TypeCheckerContext {
 
         return new TypeCheckerResult<>(mappedValues, context);
     }
+
+    public TypeInfo toTypeInfo() {
+        return new TypeInfo(namespaceFields);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfo.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfo.java
new file mode 100644
index 00000000..cfce56ef
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfo.java
@@ -0,0 +1,24 @@
+package org.shedlang.compiler.typechecker;
+
+import org.pcollections.Empty;
+import org.pcollections.PMap;
+import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.util.POrderedMap;
+
+public class TypeInfo {
+    public static final TypeInfo EMPTY = new TypeInfo(Empty.map());
+
+    private final PMap<NamespaceName, POrderedMap<String, Type>> namespaceFields;
+
+    public TypeInfo(
+        PMap<NamespaceName, POrderedMap<String, Type>> namespaceFields
+    ) {
+        this.namespaceFields = namespaceFields;
+    }
+
+    public POrderedMap<String, Type> namespaceFields(NamespaceName namespaceName) {
+        // TODO: handle missing key
+        return namespaceFields.get(namespaceName);
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/util/POrderedMap.java b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
index aabec1d7..6d273cdc 100644
--- a/src/main/java/org/shedlang/compiler/util/POrderedMap.java
+++ b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
@@ -30,6 +30,10 @@ public class POrderedMap<K, V> {
         return map.containsKey(key);
     }
 
+    public PVector<K> keys() {
+        return keys;
+    }
+
     public Collection<V> values() {
         return map.values();
     }
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaceFieldStoreTests.java b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaceFieldStoreTests.java
new file mode 100644
index 00000000..7877ed87
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCompilerNamespaceFieldStoreTests.java
@@ -0,0 +1,67 @@
+package org.shedlang.compiler.backends.llvm;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmIr;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmGetElementPtr;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmStore;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
+import org.shedlang.compiler.stackir.instructions.Instructions;
+import org.shedlang.compiler.typechecker.TypeCheckerContext;
+import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.Types;
+
+import java.util.List;
+
+import static org.shedlang.compiler.backends.llvm.ir.ast.LlvmIrMatchers.*;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.*;
+
+public class LlvmCompilerNamespaceFieldStoreTests {
+    @Test
+    public void namespaceFieldStorePopsFromStackAndUpdatesFieldOnNamespace() {
+        var namespaceName = new NamespaceName(List.of("Example", "Project"));
+        var instruction = Instructions.namespaceFieldStore(namespaceName, "y");
+        var typeInfo = TypeCheckerContext.EMPTY
+            .addNamespaceField(namespaceName, "x", Types.INT)
+            .addNamespaceField(namespaceName, "y", Types.INT)
+            .addNamespaceField(namespaceName, "z", Types.INT)
+            .toTypeInfo();
+        var context = LlvmCompilerContext.fromTypeInfo(typeInfo)
+            .push(LlvmIr.operandInt(0))
+            .push(LlvmIr.operandInt(1));
+
+        var result = LlvmCompiler.compileInstruction(instruction, context);
+
+        assertThat(result.instructions(), isSequence(
+            instanceOf(
+                LlvmGetElementPtr.class,
+                has("target", x -> x.target(), isLlvmVariableLocal("fieldPointer__y")),
+                has(
+                    "pointerType",
+                    x -> x.pointerType(),
+                    equalTo(LlvmTypes.pointer(LlvmTypes.structure(List.of(
+                        LlvmTypes.INT_64,
+                        LlvmTypes.INT_64,
+                        LlvmTypes.INT_64
+                    ))))
+                ),
+                has(
+                    "pointer",
+                    x -> x.pointer(),
+                    isLlvmVariableGlobal("shed__namespace__Example_Project")
+                ),
+                has("indexes", x -> x.indexes(), isSequence(
+                    isLlvmTypedOperand(LlvmTypes.INT_64, isLlvmOperandInt(0)),
+                    isLlvmTypedOperand(LlvmTypes.INT_32, isLlvmOperandInt(1))
+                ))
+            ),
+            instanceOf(
+                LlvmStore.class,
+                has("type", x -> x.type(), equalTo(LlvmTypes.INT_64)),
+                has("value", x -> x.value(), isLlvmOperandInt(1)),
+                has("pointer", x -> x.pointer(), isLlvmVariableLocal("fieldPointer__y"))
+            )
+        ));
+        assertThat(result.stack(), isSequence(isLlvmOperandInt(0)));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIrMatchers.java b/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIrMatchers.java
index 8582ce93..21755f6e 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIrMatchers.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIrMatchers.java
@@ -30,6 +30,13 @@ public class LlvmIrMatchers {
         );
     }
 
+    public static Matcher<LlvmTypedOperand> isLlvmTypedOperand(LlvmType type, Matcher<LlvmOperand> operand) {
+        return allOf(
+            has("type", x -> x.type(), equalTo(type)),
+            has("operand", x -> x.operand(), operand)
+        );
+    }
+
     public static Matcher<LlvmOperand> isLlvmVariableGlobal(String name) {
         return instanceOf(
             LlvmVariableGlobal.class,