Garbage collect on function exit

5 files changed, 185 insertions, 19 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 81fa758d..6a18af6b 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -601,6 +601,12 @@ public class PatioCodeGen {
                 var value = valueResult.value();
                 context = valueResult.context();
 
+                context = PatioCodeGenClosures.beforeExit(
+                    List.of(value),
+                    resume.source(),
+                    context
+                );
+
                 yield PatioCodeGenEffects.resume(value, resume.source(), context);
             }
 
@@ -609,6 +615,12 @@ public class PatioCodeGen {
                 var value = valueResult.value();
                 context = valueResult.context();
 
+                context = PatioCodeGenClosures.beforeExit(
+                    List.of(value),
+                    return_.source(),
+                    context
+                );
+
                 yield context.addInstruction(new PatioReturn(value, return_.source()));
             }
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index bc1cd8c1..c42b1be8 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -124,6 +124,21 @@ class PatioCodeGenClosures {
                 bodyContextParent
             );
 
+            var gcRootStackIndexResult = bodyContext.allocateVal(
+                PatioCodeGenLocalSpace.GC_ROOT_STACK_INDEX,
+                PatioCodeGenGc.GC_ROOT_STACK_INDEX_TYPE
+            );
+            var gcRootStackIndex = gcRootStackIndexResult.value();
+            bodyContext = gcRootStackIndexResult.context();
+
+            bodyContext = PatioCodeGenGc.getGcRootStackIndex(
+                gcRootStackIndex,
+                source,
+                bodyContext
+            );
+
+            bodyContext = bodyContext.stashGcRootStackIndex(gcRootStackIndex);
+
             bodyContext = generateBody.apply(closurePointer, bodyContext);
 
             return bodyContext;
@@ -248,6 +263,26 @@ class PatioCodeGenClosures {
         return bodyContext;
     }
 
+    static PatioCodeGenContext beforeExit(
+        List<PatioOperand> escapees,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        context = PatioCodeGenGc.setGcRootStackIndex(
+            context.stashedGcRootStackIndex(),
+            source,
+            context
+        );
+
+        for (var escapee : escapees) {
+            context = PatioCodeGenGc.pushGcRoot(escapee, source, context);
+        }
+
+        context = PatioCodeGenGc.collect(source, context);
+
+        return context;
+    }
+
     static PatioCodeGenContext call(
         PatioLocalRef target,
         PatioLocalRef closure,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 6083cf75..87d5901a 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -187,7 +187,8 @@ public record PatioCodeGenContext(
         OrderedPMap<String, PatioLocal> localsByName,
         PMap<String, PatioLocalRef> shidNameToLocal,
         PMap<String, PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> waitingForShidLocal,
-        Optional<PatioOperand> stashedEffectHandler
+        Optional<PatioOperand> stashedEffectHandler,
+        Optional<PatioOperand> stashedGcRootStackIndex
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
             P.stack(),
@@ -195,6 +196,7 @@ public record PatioCodeGenContext(
             OrderedPMap.empty(),
             P.map(),
             P.map(),
+            Optional.empty(),
             Optional.empty()
         );
 
@@ -206,7 +208,8 @@ public record PatioCodeGenContext(
                 this.localsByName,
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             ).toResult(this.temporaryStack.get(0));
         }
 
@@ -218,7 +221,8 @@ public record PatioCodeGenContext(
                 this.localsByName,
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             ).toResult(this.temporaryStack);
         }
 
@@ -230,7 +234,8 @@ public record PatioCodeGenContext(
                 this.localsByName,
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             );
         }
 
@@ -242,7 +247,8 @@ public record PatioCodeGenContext(
                 this.localsByName,
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             );
         }
 
@@ -264,7 +270,8 @@ public record PatioCodeGenContext(
                 this.localsByName.plus(fullName, createLocal.apply(fullName)),
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             ).toResult(new PatioLocalRef(fullName));
         }
 
