Composing effect handler, again

(This is a follow on to my previous question)

Suppose that in addition to ‘running’ my effectful computations, I would like to provide a way of stepping through them. I can do this by creating a second handler which just stores the current status of the computation and the continuation and pairing it with a ‘next’ function which continues the computation.

module Status = struct
  type 'st t =
    | Suspended of 'st
    | Complete
    | Failed of exn
end

module Make_state (State : sig
    type t
  end) =
struct
  type _ Effect.t += Get : State.t Effect.t | Put : State.t -> unit Effect.t

  let get () = Effect.perform Get
  let put st = Effect.perform (Put st)

  let run comp st =
    let st_ref = ref st in
    Effect.Deep.match_with
      comp
      ()
      { effc =
          (fun (type a) (eff : a Effect.t) ->
            match eff with
            | Get -> Some (fun (k : (a, _) Effect.Deep.continuation) -> Effect.Deep.continue k (!st_ref : State.t))
            | Put st ->
              st_ref := st;
              Some (fun (k : (a, _) Effect.Deep.continuation) -> Effect.Deep.continue k ())
            | _ -> None)
      ; retc = (fun res -> res, !st_ref)
      ; exnc = (fun exn -> raise exn)
      }
  ;;

  type status =
    | Got of (State.t, status Status.t) Effect.Deep.continuation
    | Put of State.t * (unit, status Status.t) Effect.Deep.continuation

  let next t st =
    match t with
    | Got k -> Effect.Deep.continue k st, st
    | Put (_, k) -> Effect.Deep.continue k (), st
  ;;

  let step comp =
    Effect.Deep.match_with
      comp
      ()
      { effc =
          (fun (type a) (eff : a Effect.t) ->
            match eff with
            | Get -> Some (fun (k : (a, _) Effect.Deep.continuation) -> Status.Suspended (Got k))
            | Put st -> Some (fun (k : (a, _) Effect.Deep.continuation) -> Status.Suspended (Put (st, k)))
            | _ -> None)
      ; retc = (fun _res -> Status.Complete)
      ; exnc = (fun exn -> Status.Failed exn)
      }
  ;;
end

As before, I then want to use the sub-computations in a larger effectful computation.

module Sub_one = struct
  module State = struct
    type t = int
  end

  module Effect = Make_state (State)

  let multiply_by factor () =
    Effect.(
      let st = get () in
      let st = st * factor in
      put st)
  ;;
end

module Sub_two = struct
  module State = struct
    type t = string
  end

  module Effect = Make_state (State)

  let append suffix () =
    Effect.(
      let st = get () in
      let st = st ^ suffix in
      put st)
  ;;
end

Then

module Effect = struct
    type _ Effect.t += Log : string -> unit Effect.t

    let run_log comp (st : State.t) (log : string list) =
      let st_ref = ref st in
      let log_ref = ref log in
      Effect.Deep.match_with
        comp
        ()
        { effc =
            (fun (type a) (eff : a Effect.t) ->
              match eff with
              | Log msg ->
                log_ref := msg :: !log_ref;
                Some (fun (k : (a, _) Effect.Deep.continuation) -> Effect.Deep.continue k ())
              | _ -> None)
        ; retc = (fun res -> res, !st_ref, !log_ref)
        ; exnc = (fun exn -> raise exn)
        }
    ;;

    let run comp (st : State.t) (log : string list) =
      run_log (fun () -> Sub_one.Effect.run (fun () -> Sub_two.Effect.run comp st.sub_two) st.sub_one) st log
    ;;
end

I can write a composed step implementation in the same way as run:

type status =
    | Logged of string * (unit, status Status.t) Effect.Deep.continuation
    | Sub_one of Sub_one.Effect.status
    | Sub_two of Sub_two.Effect.status

let lift_sub_one = function
    | Status.Complete -> Status.Complete
    | Status.Failed exn -> Status.Failed exn
    | Status.Suspended status -> Status.Suspended (Sub_one status)
    
let lift_sub_two = function
    | Status.Complete -> Status.Complete
    | Status.Failed exn -> Status.Failed exn
    | Status.Suspended status -> Status.Suspended (Sub_two status)

let next_status status (st_one, st_two, log) =
    match status with
    | Logged (msg, k) -> Effect.Deep.continue k (), (st_one, st_two, msg :: log)
    | Sub_one status ->
        let status, st_one = Sub_one.Effect.next status st_one in
        lift_sub_one status, (st_one, st_two, log)
    | Sub_two status ->
        let status, st_two = Sub_two.Effect.next status st_two in
        lift_sub_two status, (st_one, st_two, log)

let next t st =
    match t with
    | Status.Complete | Status.Failed _ -> t, st
    | Status.Suspended status -> next_status status st

let step_log comp =
      Effect.Deep.match_with
        comp
        ()
        { effc =
            (fun (type a) (eff : a Effect.t) ->
              match eff with
              | Log msg -> Some (fun (k : (a, _) Effect.Deep.continuation) -> Status.Suspended (Logged (msg, k)))
              | _ -> None)
        ; retc = (fun _res -> Status.Complete)
        ; exnc = (fun exn -> Status.Failed exn)
        }


let step comp = step_log (fun () -> Sub_one.Effect.step (fun () -> Sub_two.Effect.step comp))

However, if I did this I would end up with Completed subcomputation but have now continuation with which to continue the outer computation.

Is it possible to write such a handler or do I have to explicitly handle the effects of the subcomputations in my top-level handler?

Because your step effect handlers aren’t continuing the continuation, their return value is the Status.Suspended value returned by your effc functions. This is generally speaking valid, your issue however is, that the retc parts of the step effect handlers discard the return value of the inner computations. If you replace the retcs with Fun.id it would work, but it would also require your innermost computation to return a Status.t.

But even if you do that change, and fix all the places the Typechecker is now unhappy with you, there is another problem: If you are calling Effect.Deep.continue outside of your effect handlers, the effect handlers will not apply to the continued computation.

One way to solve this might be to have your effect handlers perform an Suspended effect before continuing their computation.