Use GC alloc when allocating on the heap

10 files changed, 173 insertions, 37 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
index cac25645..da9c0c37 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/codegen/LlvmPatioCodeGen.java
@@ -22,9 +22,13 @@ import java.util.stream.StreamSupport;
 
 public class LlvmPatioCodeGen {
     private static final LlvmFunctionDeclaration ALLOC_DECLARATION = new LlvmFunctionDeclaration(
-        "shed__runtime__alloc",
+        "shed__runtime__gc_alloc",
         LlvmTypes.pointer(LlvmTypes.INT_8),
-        List.of(new LlvmParam(LlvmTypes.INT_64, "size"))
+        List.of(
+            new LlvmParam(LlvmTypes.INT_64, "size"),
+            new LlvmParam(LlvmTypes.INT_64, "children_offset"),
+            new LlvmParam(LlvmTypes.INT_64, "children_count")
+        )
     );
 
     private static final LlvmType EXCEPTION_TAG_TYPE = LlvmTypes.INT_32;
@@ -252,9 +256,21 @@ public class LlvmPatioCodeGen {
                 var target = targetResult.value();
                 context = targetResult.context();
 
+                var offsetResult = offsetOf(
+                    patioAllocHeap.layout(),
+                    patioAllocHeap.gcChildrenOffset(),
+                    context
+                );
+                var offset = offsetResult.value();
+                context = offsetResult.context();
+
                 yield context.addInstruction(ALLOC_DECLARATION.call(
                     Optional.of(target),
-                    List.of(size)
+                    List.of(
+                        size,
+                        offset,
+                        new LlvmOperandInt(patioAllocHeap.gcChildrenCount())
+                    )
                 ));
             }
 
@@ -1213,4 +1229,36 @@ public class LlvmPatioCodeGen {
             )
             .toResult(size);
     }
