More complex counter logic

polytypic · polytypic · commit 2de717ec0b5d · 2025-11-01T11:32:28.000Z
diff --git a/lib/picos_std.sync/conditions.ml b/lib/picos_std.sync/conditions.ml
@@ -8,7 +8,7 @@ let create ?padded () = Awaitable.make ?padded 0
 let[@inline] wait t lock ~acquire ~release =
   let before = Awaitable.get t in
   release lock;
-  let lock_forbidden lock =
+  let lock_forbidden lock ~acquire =
     let fiber = Fiber.current () in
     let forbid = Fiber.exchange fiber ~forbid:true in
     match acquire lock with
@@ -19,17 +19,40 @@ let[@inline] wait t lock ~acquire ~release =
         Fiber.set fiber ~forbid;
         Printexc.raise_with_backtrace exn bt
   in
-  match if Awaitable.get t == before then Awaitable.await t before with
-  | () -> lock_forbidden lock
-  | exception exn ->
-      let bt = Printexc.get_raw_backtrace () in
-      lock_forbidden lock;
-      Printexc.raise_with_backtrace exn bt
+  let await t lock ~acquire before =
+    match if Awaitable.get t == before then Awaitable.await t before with
+    | () ->
+        let before = Awaitable.get t in
+        if before land 1 = 0 then begin
+          let after = before lor 1 in
+          if
+            (not (Awaitable.compare_and_set t before after))
+            && Awaitable.get t land 1 = 0
+          then Awaitable.signal t
+        end;
+        lock_forbidden lock ~acquire
+    | exception exn ->
+        let bt = Printexc.get_raw_backtrace () in
+        lock_forbidden lock ~acquire;
+        Printexc.raise_with_backtrace exn bt
+  in
+  if before land 1 = 0 then
+    let after = before lor 1 in
+    if Awaitable.compare_and_set t before after || Awaitable.get t == after then
+      await t lock ~acquire after
+    else lock_forbidden lock ~acquire
+  else await t lock ~acquire before
 
 let signal t =
-  Awaitable.incr t;
-  Awaitable.signal t
+  let prior = Awaitable.fetch_and_add t 2 in
+  if prior land 1 <> 0 then begin
+    if Awaitable.compare_and_set t (prior + 2) (prior + 1) then ();
+    Awaitable.signal t
+  end
 
 let broadcast t =
-  Awaitable.incr t;
-  Awaitable.broadcast t
+  let prior = Awaitable.fetch_and_add t 2 in
+  if prior land 1 <> 0 then begin
+    if Awaitable.compare_and_set t (prior + 2) (prior + 1) then ();
+    Awaitable.broadcast t
+  end
