fold_right.ml

module Seq = struct
  include Seq

  let rec fold_right : 'a Seq.t -> 'b -> ('a -> 'b Lazy.t -> 'b Lazy.t) -> 'b Lazy.t =
    fun seq acc f ->
      match seq () with
      | Seq.Nil -> lazy acc
      | Seq.Cons (a, rest) -> f a (lazy (Lazy.force (fold_right rest acc f)))

  let map : ('a -> 'b) -> 'a Seq.t -> 'b Seq.t =
    fun f seq ->
      Lazy.force (fold_right seq Seq.empty (fun a acc ->
        lazy (fun () -> cons (f a) (Lazy.force acc) ())
      ))

  (* https://stackoverflow.com/a/15885135/58808 *)
  let take : int -> 'a Seq.t -> 'a Seq.t =
    fun n seq ->
      Lazy.force (fold_right (zip seq (ints 0)) Seq.empty (fun (a, i) acc ->
        if i = n
        then lazy empty
        else lazy (cons a (Lazy.force acc))
      ))
end

let () =
  let doubled = Seq.repeat 1 |> Seq.map (fun i -> i * 2) in
  let first_3 = Seq.take 3 doubled in
  Seq.iter (fun i -> print_int i; print_newline ()) first_3