From d05f82023683940da64b87373375ba2c444686bd Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 09:32:51 +0100 Subject: [PATCH 01/11] perf(fiber): represent common continuations explicitly Represent continuations introduced by map and bind with typed frames while keeping a function frame as the escape hatch for arbitrary callbacks. Carry these continuations through effects and the scheduler without changing the public Fiber API. On a no-op @install self-build this reduced minor allocation by 14.1M words, promoted allocation by 1.1M words, the live heap by 0.6M words, and Cachegrind instructions by 0.8%. Native wall time remained neutral. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 223 ++++++++++++++++++++++--------------- src/fiber/src/lazy.ml | 4 +- src/fiber/src/mutex.ml | 4 +- src/fiber/src/mvar.ml | 8 +- src/fiber/src/pool.ml | 16 +-- src/fiber/src/scheduler.ml | 39 ++++--- src/fiber/src/stream.ml | 17 +-- src/fiber/src/svar.ml | 6 +- 8 files changed, 190 insertions(+), 127 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index a604c04a476..4daa7de6f8c 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -1,24 +1,32 @@ open Stdune -type 'a t = ('a -> eff) -> eff +type 'a t = 'a continuation -> eff + +(* Keep arbitrary callbacks as functions, but represent the frames introduced by common + fiber combinators directly. *) +and 'a continuation = + | Function : ('a -> eff) -> 'a continuation + | Effect : eff continuation + | Map : ('a -> 'b) * 'b continuation -> 'a continuation + | Bind : ('a -> 'b t) * 'b continuation -> 'a continuation and eff = - | Read_ivar : 'a ivar * ('a -> eff) -> eff - | Fill_ivar : 'a ivar * 'a * (unit -> eff) -> eff - | Suspend : ('a k -> unit) * ('a -> eff) -> eff - | Resume : 'a k * 'a * (unit -> eff) -> eff - | Get_var : 'a Var_map.Key.t * ('a -> eff) -> eff - | Set_var : 'a Var_map.Key.t * 'a * (unit -> eff) -> eff - | Update_var : 'a Var_map.Key.t * ('a -> 'a) * (unit -> eff) -> eff - | With_error_handler : (Exn_with_backtrace.t -> Nothing.t t) * (unit -> eff) -> eff - | Unwind : ('a -> eff) * 'a -> eff + | Read_ivar : 'a ivar * 'a continuation -> eff + | Fill_ivar : 'a ivar * 'a * unit continuation -> eff + | Suspend : ('a k -> unit) * 'a continuation -> eff + | Resume : 'a k * 'a * unit continuation -> eff + | Get_var : 'a Var_map.Key.t * 'a continuation -> eff + | Set_var : 'a Var_map.Key.t * 'a * unit continuation -> eff + | Update_var : 'a Var_map.Key.t * ('a -> 'a) * unit continuation -> eff + | With_error_handler : (Exn_with_backtrace.t -> Nothing.t t) * unit continuation -> eff + | Unwind : 'a continuation * 'a -> eff | Map_reduce_errors : (module Monoid with type t = 'a) * (Exn_with_backtrace.t -> 'a t) * (unit -> eff) - * (('b, 'a) result -> eff) + * ('b, 'a) result continuation -> eff - | Unwind_map_reduce : ('a -> eff) * 'a -> eff + | Unwind_map_reduce : 'a continuation * 'a -> eff | End_of_map_reduce_error_handler : (_, _) map_reduce_context' -> eff | End_of_fiber of unit | Never of unit @@ -36,7 +44,7 @@ and ('a, _) ivar_state = | Full : 'a -> ('a, [> `Full ]) ivar_state | Empty : ('a, [> `Empty ]) ivar_state | Empty_with_readers : - context * ('a -> eff) * ('a, [ `Empty ]) ivar_state + context * 'a continuation * ('a, [ `Empty ]) ivar_state -> ('a, [> `Empty ]) ivar_state and value = .. @@ -59,21 +67,31 @@ and map_reduce_context = | Map_reduce_context : (_, _) map_reduce_context' -> map_reduce_context and 'a k = - { run : 'a -> eff + { run : 'a continuation ; ctx : context } -let return x k = k x -let bind t ~f k = t (fun x -> f x k) -let map t ~f k = t (fun x -> k (f x)) +let rec continue : type a. a continuation -> a -> eff = + fun k x -> + match k with + | Function f -> f x + | Effect -> x + | Map (f, k) -> continue k (f x) + | Bind (f, k) -> f x k +;; + +let return x k = continue k x +let bind t ~f k = t (Bind (f, k)) +let map t ~f k = t (Map (f, k)) let with_error_handler f ~on_error k = - With_error_handler (on_error, fun () -> f () (fun x -> Unwind (k, x))) + With_error_handler + (on_error, Function (fun () -> f () (Function (fun x -> Unwind (k, x))))) ;; let map_reduce_errors m ~on_error f k = Map_reduce_errors - (m, on_error, (fun () -> f () (fun x -> Unwind_map_reduce (k, Ok x))), k) + (m, on_error, (fun () -> f () (Function (fun x -> Unwind_map_reduce (k, Ok x)))), k) ;; let suspend f k = Suspend (f, k) @@ -144,58 +162,70 @@ let rec nfork_array a i f = let parallel_iter_seq (seq : _ Seq.t) ~f k = match seq () with - | Nil -> k () + | Nil -> continue k () | Cons (x, seq) -> let left_over = ref 1 in let f x = - f x (fun () -> - decr left_over; - if !left_over = 0 then k () else end_of_fiber) + f + x + (Function + (fun () -> + decr left_over; + if !left_over = 0 then continue k () else end_of_fiber)) in nfork_seq left_over x seq f ;; let map_reduce_seq (seq : _ Seq.t) ~f ~empty ~combine k = match seq () with - | Nil -> k empty + | Nil -> continue k empty | Cons (x, seq) -> let current = ref empty in let running = ref 1 in let f x = - f x (fun y -> - current := combine !current y; - decr running; - if !running = 0 then k !current else end_of_fiber) + f + x + (Function + (fun y -> + current := combine !current y; + decr running; + if !running = 0 then continue k !current else end_of_fiber)) in nfork_seq running x seq f ;; let map_reduce_array a ~f ~empty ~combine k = match Array.length a with - | 0 -> k empty + | 0 -> continue k empty | len -> let current = ref empty in let running = ref len in let f x = - f x (fun y -> - current := combine !current y; - decr running; - if !running = 0 then k !current else end_of_fiber) + f + x + (Function + (fun y -> + current := combine !current y; + decr running; + if !running = 0 then continue k !current else end_of_fiber)) in nfork_array a 0 f ;; let map_reduce l ~f ~empty ~combine k = match l with - | [] -> k empty + | [] -> continue k empty | x :: l -> let current = ref empty in let running = ref (List.length l + 1) in let f x = - f x (fun y -> - current := combine !current y; - decr running; - if !running = 0 then k !current else end_of_fiber) + f + x + (Function + (fun y -> + current := combine !current y; + decr running; + if !running = 0 then continue k !current else end_of_fiber)) in nfork x l f ;; @@ -213,15 +243,17 @@ let fork_and_join fa fb k = state := Got_a a; end_of_fiber | Got_a _ -> assert false - | Got_b b -> k (a, b) + | Got_b b -> continue k (a, b) and kb b = match !state with | Nothing_yet -> state := Got_b b; end_of_fiber - | Got_a a -> k (a, b) + | Got_a a -> continue k (a, b) | Got_b _ -> assert false in + let ka = Function ka in + let kb = Function kb in match apply2 fa () ka with | End_of_fiber () -> fb () kb | eff -> Fork (eff, fun () -> fb () kb) @@ -230,26 +262,33 @@ let fork_and_join fa fb k = let fork_and_join_unit fa fb k = let state = ref Nothing_yet in match - apply2 fa () (fun () -> - match !state with - | Nothing_yet -> - state := Got_a (); - end_of_fiber - | Got_a _ -> assert false - | Got_b b -> k b) + apply2 + fa + () + (Function + (fun () -> + match !state with + | Nothing_yet -> + state := Got_a (); + end_of_fiber + | Got_a _ -> assert false + | Got_b b -> continue k b)) with | End_of_fiber () -> fb () k | eff -> Fork ( eff , fun () -> - fb () (fun b -> - match !state with - | Nothing_yet -> - state := Got_b b; - end_of_fiber - | Got_a () -> k b - | Got_b _ -> assert false) ) + fb + () + (Function + (fun b -> + match !state with + | Nothing_yet -> + state := Got_b b; + end_of_fiber + | Got_a () -> continue k b + | Got_b _ -> assert false)) ) ;; let rec length_and_rev l len acc = @@ -274,7 +313,7 @@ module Ivar = struct let read t k = match t.state with - | Full x -> k x + | Full x -> continue k x | Empty_with_readers _ | Empty -> Read_ivar (t, k) ;; @@ -282,7 +321,7 @@ module Ivar = struct match t.state with | Empty -> t.state <- Full x; - k () + continue k () | Full _ | Empty_with_readers _ -> Fill_ivar (t, x, k) ;; @@ -299,7 +338,8 @@ module Var = struct let get (key : 'a Var_map.Key.t) : 'a t = fun k -> Get_var (key, k) let set (key : 'a Var_map.Key.t) (value : 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> Set_var (key, value, fun () -> fiber () (fun x -> Unwind (k, x))) + fun k -> + Set_var (key, value, Function (fun () -> fiber () (Function (fun x -> Unwind (k, x))))) ;; let get_exn (key : 'a option Var_map.Key.t) : 'a t = @@ -309,25 +349,28 @@ module Var = struct ;; let update (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> Update_var (key, f, fun () -> fiber () (fun x -> Unwind (k, x))) + fun k -> + Update_var (key, f, Function (fun () -> fiber () (Function (fun x -> Unwind (k, x))))) ;; let get_apply (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Get_var (key, fun value -> f value x k) + fun k -> Get_var (key, Function (fun value -> f value x k)) ;; let get_apply_map (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c) (x : 'b) : 'c t = - fun k -> Get_var (key, fun value -> k (f value x)) + fun k -> Get_var (key, Function (fun value -> continue k (f value x))) ;; let set_apply (key : 'a Var_map.Key.t) (value : 'a) (f : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Set_var (key, value, fun () -> f x (fun y -> Unwind (k, y))) + fun k -> + Set_var (key, value, Function (fun () -> f x (Function (fun y -> Unwind (k, y))))) ;; let update_apply (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (g : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Update_var (key, f, fun () -> g x (fun y -> Unwind (k, y))) + fun k -> + Update_var (key, f, Function (fun () -> g x (Function (fun y -> Unwind (k, y))))) ;; include Var_map.Key @@ -337,9 +380,9 @@ let of_thunk f k = f () k let of_thunk_apply f x k = f x k module O = struct - let ( >>> ) a b k = a (fun () -> b k) - let ( >>= ) t f k = t (fun x -> f x k) - let ( >>| ) t f k = t (fun x -> k (f x)) + let ( >>> ) a b = bind a ~f:(fun () -> b) + let ( >>= ) t f = bind t ~f + let ( >>| ) t f = map t ~f let ( let+ ) = ( >>| ) let ( let* ) = ( >>= ) let ( and* ) a b = fork_and_join (fun () -> a) (fun () -> b) @@ -378,15 +421,18 @@ let sequential_iter l ~f = let parallel_iter l ~f k = match l with - | [] -> k () + | [] -> continue k () | [ x ] -> f x k | x :: l -> let len = List.length l + 1 in let left_over = ref len in let f x = - f x (fun () -> - decr left_over; - if !left_over = 0 then k () else end_of_fiber) + f + x + (Function + (fun () -> + decr left_over; + if !left_over = 0 then continue k () else end_of_fiber)) in nfork x l f ;; @@ -396,35 +442,38 @@ let parallel_array_of_list_map' x l ~f k = let left_over = ref len in let results = ref [||] in let f i x = - f x (fun y -> - let a = - match !results with - | [||] -> - let a = Array.make len y in - results := a; - a - | a -> - a.(i) <- y; - a - in - decr left_over; - if !left_over = 0 then k a else end_of_fiber) + f + x + (Function + (fun y -> + let a = + match !results with + | [||] -> + let a = Array.make len y in + results := a; + a + | a -> + a.(i) <- y; + a + in + decr left_over; + if !left_over = 0 then continue k a else end_of_fiber)) in nforki x l f ;; let parallel_array_of_list_map l ~f k = match l with - | [] -> k [||] - | [ x ] -> f x (fun x -> k [| x |]) + | [] -> continue k [||] + | [ x ] -> f x (Map ((fun x -> [| x |]), k)) | x :: l -> parallel_array_of_list_map' x l ~f k ;; let parallel_map l ~f k = match l with - | [] -> k [] - | [ x ] -> f x (fun x -> k [ x ]) - | x :: l -> parallel_array_of_list_map' x l ~f (fun a -> k (Array.to_list a)) + | [] -> continue k [] + | [ x ] -> f x (Map ((fun x -> [ x ]), k)) + | x :: l -> parallel_array_of_list_map' x l ~f (Map (Array.to_list, k)) ;; let all = sequential_map ~f:Fun.id diff --git a/src/fiber/src/lazy.ml b/src/fiber/src/lazy.ml index 5a9a9724322..4a4c7fab99d 100644 --- a/src/fiber/src/lazy.ml +++ b/src/fiber/src/lazy.ml @@ -61,7 +61,7 @@ let is_value t = ;; let force_all_unit = - let stop () = end_of_fiber in + let stop = Function (fun () -> end_of_fiber) in (* Fork all computations that haven't been forced yet. Note that this should be substantially more efficient that [parallel_map ~f:force] since we ignore computations which have already been forced. *) @@ -71,7 +71,7 @@ let force_all_unit = | Done _ | Running _ -> return () | Init f -> let v = prep t in - fun k -> fork (fun () -> (execute t v f) stop) k) + fun k -> fork (fun () -> (execute t v f) stop) (fun () -> continue k ())) in (* Wait for all computations, collecting all exceptions. *) (* CR-someday rgrinberg: use [Appendable.t] for [acc] rather than [Appendable.t option]. *) diff --git a/src/fiber/src/mutex.ml b/src/fiber/src/mutex.ml index 7678e68f5f0..aa2ec2d7b73 100644 --- a/src/fiber/src/mutex.ml +++ b/src/fiber/src/mutex.ml @@ -12,7 +12,7 @@ let lock t k = then suspend (fun k -> Queue.push t.waiters k) k else ( t.locked <- true; - k ()) + continue k ()) ;; let unlock t k = @@ -20,7 +20,7 @@ let unlock t k = match Queue.pop t.waiters with | None -> t.locked <- false; - k () + continue k () | Some next -> resume next () k ;; diff --git a/src/fiber/src/mvar.ml b/src/fiber/src/mvar.ml index 9509a36d49a..0cc3cfc4f5a 100644 --- a/src/fiber/src/mvar.ml +++ b/src/fiber/src/mvar.ml @@ -28,10 +28,10 @@ let read t k = (match Queue.pop t.writers with | None -> t.value <- None; - k v + continue k v | Some (v', w) -> t.value <- Some v'; - resume w () (fun () -> k v)) + resume w () (Map ((fun () -> v), k))) ;; let write t x k = @@ -41,6 +41,6 @@ let write t x k = (match Queue.pop t.readers with | None -> t.value <- Some x; - k () - | Some r -> resume r x (fun () -> k ())) + continue k () + | Some r -> resume r x (Map ((fun () -> ()), k))) ;; diff --git a/src/fiber/src/pool.ml b/src/fiber/src/pool.ml index 044e2cf4349..359db08e02a 100644 --- a/src/fiber/src/pool.ml +++ b/src/fiber/src/pool.ml @@ -27,8 +27,8 @@ type nonrec t = let running t k = match t.status with - | Open -> k true - | Closed -> k false + | Open -> continue k true + | Closed -> continue k false ;; let create () = { tasks = Queue.create (); runner = Awaiting_run; status = Open } @@ -39,7 +39,7 @@ let task t ~f k = | Open -> Queue.push t.tasks f; (match t.runner with - | Running | Awaiting_run -> k () + | Running | Awaiting_run -> continue k () | Awaiting_resume r -> t.runner <- Running; resume r () k) @@ -47,11 +47,11 @@ let task t ~f k = let close t k = match t.status with - | Closed -> k () + | Closed -> continue k () | Open -> t.status <- Closed; (match t.runner with - | Running | Awaiting_run -> k () + | Running | Awaiting_run -> continue k () | Awaiting_resume r -> t.runner <- Running; resume r () k) @@ -68,14 +68,14 @@ let run t k = let n = ref 1 in let done_fiber () = decr n; - if !n = 0 then k () else end_of_fiber + if !n = 0 then continue k () else end_of_fiber in let rec read t = match Queue.pop t.tasks with | None -> finish_or_suspend t | Some v -> incr n; - fork (fun () -> v () done_fiber) read_delayed + fork (fun () -> v () (Function done_fiber)) read_delayed and read_delayed () = read t and suspend_k k = (* we are suspending because we have no tasks *) @@ -84,7 +84,7 @@ let run t k = and finish_or_suspend t = match t.status with | Closed -> done_fiber () - | Open -> suspend suspend_k read_delayed + | Open -> suspend suspend_k (Function read_delayed) in read t ;; diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 1b25b9e8a2b..e776c400bf4 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -6,7 +6,7 @@ type fill = Fill : 'a ivar * 'a -> fill module Jobs = struct type t = | Empty - | Job : context * ('a -> eff) * 'a * t -> t + | Job : context * 'a continuation * 'a * t -> t | Concat : t * t -> t let concat a b = @@ -71,9 +71,9 @@ and loop2 a b = | Job (ctx, run, x, a) -> exec ctx run x (Jobs.concat a b) | Concat (a1, a2) -> loop2 a1 (Jobs.concat a2 b) -and exec : 'a. context -> ('a -> eff) -> 'a -> Jobs.t -> step' = +and exec : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = fun ctx k x jobs -> - match k x with + match continue k x with | exception exn -> let exn = Exn_with_backtrace.capture exn in exec ctx.on_error.ctx ctx.on_error.run exn jobs @@ -109,7 +109,11 @@ and exec : 'a. context -> ('a -> eff) -> 'a -> Jobs.t -> step' = in exec ctx k () jobs | With_error_handler (on_error, k) -> - let on_error = { ctx; run = (fun exn -> on_error exn Nothing.unreachable_code) } in + let on_error = + { ctx + ; run = Function (fun exn -> on_error exn (Function Nothing.unreachable_code)) + } + in let ctx = { ctx with parent = ctx; on_error } in exec ctx k () jobs | Map_reduce_errors (m, on_error, f, k) -> map_reduce_errors ctx m on_error f k jobs @@ -127,7 +131,7 @@ and exec : 'a. context -> ('a -> eff) -> 'a -> Jobs.t -> step' = | Fork (a, b) -> let (Map_reduce_context r) = ctx.map_reduce_context in r.ref_count <- r.ref_count + 1; - exec ctx Fun.id a (Job (ctx, b, (), jobs)) + exec ctx Effect a (Job (ctx, Function b, (), jobs)) | Reraise exn -> let { ctx; run } = ctx.on_error in exec ctx run exn jobs @@ -160,7 +164,7 @@ and map_reduce_errors -> (module Monoid with type t = errors) -> (Exn_with_backtrace.t -> errors t) -> (unit -> eff) - -> ((b, errors) result -> eff) + -> (b, errors) result continuation -> Jobs.t -> step' = @@ -169,10 +173,14 @@ and map_reduce_errors let on_error = { ctx ; run = - (fun exn -> - on_error exn (fun m -> - map_reduce_context.errors <- M.combine map_reduce_context.errors m; - End_of_map_reduce_error_handler map_reduce_context)) + Function + (fun exn -> + on_error + exn + (Function + (fun m -> + map_reduce_context.errors <- M.combine map_reduce_context.errors m; + End_of_map_reduce_error_handler map_reduce_context))) } in let ctx = @@ -182,7 +190,7 @@ and map_reduce_errors ; map_reduce_context = Map_reduce_context map_reduce_context } in - exec ctx f () jobs + exec ctx (Function f) () jobs ;; let repack_step (type a) (module W : Witness with type t = a) (step' : step') = @@ -208,12 +216,15 @@ let start (type a) (t : a t) = in let rec ctx = { parent = ctx - ; on_error = { ctx; run = (fun exn -> Toplevel_exception exn) } + ; on_error = { ctx; run = Function (fun exn -> Toplevel_exception exn) } ; vars = Var_map.empty ; map_reduce_context = Map_reduce_context - { k = { ctx; run = (fun _ -> assert false) }; ref_count = 1; errors = () } + { k = { ctx; run = Function (fun _ -> assert false) } + ; ref_count = 1 + ; errors = () + } } in - exec ctx t (fun x -> Done (W.X x)) Empty |> repack_step (module W) + exec ctx (Function t) (Function (fun x -> Done (W.X x))) Empty |> repack_step (module W) ;; diff --git a/src/fiber/src/stream.ml b/src/fiber/src/stream.ml index a7b0c21b687..d7b1f5097fc 100644 --- a/src/fiber/src/stream.ml +++ b/src/fiber/src/stream.ml @@ -101,20 +101,23 @@ module In = struct let parallel_iter t ~f k = let n = ref 1 in - let k () = + let done_ () = decr n; if !n = 0 then ( unlock t; - k ()) + continue k ()) else end_of_fiber in let rec loop t = - t.read () (function - | None -> k () - | Some x -> - incr n; - fork (fun () -> f x k) (fun () -> loop t)) + t.read + () + (Function + (function + | None -> done_ () + | Some x -> + incr n; + fork (fun () -> f x (Function done_)) (fun () -> loop t))) in lock t; loop t diff --git a/src/fiber/src/svar.ml b/src/fiber/src/svar.ml index 7b828b002be..17e7033e7dc 100644 --- a/src/fiber/src/svar.ml +++ b/src/fiber/src/svar.ml @@ -25,8 +25,8 @@ let create current = { current; waiters = [] } let write = let rec run_awakers final = function - | [] -> final () - | k :: ks -> resume k () (fun () -> run_awakers final ks) + | [] -> continue final () + | k :: ks -> resume k () (Function (fun () -> run_awakers final ks)) in fun t a k -> t.current <- a; @@ -35,7 +35,7 @@ let write = if f t.current then Right k else Left (k, f)) in match awake with - | [] -> k () + | [] -> continue k () | awake -> t.waiters <- List.rev sleep; run_awakers k awake From d57af41d35e6641bebd61d31db7dd2dc7f0aa584 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 09:40:41 +0100 Subject: [PATCH 02/11] perf(fiber): specialize apply and unwind continuations Represent the continuations used by Fiber variable helpers and context unwinding directly instead of wrapping them in function closures. This keeps the function escape hatch for arbitrary callbacks while making common internal control transfers smaller. On a no-op @install self-build this reduced minor allocation by 3.0M words, promoted allocation by 1.0M words, and Cachegrind instructions by 1.5% on top of explicit map and bind continuations. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 4daa7de6f8c..ea3355034ac 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -9,6 +9,10 @@ and 'a continuation = | Effect : eff continuation | Map : ('a -> 'b) * 'b continuation -> 'a continuation | Bind : ('a -> 'b t) * 'b continuation -> 'a continuation + | Apply : ('a -> 'b -> 'c t) * 'b * 'c continuation -> 'a continuation + | Apply_map : ('a -> 'b -> 'c) * 'b * 'c continuation -> 'a continuation + | Unwind_to : 'a continuation -> 'a continuation + | Unwind_map_reduce_to : ('a, 'b) result continuation -> 'a continuation and eff = | Read_ivar : 'a ivar * 'a continuation -> eff @@ -78,6 +82,10 @@ let rec continue : type a. a continuation -> a -> eff = | Effect -> x | Map (f, k) -> continue k (f x) | Bind (f, k) -> f x k + | Apply (f, y, k) -> f x y k + | Apply_map (f, y, k) -> continue k (f x y) + | Unwind_to k -> Unwind (k, x) + | Unwind_map_reduce_to k -> Unwind_map_reduce (k, Ok x) ;; let return x k = continue k x @@ -85,13 +93,11 @@ let bind t ~f k = t (Bind (f, k)) let map t ~f k = t (Map (f, k)) let with_error_handler f ~on_error k = - With_error_handler - (on_error, Function (fun () -> f () (Function (fun x -> Unwind (k, x))))) + With_error_handler (on_error, Function (fun () -> f () (Unwind_to k))) ;; let map_reduce_errors m ~on_error f k = - Map_reduce_errors - (m, on_error, (fun () -> f () (Function (fun x -> Unwind_map_reduce (k, Ok x)))), k) + Map_reduce_errors (m, on_error, (fun () -> f () (Unwind_map_reduce_to k)), k) ;; let suspend f k = Suspend (f, k) @@ -338,8 +344,7 @@ module Var = struct let get (key : 'a Var_map.Key.t) : 'a t = fun k -> Get_var (key, k) let set (key : 'a Var_map.Key.t) (value : 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> - Set_var (key, value, Function (fun () -> fiber () (Function (fun x -> Unwind (k, x))))) + fun k -> Set_var (key, value, Function (fun () -> fiber () (Unwind_to k))) ;; let get_exn (key : 'a option Var_map.Key.t) : 'a t = @@ -349,28 +354,25 @@ module Var = struct ;; let update (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> - Update_var (key, f, Function (fun () -> fiber () (Function (fun x -> Unwind (k, x))))) + fun k -> Update_var (key, f, Function (fun () -> fiber () (Unwind_to k))) ;; let get_apply (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Get_var (key, Function (fun value -> f value x k)) + fun k -> Get_var (key, Apply (f, x, k)) ;; let get_apply_map (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c) (x : 'b) : 'c t = - fun k -> Get_var (key, Function (fun value -> continue k (f value x))) + fun k -> Get_var (key, Apply_map (f, x, k)) ;; let set_apply (key : 'a Var_map.Key.t) (value : 'a) (f : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> - Set_var (key, value, Function (fun () -> f x (Function (fun y -> Unwind (k, y))))) + fun k -> Set_var (key, value, Function (fun () -> f x (Unwind_to k))) ;; let update_apply (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (g : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> - Update_var (key, f, Function (fun () -> g x (Function (fun y -> Unwind (k, y))))) + fun k -> Update_var (key, f, Function (fun () -> g x (Unwind_to k))) ;; include Var_map.Key From 072703ede6507ff2a53a0f2705e8df3e406a3cf5 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 09:44:45 +0100 Subject: [PATCH 03/11] perf(fiber): fuse short map continuation chains Represent runs of up to three adjacent map continuations with one typed frame. This shortens promoted continuation chains and reduces scheduler dispatch without changing map order or the public API. On a no-op @install self-build this reduced promoted allocation by 0.25M words, peak heap size by 0.15M words, and Cachegrind instructions by 0.3%. Minor allocation increased by 0.46M words and native wall time remained neutral. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index ea3355034ac..7cdcd971506 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -8,6 +8,8 @@ and 'a continuation = | Function : ('a -> eff) -> 'a continuation | Effect : eff continuation | Map : ('a -> 'b) * 'b continuation -> 'a continuation + | Map2 : ('a -> 'b) * ('b -> 'c) * 'c continuation -> 'a continuation + | Map3 : ('a -> 'b) * ('b -> 'c) * ('c -> 'd) * 'd continuation -> 'a continuation | Bind : ('a -> 'b t) * 'b continuation -> 'a continuation | Apply : ('a -> 'b -> 'c t) * 'b * 'c continuation -> 'a continuation | Apply_map : ('a -> 'b -> 'c) * 'b * 'c continuation -> 'a continuation @@ -81,6 +83,8 @@ let rec continue : type a. a continuation -> a -> eff = | Function f -> f x | Effect -> x | Map (f, k) -> continue k (f x) + | Map2 (f, g, k) -> continue k (g (f x)) + | Map3 (f, g, h, k) -> continue k (h (g (f x))) | Bind (f, k) -> f x k | Apply (f, y, k) -> f x y k | Apply_map (f, y, k) -> continue k (f x y) @@ -90,7 +94,13 @@ let rec continue : type a. a continuation -> a -> eff = let return x k = continue k x let bind t ~f k = t (Bind (f, k)) -let map t ~f k = t (Map (f, k)) + +let map t ~f k = + match k with + | Map (g, k) -> t (Map2 (f, g, k)) + | Map2 (g, h, k) -> t (Map3 (f, g, h, k)) + | k -> t (Map (f, k)) +;; let with_error_handler f ~on_error k = With_error_handler (on_error, Function (fun () -> f () (Unwind_to k))) From f68440ec23937d1b79c97d8fc0efd77ee40f3152 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 11:19:23 +0100 Subject: [PATCH 04/11] test(fiber): cover reusable fiber thunks Check that constructing a thunk fiber does not execute it and that running the same fiber twice executes the thunk twice. This guards the documented reuse semantics before changing the internal computation representation. Signed-off-by: Rudi Grinberg --- src/fiber/test/fiber_tests.ml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/fiber/test/fiber_tests.ml b/src/fiber/test/fiber_tests.ml index 66d73f7f866..639335fcead 100644 --- a/src/fiber/test/fiber_tests.ml +++ b/src/fiber/test/fiber_tests.ml @@ -18,6 +18,23 @@ let%expect_test "basics" = [%expect {| () |}] ;; +let%expect_test "fibers are reusable and thunks run during execution" = + let runs = ref 0 in + let fiber = + Fiber.of_thunk (fun () -> + incr runs; + Fiber.return !runs) + in + printfn "before: %d" !runs; + test int fiber; + test int fiber; + [%expect + {| + before: 0 + 1 + 2 |}] +;; + let%expect_test "collect_errors" = test (backtrace_result unit) (Fiber.collect_errors (fun () -> raise Exit)); [%expect {| Error [ { exn = "Stdlib.Exit"; backtrace = "" } ] |}] From 5565a85d16a95e320f67dc36793a44b65a85653c Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 11:47:40 +0100 Subject: [PATCH 05/11] perf(fiber): represent fiber computations explicitly Replace the internal CPS computation function with typed computation nodes and interpret computation and continuation spines directly in the scheduler. Keep user callbacks as opaque function payloads and pair internal primitive runners with explicit state instead of allocated closures. On a no-op @install self-build this reduced minor allocation by 23.8M words, promoted allocation by 1.2M words, the live heap by 1.1M words, and Cachegrind instructions by 1.1%. Alternating native measurements also showed lower wall time. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 293 ++++++++++++++++++++----------- src/fiber/src/fiber.ml | 4 +- src/fiber/src/lazy.ml | 9 +- src/fiber/src/mutex.ml | 12 +- src/fiber/src/mvar.ml | 16 +- src/fiber/src/pool.ml | 23 ++- src/fiber/src/scheduler.ml | 348 ++++++++++++++++++++++++++++++++----- src/fiber/src/stream.ml | 9 +- src/fiber/src/svar.ml | 32 ++-- 9 files changed, 563 insertions(+), 183 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 7cdcd971506..93af6bfb14c 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -1,12 +1,44 @@ open Stdune -type 'a t = 'a continuation -> eff +(* Fiber computations are explicit nodes. Callbacks supplied by callers remain functions, + while the Fiber-created computation spine is interpreted by the scheduler. *) +type _ t = + | Return_t : 'a -> 'a t + | Never_t : 'a t + | Map_t : 'a t * ('a -> 'b) -> 'b t + | Map2_t : 'a t * ('a -> 'b) * ('b -> 'c) -> 'c t + | Map3_t : 'a t * ('a -> 'b) * ('b -> 'c) * ('c -> 'd) -> 'd t + | Bind_t : 'a t * ('a -> 'b t) -> 'b t + | Thunk_t : (unit -> 'a t) -> 'a t + | Thunk_apply_t : ('a -> 'b t) * 'a -> 'b t + | With_error_handler_t : (unit -> 'a t) * (Exn_with_backtrace.t -> Nothing.t t) -> 'a t + | Map_reduce_errors_t : + (module Monoid with type t = 'a) * (Exn_with_backtrace.t -> 'a t) * (unit -> 'b t) + -> ('b, 'a) result t + | Suspend_t : ('a k -> unit) -> 'a t + | Resume_t : 'a k * 'a -> unit t + | Reraise_all_t : Exn_with_backtrace.t list -> 'a t + | Ivar_read_t : 'a ivar -> 'a t + | Ivar_fill_t : 'a ivar * 'a -> unit t + | Get_var_t : 'a Var_map.Key.t -> 'a t + | Set_var_t : 'a Var_map.Key.t * 'a * (unit -> 'b t) -> 'b t + | Update_var_t : 'a Var_map.Key.t * ('a -> 'a) * (unit -> 'b t) -> 'b t + | Get_apply_t : 'a Var_map.Key.t * ('a -> 'b -> 'c t) * 'b -> 'c t + | Get_apply_map_t : 'a Var_map.Key.t * ('a -> 'b -> 'c) * 'b -> 'c t + | Set_apply_t : 'a Var_map.Key.t * 'a * ('b -> 'c t) * 'b -> 'c t + | Update_apply_t : 'a Var_map.Key.t * ('a -> 'a) * ('b -> 'c t) * 'b -> 'c t + (* Pair internal operations with their state without allocating a closure. *) + | Primitive_t : ('a -> 'b continuation -> eff) * 'a -> 'b t + | Primitive2_t : ('a -> 'b -> 'c continuation -> eff) * 'a * 'b -> 'c t + | Primitive3_t : ('a -> 'b -> 'c -> 'd continuation -> eff) * 'a * 'b * 'c -> 'd t + | Primitive4_t : + ('a -> 'b -> 'c -> 'd -> 'e continuation -> eff) * 'a * 'b * 'c * 'd + -> 'e t (* Keep arbitrary callbacks as functions, but represent the frames introduced by common fiber combinators directly. *) and 'a continuation = | Function : ('a -> eff) -> 'a continuation - | Effect : eff continuation | Map : ('a -> 'b) * 'b continuation -> 'a continuation | Map2 : ('a -> 'b) * ('b -> 'c) * 'c continuation -> 'a continuation | Map3 : ('a -> 'b) * ('b -> 'c) * ('c -> 'd) * 'd continuation -> 'a continuation @@ -17,19 +49,26 @@ and 'a continuation = | Unwind_map_reduce_to : ('a, 'b) result continuation -> 'a continuation and eff = + | Run : 'a t * 'a continuation -> eff | Read_ivar : 'a ivar * 'a continuation -> eff | Fill_ivar : 'a ivar * 'a * unit continuation -> eff | Suspend : ('a k -> unit) * 'a continuation -> eff | Resume : 'a k * 'a * unit continuation -> eff | Get_var : 'a Var_map.Key.t * 'a continuation -> eff - | Set_var : 'a Var_map.Key.t * 'a * unit continuation -> eff - | Update_var : 'a Var_map.Key.t * ('a -> 'a) * unit continuation -> eff - | With_error_handler : (Exn_with_backtrace.t -> Nothing.t t) * unit continuation -> eff + | Set_var : 'a Var_map.Key.t * 'a * (unit -> 'b t) * 'b continuation -> eff + | Update_var : 'a Var_map.Key.t * ('a -> 'a) * (unit -> 'b t) * 'b continuation -> eff + | Set_var_apply : 'a Var_map.Key.t * 'a * ('b -> 'c t) * 'b * 'c continuation -> eff + | Update_var_apply : + 'a Var_map.Key.t * ('a -> 'a) * ('b -> 'c t) * 'b * 'c continuation + -> eff + | With_error_handler : + (Exn_with_backtrace.t -> Nothing.t t) * (unit -> 'a t) * 'a continuation + -> eff | Unwind : 'a continuation * 'a -> eff | Map_reduce_errors : (module Monoid with type t = 'a) * (Exn_with_backtrace.t -> 'a t) - * (unit -> eff) + * (unit -> 'b t) * ('b, 'a) result continuation -> eff | Unwind_map_reduce : 'a continuation * 'a -> eff @@ -81,39 +120,32 @@ let rec continue : type a. a continuation -> a -> eff = fun k x -> match k with | Function f -> f x - | Effect -> x | Map (f, k) -> continue k (f x) | Map2 (f, g, k) -> continue k (g (f x)) | Map3 (f, g, h, k) -> continue k (h (g (f x))) - | Bind (f, k) -> f x k - | Apply (f, y, k) -> f x y k + | Bind (f, k) -> Run (f x, k) + | Apply (f, y, k) -> Run (f x y, k) | Apply_map (f, y, k) -> continue k (f x y) | Unwind_to k -> Unwind (k, x) | Unwind_map_reduce_to k -> Unwind_map_reduce (k, Ok x) ;; -let return x k = continue k x -let bind t ~f k = t (Bind (f, k)) - -let map t ~f k = - match k with - | Map (g, k) -> t (Map2 (f, g, k)) - | Map2 (g, h, k) -> t (Map3 (f, g, h, k)) - | k -> t (Map (f, k)) -;; - -let with_error_handler f ~on_error k = - With_error_handler (on_error, Function (fun () -> f () (Unwind_to k))) -;; +let return x = Return_t x +let bind t ~f = Bind_t (t, f) -let map_reduce_errors m ~on_error f k = - Map_reduce_errors (m, on_error, (fun () -> f () (Unwind_map_reduce_to k)), k) +let map t ~f = + match t with + | Map_t (t, g) -> Map2_t (t, g, f) + | Map2_t (t, g, h) -> Map3_t (t, g, h, f) + | t -> Map_t (t, f) ;; -let suspend f k = Suspend (f, k) -let resume suspended x k = Resume (suspended, x, k) +let with_error_handler f ~on_error = With_error_handler_t (f, on_error) +let map_reduce_errors m ~on_error f = Map_reduce_errors_t (m, on_error, f) +let suspend f = Suspend_t f +let resume suspended x = Resume_t (suspended, x) let end_of_fiber = End_of_fiber () -let never _k = Never () +let never = Never_t let apply f x = try f x with @@ -129,6 +161,64 @@ let apply2 f x y = Reraise exn ;; +let rec eval : type a. a t -> a continuation -> eff = + fun t k -> + match t with + | Return_t x -> continue k x + | Never_t -> Never () + | Map_t (t, f) -> eval t (Map (f, k)) + | Map2_t (t, f, g) -> eval t (Map2 (f, g, k)) + | Map3_t (t, f, g, h) -> eval t (Map3 (f, g, h, k)) + | Bind_t (t, f) -> eval t (Bind (f, k)) + | Thunk_t f -> eval (f ()) k + | Thunk_apply_t (f, x) -> eval (f x) k + | With_error_handler_t (f, on_error) -> With_error_handler (on_error, f, k) + | Map_reduce_errors_t (m, on_error, f) -> Map_reduce_errors (m, on_error, f, k) + | Suspend_t f -> Suspend (f, k) + | Resume_t (suspended, x) -> Resume (suspended, x, k) + | Reraise_all_t exns -> + (match exns with + | [] -> Never () + | [ exn ] -> Exn_with_backtrace.reraise exn + | _ -> Reraise_all exns) + | Ivar_read_t ivar -> + (match ivar.state with + | Full x -> continue k x + | Empty_with_readers _ | Empty -> Read_ivar (ivar, k)) + | Ivar_fill_t (ivar, x) -> + (match ivar.state with + | Empty -> + ivar.state <- Full x; + continue k () + | Full _ | Empty_with_readers _ -> Fill_ivar (ivar, x, k)) + | Get_var_t key -> Get_var (key, k) + | Set_var_t (key, value, f) -> Set_var (key, value, f, k) + | Update_var_t (key, f, body) -> Update_var (key, f, body, k) + | Get_apply_t (key, f, x) -> Get_var (key, Apply (f, x, k)) + | Get_apply_map_t (key, f, x) -> Get_var (key, Apply_map (f, x, k)) + | Set_apply_t (key, value, f, x) -> Set_var_apply (key, value, f, x, k) + | Update_apply_t (key, f, body, x) -> Update_var_apply (key, f, body, x, k) + | Primitive_t (f, x) -> f x k + | Primitive2_t (f, x, y) -> f x y k + | Primitive3_t (f, x, y, z) -> f x y z k + | Primitive4_t (f, w, x, y, z) -> f w x y z k +;; + +let primitive f x = Primitive_t (f, x) +let primitive2 f x y = Primitive2_t (f, x, y) +let primitive3 f x y z = Primitive3_t (f, x, y, z) +let primitive4 f w x y z = Primitive4_t (f, w, x, y, z) + +let apply_t f x k = + try eval (f x) k with + | exn -> Reraise (Exn_with_backtrace.capture exn) +;; + +let apply_t2 f x y k = + try eval (f x y) k with + | exn -> Reraise (Exn_with_backtrace.capture exn) +;; + let[@inline always] fork a b = match apply a () with | End_of_fiber () -> b () @@ -176,30 +266,34 @@ let rec nfork_array a i f = | eff -> Fork (eff, fun () -> nfork_array a (i + 1) f)) ;; -let parallel_iter_seq (seq : _ Seq.t) ~f k = +let run_parallel_iter_seq (seq : _ Seq.t) f k = match seq () with | Nil -> continue k () | Cons (x, seq) -> let left_over = ref 1 in - let f x = - f + let f' x = + apply_t + f x (Function (fun () -> decr left_over; if !left_over = 0 then continue k () else end_of_fiber)) in - nfork_seq left_over x seq f + nfork_seq left_over x seq f' ;; -let map_reduce_seq (seq : _ Seq.t) ~f ~empty ~combine k = +let parallel_iter_seq seq ~f = primitive2 run_parallel_iter_seq seq f + +let run_map_reduce_seq (seq : _ Seq.t) f empty combine k = match seq () with | Nil -> continue k empty | Cons (x, seq) -> let current = ref empty in let running = ref 1 in - let f x = - f + let f' x = + apply_t + f x (Function (fun y -> @@ -207,17 +301,22 @@ let map_reduce_seq (seq : _ Seq.t) ~f ~empty ~combine k = decr running; if !running = 0 then continue k !current else end_of_fiber)) in - nfork_seq running x seq f + nfork_seq running x seq f' ;; -let map_reduce_array a ~f ~empty ~combine k = +let map_reduce_seq seq ~f ~empty ~combine = + primitive4 run_map_reduce_seq seq f empty combine +;; + +let run_map_reduce_array a f empty combine k = match Array.length a with | 0 -> continue k empty | len -> let current = ref empty in let running = ref len in - let f x = - f + let f' x = + apply_t + f x (Function (fun y -> @@ -225,17 +324,22 @@ let map_reduce_array a ~f ~empty ~combine k = decr running; if !running = 0 then continue k !current else end_of_fiber)) in - nfork_array a 0 f + nfork_array a 0 f' ;; -let map_reduce l ~f ~empty ~combine k = +let map_reduce_array a ~f ~empty ~combine = + primitive4 run_map_reduce_array a f empty combine +;; + +let run_map_reduce l f empty combine k = match l with | [] -> continue k empty | x :: l -> let current = ref empty in let running = ref (List.length l + 1) in - let f x = - f + let f' x = + apply_t + f x (Function (fun y -> @@ -243,15 +347,17 @@ let map_reduce l ~f ~empty ~combine k = decr running; if !running = 0 then continue k !current else end_of_fiber)) in - nfork x l f + nfork x l f' ;; +let map_reduce l ~f ~empty ~combine = primitive4 run_map_reduce l f empty combine + type ('a, 'b) fork_and_join_state = | Nothing_yet | Got_a of 'a | Got_b of 'b -let fork_and_join fa fb k = +let run_fork_and_join fa fb k = let state = ref Nothing_yet in let ka a = match !state with @@ -270,15 +376,17 @@ let fork_and_join fa fb k = in let ka = Function ka in let kb = Function kb in - match apply2 fa () ka with - | End_of_fiber () -> fb () kb - | eff -> Fork (eff, fun () -> fb () kb) + match apply_t fa () ka with + | End_of_fiber () -> apply_t fb () kb + | eff -> Fork (eff, fun () -> apply_t fb () kb) ;; -let fork_and_join_unit fa fb k = +let fork_and_join fa fb = primitive2 run_fork_and_join fa fb + +let run_fork_and_join_unit fa fb k = let state = ref Nothing_yet in match - apply2 + apply_t fa () (Function @@ -290,12 +398,13 @@ let fork_and_join_unit fa fb k = | Got_a _ -> assert false | Got_b b -> continue k b)) with - | End_of_fiber () -> fb () k + | End_of_fiber () -> apply_t fb () k | eff -> Fork ( eff , fun () -> - fb + apply_t + fb () (Function (fun b -> @@ -307,6 +416,8 @@ let fork_and_join_unit fa fb k = | Got_b _ -> assert false)) ) ;; +let fork_and_join_unit fa fb = primitive2 run_fork_and_join_unit fa fb + let rec length_and_rev l len acc = match l with | [] -> len, acc @@ -314,33 +425,14 @@ let rec length_and_rev l len acc = ;; let length_and_rev l = length_and_rev l 0 [] - -let reraise_all l _k = - match l with - | [] -> Never () - | [ exn ] -> Exn_with_backtrace.reraise exn - | _ -> Reraise_all l -;; +let reraise_all l = Reraise_all_t l module Ivar = struct type 'a t = 'a ivar let create () = { state = Empty } - - let read t k = - match t.state with - | Full x -> continue k x - | Empty_with_readers _ | Empty -> Read_ivar (t, k) - ;; - - let fill t x k = - match t.state with - | Empty -> - t.state <- Full x; - continue k () - | Full _ | Empty_with_readers _ -> Fill_ivar (t, x, k) - ;; - + let read t = Ivar_read_t t + let fill t x = Ivar_fill_t (t, x) let create_full a = { state = Full a } let peek t = @@ -351,10 +443,10 @@ module Ivar = struct end module Var = struct - let get (key : 'a Var_map.Key.t) : 'a t = fun k -> Get_var (key, k) + let get (key : 'a Var_map.Key.t) : 'a t = Get_var_t key let set (key : 'a Var_map.Key.t) (value : 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> Set_var (key, value, Function (fun () -> fiber () (Unwind_to k))) + Set_var_t (key, value, fiber) ;; let get_exn (key : 'a option Var_map.Key.t) : 'a t = @@ -364,32 +456,32 @@ module Var = struct ;; let update (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (fiber : unit -> 'b t) : 'b t = - fun k -> Update_var (key, f, Function (fun () -> fiber () (Unwind_to k))) + Update_var_t (key, f, fiber) ;; let get_apply (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Get_var (key, Apply (f, x, k)) + Get_apply_t (key, f, x) ;; let get_apply_map (key : 'a Var_map.Key.t) (f : 'a -> 'b -> 'c) (x : 'b) : 'c t = - fun k -> Get_var (key, Apply_map (f, x, k)) + Get_apply_map_t (key, f, x) ;; let set_apply (key : 'a Var_map.Key.t) (value : 'a) (f : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Set_var (key, value, Function (fun () -> f x (Unwind_to k))) + Set_apply_t (key, value, f, x) ;; let update_apply (key : 'a Var_map.Key.t) ~(f : 'a -> 'a) (g : 'b -> 'c t) (x : 'b) : 'c t = - fun k -> Update_var (key, f, Function (fun () -> g x (Unwind_to k))) + Update_apply_t (key, f, g, x) ;; include Var_map.Key end -let of_thunk f k = f () k -let of_thunk_apply f x k = f x k +let of_thunk f = Thunk_t f +let of_thunk_apply f x = Thunk_apply_t (f, x) module O = struct let ( >>> ) a b = bind a ~f:(fun () -> b) @@ -431,30 +523,34 @@ let sequential_iter l ~f = loop l ;; -let parallel_iter l ~f k = +let run_parallel_iter l f k = match l with | [] -> continue k () - | [ x ] -> f x k + | [ x ] -> apply_t f x k | x :: l -> let len = List.length l + 1 in let left_over = ref len in - let f x = - f + let f' x = + apply_t + f x (Function (fun () -> decr left_over; if !left_over = 0 then continue k () else end_of_fiber)) in - nfork x l f + nfork x l f' ;; -let parallel_array_of_list_map' x l ~f k = +let parallel_iter l ~f = primitive2 run_parallel_iter l f + +let run_parallel_array_of_list_map' x l f k = let len = List.length l + 1 in let left_over = ref len in let results = ref [||] in - let f i x = - f + let f' i x = + apply_t + f x (Function (fun y -> @@ -471,23 +567,26 @@ let parallel_array_of_list_map' x l ~f k = decr left_over; if !left_over = 0 then continue k a else end_of_fiber)) in - nforki x l f + nforki x l f' ;; -let parallel_array_of_list_map l ~f k = +let run_parallel_array_of_list_map l f k = match l with | [] -> continue k [||] - | [ x ] -> f x (Map ((fun x -> [| x |]), k)) - | x :: l -> parallel_array_of_list_map' x l ~f k + | [ x ] -> apply_t f x (Map ((fun x -> [| x |]), k)) + | x :: l -> run_parallel_array_of_list_map' x l f k ;; -let parallel_map l ~f k = +let parallel_array_of_list_map l ~f = primitive2 run_parallel_array_of_list_map l f + +let run_parallel_map l f k = match l with | [] -> continue k [] - | [ x ] -> f x (Map ((fun x -> [ x ]), k)) - | x :: l -> parallel_array_of_list_map' x l ~f (Map (Array.to_list, k)) + | [ x ] -> apply_t f x (Map ((fun x -> [ x ]), k)) + | x :: l -> run_parallel_array_of_list_map' x l f (Map (Array.to_list, k)) ;; +let parallel_map l ~f = primitive2 run_parallel_map l f let all = sequential_map ~f:Fun.id let all_concurrently = parallel_map ~f:Fun.id let all_concurrently_unit l = parallel_iter l ~f:Fun.id diff --git a/src/fiber/src/fiber.ml b/src/fiber/src/fiber.ml index b0cb0543300..22eea5f6e1e 100644 --- a/src/fiber/src/fiber.ml +++ b/src/fiber/src/fiber.ml @@ -25,8 +25,8 @@ type fill = Scheduler.fill = Fill : 'a ivar * 'a -> fill module Expert = struct type nonrec 'a k = 'a k - let suspend f k = suspend f k - let resume a x k = resume a x k + let suspend = suspend + let resume = resume end module Temp = Stdune.Temp.Monad (struct diff --git a/src/fiber/src/lazy.ml b/src/fiber/src/lazy.ml index 4a4c7fab99d..b34c704a720 100644 --- a/src/fiber/src/lazy.ml +++ b/src/fiber/src/lazy.ml @@ -60,8 +60,13 @@ let is_value t = | Running _ | Init _ -> false ;; +let stop = Function (fun () -> end_of_fiber) + +let run_force t v f k = + fork (fun () -> eval (execute t v f) stop) (fun () -> continue k ()) +;; + let force_all_unit = - let stop = Function (fun () -> end_of_fiber) in (* Fork all computations that haven't been forced yet. Note that this should be substantially more efficient that [parallel_map ~f:force] since we ignore computations which have already been forced. *) @@ -71,7 +76,7 @@ let force_all_unit = | Done _ | Running _ -> return () | Init f -> let v = prep t in - fun k -> fork (fun () -> (execute t v f) stop) (fun () -> continue k ())) + primitive3 run_force t v f) in (* Wait for all computations, collecting all exceptions. *) (* CR-someday rgrinberg: use [Appendable.t] for [acc] rather than [Appendable.t option]. *) diff --git a/src/fiber/src/mutex.ml b/src/fiber/src/mutex.ml index aa2ec2d7b73..828adac5985 100644 --- a/src/fiber/src/mutex.ml +++ b/src/fiber/src/mutex.ml @@ -7,23 +7,27 @@ type t = ; mutable waiters : unit k Queue.t } -let lock t k = +let run_lock t k = if t.locked - then suspend (fun k -> Queue.push t.waiters k) k + then Suspend ((fun k -> Queue.push t.waiters k), k) else ( t.locked <- true; continue k ()) ;; -let unlock t k = +let lock t = primitive run_lock t + +let run_unlock t k = assert t.locked; match Queue.pop t.waiters with | None -> t.locked <- false; continue k () - | Some next -> resume next () k + | Some next -> Resume (next, (), k) ;; +let unlock t = primitive run_unlock t + let with_lock t ~f = let* () = lock t in finalize f ~finally:(fun () -> unlock t) diff --git a/src/fiber/src/mvar.ml b/src/fiber/src/mvar.ml index 0cc3cfc4f5a..e264d495fd3 100644 --- a/src/fiber/src/mvar.ml +++ b/src/fiber/src/mvar.ml @@ -21,9 +21,9 @@ let create_full x = { value = Some x; writers = Queue.create (); readers = Queue.create () } ;; -let read t k = +let run_read t k = match t.value with - | None -> suspend (fun k -> Queue.push t.readers k) k + | None -> Suspend ((fun k -> Queue.push t.readers k), k) | Some v -> (match Queue.pop t.writers with | None -> @@ -31,16 +31,20 @@ let read t k = continue k v | Some (v', w) -> t.value <- Some v'; - resume w () (Map ((fun () -> v), k))) + Resume (w, (), Map ((fun () -> v), k))) ;; -let write t x k = +let read t = primitive run_read t + +let run_write t x k = match t.value with - | Some _ -> suspend (fun k -> Queue.push t.writers (x, k)) k + | Some _ -> Suspend ((fun k -> Queue.push t.writers (x, k)), k) | None -> (match Queue.pop t.readers with | None -> t.value <- Some x; continue k () - | Some r -> resume r x (Map ((fun () -> ()), k))) + | Some r -> Resume (r, x, Map ((fun () -> ()), k))) ;; + +let write t x = primitive2 run_write t x diff --git a/src/fiber/src/pool.ml b/src/fiber/src/pool.ml index 359db08e02a..a838de928f1 100644 --- a/src/fiber/src/pool.ml +++ b/src/fiber/src/pool.ml @@ -25,15 +25,16 @@ type nonrec t = ; mutable status : status } -let running t k = +let run_running t k = match t.status with | Open -> continue k true | Closed -> continue k false ;; +let running t = primitive run_running t let create () = { tasks = Queue.create (); runner = Awaiting_run; status = Open } -let task t ~f k = +let run_task t f k = match t.status with | Closed -> Code_error.raise "pool is closed. new tasks may not be submitted" [] | Open -> @@ -42,10 +43,12 @@ let task t ~f k = | Running | Awaiting_run -> continue k () | Awaiting_resume r -> t.runner <- Running; - resume r () k) + Resume (r, (), k)) ;; -let close t k = +let task t ~f = primitive2 run_task t f + +let run_close t k = match t.status with | Closed -> continue k () | Open -> @@ -54,10 +57,12 @@ let close t k = | Running | Awaiting_run -> continue k () | Awaiting_resume r -> t.runner <- Running; - resume r () k) + Resume (r, (), k)) ;; -let run t k = +let close t = primitive run_close t + +let run_pool t k = match t.runner with | Awaiting_resume _ | Running -> Code_error.raise "Fiber.Pool.run: concurent calls to run aren't allowed" [] @@ -75,7 +80,7 @@ let run t k = | None -> finish_or_suspend t | Some v -> incr n; - fork (fun () -> v () (Function done_fiber)) read_delayed + fork (fun () -> apply_t v () (Function done_fiber)) read_delayed and read_delayed () = read t and suspend_k k = (* we are suspending because we have no tasks *) @@ -84,11 +89,13 @@ let run t k = and finish_or_suspend t = match t.status with | Closed -> done_fiber () - | Open -> suspend suspend_k (Function read_delayed) + | Open -> Suspend (suspend_k, Function read_delayed) in read t ;; +let run t = primitive run_pool t + let with_ f = of_thunk (fun () -> let pool = create () in diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index e776c400bf4..31619ac5dd2 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -60,6 +60,14 @@ type 'a step = | Done of 'a | Stalled of 'a stalled +let update_var ctx key f = + (* CR-someday rgrinberg: If [vars = ctx.vars], we could elide the re-allocation of + [ctx] here. This doesn't seem important for us at the moment though because all + existing call sites do change the value of the variable. *) + let vars = Var_map.update ctx.vars ~f key in + { ctx with parent = ctx; vars } +;; + let rec loop : Jobs.t -> step' = function | Empty -> Stalled | Job (ctx, run, x, jobs) -> exec ctx run x jobs @@ -71,12 +79,94 @@ and loop2 a b = | Job (ctx, run, x, a) -> exec ctx run x (Jobs.concat a b) | Concat (a1, a2) -> loop2 a1 (Jobs.concat a2 b) -and exec : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = +and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = fun ctx k x jobs -> - match continue k x with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + match k with + | Function f -> exec_function ctx f x jobs + | Map (f, k) -> exec_map ctx f x k jobs + | Map2 (f, g, k) -> exec_map2 ctx f g x k jobs + | Map3 (f, g, h, k) -> exec_map3 ctx f g h x k jobs + | Bind (f, k) -> exec_fiber_apply ctx f x k jobs + | Apply (f, y, k) -> exec_fiber_apply2 ctx f x y k jobs + | Apply_map (f, y, k) -> exec_apply_map ctx f x y k jobs + | Unwind_to k -> exec ctx.parent k x jobs + | Unwind_map_reduce_to k -> unwind_map_reduce ctx k (Ok x) jobs + +and exec_function : 'a. context -> ('a -> eff) -> 'a -> Jobs.t -> step' = + fun ctx f x jobs -> + match f x with + | exception exn -> handle_exception ctx exn jobs + | eff -> exec_effect ctx eff jobs + +and exec_map : 'a 'b. context -> ('a -> 'b) -> 'a -> 'b continuation -> Jobs.t -> step' = + fun ctx f x k jobs -> + match f x with + | exception exn -> handle_exception ctx exn jobs + | y -> exec ctx k y jobs + +and exec_map2 + : 'a 'b 'c. + context + -> ('a -> 'b) + -> ('b -> 'c) + -> 'a + -> 'c continuation + -> Jobs.t + -> step' + = + fun ctx f g x k jobs -> + match g (f x) with + | exception exn -> handle_exception ctx exn jobs + | y -> exec ctx k y jobs + +and exec_map3 + : 'a 'b 'c 'd. + context + -> ('a -> 'b) + -> ('b -> 'c) + -> ('c -> 'd) + -> 'a + -> 'd continuation + -> Jobs.t + -> step' + = + fun ctx f g h x k jobs -> + match h (g (f x)) with + | exception exn -> handle_exception ctx exn jobs + | y -> exec ctx k y jobs + +and exec_apply_map + : 'a 'b 'c. + context + -> ('a -> 'b -> 'c) + -> 'a + -> 'b + -> 'c continuation + -> Jobs.t + -> step' + = + fun ctx f x y k jobs -> + match f x y with + | exception exn -> handle_exception ctx exn jobs + | z -> exec ctx k z jobs + +and handle_exception ctx exn jobs = + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + +and unwind_map_reduce + : 'a 'b. context -> ('a, 'b) result continuation -> ('a, 'b) result -> Jobs.t -> step' + = + fun ctx k x jobs -> + let (Map_reduce_context r) = ctx.map_reduce_context in + let ref_count = r.ref_count - 1 in + r.ref_count <- ref_count; + assert (ref_count = 0); + exec ctx.parent k x jobs + +and exec_effect ctx eff jobs = + match eff with + | Run (t, k) -> exec_fiber ctx t k jobs | Done v -> Done v | Toplevel_exception exn -> Exn_with_backtrace.reraise exn | Unwind (k, x) -> exec ctx.parent k x jobs @@ -96,26 +186,19 @@ and exec : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = | Resume (suspended, x, k) -> exec ctx k () (Jobs.concat jobs (Job (suspended.ctx, suspended.run, x, Empty))) | Get_var (key, k) -> exec ctx k (Var_map.get ctx.vars key) jobs - | Set_var (key, x, k) -> + | Set_var (key, x, f, k) -> let ctx = { ctx with parent = ctx; vars = Var_map.set ctx.vars key x } in - exec ctx k () jobs - | Update_var (key, f, k) -> - let ctx = - (* CR-someday rgrinberg: If [vars = ctx.vars], we could elide the re-allocation of - [ctx] here. This doesn't seem important for us at the moment though because all - existing call sites do change the value of the variable. *) - let vars = Var_map.update ctx.vars ~f key in - { ctx with parent = ctx; vars } - in - exec ctx k () jobs - | With_error_handler (on_error, k) -> - let on_error = - { ctx - ; run = Function (fun exn -> on_error exn (Function Nothing.unreachable_code)) - } - in - let ctx = { ctx with parent = ctx; on_error } in - exec ctx k () jobs + exec_fiber_thunk ctx f (Unwind_to k) jobs + | Update_var (key, f, body, k) -> + let ctx = update_var ctx key f in + exec_fiber_thunk ctx body (Unwind_to k) jobs + | Set_var_apply (key, x, f, y, k) -> + let ctx = { ctx with parent = ctx; vars = Var_map.set ctx.vars key x } in + exec_fiber_apply ctx f y (Unwind_to k) jobs + | Update_var_apply (key, f, body, x, k) -> + let ctx = update_var ctx key f in + exec_fiber_apply ctx body x (Unwind_to k) jobs + | With_error_handler (on_error, f, k) -> with_error_handler ctx on_error f k jobs | Map_reduce_errors (m, on_error, f, k) -> map_reduce_errors ctx m on_error f k jobs | End_of_fiber () -> let (Map_reduce_context r) = ctx.map_reduce_context in @@ -131,22 +214,195 @@ and exec : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = | Fork (a, b) -> let (Map_reduce_context r) = ctx.map_reduce_context in r.ref_count <- r.ref_count + 1; - exec ctx Effect a (Job (ctx, Function b, (), jobs)) + exec_effect ctx a (Job (ctx, Function b, (), jobs)) | Reraise exn -> let { ctx; run } = ctx.on_error in exec ctx run exn jobs - | Reraise_all exns -> - (match length_and_rev exns with - | 0, _ -> loop jobs - | n, exns -> - let (Map_reduce_context r) = ctx.map_reduce_context in - r.ref_count <- r.ref_count + (n - 1); - let { ctx; run } = ctx.on_error in - let jobs = - List.fold_left exns ~init:jobs ~f:(fun jobs exn -> - Jobs.Job (ctx, run, exn, jobs)) - in - loop jobs) + | Reraise_all exns -> reraise_all ctx exns jobs + +and with_error_handler + : 'a. + context + -> (Exn_with_backtrace.t -> Nothing.t t) + -> (unit -> 'a t) + -> 'a continuation + -> Jobs.t + -> step' + = + fun ctx on_error f k jobs -> + let on_error = + { ctx + ; run = Function (fun exn -> Run (on_error exn, Function Nothing.unreachable_code)) + } + in + let ctx = { ctx with parent = ctx; on_error } in + exec_fiber_thunk ctx f (Unwind_to k) jobs + +and reraise_all ctx exns jobs = + match length_and_rev exns with + | 0, _ -> loop jobs + | n, exns -> + let (Map_reduce_context r) = ctx.map_reduce_context in + r.ref_count <- r.ref_count + (n - 1); + let { ctx; run } = ctx.on_error in + let jobs = + List.fold_left exns ~init:jobs ~f:(fun jobs exn -> Jobs.Job (ctx, run, exn, jobs)) + in + loop jobs + +and exec_fiber : type a. context -> a t -> a continuation -> Jobs.t -> step' = + fun ctx t k jobs -> + match t with + | Return_t x -> exec ctx k x jobs + | Never_t -> loop jobs + | Map_t (t, f) -> exec_fiber ctx t (Map (f, k)) jobs + | Map2_t (t, f, g) -> exec_fiber ctx t (Map2 (f, g, k)) jobs + | Map3_t (t, f, g, h) -> exec_fiber ctx t (Map3 (f, g, h, k)) jobs + | Bind_t (t, f) -> exec_fiber ctx t (Bind (f, k)) jobs + | Thunk_t f -> exec_fiber_thunk ctx f k jobs + | Thunk_apply_t (f, x) -> exec_fiber_apply ctx f x k jobs + | With_error_handler_t (f, on_error) -> with_error_handler ctx on_error f k jobs + | Map_reduce_errors_t (m, on_error, f) -> map_reduce_errors ctx m on_error f k jobs + | Suspend_t f -> + let k = { ctx; run = k } in + f k; + loop jobs + | Resume_t (suspended, x) -> + exec ctx k () (Jobs.concat jobs (Job (suspended.ctx, suspended.run, x, Empty))) + | Reraise_all_t exns -> reraise_all ctx exns jobs + | Ivar_read_t ivar -> + (match ivar.state with + | (Empty | Empty_with_readers _) as readers -> + ivar.state <- Empty_with_readers (ctx, k, readers); + loop jobs + | Full x -> exec ctx k x jobs) + | Ivar_fill_t (ivar, x) -> + let jobs = Jobs.concat jobs (Jobs.fill_ivar ivar x Empty) in + exec ctx k () jobs + | Get_var_t key -> exec ctx k (Var_map.get ctx.vars key) jobs + | Set_var_t (key, x, f) -> + let ctx = { ctx with parent = ctx; vars = Var_map.set ctx.vars key x } in + exec_fiber_thunk ctx f (Unwind_to k) jobs + | Update_var_t (key, f, body) -> + let ctx = update_var ctx key f in + exec_fiber_thunk ctx body (Unwind_to k) jobs + | Get_apply_t (key, f, x) -> exec ctx (Apply (f, x, k)) (Var_map.get ctx.vars key) jobs + | Get_apply_map_t (key, f, x) -> + exec ctx (Apply_map (f, x, k)) (Var_map.get ctx.vars key) jobs + | Set_apply_t (key, value, f, x) -> + let ctx = { ctx with parent = ctx; vars = Var_map.set ctx.vars key value } in + exec_fiber_apply ctx f x (Unwind_to k) jobs + | Update_apply_t (key, f, body, x) -> + let ctx = update_var ctx key f in + exec_fiber_apply ctx body x (Unwind_to k) jobs + | Primitive_t (f, x) -> exec_primitive1 ctx f x k jobs + | Primitive2_t (f, x, y) -> exec_primitive2 ctx f x y k jobs + | Primitive3_t (f, x, y, z) -> exec_primitive3 ctx f x y z k jobs + | Primitive4_t (f, w, x, y, z) -> exec_primitive4 ctx f w x y z k jobs + +and exec_primitive1 + : 'a 'b. + context + -> ('a -> 'b continuation -> eff) + -> 'a + -> 'b continuation + -> Jobs.t + -> step' + = + fun ctx f x k jobs -> + match f x k with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | eff -> exec_effect ctx eff jobs + +and exec_primitive2 + : 'a 'b 'c. + context + -> ('a -> 'b -> 'c continuation -> eff) + -> 'a + -> 'b + -> 'c continuation + -> Jobs.t + -> step' + = + fun ctx f x y k jobs -> + match f x y k with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | eff -> exec_effect ctx eff jobs + +and exec_primitive3 + : 'a 'b 'c 'd. + context + -> ('a -> 'b -> 'c -> 'd continuation -> eff) + -> 'a + -> 'b + -> 'c + -> 'd continuation + -> Jobs.t + -> step' + = + fun ctx f x y z k jobs -> + match f x y z k with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | eff -> exec_effect ctx eff jobs + +and exec_primitive4 + : 'a 'b 'c 'd 'e. + context + -> ('a -> 'b -> 'c -> 'd -> 'e continuation -> eff) + -> 'a + -> 'b + -> 'c + -> 'd + -> 'e continuation + -> Jobs.t + -> step' + = + fun ctx f w x y z k jobs -> + match f w x y z k with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | eff -> exec_effect ctx eff jobs + +and exec_fiber_thunk : 'a. context -> (unit -> 'a t) -> 'a continuation -> Jobs.t -> step' + = + fun ctx f k jobs -> + match f () with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | t -> exec_fiber ctx t k jobs + +and exec_fiber_apply + : 'a 'b. context -> ('a -> 'b t) -> 'a -> 'b continuation -> Jobs.t -> step' + = + fun ctx f x k jobs -> + match f x with + | exception exn -> + let exn = Exn_with_backtrace.capture exn in + exec ctx.on_error.ctx ctx.on_error.run exn jobs + | t -> exec_fiber ctx t k jobs + +and exec_fiber_apply2 + : 'a 'b 'c. + context + -> ('a -> 'b -> 'c t) + -> 'a + -> 'b + -> 'c continuation + -> Jobs.t + -> step' + = + fun ctx f x y k jobs -> + match f x y with + | exception exn -> handle_exception ctx exn jobs + | t -> exec_fiber ctx t k jobs and deref : 'a 'b. ('a, 'b) map_reduce_context' -> Jobs.t -> step' = fun r jobs -> @@ -163,7 +419,7 @@ and map_reduce_errors context -> (module Monoid with type t = errors) -> (Exn_with_backtrace.t -> errors t) - -> (unit -> eff) + -> (unit -> b t) -> (b, errors) result continuation -> Jobs.t -> step' @@ -175,12 +431,12 @@ and map_reduce_errors ; run = Function (fun exn -> - on_error - exn - (Function - (fun m -> - map_reduce_context.errors <- M.combine map_reduce_context.errors m; - End_of_map_reduce_error_handler map_reduce_context))) + Run + ( on_error exn + , Function + (fun m -> + map_reduce_context.errors <- M.combine map_reduce_context.errors m; + End_of_map_reduce_error_handler map_reduce_context) )) } in let ctx = @@ -190,7 +446,7 @@ and map_reduce_errors ; map_reduce_context = Map_reduce_context map_reduce_context } in - exec ctx (Function f) () jobs + exec_fiber_thunk ctx f (Unwind_map_reduce_to k) jobs ;; let repack_step (type a) (module W : Witness with type t = a) (step' : step') = @@ -226,5 +482,5 @@ let start (type a) (t : a t) = } } in - exec ctx (Function t) (Function (fun x -> Done (W.X x))) Empty |> repack_step (module W) + exec_fiber ctx t (Function (fun x -> Done (W.X x))) Empty |> repack_step (module W) ;; diff --git a/src/fiber/src/stream.ml b/src/fiber/src/stream.ml index d7b1f5097fc..a68973161d3 100644 --- a/src/fiber/src/stream.ml +++ b/src/fiber/src/stream.ml @@ -99,7 +99,7 @@ module In = struct loop t ~f ;; - let parallel_iter t ~f k = + let run_parallel_iter t f k = let n = ref 1 in let done_ () = decr n; @@ -110,18 +110,21 @@ module In = struct else end_of_fiber in let rec loop t = - t.read + apply_t + t.read () (Function (function | None -> done_ () | Some x -> incr n; - fork (fun () -> f x (Function done_)) (fun () -> loop t))) + fork (fun () -> apply_t f x (Function done_)) (fun () -> loop t))) in lock t; loop t ;; + + let parallel_iter t ~f = primitive2 run_parallel_iter t f end module Out = struct diff --git a/src/fiber/src/svar.ml b/src/fiber/src/svar.ml index 17e7033e7dc..8eb0e233206 100644 --- a/src/fiber/src/svar.ml +++ b/src/fiber/src/svar.ml @@ -23,20 +23,22 @@ let wait = let create current = { current; waiters = [] } -let write = - let rec run_awakers final = function - | [] -> continue final () - | k :: ks -> resume k () (Function (fun () -> run_awakers final ks)) +let rec run_awakers final = function + | [] -> continue final () + | k :: ks -> Resume (k, (), Function (fun () -> run_awakers final ks)) +;; + +let run_write t a k = + t.current <- a; + let sleep, awake = + List.rev_partition_map t.waiters ~f:(fun (k, f) -> + if f t.current then Right k else Left (k, f)) in - fun t a k -> - t.current <- a; - let sleep, awake = - List.rev_partition_map t.waiters ~f:(fun (k, f) -> - if f t.current then Right k else Left (k, f)) - in - match awake with - | [] -> continue k () - | awake -> - t.waiters <- List.rev sleep; - run_awakers k awake + match awake with + | [] -> continue k () + | awake -> + t.waiters <- List.rev sleep; + run_awakers k awake ;; + +let write t a = primitive2 run_write t a From 196b1230faf195d7eee3a21154b78efcde486a25 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 12:10:35 +0100 Subject: [PATCH 06/11] perf(fiber): represent deferred scheduler work explicitly Store deferred fork loops and fiber applications as typed work items instead of allocating thunk closures and wrapping them in scheduler continuations. Keep a function-backed work item for module-local loops that still require an opaque escape hatch. On a no-op @install self-build this reduced minor allocation by 1.0M words and slightly reduced promoted allocation and Cachegrind instructions. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 67 +++++++++++++++++++++++--------------- src/fiber/src/lazy.ml | 4 ++- src/fiber/src/pool.ml | 4 ++- src/fiber/src/scheduler.ml | 10 +++++- src/fiber/src/stream.ml | 4 ++- 5 files changed, 58 insertions(+), 31 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 93af6bfb14c..93d425986f9 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -77,12 +77,22 @@ and eff = | Never of unit (* Add a dummy unit argument to [End_of_fiber] and [Never] so that all constructors are boxed, which removes a branch in the pattern match. *) - | Fork : eff * (unit -> eff) -> eff + | Fork : eff * work -> eff | Reraise : Exn_with_backtrace.t -> eff | Reraise_all : Exn_with_backtrace.t list -> eff | Toplevel_exception : Exn_with_backtrace.t -> eff | Done of value +and work = + | Function_work of (unit -> eff) + | Nfork_work : 'a * 'a list * ('a -> eff) -> work + | Nforki_work : int * 'a * 'a list * (int -> 'a -> eff) -> work + | Nfork_seq_work : int ref * 'a * 'a Seq.t * ('a -> eff) -> work + | Nfork_array_work : 'a array * int * ('a -> eff) -> work + | Apply_thunk_work : (unit -> 'a t) * 'a continuation -> work + | Eval_work : 'a t * 'a continuation -> work + | Continue_work : unit continuation -> work + and 'a ivar = { mutable state : ('a, [ `Full | `Empty ]) ivar_state } and ('a, _) ivar_state = @@ -219,12 +229,6 @@ let apply_t2 f x y k = | exn -> Reraise (Exn_with_backtrace.capture exn) ;; -let[@inline always] fork a b = - match apply a () with - | End_of_fiber () -> b () - | eff -> Fork (eff, b) -;; - let rec nfork x l f = match l with | [] -> f x @@ -233,19 +237,19 @@ let rec nfork x l f = not getting rid of the closures. *) (match apply f x with | End_of_fiber () -> nfork y l f - | eff -> Fork (eff, fun () -> nfork y l f)) + | eff -> Fork (eff, Nfork_work (y, l, f))) ;; -let rec nforki i x l f = +let rec nforki_from i x l f = match l with | [] -> f i x | y :: l -> (match apply2 f i x with - | End_of_fiber () -> nforki (i + 1) y l f - | eff -> Fork (eff, fun () -> nforki (i + 1) y l f)) + | End_of_fiber () -> nforki_from (i + 1) y l f + | eff -> Fork (eff, Nforki_work (i + 1, y, l, f))) ;; -let nforki x l f = nforki 0 x l f +let nforki x l f = nforki_from 0 x l f let rec nfork_seq left_over x (seq : _ Seq.t) f = match seq () with @@ -254,7 +258,7 @@ let rec nfork_seq left_over x (seq : _ Seq.t) f = incr left_over; (match apply f x with | End_of_fiber () -> nfork_seq left_over y seq f - | eff -> Fork (eff, fun () -> nfork_seq left_over y seq f)) + | eff -> Fork (eff, Nfork_seq_work (left_over, y, seq, f))) ;; let rec nfork_array a i f = @@ -263,7 +267,18 @@ let rec nfork_array a i f = else ( match apply f a.(i) with | End_of_fiber () -> nfork_array a (i + 1) f - | eff -> Fork (eff, fun () -> nfork_array a (i + 1) f)) + | eff -> Fork (eff, Nfork_array_work (a, i + 1, f))) +;; + +let run_work = function + | Function_work f -> f () + | Nfork_work (x, l, f) -> nfork x l f + | Nforki_work (i, x, l, f) -> nforki_from i x l f + | Nfork_seq_work (left_over, x, seq, f) -> nfork_seq left_over x seq f + | Nfork_array_work (a, i, f) -> nfork_array a i f + | Apply_thunk_work (f, k) -> apply_t f () k + | Eval_work (t, k) -> eval t k + | Continue_work k -> continue k () ;; let run_parallel_iter_seq (seq : _ Seq.t) f k = @@ -378,7 +393,7 @@ let run_fork_and_join fa fb k = let kb = Function kb in match apply_t fa () ka with | End_of_fiber () -> apply_t fb () kb - | eff -> Fork (eff, fun () -> apply_t fb () kb) + | eff -> Fork (eff, Apply_thunk_work (fb, kb)) ;; let fork_and_join fa fb = primitive2 run_fork_and_join fa fb @@ -402,18 +417,16 @@ let run_fork_and_join_unit fa fb k = | eff -> Fork ( eff - , fun () -> - apply_t - fb - () - (Function - (fun b -> - match !state with - | Nothing_yet -> - state := Got_b b; - end_of_fiber - | Got_a () -> continue k b - | Got_b _ -> assert false)) ) + , Apply_thunk_work + ( fb + , Function + (fun b -> + match !state with + | Nothing_yet -> + state := Got_b b; + end_of_fiber + | Got_a () -> continue k b + | Got_b _ -> assert false) ) ) ;; let fork_and_join_unit fa fb = primitive2 run_fork_and_join_unit fa fb diff --git a/src/fiber/src/lazy.ml b/src/fiber/src/lazy.ml index b34c704a720..7144dec1ac9 100644 --- a/src/fiber/src/lazy.ml +++ b/src/fiber/src/lazy.ml @@ -63,7 +63,9 @@ let is_value t = let stop = Function (fun () -> end_of_fiber) let run_force t v f k = - fork (fun () -> eval (execute t v f) stop) (fun () -> continue k ()) + match eval (execute t v f) stop with + | End_of_fiber () -> continue k () + | eff -> Fork (eff, Continue_work k) ;; let force_all_unit = diff --git a/src/fiber/src/pool.ml b/src/fiber/src/pool.ml index a838de928f1..98e6a93a052 100644 --- a/src/fiber/src/pool.ml +++ b/src/fiber/src/pool.ml @@ -80,7 +80,9 @@ let run_pool t k = | None -> finish_or_suspend t | Some v -> incr n; - fork (fun () -> apply_t v () (Function done_fiber)) read_delayed + (match apply_t v () (Function done_fiber) with + | End_of_fiber () -> read_delayed () + | eff -> Fork (eff, Function_work read_delayed)) and read_delayed () = read t and suspend_k k = (* we are suspending because we have no tasks *) diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 31619ac5dd2..82476d03c77 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -7,6 +7,7 @@ module Jobs = struct type t = | Empty | Job : context * 'a continuation * 'a * t -> t + | Work of context * work * t | Concat : t * t -> t let concat a b = @@ -71,12 +72,14 @@ let update_var ctx key f = let rec loop : Jobs.t -> step' = function | Empty -> Stalled | Job (ctx, run, x, jobs) -> exec ctx run x jobs + | Work (ctx, work, jobs) -> exec_work ctx work jobs | Concat (a, b) -> loop2 a b and loop2 a b = match a with | Empty -> loop b | Job (ctx, run, x, a) -> exec ctx run x (Jobs.concat a b) + | Work (ctx, work, a) -> exec_work ctx work (Jobs.concat a b) | Concat (a1, a2) -> loop2 a1 (Jobs.concat a2 b) and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = @@ -214,12 +217,17 @@ and exec_effect ctx eff jobs = | Fork (a, b) -> let (Map_reduce_context r) = ctx.map_reduce_context in r.ref_count <- r.ref_count + 1; - exec_effect ctx a (Job (ctx, Function b, (), jobs)) + exec_effect ctx a (Work (ctx, b, jobs)) | Reraise exn -> let { ctx; run } = ctx.on_error in exec ctx run exn jobs | Reraise_all exns -> reraise_all ctx exns jobs +and exec_work ctx work jobs = + match run_work work with + | exception exn -> handle_exception ctx exn jobs + | eff -> exec_effect ctx eff jobs + and with_error_handler : 'a. context diff --git a/src/fiber/src/stream.ml b/src/fiber/src/stream.ml index a68973161d3..2fd36402593 100644 --- a/src/fiber/src/stream.ml +++ b/src/fiber/src/stream.ml @@ -118,7 +118,9 @@ module In = struct | None -> done_ () | Some x -> incr n; - fork (fun () -> apply_t f x (Function done_)) (fun () -> loop t))) + (match apply_t f x (Function done_) with + | End_of_fiber () -> loop t + | eff -> Fork (eff, Function_work (fun () -> loop t))))) in lock t; loop t From 9bfaa981c07ef0252e4bd24bb37d8e893d34c5e2 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 12:19:28 +0100 Subject: [PATCH 07/11] perf(fiber): specialize parallel completion frames Replace per-child completion closures in parallel iteration, map-reduce, array mapping, fork-and-join, and Svar wakeups with typed continuation frames. Share one completion frame where all children perform the same update. On a no-op @install self-build this reduced minor allocation by 1.0M words, promoted allocation by 0.2M words, and Cachegrind instructions by 0.1%. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 194 +++++++++++++++---------------------- src/fiber/src/lazy.ml | 2 +- src/fiber/src/scheduler.ml | 13 +++ src/fiber/src/svar.ml | 7 +- 4 files changed, 93 insertions(+), 123 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 93d425986f9..2ea5dde7777 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -1,5 +1,10 @@ open Stdune +type ('a, 'b) fork_and_join_state = + | Nothing_yet + | Got_a of 'a + | Got_b of 'b + (* Fiber computations are explicit nodes. Callbacks supplied by callers remain functions, while the Fiber-created computation spine is interpreted by the scheduler. *) type _ t = @@ -47,6 +52,21 @@ and 'a continuation = | Apply_map : ('a -> 'b -> 'c) * 'b * 'c continuation -> 'a continuation | Unwind_to : 'a continuation -> 'a continuation | Unwind_map_reduce_to : ('a, 'b) result continuation -> 'a continuation + | End : unit continuation + | Parallel_unit_complete : int ref * unit continuation -> unit continuation + | Map_reduce_complete : + 'a ref * int ref * ('a -> 'a -> 'a) * 'a continuation + -> 'a continuation + | Array_map_complete : + 'a array ref * int * int * int ref * 'a array continuation + -> 'a continuation + | Fork_join_left : + ('a, 'b) fork_and_join_state ref * ('a * 'b) continuation + -> 'a continuation + | Fork_join_right : + ('a, 'b) fork_and_join_state ref * ('a * 'b) continuation + -> 'b continuation + | Resume_many : unit k list * unit continuation -> unit continuation and eff = | Run : 'a t * 'a continuation -> eff @@ -138,6 +158,45 @@ let rec continue : type a. a continuation -> a -> eff = | Apply_map (f, y, k) -> continue k (f x y) | Unwind_to k -> Unwind (k, x) | Unwind_map_reduce_to k -> Unwind_map_reduce (k, Ok x) + | End -> End_of_fiber () + | Parallel_unit_complete (running, k) -> + decr running; + if !running = 0 then continue k () else End_of_fiber () + | Map_reduce_complete (current, running, combine, k) -> + current := combine !current x; + decr running; + if !running = 0 then continue k !current else End_of_fiber () + | Array_map_complete (results, len, i, running, k) -> + let a = + match !results with + | [||] -> + let a = Array.make len x in + results := a; + a + | a -> + a.(i) <- x; + a + in + decr running; + if !running = 0 then continue k a else End_of_fiber () + | Fork_join_left (state, k) -> + (match !state with + | Nothing_yet -> + state := Got_a x; + End_of_fiber () + | Got_a _ -> assert false + | Got_b b -> continue k (x, b)) + | Fork_join_right (state, k) -> + (match !state with + | Nothing_yet -> + state := Got_b x; + End_of_fiber () + | Got_a a -> continue k (a, x) + | Got_b _ -> assert false) + | Resume_many (suspended, k) -> + (match suspended with + | [] -> continue k () + | suspended :: rest -> Resume (suspended, (), Resume_many (rest, k))) ;; let return x = Return_t x @@ -286,15 +345,8 @@ let run_parallel_iter_seq (seq : _ Seq.t) f k = | Nil -> continue k () | Cons (x, seq) -> let left_over = ref 1 in - let f' x = - apply_t - f - x - (Function - (fun () -> - decr left_over; - if !left_over = 0 then continue k () else end_of_fiber)) - in + let complete = Parallel_unit_complete (left_over, k) in + let f' x = apply_t f x complete in nfork_seq left_over x seq f' ;; @@ -306,16 +358,8 @@ let run_map_reduce_seq (seq : _ Seq.t) f empty combine k = | Cons (x, seq) -> let current = ref empty in let running = ref 1 in - let f' x = - apply_t - f - x - (Function - (fun y -> - current := combine !current y; - decr running; - if !running = 0 then continue k !current else end_of_fiber)) - in + let complete = Map_reduce_complete (current, running, combine, k) in + let f' x = apply_t f x complete in nfork_seq running x seq f' ;; @@ -329,16 +373,8 @@ let run_map_reduce_array a f empty combine k = | len -> let current = ref empty in let running = ref len in - let f' x = - apply_t - f - x - (Function - (fun y -> - current := combine !current y; - decr running; - if !running = 0 then continue k !current else end_of_fiber)) - in + let complete = Map_reduce_complete (current, running, combine, k) in + let f' x = apply_t f x complete in nfork_array a 0 f' ;; @@ -352,45 +388,17 @@ let run_map_reduce l f empty combine k = | x :: l -> let current = ref empty in let running = ref (List.length l + 1) in - let f' x = - apply_t - f - x - (Function - (fun y -> - current := combine !current y; - decr running; - if !running = 0 then continue k !current else end_of_fiber)) - in + let complete = Map_reduce_complete (current, running, combine, k) in + let f' x = apply_t f x complete in nfork x l f' ;; let map_reduce l ~f ~empty ~combine = primitive4 run_map_reduce l f empty combine -type ('a, 'b) fork_and_join_state = - | Nothing_yet - | Got_a of 'a - | Got_b of 'b - let run_fork_and_join fa fb k = let state = ref Nothing_yet in - let ka a = - match !state with - | Nothing_yet -> - state := Got_a a; - end_of_fiber - | Got_a _ -> assert false - | Got_b b -> continue k (a, b) - and kb b = - match !state with - | Nothing_yet -> - state := Got_b b; - end_of_fiber - | Got_a a -> continue k (a, b) - | Got_b _ -> assert false - in - let ka = Function ka in - let kb = Function kb in + let ka = Fork_join_left (state, k) in + let kb = Fork_join_right (state, k) in match apply_t fa () ka with | End_of_fiber () -> apply_t fb () kb | eff -> Fork (eff, Apply_thunk_work (fb, kb)) @@ -400,33 +408,12 @@ let fork_and_join fa fb = primitive2 run_fork_and_join fa fb let run_fork_and_join_unit fa fb k = let state = ref Nothing_yet in - match - apply_t - fa - () - (Function - (fun () -> - match !state with - | Nothing_yet -> - state := Got_a (); - end_of_fiber - | Got_a _ -> assert false - | Got_b b -> continue k b)) - with - | End_of_fiber () -> apply_t fb () k - | eff -> - Fork - ( eff - , Apply_thunk_work - ( fb - , Function - (fun b -> - match !state with - | Nothing_yet -> - state := Got_b b; - end_of_fiber - | Got_a () -> continue k b - | Got_b _ -> assert false) ) ) + let pair = Map (snd, k) in + let ka = Fork_join_left (state, pair) in + let kb = Fork_join_right (state, pair) in + match apply_t fa () ka with + | End_of_fiber () -> apply_t fb () kb + | eff -> Fork (eff, Apply_thunk_work (fb, kb)) ;; let fork_and_join_unit fa fb = primitive2 run_fork_and_join_unit fa fb @@ -543,15 +530,8 @@ let run_parallel_iter l f k = | x :: l -> let len = List.length l + 1 in let left_over = ref len in - let f' x = - apply_t - f - x - (Function - (fun () -> - decr left_over; - if !left_over = 0 then continue k () else end_of_fiber)) - in + let complete = Parallel_unit_complete (left_over, k) in + let f' x = apply_t f x complete in nfork x l f' ;; @@ -561,25 +541,7 @@ let run_parallel_array_of_list_map' x l f k = let len = List.length l + 1 in let left_over = ref len in let results = ref [||] in - let f' i x = - apply_t - f - x - (Function - (fun y -> - let a = - match !results with - | [||] -> - let a = Array.make len y in - results := a; - a - | a -> - a.(i) <- y; - a - in - decr left_over; - if !left_over = 0 then continue k a else end_of_fiber)) - in + let f' i x = apply_t f x (Array_map_complete (results, len, i, left_over, k)) in nforki x l f' ;; diff --git a/src/fiber/src/lazy.ml b/src/fiber/src/lazy.ml index 7144dec1ac9..81a7899e9f5 100644 --- a/src/fiber/src/lazy.ml +++ b/src/fiber/src/lazy.ml @@ -60,7 +60,7 @@ let is_value t = | Running _ | Init _ -> false ;; -let stop = Function (fun () -> end_of_fiber) +let stop = End let run_force t v f k = match eval (execute t v f) stop with diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 82476d03c77..1daddc5d1de 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -94,6 +94,19 @@ and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = | Apply_map (f, y, k) -> exec_apply_map ctx f x y k jobs | Unwind_to k -> exec ctx.parent k x jobs | Unwind_map_reduce_to k -> unwind_map_reduce ctx k (Ok x) jobs + | End as k -> exec_core_continuation ctx k x jobs + | Parallel_unit_complete _ as k -> exec_core_continuation ctx k x jobs + | Map_reduce_complete _ as k -> exec_core_continuation ctx k x jobs + | Array_map_complete _ as k -> exec_core_continuation ctx k x jobs + | Fork_join_left _ as k -> exec_core_continuation ctx k x jobs + | Fork_join_right _ as k -> exec_core_continuation ctx k x jobs + | Resume_many _ as k -> exec_core_continuation ctx k x jobs + +and exec_core_continuation : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = + fun ctx k x jobs -> + match continue k x with + | exception exn -> handle_exception ctx exn jobs + | eff -> exec_effect ctx eff jobs and exec_function : 'a. context -> ('a -> eff) -> 'a -> Jobs.t -> step' = fun ctx f x jobs -> diff --git a/src/fiber/src/svar.ml b/src/fiber/src/svar.ml index 8eb0e233206..8c5a40a21ec 100644 --- a/src/fiber/src/svar.ml +++ b/src/fiber/src/svar.ml @@ -23,11 +23,6 @@ let wait = let create current = { current; waiters = [] } -let rec run_awakers final = function - | [] -> continue final () - | k :: ks -> Resume (k, (), Function (fun () -> run_awakers final ks)) -;; - let run_write t a k = t.current <- a; let sleep, awake = @@ -38,7 +33,7 @@ let run_write t a k = | [] -> continue k () | awake -> t.waiters <- List.rev sleep; - run_awakers k awake + continue (Resume_many (awake, k)) () ;; let write t a = primitive2 run_write t a From 8d87a4a8aea0be8dffd4bd8390839532fa270806 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 12:25:57 +0100 Subject: [PATCH 08/11] perf(fiber): specialize scheduler error continuations Represent top-level, scoped, and map-reduce error continuations directly. This removes closure wrappers from error contexts and accumulates map-reduce errors through typed frames while retaining user-provided handlers as opaque callbacks. On a no-op @install self-build this reduced minor allocation by 0.43M words and promoted allocation by 0.25M words, with a small Cachegrind instruction reduction. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 23 +++++++++++++++++++++++ src/fiber/src/scheduler.ml | 38 ++++++++++++++++---------------------- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 2ea5dde7777..7c1fd5ccb75 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -67,6 +67,20 @@ and 'a continuation = ('a, 'b) fork_and_join_state ref * ('a * 'b) continuation -> 'b continuation | Resume_many : unit k list * unit continuation -> unit continuation + | Top_level_error : Exn_with_backtrace.t continuation + | Unreachable : Nothing.t continuation + | Never_called : 'a continuation + | Handle_error : + (Exn_with_backtrace.t -> Nothing.t t) + -> Exn_with_backtrace.t continuation + | Map_reduce_error : + (Exn_with_backtrace.t -> 'errors t) + * ('result, 'errors) map_reduce_context' + * ('errors -> 'errors -> 'errors) + -> Exn_with_backtrace.t continuation + | Accumulate_error : + ('result, 'errors) map_reduce_context' * ('errors -> 'errors -> 'errors) + -> 'errors continuation and eff = | Run : 'a t * 'a continuation -> eff @@ -197,6 +211,15 @@ let rec continue : type a. a continuation -> a -> eff = (match suspended with | [] -> continue k () | suspended :: rest -> Resume (suspended, (), Resume_many (rest, k))) + | Top_level_error -> Toplevel_exception x + | Unreachable -> Nothing.unreachable_code x + | Never_called -> assert false + | Handle_error f -> Run (f x, Unreachable) + | Map_reduce_error (f, map_reduce_context, combine) -> + Run (f x, Accumulate_error (map_reduce_context, combine)) + | Accumulate_error (map_reduce_context, combine) -> + map_reduce_context.errors <- combine map_reduce_context.errors x; + End_of_map_reduce_error_handler map_reduce_context ;; let return x = Return_t x diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 1daddc5d1de..6b6060edcad 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -101,6 +101,18 @@ and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = | Fork_join_left _ as k -> exec_core_continuation ctx k x jobs | Fork_join_right _ as k -> exec_core_continuation ctx k x jobs | Resume_many _ as k -> exec_core_continuation ctx k x jobs + | Top_level_error -> Exn_with_backtrace.reraise x + | Unreachable -> Nothing.unreachable_code x + | Never_called -> assert false + | Handle_error f -> exec_fiber_apply ctx f x Unreachable jobs + | Map_reduce_error (f, map_reduce_context, combine) -> + exec_fiber_apply ctx f x (Accumulate_error (map_reduce_context, combine)) jobs + | Accumulate_error (map_reduce_context, combine) -> + (match combine map_reduce_context.errors x with + | exception exn -> handle_exception ctx exn jobs + | errors -> + map_reduce_context.errors <- errors; + deref map_reduce_context jobs) and exec_core_continuation : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = fun ctx k x jobs -> @@ -251,11 +263,7 @@ and with_error_handler -> step' = fun ctx on_error f k jobs -> - let on_error = - { ctx - ; run = Function (fun exn -> Run (on_error exn, Function Nothing.unreachable_code)) - } - in + let on_error = { ctx; run = Handle_error on_error } in let ctx = { ctx with parent = ctx; on_error } in exec_fiber_thunk ctx f (Unwind_to k) jobs @@ -448,17 +456,7 @@ and map_reduce_errors fun ctx (module M : Monoid with type t = errors) on_error f k jobs -> let map_reduce_context = { k = { ctx; run = k }; ref_count = 1; errors = M.empty } in let on_error = - { ctx - ; run = - Function - (fun exn -> - Run - ( on_error exn - , Function - (fun m -> - map_reduce_context.errors <- M.combine map_reduce_context.errors m; - End_of_map_reduce_error_handler map_reduce_context) )) - } + { ctx; run = Map_reduce_error (on_error, map_reduce_context, M.combine) } in let ctx = { ctx with @@ -493,14 +491,10 @@ let start (type a) (t : a t) = in let rec ctx = { parent = ctx - ; on_error = { ctx; run = Function (fun exn -> Toplevel_exception exn) } + ; on_error = { ctx; run = Top_level_error } ; vars = Var_map.empty ; map_reduce_context = - Map_reduce_context - { k = { ctx; run = Function (fun _ -> assert false) } - ; ref_count = 1 - ; errors = () - } + Map_reduce_context { k = { ctx; run = Never_called }; ref_count = 1; errors = () } } in exec_fiber ctx t (Function (fun x -> Done (W.X x))) Empty |> repack_step (module W) From 8f069422c529f11fb82374977f5105644ac74bc2 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 14:36:46 +0100 Subject: [PATCH 09/11] perf(fiber): compact scheduler error handlers Separate error handlers from the hot continuation representation and store map-reduce return state, handlers, and combiners in one context. Derive the active collection handler from that context and queue captured errors directly instead of reconstructing continuation jobs. Across no-op self-build targets this reduced minor allocation by 0.55M to 0.89M words and promoted allocation by 0.37M to 0.63M words. Cachegrind instructions fell by 0.8% to 1.4%, with neutral to improved native wall time. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 33 ++++++--------- src/fiber/src/scheduler.ml | 85 +++++++++++++++++++++----------------- 2 files changed, 60 insertions(+), 58 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index 7c1fd5ccb75..c06c6fc3512 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -67,20 +67,9 @@ and 'a continuation = ('a, 'b) fork_and_join_state ref * ('a * 'b) continuation -> 'b continuation | Resume_many : unit k list * unit continuation -> unit continuation - | Top_level_error : Exn_with_backtrace.t continuation | Unreachable : Nothing.t continuation | Never_called : 'a continuation - | Handle_error : - (Exn_with_backtrace.t -> Nothing.t t) - -> Exn_with_backtrace.t continuation - | Map_reduce_error : - (Exn_with_backtrace.t -> 'errors t) - * ('result, 'errors) map_reduce_context' - * ('errors -> 'errors -> 'errors) - -> Exn_with_backtrace.t continuation - | Accumulate_error : - ('result, 'errors) map_reduce_context' * ('errors -> 'errors -> 'errors) - -> 'errors continuation + | Accumulate_error : ('result, 'errors) map_reduce_context' -> 'errors continuation and eff = | Run : 'a t * 'a continuation -> eff @@ -140,13 +129,21 @@ and value = .. and context = { parent : context - ; on_error : Exn_with_backtrace.t k + ; on_error : error_handler ; vars : Var_map.t ; map_reduce_context : map_reduce_context } +and error_handler = + | Top_level_error + | Handle_error of context * (Exn_with_backtrace.t -> Nothing.t t) + | Collect_errors + and ('a, 'b) map_reduce_context' = - { k : ('a, 'b) result k + { ctx : context + ; k : ('a, 'b) result continuation + ; on_error : Exn_with_backtrace.t -> 'b t + ; combine : 'b -> 'b -> 'b ; mutable ref_count : int ; mutable errors : 'b } @@ -211,14 +208,10 @@ let rec continue : type a. a continuation -> a -> eff = (match suspended with | [] -> continue k () | suspended :: rest -> Resume (suspended, (), Resume_many (rest, k))) - | Top_level_error -> Toplevel_exception x | Unreachable -> Nothing.unreachable_code x | Never_called -> assert false - | Handle_error f -> Run (f x, Unreachable) - | Map_reduce_error (f, map_reduce_context, combine) -> - Run (f x, Accumulate_error (map_reduce_context, combine)) - | Accumulate_error (map_reduce_context, combine) -> - map_reduce_context.errors <- combine map_reduce_context.errors x; + | Accumulate_error map_reduce_context -> + map_reduce_context.errors <- map_reduce_context.combine map_reduce_context.errors x; End_of_map_reduce_error_handler map_reduce_context ;; diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 6b6060edcad..178a008407a 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -7,6 +7,7 @@ module Jobs = struct type t = | Empty | Job : context * 'a continuation * 'a * t -> t + | Error of context * Exn_with_backtrace.t * t | Work of context * work * t | Concat : t * t -> t @@ -72,6 +73,7 @@ let update_var ctx key f = let rec loop : Jobs.t -> step' = function | Empty -> Stalled | Job (ctx, run, x, jobs) -> exec ctx run x jobs + | Error (ctx, exn, jobs) -> handle_captured_exception ctx exn jobs | Work (ctx, work, jobs) -> exec_work ctx work jobs | Concat (a, b) -> loop2 a b @@ -79,6 +81,7 @@ and loop2 a b = match a with | Empty -> loop b | Job (ctx, run, x, a) -> exec ctx run x (Jobs.concat a b) + | Error (ctx, exn, a) -> handle_captured_exception ctx exn (Jobs.concat a b) | Work (ctx, work, a) -> exec_work ctx work (Jobs.concat a b) | Concat (a1, a2) -> loop2 a1 (Jobs.concat a2 b) @@ -101,14 +104,10 @@ and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = | Fork_join_left _ as k -> exec_core_continuation ctx k x jobs | Fork_join_right _ as k -> exec_core_continuation ctx k x jobs | Resume_many _ as k -> exec_core_continuation ctx k x jobs - | Top_level_error -> Exn_with_backtrace.reraise x | Unreachable -> Nothing.unreachable_code x | Never_called -> assert false - | Handle_error f -> exec_fiber_apply ctx f x Unreachable jobs - | Map_reduce_error (f, map_reduce_context, combine) -> - exec_fiber_apply ctx f x (Accumulate_error (map_reduce_context, combine)) jobs - | Accumulate_error (map_reduce_context, combine) -> - (match combine map_reduce_context.errors x with + | Accumulate_error map_reduce_context -> + (match map_reduce_context.combine map_reduce_context.errors x with | exception exn -> handle_exception ctx exn jobs | errors -> map_reduce_context.errors <- errors; @@ -180,7 +179,20 @@ and exec_apply_map and handle_exception ctx exn jobs = let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + handle_captured_exception ctx exn jobs + +and handle_captured_exception ctx exn jobs = + match ctx.on_error with + | Top_level_error -> Exn_with_backtrace.reraise exn + | Handle_error (handler_ctx, f) -> exec_fiber_apply handler_ctx f exn Unreachable jobs + | Collect_errors -> + let (Map_reduce_context map_reduce_context) = ctx.map_reduce_context in + exec_fiber_apply + map_reduce_context.ctx + map_reduce_context.on_error + exn + (Accumulate_error map_reduce_context) + jobs and unwind_map_reduce : 'a 'b. context -> ('a, 'b) result continuation -> ('a, 'b) result -> Jobs.t -> step' @@ -243,9 +255,7 @@ and exec_effect ctx eff jobs = let (Map_reduce_context r) = ctx.map_reduce_context in r.ref_count <- r.ref_count + 1; exec_effect ctx a (Work (ctx, b, jobs)) - | Reraise exn -> - let { ctx; run } = ctx.on_error in - exec ctx run exn jobs + | Reraise exn -> handle_captured_exception ctx exn jobs | Reraise_all exns -> reraise_all ctx exns jobs and exec_work ctx work jobs = @@ -263,7 +273,7 @@ and with_error_handler -> step' = fun ctx on_error f k jobs -> - let on_error = { ctx; run = Handle_error on_error } in + let on_error = Handle_error (ctx, on_error) in let ctx = { ctx with parent = ctx; on_error } in exec_fiber_thunk ctx f (Unwind_to k) jobs @@ -273,9 +283,8 @@ and reraise_all ctx exns jobs = | n, exns -> let (Map_reduce_context r) = ctx.map_reduce_context in r.ref_count <- r.ref_count + (n - 1); - let { ctx; run } = ctx.on_error in let jobs = - List.fold_left exns ~init:jobs ~f:(fun jobs exn -> Jobs.Job (ctx, run, exn, jobs)) + List.fold_left exns ~init:jobs ~f:(fun jobs exn -> Jobs.Error (ctx, exn, jobs)) in loop jobs @@ -340,9 +349,7 @@ and exec_primitive1 = fun ctx f x k jobs -> match f x k with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | eff -> exec_effect ctx eff jobs and exec_primitive2 @@ -357,9 +364,7 @@ and exec_primitive2 = fun ctx f x y k jobs -> match f x y k with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | eff -> exec_effect ctx eff jobs and exec_primitive3 @@ -375,9 +380,7 @@ and exec_primitive3 = fun ctx f x y z k jobs -> match f x y z k with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | eff -> exec_effect ctx eff jobs and exec_primitive4 @@ -394,18 +397,14 @@ and exec_primitive4 = fun ctx f w x y z k jobs -> match f w x y z k with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | eff -> exec_effect ctx eff jobs and exec_fiber_thunk : 'a. context -> (unit -> 'a t) -> 'a continuation -> Jobs.t -> step' = fun ctx f k jobs -> match f () with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | t -> exec_fiber ctx t k jobs and exec_fiber_apply @@ -413,9 +412,7 @@ and exec_fiber_apply = fun ctx f x k jobs -> match f x with - | exception exn -> - let exn = Exn_with_backtrace.capture exn in - exec ctx.on_error.ctx ctx.on_error.run exn jobs + | exception exn -> handle_exception ctx exn jobs | t -> exec_fiber ctx t k jobs and exec_fiber_apply2 @@ -438,7 +435,7 @@ and deref : 'a 'b. ('a, 'b) map_reduce_context' -> Jobs.t -> step' = let ref_count = r.ref_count - 1 in r.ref_count <- ref_count; match ref_count with - | 0 -> exec r.k.ctx r.k.run (Error r.errors) jobs + | 0 -> exec r.ctx r.k (Error r.errors) jobs | _ -> assert (ref_count > 0); loop jobs @@ -454,14 +451,13 @@ and map_reduce_errors -> step' = fun ctx (module M : Monoid with type t = errors) on_error f k jobs -> - let map_reduce_context = { k = { ctx; run = k }; ref_count = 1; errors = M.empty } in - let on_error = - { ctx; run = Map_reduce_error (on_error, map_reduce_context, M.combine) } + let map_reduce_context = + { ctx; k; on_error; combine = M.combine; ref_count = 1; errors = M.empty } in let ctx = { ctx with parent = ctx - ; on_error + ; on_error = Collect_errors ; map_reduce_context = Map_reduce_context map_reduce_context } in @@ -483,6 +479,12 @@ let advance (type a) (module W : Witness with type t = a) fill : a step = fill |> Jobs.exec_fills |> loop |> repack_step (module W) ;; +let never_handle_error _ = + Code_error.raise "Fiber scheduler dummy error handler called" [] +;; + +let combine_unit () () = () + let start (type a) (t : a t) = let module W = struct type t = a @@ -491,10 +493,17 @@ let start (type a) (t : a t) = in let rec ctx = { parent = ctx - ; on_error = { ctx; run = Top_level_error } + ; on_error = Top_level_error ; vars = Var_map.empty ; map_reduce_context = - Map_reduce_context { k = { ctx; run = Never_called }; ref_count = 1; errors = () } + Map_reduce_context + { ctx + ; k = Never_called + ; on_error = never_handle_error + ; combine = combine_unit + ; ref_count = 1 + ; errors = () + } } in exec_fiber ctx t (Function (fun x -> Done (W.X x))) Empty |> repack_step (module W) From c87782710ccf0600c6bf760053e29c3412fe82b4 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 15:14:13 +0100 Subject: [PATCH 10/11] perf(fiber): share parallel array completion state Store parallel array results, counters, and the final continuation once per map rather than repeating them in every child completion frame. Child frames now contain only the shared state and their array index. Across no-op self-build targets this reduced minor allocation by 0.32M to 0.58M words and promoted allocation by 0.05M to 0.12M words. Cachegrind instructions fell by 0.04% to 0.10%, with neutral native wall time. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index c06c6fc3512..e9270020e96 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -57,9 +57,7 @@ and 'a continuation = | Map_reduce_complete : 'a ref * int ref * ('a -> 'a -> 'a) * 'a continuation -> 'a continuation - | Array_map_complete : - 'a array ref * int * int * int ref * 'a array continuation - -> 'a continuation + | Array_map_complete : 'a array_map_state * int -> 'a continuation | Fork_join_left : ('a, 'b) fork_and_join_state ref * ('a * 'b) continuation -> 'a continuation @@ -71,6 +69,13 @@ and 'a continuation = | Never_called : 'a continuation | Accumulate_error : ('result, 'errors) map_reduce_context' -> 'errors continuation +and 'a array_map_state = + { results : 'a array ref + ; len : int + ; running : int ref + ; k : 'a array continuation + } + and eff = | Run : 'a t * 'a continuation -> eff | Read_ivar : 'a ivar * 'a continuation -> eff @@ -177,7 +182,8 @@ let rec continue : type a. a continuation -> a -> eff = current := combine !current x; decr running; if !running = 0 then continue k !current else End_of_fiber () - | Array_map_complete (results, len, i, running, k) -> + | Array_map_complete (state, i) -> + let { results; len; running; k } = state in let a = match !results with | [||] -> @@ -555,9 +561,8 @@ let parallel_iter l ~f = primitive2 run_parallel_iter l f let run_parallel_array_of_list_map' x l f k = let len = List.length l + 1 in - let left_over = ref len in - let results = ref [||] in - let f' i x = apply_t f x (Array_map_complete (results, len, i, left_over, k)) in + let state = { results = ref [||]; len; running = ref len; k } in + let f' i x = apply_t f x (Array_map_complete (state, i)) in nforki x l f' ;; From cdd448eb200e5aba7973be9e46f849c34c1799e8 Mon Sep 17 00:00:00 2001 From: Rudi Grinberg Date: Fri, 7 Aug 2026 17:22:53 +0100 Subject: [PATCH 11/11] perf(fiber): represent collect_errors directly Use a dedicated computation and completion frame for [collect_errors] instead of constructing a generic map-reduce node followed by a map node. Caller callbacks remain opaque and error collection still uses the existing scope machinery. Across no-op self-build targets this reduced minor allocation by 0.19M to 0.33M words and promoted allocation by 0.02M to 0.04M words. Cachegrind and native wall time remained effectively neutral. Signed-off-by: Rudi Grinberg --- src/fiber/src/core.ml | 27 ++++++++++++++++----------- src/fiber/src/scheduler.ml | 11 +++++++++++ 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/fiber/src/core.ml b/src/fiber/src/core.ml index e9270020e96..f10b32ca5ce 100644 --- a/src/fiber/src/core.ml +++ b/src/fiber/src/core.ml @@ -20,6 +20,7 @@ type _ t = | Map_reduce_errors_t : (module Monoid with type t = 'a) * (Exn_with_backtrace.t -> 'a t) * (unit -> 'b t) -> ('b, 'a) result t + | Collect_errors_t : (unit -> 'a t) -> ('a, Exn_with_backtrace.t list) result t | Suspend_t : ('a k -> unit) -> 'a t | Resume_t : 'a k * 'a -> unit t | Reraise_all_t : Exn_with_backtrace.t list -> 'a t @@ -68,6 +69,9 @@ and 'a continuation = | Unreachable : Nothing.t continuation | Never_called : 'a continuation | Accumulate_error : ('result, 'errors) map_reduce_context' -> 'errors continuation + | Collect_errors_complete : + ('a, Exn_with_backtrace.t list) result continuation + -> ('a, Exn_with_backtrace.t Appendable_list.t) result continuation and 'a array_map_state = { results : 'a array ref @@ -99,6 +103,9 @@ and eff = * (unit -> 'b t) * ('b, 'a) result continuation -> eff + | Run_collect_errors : + (unit -> 'a t) * ('a, Exn_with_backtrace.t list) result continuation + -> eff | Unwind_map_reduce : 'a continuation * 'a -> eff | End_of_map_reduce_error_handler : (_, _) map_reduce_context' -> eff | End_of_fiber of unit @@ -219,6 +226,12 @@ let rec continue : type a. a continuation -> a -> eff = | Accumulate_error map_reduce_context -> map_reduce_context.errors <- map_reduce_context.combine map_reduce_context.errors x; End_of_map_reduce_error_handler map_reduce_context + | Collect_errors_complete k -> + continue + k + (match x with + | Ok x -> Ok x + | Error errors -> Error (Appendable_list.to_list errors)) ;; let return x = Return_t x @@ -265,6 +278,7 @@ let rec eval : type a. a t -> a continuation -> eff = | Thunk_apply_t (f, x) -> eval (f x) k | With_error_handler_t (f, on_error) -> With_error_handler (on_error, f, k) | Map_reduce_errors_t (m, on_error, f) -> Map_reduce_errors (m, on_error, f, k) + | Collect_errors_t f -> Run_collect_errors (f, k) | Suspend_t f -> Suspend (f, k) | Resume_t (suspended, x) -> Resume (suspended, x, k) | Reraise_all_t exns -> @@ -638,17 +652,8 @@ let rec repeat_while : 'a. f:('a -> 'a option t) -> init:'a -> unit t = module Exns = Monoid.Appendable_list (Exn_with_backtrace) -let collect_errors f = - let+ res = - map_reduce_errors - (module Exns) - f - ~on_error:(fun e -> return (Appendable_list.singleton e)) - in - match res with - | Ok x -> Ok x - | Error l -> Error (Appendable_list.to_list l) -;; +let collect_error e = return (Appendable_list.singleton e) +let collect_errors f = Collect_errors_t f let finalize f ~finally = let* res1 = collect_errors f in diff --git a/src/fiber/src/scheduler.ml b/src/fiber/src/scheduler.ml index 178a008407a..cc0e4b9eb64 100644 --- a/src/fiber/src/scheduler.ml +++ b/src/fiber/src/scheduler.ml @@ -112,6 +112,13 @@ and exec : type a. context -> a continuation -> a -> Jobs.t -> step' = | errors -> map_reduce_context.errors <- errors; deref map_reduce_context jobs) + | Collect_errors_complete k -> + let x = + match x with + | Ok x -> Ok x + | Error errors -> Error (Appendable_list.to_list errors) + in + exec ctx k x jobs and exec_core_continuation : 'a. context -> 'a continuation -> 'a -> Jobs.t -> step' = fun ctx k x jobs -> @@ -240,6 +247,8 @@ and exec_effect ctx eff jobs = exec_fiber_apply ctx body x (Unwind_to k) jobs | With_error_handler (on_error, f, k) -> with_error_handler ctx on_error f k jobs | Map_reduce_errors (m, on_error, f, k) -> map_reduce_errors ctx m on_error f k jobs + | Run_collect_errors (f, k) -> + map_reduce_errors ctx (module Exns) collect_error f (Collect_errors_complete k) jobs | End_of_fiber () -> let (Map_reduce_context r) = ctx.map_reduce_context in deref r jobs @@ -301,6 +310,8 @@ and exec_fiber : type a. context -> a t -> a continuation -> Jobs.t -> step' = | Thunk_apply_t (f, x) -> exec_fiber_apply ctx f x k jobs | With_error_handler_t (f, on_error) -> with_error_handler ctx on_error f k jobs | Map_reduce_errors_t (m, on_error, f) -> map_reduce_errors ctx m on_error f k jobs + | Collect_errors_t f -> + map_reduce_errors ctx (module Exns) collect_error f (Collect_errors_complete k) jobs | Suspend_t f -> let k = { ctx; run = k } in f k;