feat(runtime): Prevent stack overflow while collecting large lists (#2248)

spotandjake · ospencer · web-flow · commit 097ae7d403aa · 2025-02-17T04:30:59.000Z
Co-authored-by: Oscar Spencer &lt;oscar.spen@gmail.com&gt;
diff --git a/compiler/test/suites/basic_functionality.re b/compiler/test/suites/basic_functionality.re
@@ -377,6 +377,6 @@ describe("basic functionality", ({test, testSkip}) => {
     ~config_fn=smallestFileConfig,
     "smallest_grain_program",
     "",
-    6540,
+    6494,
   );
 });
diff --git a/compiler/test/suites/gc.re b/compiler/test/suites/gc.re
@@ -139,7 +139,7 @@ describe("garbage collection", ({test, testSkip}) => {
     "OK\n",
   );
   assertRunGC(
-    "long_lists",
+    "medium_lists",
     350,
     {|
     from "list" include List
@@ -168,6 +168,20 @@ describe("garbage collection", ({test, testSkip}) => {
     |},
     "true\n",
   );
+  assertRun(
+    "large_lists",
+    {|
+    let rec make_list = (n, acc) => {
+      if (n == 0) {
+        acc
+      } else {
+        make_list(n - 1, [1, ...acc])
+      }
+    }
+    assert make_list(500_000, []) != []
+    |},
+    "",
+  );
   assertFileRun("memory_grow1", "memoryGrow", "1000000000000\n");
   assertMemoryLimitedFileRun(
     "loop_memory_reclaim",
diff --git a/stdlib/runtime/gc.gr b/stdlib/runtime/gc.gr
@@ -25,6 +25,7 @@ from "runtime/unsafe/panic" include Panic
 from "runtime/unsafe/wasmi32" include WasmI32
 use WasmI32.{ (+), (-), (*), (&), (==), (!=) }
 
+primitive (!) = "@not"
 primitive (&&) = "@and"
 primitive (||) = "@or"
 primitive ignore = "@ignore"
@@ -84,24 +85,37 @@ let rec decRef = (userPtr: WasmI32, ignoreZeros: Bool) => {
     // }
 
     if (WasmI32.eqz(refCount)) {
-      if (ignoreZeros) {
-        userPtr
-      } else {
+      if (!ignoreZeros) {
         throwDecRefError()
       }
     } else {
       let refCount = refCount - 1n
       setRefCount(userPtr, refCount)
 
       if (WasmI32.eqz(refCount)) {
-        decRefChildren(userPtr)
+        /*
+         * Note: We call free before decRefChildren to allow for a tail call.
+         * This is okay because no allocations occur while we traverse the
+         * structure and free does not mangle the data.
+         */
         free(userPtr)
+        decRefChildren(userPtr)
       }
-
-      userPtr
     }
-  } else {
-    userPtr
+  }
+}
+and decRefChildrenHelp = (
+  userPtr: WasmI32,
+  arityOffset: WasmI32,
+  offset: WasmI32,
+) => {
+  let arity = WasmI32.load(userPtr, arityOffset)
+  if (arity != 0n) {
+    let maxOffset = (arity - 1n) * 4n
+    for (let mut i = 0n; WasmI32.ltU(i, maxOffset); i += 4n) {
+      decRef(WasmI32.load(userPtr + i, offset), false)
+    }
+    decRef(WasmI32.load(userPtr + maxOffset, offset), false)
   }
 }
 and decRefChildren = (userPtr: WasmI32) => {
@@ -110,43 +124,25 @@ and decRefChildren = (userPtr: WasmI32) => {
       let tag = WasmI32.load(userPtr, 4n)
       if (userPtr == Tags._GRAIN_RATIONAL_BOXED_NUM_TAG) {
         // decRef underlying BigInts
-        ignore(decRef(WasmI32.load(userPtr, 8n), false))
-        ignore(decRef(WasmI32.load(userPtr, 12n), false))
+        decRef(WasmI32.load(userPtr, 8n), false)
+        decRef(WasmI32.load(userPtr, 12n), false)
       }
     },
     t when t == Tags._GRAIN_ADT_HEAP_TAG => {
-      let arity = WasmI32.load(userPtr, 16n)
-      let maxOffset = arity * 4n
-      for (let mut i = 0n; WasmI32.ltU(i, maxOffset); i += 4n) {
-        ignore(decRef(WasmI32.load(userPtr + i, 20n), false))
-      }
+      decRefChildrenHelp(userPtr, 16n, 20n)
     },
-    t when t == Tags._GRAIN_RECORD_HEAP_TAG => {
-      let arity = WasmI32.load(userPtr, 12n)
-      let maxOffset = arity * 4n
-      for (let mut i = 0n; WasmI32.ltU(i, maxOffset); i += 4n) {
-        ignore(decRef(WasmI32.load(userPtr + i, 16n), false))
-      }
+    t when t == Tags._GRAIN_RECORD_HEAP_TAG || t == Tags._GRAIN_LAMBDA_HEAP_TAG => {
+      decRefChildrenHelp(userPtr, 12n, 16n)
     },
     t when t == Tags._GRAIN_ARRAY_HEAP_TAG || t == Tags._GRAIN_TUPLE_HEAP_TAG => {
-      let arity = WasmI32.load(userPtr, 4n)
-      let maxOffset = arity * 4n
-      for (let mut i = 0n; WasmI32.ltU(i, maxOffset); i += 4n) {
-        ignore(decRef(WasmI32.load(userPtr + i, 8n), false))
-      }
-    },
-    t when t == Tags._GRAIN_LAMBDA_HEAP_TAG => {
-      let arity = WasmI32.load(userPtr, 12n)
-      let maxOffset = arity * 4n
-      for (let mut i = 0n; WasmI32.ltU(i, maxOffset); i += 4n) {
-        ignore(decRef(WasmI32.load(userPtr + i, 16n), false))
-      }
-    },
-    _ => {
-      // No travelsal necessary for other tags
-      void
+      decRefChildrenHelp(userPtr, 4n, 8n)
     },
+    // No traversal necessary for other tags
+    _ => void,
   }
 }
 
-provide let decRef = userPtr => decRef(userPtr, false)
+provide let decRef = userPtr => {
+  decRef(userPtr, false)
+  userPtr
+}
