Skip to content

Commit 7692f14

Browse files
committed
And more tweaks to improve code size and performance
1 parent 2c5c800 commit 7692f14

File tree

4 files changed

+306
-268
lines changed

4 files changed

+306
-268
lines changed

lib/picos/picos.ocaml5.ml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ module Computation = struct
219219
let trigger = Trigger.from_action x y action in
220220
Atomic.make (S (Continue { balance_and_mode; triggers = [ trigger ] }))
221221

222-
let is_canceled t =
222+
let[@inline] is_canceled t =
223223
match Atomic.get t with
224224
| S (Canceled { tx; _ }) -> tx == Stopped
225225
| S (Returned _) | S (Continue _) -> false
@@ -460,13 +460,13 @@ module Fiber = struct
460460

461461
let has_forbidden (Fiber r : t) = r.forbid
462462

463-
let is_canceled (Fiber r : t) =
463+
let[@inline] is_canceled (Fiber r : t) =
464464
(not r.forbid)
465465
&&
466466
let (Packed computation) = r.packed in
467467
Computation.is_canceled computation
468468

469-
let canceled (Fiber r : t) =
469+
let[@inline] canceled (Fiber r : t) =
470470
if r.forbid then None
471471
else
472472
let (Packed computation) = r.packed in
Lines changed: 110 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
open Picos
22

3-
let[@inline never] quota_non_positive () = invalid_arg "quota must be positive"
3+
let[@inline never] quota_non_positive _ = invalid_arg "quota must be positive"
44

