Generate main function using Patio code gen in llvm-patio backend

5 files changed, 67 insertions, 10 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
index 86c630fe..7cc4f85b 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmBackend.java
@@ -2,13 +2,22 @@ package org.shedlang.compiler.backends.llvm;
 
 import org.shedlang.compiler.backends.Backend;
 import org.shedlang.compiler.backends.llvm.ir.LlvmIrSerializer;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmBasicBlock;
+import org.shedlang.compiler.backends.llvm.ir.ast.LlvmFunctionDefinition;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmModule;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmCall;
+import org.shedlang.compiler.backends.llvm.ir.ast.instructions.LlvmReturn;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableGlobal;
+import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
+import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
 import org.shedlang.compiler.config.ProjectConfig;
 import org.shedlang.compiler.filesystem.ShedCompilerFiles;
 import org.shedlang.compiler.patio.PatioCodeGen;
 import org.shedlang.compiler.patio.PatioCodeGenContext;
+import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.stackir.StackIrNamespace;
 import org.shedlang.compiler.typechecker.TypeInfo;
+import org.shedlang.compiler.types.NamespaceName;
 
 import java.io.BufferedWriter;
 import java.io.FileWriter;
@@ -18,6 +27,7 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 import static org.shedlang.compiler.backends.ProcessRunner.runInheritIo;
 
@@ -27,6 +37,11 @@ public class LlvmBackend implements Backend {
         public LlvmModule generateNamespace(StackIrNamespace stackIrNamespace, TypeInfo typeInfo) {
             return LlvmCodeGen.generateNamespace(stackIrNamespace, typeInfo);
         }
+
+        @Override
+        public LlvmModule generateMain(NamespaceName main, TypeInfo typeInfo) {
+            return LlvmCodeGen.generateMain(main, typeInfo);
+        }
     });
 
     public static final LlvmBackend PATIO = new LlvmBackend("llvm-patio", new CodeGen() {
@@ -39,6 +54,44 @@ public class LlvmBackend implements Backend {
 
             return LlvmPatioCodeGen.generateCompilationUnit(patioCompilationUnit);
         }
+
+        @Override
+        public LlvmModule generateMain(NamespaceName main, TypeInfo typeInfo) {
+            // TODO: use more appropriate source
+            var source = Sources.NULL;
+
+            var context = PatioCodeGenContext.forTypeInfo(typeInfo);
+
+            var generateMainResult = PatioCodeGen.generateMain(main, source, context);
+            var mainFunctionName = generateMainResult.value();
+            context = generateMainResult.context();
+
+            var patioCompilationUnit = context.toCompilationUnit(source);
+
+            var llvmModule = LlvmPatioCodeGen.generateCompilationUnit(patioCompilationUnit);
+
+            llvmModule = llvmModule.addTopLevelEntity(new LlvmFunctionDefinition(
+                "main",
+                LlvmTypes.INT_32,
+                List.of(),
+                List.of(
+                    new LlvmBasicBlock(
+                        Optional.empty(),
+                        List.of(
+                            new LlvmCall(
+                                new LlvmVariableLocal("exitCode"),
+                                LlvmTypes.INT_32,
+                                new LlvmVariableGlobal(mainFunctionName),
+                                List.of()
+                            )
+                        ),
+                        new LlvmReturn(LlvmTypes.INT_32, new LlvmVariableLocal("exitCode"))
+                    )
+                )
+            ));
+
+            return llvmModule;
+        }
     });
 
     interface CodeGen {
@@ -46,6 +99,8 @@ public class LlvmBackend implements Backend {
             StackIrNamespace stackIrNamespace,
             TypeInfo typeInfo
         );
+
+        LlvmModule generateMain(NamespaceName main, TypeInfo typeInfo);
     }
 
     private final String name;
@@ -76,7 +131,7 @@ public class LlvmBackend implements Backend {
         }
 
         for (var bin : config.bins()) {
-            var mainModule = LlvmCodeGen.generateMain(config.main(), typeInfo);
+            var mainModule = this.codeGen.generateMain(config.main(), typeInfo);
             var mainObjectPath = compileLlvmModule(mainModule, targetDir.resolve("bins"), bin.name());
 
             var gccArgs = new ArrayList<String>();
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmModule.java b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmModule.java
index f4ea9dc9..0a899aa3 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmModule.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmModule.java
@@ -1,10 +1,11 @@
 package org.shedlang.compiler.backends.llvm.ir.ast;
 
-import org.shedlang.compiler.sources.Source;
-
-import java.util.List;
+import org.pcollections.PVector;
 
 public record LlvmModule(
-    List<LlvmTopLevelEntity> body
+    PVector<LlvmTopLevelEntity> body
 ) {
+    public LlvmModule addTopLevelEntity(LlvmTopLevelEntity topLevelEntity) {
+        return new LlvmModule(this.body.plus(topLevelEntity));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/util/POrderedMap.java b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
index 9363d42e..3ab74270 100644
--- a/src/main/java/org/shedlang/compiler/util/POrderedMap.java
+++ b/src/main/java/org/shedlang/compiler/util/POrderedMap.java
@@ -42,10 +42,10 @@ public class POrderedMap<K, V> {
         return keys;
     }
 
-    public List<V> toValues() {
+    public PVector<V> toValues() {
         return keys.stream()
             .map(key -> map.get(key))
-            .toList();
+            .collect(PCollectors.toVector());
     }
 
     public int size() {
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/ir/LlvmIrSerializerModuleTests.java b/src/test/java/org/shedlang/compiler/backends/llvm/ir/LlvmIrSerializerModuleTests.java
index a66339d0..0c55d6d2 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/ir/LlvmIrSerializerModuleTests.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/ir/LlvmIrSerializerModuleTests.java
@@ -4,9 +4,9 @@ import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmGlobalDefinition;
 import org.shedlang.compiler.backends.llvm.ir.ast.LlvmIr;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmTypes;
+import org.shedlang.compiler.util.P;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Optional;
 
 import static org.shedlang.compiler.backends.llvm.ir.LlvmIrSerializerTesting.serializeToString;
@@ -16,7 +16,7 @@ import static org.zwobble.precisely.Matchers.equalTo;
 public class LlvmIrSerializerModuleTests {
     @Test
     public void canSerializeModules() throws IOException {
-        var module = LlvmIr.module(List.of(
+        var module = LlvmIr.module(P.vector(
             new LlvmGlobalDefinition(
                 "x",
                 LlvmTypes.INT_64,
diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIr.java b/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIr.java
index 8bd93a6f..49a653c5 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIr.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/ir/ast/LlvmIr.java
@@ -1,5 +1,6 @@
 package org.shedlang.compiler.backends.llvm.ir.ast;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.ir.ast.instructions.*;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.*;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
@@ -97,7 +98,7 @@ public class LlvmIr {
         return new LlvmLoad(target, type, pointer);
     }
 
-    public static LlvmModule module(List<LlvmTopLevelEntity> topLevelEntities) {
+    public static LlvmModule module(PVector<LlvmTopLevelEntity> topLevelEntities) {
         return new LlvmModule(topLevelEntities);
     }