composable generators for `'a gen`

gen_coroutine.ml

type 'a gen = unit -> 'a option

module Co : sig
  type 'a t (* yields ['a] *)

  val make : 'a t -> 'a gen

  val return : 'a -> 'a t

  val yield : 'a -> 'a t -> 'a t

  val append : 'a t -> 'a t -> 'a t

  val flat_map : ('a -> 'b t) -> 'a t -> 'b t

  module Infix : sig
    val (<+>) : 'a -> 'a t -> 'a t
    val (@@@) : 'a t -> 'a t -> 'a t
    val (>>=) : 'a t -> ('a -> 'b t) -> 'b t
  end

  include module type of Infix

  val of_list : 'a list -> 'a t

  val of_gen : 'a gen -> 'a t

  val yield_from_list : 'a list -> ('a -> 'b t) -> 'b t

  val yield_from_gen : 'a gen -> ('a -> 'b t) -> 'b t
end = struct
  type 'a kont =
    | K_nil
    | K_return of 'a * 'a kont
    | K_tailcall of 'a t * 'a kont
    | K_bind : {
        f: 'a -> 'b t;
        st: 'a state; (* used to generate new ['a] *)
        k: 'b kont;
      } -> 'b kont
    | K_list of {
        mutable l: 'a list;
        k: 'a kont;
      }
    | K_gen of 'a gen * 'a kont

  (* yield some ['a] by adding them to the continuation stack *)
  and 'a t = 'a kont -> 'a kont

  and 'a state = {
    mutable k : 'a kont;
  }

  (* generator *)
  let rec next_k
    : type a. a state -> a kont -> a option
    = fun st k -> match k with
      | K_nil -> st.k <- K_nil; None
      | K_tailcall (co,k) -> next_k st (co k)
      | K_return (x, k') -> st.k <- k'; Some x
      | K_bind ({f; st=st'; _} as r) ->
        begin match next_k st' st'.k with
          | None -> next_k st r.k
          | Some x -> next_k st (f x k)
        end
      | K_gen (g,k') ->
        begin match g () with
          | None -> next_k st k'
          | Some _ as res -> st.k <- k; res (* yield *)
        end
      | K_list r ->
        begin match r.l with
          | [] -> next_k st r.k
          | x :: tl -> r.l <- tl; st.k <- k; Some x
        end

  (* gen *)
  let next (st:'a state) () : 'a option =
    next_k st st.k

  let create_ (x:'a t) : 'a state = {k=K_tailcall (x, K_nil)}

  (* main entry point *)
  let make (co:'a t) : 'a gen =
    next (create_ co)

  let[@inline] append (x:'a t) (y:'a t) : 'a t =
    fun k-> x (y k)

  let[@inline] return (x:'a) : 'a t = fun k -> K_return (x, k)

  let[@inline] yield (x:'a) (co:'a t) : 'a t = fun k -> K_return (x, co k)

  let[@inline] flat_map (f: 'a -> 'b t) (x:'a t) : 'b t =
    fun k -> K_bind {f; st=create_ x; k}

  let[@inline] of_gen g : _ t = fun k -> K_gen (g,k)

  let[@inline] of_list  l : _ t = fun k -> K_list {l;k}

  module Infix = struct
    let (<+>) = yield
    let (@@@) = append
    let[@inline] (>>=) x f = flat_map f x
  end
  include Infix

  let[@inline] yield_from_gen (g:'a gen) (f:'a -> 'b t) : 'b t =
    flat_map f (of_gen g)

  let[@inline] yield_from_list (l:'a list) (f:'a -> 'b t) : 'b t =
    flat_map f (of_list l)
end

type 'a tree = L of 'a | N of 'a * 'a tree list | B of 'a tree * 'a tree

let gen_of_t (t:'a tree) : 'a gen =
  let open Co.Infix in
  let rec from_tree (t:'a tree) : 'a Co.t = match t with
    | L x -> Co.return x
    | B (x,y) -> from_tree x @@@ from_tree y
    | N (i, l) ->
      i <+> Co.yield_from_list l from_tree
  in
  Co.make @@ from_tree t

let gen_of_t' (t:'a tree) : 'a gen =
  let open Gen.Infix in
  let rec aux = function
    | L x -> Gen.return x
    | B (a,b) -> Gen.append (aux a) (aux b)
    | N (i,l) -> Gen.append (Gen.return i) (Gen.of_list l |> Gen.flat_map aux)
  in
  aux t

(* make a tree *)
let rec mk_tree (n:int) =
  if n <= 0 then L n
  else if n mod 10 = 0 then B(mk_tree (n-1), mk_tree (n-1))
  else N(n, mk_tree_l (n-1))
and mk_tree_l = function
  | 0 -> []
  | 1 -> [mk_tree 1]
  | n -> mk_tree n :: mk_tree_l (n-1)
;;

#require "gen";;
#require "unix";;

let with_time f x =
  let start = Unix.gettimeofday() in
  let r = f x in
  Printf.printf "time: %.3fs\n%!" (Unix.gettimeofday() -. start);
  r

let t20 = mk_tree 20 ;;
let n20 = with_time Gen.length @@ gen_of_t t20;;
let n20' = with_time Gen.length @@ gen_of_t' t20;;

let t25 = mk_tree 25 ;;
let n25 = with_time Gen.length @@ gen_of_t t25;;
let n25' = with_time Gen.length @@ gen_of_t' t25;;