55
type ready =
66
| Spawn of Fiber.t * (Fiber.t -> unit)
@@ -17,16 +17,16 @@ type t = {
1717
needs_wakeup : bool Atomic.t;
1818
mutex : Mutex.t;
1919
condition : Condition.t;
20-
resume :
20+
mutable resume :
2121
Trigger.t ->
2222
Fiber.t ->
2323
((exn * Printexc.raw_backtrace) option, unit) Effect.Deep.continuation ->
2424
unit;
25-
current : ((Fiber.t, unit) Effect.Deep.continuation -> unit) option;
26-
yield : ((unit, unit) Effect.Deep.continuation -> unit) option;
27-
return : ((unit, unit) Effect.Deep.continuation -> unit) option;
28-
discontinue : ((unit, unit) Effect.Deep.continuation -> unit) option;
29-
handler : (unit, unit) Effect.Deep.handler;
25+
mutable current : ((Fiber.t, unit) Effect.Deep.continuation -> unit) option;
26+
mutable yield : ((unit, unit) Effect.Deep.continuation -> unit) option;
27+
mutable return : ((unit, unit) Effect.Deep.continuation -> unit) option;
28+
mutable discontinue : ((unit, unit) Effect.Deep.continuation -> unit) option;
29+
mutable handler : (unit, unit) Effect.Deep.handler;
3030
quota : int;
3131
mutable fiber : Fiber.Maybe.t;
3232
mutable remaining_quota : int;
@@ -67,48 +67,112 @@ let rec next t =
6767
next t
6868
end
6969

70-
let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
71-
let quota =
72-
match quota with
73-
| None -> Int.max_int
74-
| Some quota ->
75-
if quota <= 0 then quota_non_positive ();
76-
quota
77-
in
70+
let run_fiber ?quota ?fatal_exn_handler fiber main =
7871
Select.check_configured ();
79-
let ready = Mpscq.create ~padded:true ()
80-
and needs_wakeup = Atomic.make false |> Multicore_magic.copy_as_padded
81-
and mutex = Mutex.create ()
82-
and condition = Condition.create () in
83-
let rec t =
72+
let t =
73+
let quota =
74+
match quota with
75+
| None -> Int.max_int
76+
| Some quota -> if quota <= 0 then quota_non_positive quota else quota
77+
in
8478
{
85-
ready;
86-
fiber = Fiber.Maybe.of_fiber fiber;
87-
needs_wakeup;
88-
mutex;
89-
condition;
90-
resume;
91-
current;
92-
yield;
93-
return;
94-
discontinue;
95-
handler;
79+
ready = Mpscq.create ~padded:true ();
80+
needs_wakeup = Atomic.make false |> Multicore_magic.copy_as_padded;
81+
mutex = Mutex.create ();
82+
condition = Condition.create ();
83+
resume = Obj.magic ();
84+
current = Obj.magic ();
85+
yield = Obj.magic ();
86+
return = Obj.magic ();
87+
discontinue = Obj.magic ();
88+
handler = Obj.magic ();
9689
quota;
90+
fiber = Fiber.Maybe.of_fiber fiber;
9791
remaining_quota = quota;
9892
num_alive_fibers = 1;
9993
}
100-
and current =
94+
in
95+
t.handler <-
96+
{
97+
exnc = (match fatal_exn_handler with None -> raise | Some exnc -> exnc);
98+
effc =
99+
(fun (type a) (e : a Effect.t) :
100+
((a, _) Effect.Deep.continuation -> _) option ->
101+
match e with
102+
| Fiber.Current -> t.current
103+
| Fiber.Spawn r ->
104+
let fiber = Fiber.Maybe.to_fiber t.fiber in
105+
if Fiber.is_canceled fiber then t.discontinue
106+
else begin
107+
t.num_alive_fibers <- t.num_alive_fibers + 1;
108+
Mpscq.push t.ready (Spawn (r.fiber, r.main));
109+
t.return
110+
end
111+
| Fiber.Yield -> t.yield
112+
| Computation.Cancel_after r -> begin
113+
let fiber = Fiber.Maybe.to_fiber t.fiber in
114+
if Fiber.is_canceled fiber then t.discontinue
115+
else
116+
match
117+
Select.cancel_after r.computation ~seconds:r.seconds r.exn
118+
r.bt
119+
with
120+
| () -> t.return
121+
| exception exn ->
122+
let bt = Printexc.get_raw_backtrace () in
123+
Some
124+
(fun k -> Effect.Deep.discontinue_with_backtrace k exn bt)
125+
end
126+
| Trigger.Await trigger ->
127+
Some
128+
(fun k ->
129+
let fiber = Fiber.Maybe.to_fiber t.fiber in
130+
if Fiber.try_suspend fiber trigger fiber k t.resume then
131+
next t
132+
else
133+
let remaining_quota = t.remaining_quota - 1 in
134+
if 0 < remaining_quota then begin
135+
t.remaining_quota <- remaining_quota;
136+
Fiber.resume fiber k
137+
end
138+
else begin
139+
Mpscq.push t.ready (Resume (fiber, k));
140+
next t
141+
end)
142+
| _ -> None);
143+
retc =
144+
(fun () ->
145+
t.num_alive_fibers <- t.num_alive_fibers - 1;
146+
next t);
147+
};
148+
t.resume <-
149+
(fun trigger fiber k ->
150+
let resume = Resume (fiber, k) in
151+
if Fiber.unsuspend fiber trigger then Mpscq.push t.ready resume
152+
else Mpscq.push_head t.ready resume;
153+
if
154+
Atomic.get t.needs_wakeup
155+
&& Atomic.compare_and_set t.needs_wakeup true false
156+
then begin
157+
begin
158+
match Mutex.lock t.mutex with
159+
| () -> Mutex.unlock t.mutex
160+
| exception Sys_error _ -> ()
161+
end;
162+
Condition.broadcast t.condition
163+
end);
164+
t.current <-
101165
Some
102166
(fun k ->
103167
let fiber = Fiber.Maybe.to_fiber t.fiber in
104-
Effect.Deep.continue k fiber)
105-
and yield =
168+
Effect.Deep.continue k fiber);
169+
t.yield <-
106170
Some
107171
(fun k ->
108172
let fiber = Fiber.Maybe.to_fiber t.fiber in
109173
Mpscq.push t.ready (Continue (fiber, k));
110-
next t)
111-
and return =
174+
next t);
175+
t.return <-
112176
Some
113177
(fun k ->
114178
let remaining_quota = t.remaining_quota - 1 in
@@ -119,78 +183,21 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
119183
else begin
120184
Mpscq.push t.ready (Return (Fiber.Maybe.to_fiber t.fiber, k));
121185
next t
122-
end)
123-
and discontinue =
186+
end);
187+
t.discontinue <-
124188
Some
125189
(fun k ->
126190
let fiber = Fiber.Maybe.to_fiber t.fiber in
127-
Fiber.continue fiber k ())
128-
and handler = { retc; exnc; effc }
129-
and[@alert "-handler"] effc :
130-
type a. a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option =
131-
function
132-
| Fiber.Current -> t.current
133-
| Fiber.Spawn r ->
134-
let fiber = Fiber.Maybe.to_fiber t.fiber in
135-
if Fiber.is_canceled fiber then t.discontinue
136-
else begin
137-
t.num_alive_fibers <- t.num_alive_fibers + 1;
138-
Mpscq.push t.ready (Spawn (r.fiber, r.main));
139-
t.return
140-
end
141-
| Fiber.Yield -> t.yield
142-
| Computation.Cancel_after r -> begin
143-
let fiber = Fiber.Maybe.to_fiber t.fiber in
144-
if Fiber.is_canceled fiber then t.discontinue
145-
else
146-
match
147-
Select.cancel_after r.computation ~seconds:r.seconds r.exn r.bt
148-
with
149-
| () -> t.return
150-
| exception exn ->
151-
let bt = Printexc.get_raw_backtrace () in
152-
Some (fun k -> Effect.Deep.discontinue_with_backtrace k exn bt)
153-
end
154-
| Trigger.Await trigger ->
155-
Some
156-
(fun k ->
157-
let fiber = Fiber.Maybe.to_fiber t.fiber in
158-
if Fiber.try_suspend fiber trigger fiber k t.resume then next t
159-
else
160-
let remaining_quota = t.remaining_quota - 1 in
161-
if 0 < remaining_quota then begin
162-
t.remaining_quota <- remaining_quota;
163-
Fiber.resume fiber k
164-
end
165-
else begin
166-
Mpscq.push t.ready (Resume (fiber, k));
167-
next t
168-
end)
169-
| _ -> None
170-
and retc () =
171-
t.num_alive_fibers <- t.num_alive_fibers - 1;
172-
next t
173-
and resume trigger fiber k =
174-
let resume = Resume (fiber, k) in
175-
if Fiber.unsuspend fiber trigger then Mpscq.push t.ready resume
176-
else Mpscq.push_head t.ready resume;
177-
if
178-
Atomic.get t.needs_wakeup
179-
&& Atomic.compare_and_set t.needs_wakeup true false
180-
then begin
181-
begin
182-
match Mutex.lock t.mutex with
183-
| () -> Mutex.unlock t.mutex
184-
| exception Sys_error _ -> ()
185-
end;
186-
Condition.broadcast t.condition
187-
end
188-
in
191+
Fiber.continue fiber k ());
189192
Effect.Deep.match_with main fiber t.handler
190193

191-
let run ?quota ?fatal_exn_handler ?(forbid = false) main =
194+
let[@inline never] run ?quota ?fatal_exn_handler fiber main computation =
195+
run_fiber ?quota ?fatal_exn_handler fiber main;
196+
Computation.peek_exn computation
197+
198+
let run ?quota ?fatal_exn_handler ?forbid main =
199+
let forbid = match forbid with None -> false | Some forbid -> forbid in
192200
let computation = Computation.create ~mode:`LIFO () in
193201
let fiber = Fiber.create ~forbid computation in
194202
let main _ = Computation.capture computation main () in
195-
run_fiber ?quota ?fatal_exn_handler fiber main;
196-
Computation.peek_exn computation
203+
run ?quota ?fatal_exn_handler fiber main computation

0 commit comments

Comments
 (0)