From e214bab5dd0322e102adbf59d0bd9c651fe32376 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Wed, 2 Oct 2024 02:19:10 +0300
Subject: [PATCH] Tweaks to reduce space usage

Turns out OCaml has space leaks related to use of closures.
---
 lib/picos_std.structured/run.ml | 50 ++++++++++++++++-----------------
 1 file changed, 24 insertions(+), 26 deletions(-)

diff --git a/lib/picos_std.structured/run.ml b/lib/picos_std.structured/run.ml
index bc9122f5..79ebe292 100644
--- a/lib/picos_std.structured/run.ml
+++ b/lib/picos_std.structured/run.ml
@@ -1,39 +1,37 @@
 open Picos
 
-let wrap_all t main _ =
-  if Bundle.is_running t then begin
-    try main () with exn -> Bundle.error t exn (Printexc.get_raw_backtrace ())
-  end;
-  Bundle.decr t
+let[@inline never] wrap_all t main =
+  match main () with
+  | () -> Bundle.decr t
+  | exception exn -> Bundle.raised exn t
 
-let wrap_any t main _ =
-  if Bundle.is_running t then begin
-    match main () with
-    | () -> Bundle.terminate t
-    | exception exn -> Bundle.error t exn (Printexc.get_raw_backtrace ())
-  end;
-  Bundle.decr t
+let[@inline never] wrap_any t main =
+  match main () with
+  | () ->
+      Bundle.terminate t;
+      Bundle.decr t
+  | exception exn -> Bundle.raised exn t
 
-let rec spawn (Bundle r as t : Bundle.t) wrap = function
+let rec spawn (Bundle r as t : Bundle.t) ~all = function
   | [] -> ()
   | [ main ] ->
       Bundle.unsafe_incr t;
-      let unused_fake_fiber = Obj.magic () in
-      wrap t main unused_fake_fiber
+      if Bundle.is_running t then
+        if all then wrap_all t main else wrap_any t main
+      else Bundle.decr t
   | main :: mains ->
       Bundle.unsafe_incr t;
       let fiber = Fiber.create_packed ~forbid:false r.bundle in
       (* Note that [Fiber.spawn] checks the cancelation status of the bundle. *)
-      Fiber.spawn fiber (wrap t main);
-      spawn t wrap mains
+      Fiber.spawn fiber (fun _ ->
+          if Bundle.is_running t then
+            if all then wrap_all t main else wrap_any t main
+          else Bundle.decr t);
+      spawn t ~all mains
 
-let run actions wrap =
-  Bundle.join_after @@ fun (Bundle _ as t : Bundle.t) ->
-  try spawn t wrap actions
-  with exn ->
-    let bt = Printexc.get_raw_backtrace () in
-    Bundle.decr t;
-    Bundle.error t exn bt
+let run actions ~all =
+  Bundle.join_after @@ fun (t : Bundle.t) ->
+  try spawn t ~all actions with exn -> Bundle.raised exn t
 
-let all actions = run actions wrap_all
-let any actions = run actions wrap_any
+let all actions = run actions ~all:true
+let any actions = run actions ~all:false