coiterators.ml

module type Iter = sig
  type 'a coll
  type 'a t

  val iter : 'a coll -> 'a t
  val next : 'a t -> ('a * 'a t) option
end

module IterLemmas(I : Iter) = struct
  let fold (x: 'a) (f: 'a -> 'b -> 'a) (coll: 'b I.coll) =
    let rec go x f it =
      match I.next it with
      | None -> x
      | Some (hd, tl) -> go (f x hd) f tl
    in go x f (I.iter coll)
end

module type Builder = sig
  type 'a coll
  type 'a t

  val coll : 'a t -> 'a coll
  val prev : ('a * 'a t) option -> 'a t
end

module type Monad = sig
  type 'a t
  val pure : 'a -> 'a t
  val bind : 'a t -> ('a -> 'b t) -> 'b t
end

module MonadLemmas(M: Monad) = struct
  include M
  let (let*) mx f = M.bind mx f
end

(* Iter + Builder -> Monad *)
module IterBuilder2Monad(I : Iter)(B : Builder with type 'a coll = 'a I.coll)
        : Monad with type 'a t = 'a I.coll = struct
  type 'a t = 'a I.coll

  let empty () = B.prev None

  let pure x = B.coll (B.prev (Some (x, empty ())))

  open IterLemmas(I)

  let bind mx (f : 'a -> 'b t) =
    fold (empty ()) (fun builder elt ->
        fold builder (fun builder elt ->
            B.prev (Some (elt, builder))
          ) (f elt)
    ) mx
    |> B.coll
end

module ListI : Iter with type 'a coll = 'a list = struct
  type 'a coll = 'a list
  type 'a t = 'a list

  let iter l = l
  let next l = match l with
    | [] -> None
    | x::xs -> Some (x, xs)
end

module ListB : Builder with type 'a coll = 'a list = struct
  type 'a coll = 'a list
  type 'a t = 'a list

  let coll l = List.rev l
  let prev = function
    | None -> []
    | Some (x, xs) -> x::xs
end

module ListMonad = IterBuilder2Monad(ListI)(ListB)

let () =
  let l = (
      let open MonadLemmas(ListMonad) in
      let* x = [1; 2; 3; 4] in
      let* y = [5; 6; 7; 8] in
      pure (x * y)
    ) in
  List.iter (fun d -> Format.printf "%d, " d) l