loop_with_yield.ml

open Effect
open Effect.Shallow

type 'a t += Recur : 'a -> 'b t | Yield : 'a -> unit t | Return : 'a -> 'a t

let recur a = perform (Recur a)
let yield a = perform (Yield a)
let return a = perform (Return a)

type ('a, 'b) state = Finished | NotStarted | Running of ('a, 'b) continuation

let loop : 'a -> ('a -> 'b) -> unit -> 'b option =
 fun init fn ->
  let state' = ref NotStarted in
  let rec helper : 'a -> unit -> 'b option =
   fun input () ->
    match !state' with
    | NotStarted ->
        state' := Running (fiber fn);
        helper init ()
    | Finished -> None
    | Running k ->
        continue_with k input
          {
            retc =
              (fun _ ->
                state' := Finished;
                helper input ());
            exnc = (fun e -> raise e);
            effc =
              (fun (type b) (eff : b t) ->
                match eff with
                | Recur value ->
                    Some
                      (fun (_ : (b, _) continuation) ->
                        state' := Running (fiber fn);
                        helper (Obj.magic value) ())
                | Yield value ->
                    Some
                      (fun (k : (b, _) continuation) ->
                        state' := Running (Obj.magic k);
                        Some (Obj.magic value))
                | Return value ->
                    Some
                      (fun (_ : (b, _) continuation) ->
                        state' := Finished;
                        Some (Obj.magic value))
                | _ -> None);
          }
  in
  helper init

let loopwith init fn = loop init fn ()

let rec runloop f loop =
  match loop () with
  | Some x ->
      f x;
      runloop f loop
  | None -> ()

let compose f g x = f (g x)

let fact i =
  loopwith (i, 1) (fun (iter, acc) ->
      if iter > 0 then recur (iter - 1, iter * acc, 1) else return acc)

let print_optional = function
  | Some x -> print_endline (string_of_int x)
  | None -> print_endline "None"

let count n =
  let counter =
    loop 0 (fun i ->
        if i < n then (
          yield i;
          recur (i + 1))
        else return i)
  in
  runloop (compose print_endline string_of_int) counter

let () =
  print_optional (fact 5);
  count 5