Tweaks to reduce space usage

Author: polytypic

lib/picos_std.structured/bundle.ml

@@ -150,6 +150,11 @@ let[@inline never] returned value child t canceler =
   Computation.return child value;
   finish t canceler
 
+let[@inline never] plug t thunk child canceler =
+  match thunk () with
+  | value -> returned value child t canceler
+  | exception exn -> raised exn child t canceler
+
 let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
   (* The sequence of operations below ensures that nothing is leaked. *)
   incr t Backoff.default;
@@ -160,19 +165,10 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
     let canceler = Computation.attach_canceler ~from:bundle ~into:child in
     let main =
       match pass with
-      | FLS -> begin
+      | FLS ->
           Fiber.FLS.set fiber flock_key t;
-          fun fiber ->
-            match thunk () with
-            | value -> returned value child (get_flock fiber) canceler
-            | exception exn -> raised exn child (get_flock fiber) canceler
-        end
-      | Arg -> begin
-          fun _ ->
-            match thunk () with
-            | value -> returned value child t canceler
-            | exception exn -> raised exn child t canceler
-        end
+          fun fiber -> plug (get_flock fiber) thunk child canceler
+      | Arg -> fun _ -> plug t thunk child canceler
     in
     Fiber.spawn fiber main;
     child
@@ -183,36 +179,24 @@ let fork_as_promise_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
     decr t;
     raise canceled_exn
 
-let[@inline never] raised_flock exn fiber =
-  let t = get_flock fiber in
-  let bt = Printexc.get_raw_backtrace () in
-  error t exn bt;
-  decr t
-
-let[@inline never] raised_bundle exn t =
+let[@inline never] raised exn t =
   error t exn (Printexc.get_raw_backtrace ());
   decr t
 
+let[@inline never] plug t thunk =
+  match thunk () with () -> decr t | exception exn -> raised exn t
+
 let fork_pass (type a) (Bundle r as t : t) thunk (pass : a pass) =
   (* The sequence of operations below ensures that nothing is leaked. *)
   incr t Backoff.default;
   try
     let fiber = Fiber.create_packed ~forbid:false r.bundle in
     let main =
       match pass with
-      | FLS -> begin
+      | FLS ->
           Fiber.FLS.set fiber flock_key t;
-          fun fiber ->
-            match thunk () with
-            | () -> decr (get_flock fiber)
-            | exception exn -> raised_flock exn fiber
-        end
-      | Arg -> begin
-          fun _ ->
-            match thunk () with
-            | () -> decr t
-            | exception exn -> raised_bundle exn t
-        end
+          fun fiber -> plug (get_flock fiber) thunk
+      | Arg -> fun _ -> plug t thunk
     in
     Fiber.spawn fiber main
   with canceled_exn ->

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