feat(State): add try_with that overrides handlers (#29)

Co-authored-by: favonia <[email protected]>

src/State.ml

@@ -5,6 +5,7 @@ sig
   val set : state -> unit
   val modify : (state -> state) -> unit
   val run : init:state -> (unit -> 'a) -> 'a
+  val try_with : ?get:(unit -> state) -> ?set:(state -> unit) -> (unit -> 'a) -> 'a
   val register_printer : ([`Get | `Set of state] -> string option) -> unit
 end
 
@@ -29,6 +30,17 @@ struct
               st := v; continue k ()
             | _ -> None }
 
+  let try_with ?(get=get) ?(set=set) f =
+    let open Effect.Deep in
+    try_with f ()
+      { effc = fun (type a) (eff : a Effect.t) ->
+            match eff with
+            | Get -> Option.some @@ fun (k : (a, _) continuation) ->
+              continue k (get ())
+            | Set v -> Option.some @@ fun (k : (a, _) continuation) ->
+              set v; continue k ()
+            | _ -> None }
+
   let modify f = set @@ f @@ get ()
 
   let register_printer f = Printexc.register_printer @@ function

src/State.mli

@@ -30,6 +30,9 @@ sig
   val run : init:state -> (unit -> 'a) -> 'a
   (** [run ~init t] runs the thunk [t] which may perform state effects. The initial state is [init]. *)
 
+  val try_with : ?get:(unit -> state) -> ?set:(state -> unit) -> (unit -> 'a) -> 'a
+  (** [try_with ~get ~set t] runs the thunk [t] which may perform state effects, handling these effects with [get] and [set] (which may perform effects from some other module). The default handlers re-perform the effects. *)
+
   val register_printer : ([`Get | `Set of state] -> string option) -> unit
   (** [register_printer p] registers a printer [p] via {!val:Printexc.register_printer} to convert unhandled internal effects into strings for the OCaml runtime system to display. Ideally, all internal effects should have been handled by {!val:run} and there is no need to use this function, but when it is not the case, this function can be helpful for debugging. The functor {!module:Make} always registers a simple printer to suggest using {!val:run}, but you can register new ones to override it. The return type of the printer [p] should return [Some s] where [s] is the resulting string, or [None] if it chooses not to convert a particular effect. The registered printers are tried in reverse order until one of them returns [Some s] for some [s]; that is, the last registered printer is tried first. Note that this function is a wrapper of {!val:Printexc.register_printer} and all the registered printers (via this function or {!val:Printexc.register_printer}) are put into the same list.
 

test/TestState.ml

@@ -19,6 +19,7 @@ struct
   let set s = U.perform @@ StateMonad.set s
   let modify f = U.perform @@ StateMonad.modify f
   let run ~init f = fst @@ U.run f init
+  let try_with ?get:_ ?set:_ _f = failwith "state monad can't try_with"
   let register_printer _ = ()
 end