Extract functions for saving and restoring GC root stack

2 files changed, 43 insertions, 27 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index c42b1be8..593d76fe 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -124,20 +124,7 @@ 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 = PatioCodeGenGc.saveGcRootStack(source, bodyContext);
 
             bodyContext = generateBody.apply(closurePointer, bodyContext);
 
@@ -268,19 +255,7 @@ class PatioCodeGenClosures {
         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;
+        return PatioCodeGenGc.restoreGcRootStack(escapees, source, context);
     }
 
     static PatioCodeGenContext call(
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
index 525ab89a..530472fc 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenGc.java
@@ -164,4 +164,45 @@ public class PatioCodeGenGc {
                 source
             ));
     }
+
+    // == Save/Restore ==
+
+    static PatioCodeGenContext saveGcRootStack(Source source, PatioCodeGenContext context) {
+        var gcRootStackIndexResult = context.allocateVal(
+            PatioCodeGenLocalSpace.GC_ROOT_STACK_INDEX,
+            PatioCodeGenGc.GC_ROOT_STACK_INDEX_TYPE
+        );
+        var gcRootStackIndex = gcRootStackIndexResult.value();
+        context = gcRootStackIndexResult.context();
+
+        context = PatioCodeGenGc.getGcRootStackIndex(
+            gcRootStackIndex,
+            source,
+            context
+        );
+
+        context = context.stashGcRootStackIndex(gcRootStackIndex);
+
+        return context;
+    }
+
+    static PatioCodeGenContext restoreGcRootStack(
+        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;
+    }
 }