From dba25e6c00ee3d8f23ed54f79e139eeff47fca24 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 4 Dec 2024 23:21:05 -0500
Subject: [PATCH] better CPS

---
 src/cps/nanoev_cps.ml | 91 +++++++++++++++++++++++--------------------
 1 file changed, 48 insertions(+), 43 deletions(-)

diff --git a/src/cps/nanoev_cps.ml b/src/cps/nanoev_cps.ml
index 4e7aa7e..c365e0c 100644
--- a/src/cps/nanoev_cps.ml
+++ b/src/cps/nanoev_cps.ml
@@ -1,74 +1,79 @@
-type ('st, 'action, 'input, 'a, 'err) t = {
+type 'a iter = ('a -> unit) -> unit
+type 'action action_queue = { aq: 'action Queue.t } [@@unboxed]
+type void = |
+
+type ('action, 'input, 'a, 'err) t = {
   inputs: 'input Queue.t;
   actions: 'action Queue.t;
-  mutable next: ('st, 'action, 'input, 'a, 'err) next;
+  mutable next: ('action, 'input, 'a, 'err) next;
 }
 
-and ('st, 'action, 'input, 'a, 'err) next =
+and ('action, 'input, 'a, 'err) next =
   | N_return of 'a
   | N_fail of 'err
-  | N_run of ('st -> ('st, 'action, 'input, 'a, 'err) transition)
-  | N_await of ('st -> 'input -> ('st, 'action, 'input, 'a, 'err) transition)
+  | N_run of (unit -> ('action, 'input, 'a, 'err) next)
+  | N_await of ('input -> ('action, 'input, 'a, 'err) next)
 
-and ('st, 'action, 'input, 'a, 'err) transition =
-  | T_return of 'a
-  | T_fail of 'err
-  | T_perform of 'action * ('st, 'action, 'input, 'a, 'err) next
-  | T_await of ('st -> 'input -> ('st, 'action, 'input, 'a, 'err) transition)
-
-type ('a, 'err) res =
+type (+'a, +'err) res =
   | Working
-  | Return of 'a
-  | Fail of 'err
+  | Ok of 'a
+  | Error of 'err
 
 (**/*)
 
-let return x : _ transition = T_return x
-let fail err : _ transition = T_fail err
-let perform act f : _ transition = T_perform (act, N_run f)
-let perform_and_await act f : _ transition = T_perform (act, N_await f)
-let await f : _ transition = T_await f
+let return x : _ next = N_return x
+let fail err : _ next = N_fail err
+let perform ~(aq : _ action_queue) act : unit = Queue.push act aq.aq
+let await f : _ next = N_await f
 
 (**/*)
 
 let create run : _ t =
-  { actions = Queue.create (); inputs = Queue.create (); next = N_run run }
+  let actions = Queue.create () in
+  let run () = run ~aq:{ aq = actions } () in
+  { actions; inputs = Queue.create (); next = N_run run }
 
 let create_rec run : _ t =
-  let rec loop st = run st ~recurse:loop in
-  create loop
+  let actions = Queue.create () in
+  let rec loop () = run ~aq:{ aq = actions } ~recurse:loop in
+  { actions; inputs = Queue.create (); next = N_run loop }
 
 let pop_action (self : _ t) : _ option = Queue.take_opt self.actions
 
+let pop_actions (self : _ t) : 'action iter =
+  if Queue.is_empty self.actions then
+    ignore
+  else (
+    let q = Queue.create () in
+    Queue.transfer self.actions q;
+    fun yield -> Queue.iter yield q
+  )
+
+let transfer_actions (st1 : _ t) (st2 : _ t) =
+  Queue.transfer st1.actions st2.actions
+
 (** Add external input to the state machine *)
 let add_input self i : unit = Queue.push i self.inputs
 
 (** Perform some work *)
-let work self st : _ res =
-  let rec loop () =
-    match self.next with
-    | N_return x -> Return x
-    | N_fail err -> Fail err
-    | N_run f -> transition (f st)
+let work ?transfer_to_aq self : _ res =
+  let rec loop (next : _ next) =
+    match next with
+    | N_return x ->
+      self.next <- next;
+      Ok x
+    | N_fail err ->
+      self.next <- next;
+      Error err
+    | N_run f -> f () |> loop
     | N_await f -> run_on_next_input f
-  and transition tr =
-    match tr with
-    | T_return x ->
-      self.next <- N_return x;
-      Return x
-    | T_fail err ->
-      self.next <- N_fail err;
-      Fail err
-    | T_perform (a, n) ->
-      Queue.push a self.actions;
-      self.next <- n;
-      loop ()
-    | T_await f -> run_on_next_input f
   and run_on_next_input f =
     match Queue.take self.inputs with
     | exception Queue.Empty ->
       self.next <- N_await f;
       Working
-    | input -> transition (f st input)
+    | input -> loop (f input)
   in
-  loop ()
+  let r = loop self.next in
+  Option.iter (fun aq -> Queue.transfer self.actions aq.aq) transfer_to_aq;
+  r