@@ -299,7 +306,8 @@ public record PatioCodeGenContext(
                 // TODO: the mapping is redundant given we fix the mapping
                 this.shidNameToLocal.plus(name, localRef),
                 this.waitingForShidLocal.plus(name, P.vector()),
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             );
 
             return context.toResult(localRef);
@@ -313,7 +321,8 @@ public record PatioCodeGenContext(
                 this.localsByName,
                 this.shidNameToLocal,
                 this.waitingForShidLocal.minus(name),
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             );
 
             return context.toResult(this.waitingForShidLocal.getOrDefault(name, P.vector()));
@@ -363,7 +372,8 @@ public record PatioCodeGenContext(
                     name,
                     this.waitingForShidLocal.getOrDefault(name, P.vector()).plus(func)
                 ),
-                this.stashedEffectHandler
+                this.stashedEffectHandler,
+                this.stashedGcRootStackIndex
             );
         }
 
@@ -376,7 +386,8 @@ public record PatioCodeGenContext(
                 this.shidNameToLocal,
                 this.waitingForShidLocal,
                 // TODO: check current value is empty
-                Optional.of(effectHandler)
+                Optional.of(effectHandler),
+                this.stashedGcRootStackIndex
             );
         }
 
@@ -387,6 +398,27 @@ public record PatioCodeGenContext(
                 .orElseThrow(() -> new RuntimeException("No stashed effect handler"));
         }
 
+        @CheckReturnValue
+        public FunctionContext stashGcRootStackIndex(PatioOperand gcRootStackIndex) {
+            return new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.shidNameToLocal,
+                this.waitingForShidLocal,
+                this.stashedEffectHandler,
+                // TODO: check current value is empty
+                Optional.of(gcRootStackIndex)
+            );
+        }
+
+        @CheckReturnValue
+        public PatioOperand stashedGcRootStackIndexOrThrow() {
+            // TODO: better error
+            return this.stashedGcRootStackIndex
+                .orElseThrow(() -> new RuntimeException("No stashed GC root stack index"));
+        }
+
         private <T> Result<T> toResult(T value) {
             return new Result<>(value, this);
         }
@@ -691,7 +723,8 @@ public record PatioCodeGenContext(
                 this.functionContext.localsByName,
                 this.functionContext.shidNameToLocal,
                 this.functionContext.waitingForShidLocal,
-                this.functionContext.stashedEffectHandler
+                this.functionContext.stashedEffectHandler,
+                this.functionContext.stashedGcRootStackIndex
             )
         );
     }
@@ -712,7 +745,8 @@ public record PatioCodeGenContext(
                 this.functionContext.localsByName,
                 this.functionContext.shidNameToLocal,
                 this.functionContext.waitingForShidLocal,
-                this.functionContext.stashedEffectHandler
+                this.functionContext.stashedEffectHandler,
+                this.functionContext.stashedGcRootStackIndex
             )
         ).toResult(this.functionContext.instructions);
     }
@@ -734,6 +768,16 @@ public record PatioCodeGenContext(
     }
 
     @CheckReturnValue
+    public PatioCodeGenContext stashGcRootStackIndex(PatioOperand gcRootStackIndex) {
+        return withFunctionContext(this.functionContext.stashGcRootStackIndex(gcRootStackIndex));
+    }
+
+    @CheckReturnValue
+    public PatioOperand stashedGcRootStackIndex() {
+        return this.functionContext.stashedGcRootStackIndexOrThrow();
+    }
+
+    @CheckReturnValue
     private PatioCodeGenContext withCompilationUnitContext(CompilationUnitContext compilationUnitContext) {
         return new PatioCodeGenContext(this.globalContext, compilationUnitContext, this.blockStack, this.functionContext);
     }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