+
+    private static LlvmPatioCodeGenFunctionContext.Result<LlvmOperand> offsetOf(
+        PatioDataLayout patioDataLayout,
+        List<Integer> patioOffset,
+        LlvmPatioCodeGenFunctionContext context
+    ) {
+        var offsetPtrResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.OFFSET_PTR);
+        var offsetPtr = offsetPtrResult.value();
+        context = offsetPtrResult.context();
+
+        var offsetResult = context.allocateLocalName(LlvmCodeGenLocalVariableSpace.OFFSET);
+        var offset = offsetResult.value();
+        context = offsetResult.context();
+
+        return context
+            .addInstruction(
+                new LlvmGetElementPtr(
+                    offsetPtr,
+                    LlvmTypes.pointer(patioDataLayoutToLlvmType(patioDataLayout)),
+                    new LlvmOperandNull(),
+                    patioDataOffsetsToLlvmOffsets(patioDataLayout, patioOffset)
+                )
+            )
+            .addInstruction(
+                new LlvmPtrToInt(
+                    offset,
+                    offsetPtr,
+                    LlvmTypes.INT_64
+                )
+            )
+            .toResult(offset);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
index 486e5866..c1cf0fae 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmPatioCodeGen.java
@@ -21,7 +21,8 @@ import java.util.ArrayList;
 import java.util.List;
 
 public class WasmPatioCodeGen {
-    public static final String ALLOC_NAME = "shed__runtime__alloc";
+    public static final String ALLOC_NAME = "shed__runtime__gc_alloc";
+    public static final String STACK_ALLOC_NAME = "malloc";
 
     public static WasmModule generateCompilationUnit(
         PatioCompilationUnit patioCompilationUnit
@@ -31,7 +32,22 @@ public class WasmPatioCodeGen {
             ALLOC_NAME,
             new WasmImportDescriptor.Func(
                 ALLOC_NAME,
-                new WasmFuncType(List.of(WasmValueType.I32), List.of(WasmValueType.I32))
+                new WasmFuncType(
+                    List.of(WasmValueType.I32, WasmValueType.I32, WasmValueType.I32),
+                    List.of(WasmValueType.I32)
+                )
+            )
+        ));
+
+        module = module.addImport(new WasmImport(
+            "env",
+            STACK_ALLOC_NAME,
+            new WasmImportDescriptor.Func(
+                STACK_ALLOC_NAME,
+                new WasmFuncType(
+                    List.of(WasmValueType.I32),
+                    List.of(WasmValueType.I32)
+                )
             )
         ));
 
@@ -138,6 +154,11 @@ public class WasmPatioCodeGen {
             case PatioAllocHeap patioAllocHeap -> {
                 context = context
                     .addInstruction(new WasmI32Const(sizeOf(patioAllocHeap.layout())))
+                    .addInstruction(new WasmI32Const(offsetOf(
+                        patioAllocHeap.layout(),
+                        patioAllocHeap.gcChildrenOffset()
+                    )))
+                    .addInstruction(new WasmI32Const(patioAllocHeap.gcChildrenCount()))
                     .addInstruction(new WasmCall(new WasmFuncName(ALLOC_NAME)))
                     .addInstruction(new WasmLocalSet(context.wasmLocalIndex(patioAllocHeap.target())));
 
@@ -147,7 +168,7 @@ public class WasmPatioCodeGen {
             case PatioAllocStack patioAllocStack -> {
                 context = context
                     .addInstruction(new WasmI32Const(sizeOf(patioAllocStack.layout())))
-                    .addInstruction(new WasmCall(new WasmFuncName(ALLOC_NAME)))
+                    .addInstruction(new WasmCall(new WasmFuncName(STACK_ALLOC_NAME)))
                     .addInstruction(new WasmLocalSet(context.wasmLocalIndex(patioAllocStack.target())));
 
                 yield context;
@@ -520,7 +541,7 @@ public class WasmPatioCodeGen {
         };
     }
 
-    private static WasmMemoryArg wasmMemoryArg(
+    private static int offsetOf(
         PatioDataLayout patioDataLayout,
         List<Integer> patioOffsets
     ) {
@@ -556,6 +577,15 @@ public class WasmPatioCodeGen {
             }
         }
 
+        return wasmOffset;
+    }
+
+    private static WasmMemoryArg wasmMemoryArg(
+        PatioDataLayout patioDataLayout,
+        List<Integer> patioOffsets
+    ) {
+        var wasmOffset = offsetOf(patioDataLayout, patioOffsets);
+
         // TODO: alignment
         return new WasmMemoryArg(wasmOffset, WasmAlignment.NONE);
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenAlloc.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenAlloc.java
deleted file mode 100644
index e90592b9..00000000
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenAlloc.java
+++ /dev/null
@@ -1,20 +0,0 @@
-package org.shedlang.compiler.patio;
-
-import org.shedlang.compiler.patio.ast.instructions.PatioAllocHeap;
-import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
-import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
-import org.shedlang.compiler.sources.Source;
-
-public class PatioCodeGenAlloc {
-    private PatioCodeGenAlloc() {
-    }
-
-    public static PatioCodeGenContext alloc(
-        PatioLocalRef target,
-        PatioDataLayout layout,
-        Source source,
-        PatioCodeGenContext context
-    ) {
-        return context.addInstruction(new PatioAllocHeap(target, layout, source));
-    }
-}
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index c5a0096d..bc1cd8c1 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -66,6 +66,9 @@ class PatioCodeGenClosures {
         return PatioCodeGenGc.gcObjectOffset(List.of(2, freeVariable.index()));
     }
 
+    private static final List<Integer> GC_CHILDREN_OFFSET =
+        PatioCodeGenGc.gcObjectOffset(List.of(2));
+
     static PatioCodeGenContext defineParams(
         List<ShidDefineFunction.Param> positionalParams,
         List<ShidDefineFunction.Param> namedParams,
@@ -158,7 +161,14 @@ class PatioCodeGenClosures {
         var closurePointer = closurePointerResult.value();
         context = closurePointerResult.context();
 
-        context = PatioCodeGenAlloc.alloc(closurePointer, layout, source, context);
+        context = PatioCodeGenGc.alloc(
+            closurePointer,
+            layout,
+            GC_CHILDREN_OFFSET,
+            freeVariables.size(),
+            source,
+            context
+        );
 
         context = context.addInstruction(new PatioMemoryStore(
             closurePointer,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
index 1c763f17..29618c99 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -60,6 +60,9 @@ public class PatioCodeGenEffects {
         return PatioCodeGenGc.gcObjectOffset(List.of(methodIndex));
     }
 
+    private static final List<Integer> USER_DEFINED_EFFECT_GC_CHILDREN_OFFSET =
+        PatioCodeGenGc.gcObjectOffset(List.of());
+
     private static final PatioFunctionDeclaration FIND_HANDLER_METHOD_DECLARATION = new PatioFunctionDeclaration(
         "shed__runtime__effects_find_handler_method",
         List.of(
@@ -137,14 +140,17 @@ public class PatioCodeGenEffects {
 
         var effectLayout = userDefinedEffectLayout(effect, context);
 
-        context = PatioCodeGenAlloc.alloc(
+        var effectInfo = context.userDefinedEffectInfo(effect);
+
+        context = PatioCodeGenGc.alloc(
             effectPointer,
             effectLayout,
+            USER_DEFINED_EFFECT_GC_CHILDREN_OFFSET,
+            effectInfo.methodCount(),
             source,
             context
         );
 
-        var effectInfo = context.userDefinedEffectInfo(effect);
         for (var method : effectInfo.methods()) {
             var effectMethodResult = context.allocateVal(
                 PatioCodeGenLocalSpace.USER_DEFINED_EFFECT_METHOD,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
index 144cf49b..cd036368 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
@@ -2,7 +2,9 @@ package org.shedlang.compiler.patio;
 
 import org.shedlang.compiler.patio.ast.PatioFunctionDeclaration;
 import org.shedlang.compiler.patio.ast.PatioLocal;
+import org.shedlang.compiler.patio.ast.instructions.PatioAllocHeap;
 import org.shedlang.compiler.patio.ast.instructions.PatioCall;
+import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
 import org.shedlang.compiler.patio.ast.operands.PatioOperand;
 import org.shedlang.compiler.patio.ast.types.*;
 import org.shedlang.compiler.sources.Source;
@@ -70,4 +72,46 @@ public class PatioCodeGenGc {
                 source
             ));
     }
+
+    // == Allocation ==
+
+    public static PatioCodeGenContext alloc(
+        PatioLocalRef target,
+        PatioDataLayout layout,
+        List<Integer> childrenOffset,
+        int childrenCount,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        return context.addInstruction(new PatioAllocHeap(
+            target,
+            layout,
+            childrenOffset,
+            childrenCount,
+            source
+        ));
+    }
+
+    // == Collection  ==
+
+    private static final PatioFunctionDeclaration COLLECT = new PatioFunctionDeclaration(
+        "shed__runtime__gc_collect",
+        List.of(),
+        PatioTypes.VOID,
+        Sources.RUNTIME
+    );
+
+    public static PatioCodeGenContext collect(
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        return context
+            .ensureCompilationUnitElement(COLLECT)
+            .addInstruction(new PatioCall(
+                Optional.empty(),
+                COLLECT.name(),
+                List.of(),
+                source
+            ));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenStructs.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenStructs.java
index 36274023..c02f0998 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenStructs.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenStructs.java
@@ -45,9 +45,11 @@ public class PatioCodeGenStructs {
 
                 var structDataLayout = structDataLayout(type, context);
 
-                bodyContext = PatioCodeGenAlloc.alloc(
+                bodyContext = PatioCodeGenGc.alloc(
                     valuePointer,
                     structDataLayout,
+                    GC_CHILDREN_OFFSET,
+                    fields.size(),
                     source,
                     bodyContext
                 );
@@ -192,4 +194,7 @@ public class PatioCodeGenStructs {
     private static List<Integer> fieldOffset(int fieldIndex) {
         return PatioCodeGenGc.gcObjectOffset(List.of(1, fieldIndex));
     }
+
+    private static final List<Integer> GC_CHILDREN_OFFSET =
+        PatioCodeGenGc.gcObjectOffset(List.of(1));
 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTuples.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTuples.java
index 6759506d..a99689ff 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTuples.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenTuples.java
@@ -47,9 +47,11 @@ public class PatioCodeGenTuples {
 
         var tupleLayout = tupleDataLayout(elements.size());
 
-        context = PatioCodeGenAlloc.alloc(
+        context = PatioCodeGenGc.alloc(
             tuplePointer,
             tupleLayout,
+            GC_CHILDREN_OFFSET,
+            elements.size(),
             source,
             context
         );
@@ -114,14 +116,17 @@ public class PatioCodeGenTuples {
         return context.toResult(target);
     }
 
-    private static List<Integer> elementOffset(int elementIndex) {
-        return PatioCodeGenGc.gcObjectOffset(List.of(elementIndex));
-    }
-
     private static PatioDataLayout tupleDataLayout(int elementCount) {
         return PatioCodeGenGc.gcObjectDataLayout(PatioDataLayouts.struct(Lists.repeat(
             PatioDataLayouts.INT,
             elementCount
         )));
     }
+
+    private static List<Integer> elementOffset(int elementIndex) {
+        return PatioCodeGenGc.gcObjectOffset(List.of(elementIndex));
+    }
+
+    private static final List<Integer> GC_CHILDREN_OFFSET =
+        PatioCodeGenGc.gcObjectOffset(List.of());
 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenVars.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenVars.java
index c87d52b4..72218531 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenVars.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenVars.java
@@ -19,6 +19,8 @@ public class PatioCodeGenVars {
         PatioCodeGenGc.gcObjectDataLayout(PatioDataLayouts.fromType(PatioCodeGenTypes.VALUE_TYPE));
 
     private static final List<Integer> VALUE_OFFSET = PatioCodeGenGc.gcObjectOffset(List.of());
+    private static final List<Integer> GC_CHILDREN_OFFSET = PatioCodeGenGc.gcObjectOffset(List.of());
+    private static final int GC_CHILDREN_COUNT = 1;
 
     public static PatioCodeGenContext alloc(
         String shidName,
@@ -32,9 +34,11 @@ public class PatioCodeGenVars {
         var patioLocal = patioLocalResult.value();
         context = patioLocalResult.context();
 
-        context = PatioCodeGenAlloc.alloc(
+        context = PatioCodeGenGc.alloc(
             patioLocal,
             VAR_DATA_LAYOUT,
+            GC_CHILDREN_OFFSET,
+            GC_CHILDREN_COUNT,
             source,
             context
         );
diff --git a/src/main/java/org/shedlang/compiler/patio/ast/instructions/PatioAllocHeap.java b/src/main/java/org/shedlang/compiler/patio/ast/instructions/PatioAllocHeap.java
index aee64781..fa045fb5 100644
--- a/src/main/java/org/shedlang/compiler/patio/ast/instructions/PatioAllocHeap.java
+++ b/src/main/java/org/shedlang/compiler/patio/ast/instructions/PatioAllocHeap.java
@@ -4,6 +4,8 @@ import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
 import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
 import org.shedlang.compiler.sources.Source;
 
+import java.util.List;
+
 /**
  * Allocate space on the heap for the given data layout, and store the pointer
  * to the allocated memory in the target local.
@@ -11,6 +13,8 @@ import org.shedlang.compiler.sources.Source;
 public record PatioAllocHeap(
     PatioLocalRef target,
     PatioDataLayout layout,
+    List<Integer> gcChildrenOffset,
+    int gcChildrenCount,
     Source source
 ) implements PatioInstruction {
 }