index 23ee6cac..8d9c3ccd 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
@@ -15,13 +15,6 @@ import java.util.List;
 import java.util.Optional;
 
 public class PatioCodeGenGc {
-    private static final PatioFunctionDeclaration ADD_PERM_GC_ROOT = new PatioFunctionDeclaration(
-        "shed__runtime__gc_add_perm_root",
-        List.of(PatioLocal.param("gc_root", PatioCodeGenTypes.VALUE_TYPE)),
-        PatioTypes.VOID,
-        Sources.RUNTIME
-    );
-
     private PatioCodeGenGc() {
     }
 
@@ -59,6 +52,87 @@ public class PatioCodeGenGc {
         ));
     }
 
+    // == Roots ==
+
+    // == Root stack ==
+
+    public static final PatioType GC_ROOT_STACK_INDEX_TYPE = PatioTypes.INT;
+
+    private static final PatioFunctionDeclaration GET_GC_ROOT_STACK_INDEX = new PatioFunctionDeclaration(
+        "shed__runtime__gc_get_root_stack_index",
+        List.of(),
+        GC_ROOT_STACK_INDEX_TYPE,
+        Sources.RUNTIME
+    );
+
+    public static PatioCodeGenContext getGcRootStackIndex(
+        PatioLocalRef target,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        return context
+            .ensureCompilationUnitElement(GET_GC_ROOT_STACK_INDEX)
+            .addInstruction(new PatioCall(
+                Optional.of(target),
+                GET_GC_ROOT_STACK_INDEX.name(),
+                List.of(),
+                source
+            ));
+    }
+
+    private static final PatioFunctionDeclaration SET_GC_ROOT_STACK_INDEX = new PatioFunctionDeclaration(
+        "shed__runtime__gc_set_root_stack_index",
+        List.of(PatioLocal.param("gc_root_stack_index", GC_ROOT_STACK_INDEX_TYPE)),
+        PatioTypes.VOID,
+        Sources.RUNTIME
+    );
+
+    public static PatioCodeGenContext setGcRootStackIndex(
+        PatioOperand gcRootIndex,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        return context
+            .ensureCompilationUnitElement(SET_GC_ROOT_STACK_INDEX)
+            .addInstruction(new PatioCall(
+                Optional.empty(),
+                SET_GC_ROOT_STACK_INDEX.name(),
+                List.of(gcRootIndex),
+                source
+            ));
+    }
+
+    private static final PatioFunctionDeclaration PUSH_GC_ROOT = new PatioFunctionDeclaration(
+        "shed__runtime__gc_push_root",
+        List.of(PatioLocal.param("gc_root", PatioCodeGenTypes.VALUE_TYPE)),
+        PatioTypes.VOID,
+        Sources.RUNTIME
+    );
+
+    public static PatioCodeGenContext pushGcRoot(
+        PatioOperand escapee,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        return context
+            .ensureCompilationUnitElement(PUSH_GC_ROOT)
+            .addInstruction(new PatioCall(
+                Optional.empty(),
+                PUSH_GC_ROOT.name(),
+                List.of(escapee),
+                source
+            ));
+    }
+
+    // == Perm roots ==
+
+    private static final PatioFunctionDeclaration ADD_PERM_GC_ROOT = new PatioFunctionDeclaration(
+        "shed__runtime__gc_add_perm_root",
+        List.of(PatioLocal.param("gc_root", PatioCodeGenTypes.VALUE_TYPE)),
+        PatioTypes.VOID,
+        Sources.RUNTIME
+    );
+
     public static PatioCodeGenContext addPermGcRoot(
         PatioOperand gcRoot,
         Source source,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
index fcaa54ea..54337eb4 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
@@ -12,6 +12,7 @@ public enum PatioCodeGenLocalSpace {
     EFFECT_HANDLER_CURRENT("effectHandlerCurrent"),
     EFFECT_HANDLER_METHOD("effectHandlerMethod"),
     EFFECT_HANDLER_METHOD_POINTER("effectHandlerMethodPointer"),
+    GC_ROOT_STACK_INDEX("gcRootStackIndex"),
     INT("int"),
     MAIN("main"),
     NAMESPACE_IS_INITED_POINTER("namespaceIsInitedPointer"),