Source file sequence.ml

open! Import
open Container_intf.Export

module Array = Array0
module List  = List0

module Step = struct
  (* 'a is an item in the sequence, 's is the state that will produce the remainder of
     the sequence *)
  type ('a,'s) t =
    | Done
    | Skip of 's
    | Yield of 'a * 's
  [@@deriving_inline sexp_of]
  let sexp_of_t : type a s.
    (a -> Ppx_sexp_conv_lib.Sexp.t) ->
    (s -> Ppx_sexp_conv_lib.Sexp.t) -> (a, s) t -> Ppx_sexp_conv_lib.Sexp.t
    =
    fun _of_a ->
    fun _of_s ->
      function
      | Done -> Ppx_sexp_conv_lib.Sexp.Atom "Done"
      | Skip v0 ->
        let v0 = _of_s v0 in
        Ppx_sexp_conv_lib.Sexp.List
          [Ppx_sexp_conv_lib.Sexp.Atom "Skip"; v0]
      | Yield (v0, v1) ->
        let v0 = _of_a v0
        and v1 = _of_s v1 in
        Ppx_sexp_conv_lib.Sexp.List
          [Ppx_sexp_conv_lib.Sexp.Atom "Yield"; v0; v1]
  [@@@end]
end

open Step

(* 'a is an item in the sequence, 's is the state that will produce the remainder of the
   sequence *)
type +_ t =
  | Sequence : 's * ('s -> ('a,'s) Step.t) -> 'a t

type 'a sequence = 'a t

module Expert = struct
  let next_step (Sequence (s, f)) =
    match f s with
    | Done         -> Done
    | Skip      s  -> Skip     (Sequence (s, f))
    | Yield (a, s) -> Yield (a, Sequence (s, f))
  ;;

  let delayed_fold_step s ~init ~f ~finish =
    let rec loop s next finish f acc =
      match next s with
      | Done         -> finish acc
      | Skip      s  -> f acc  None    ~k:(loop s next finish f)
      | Yield (a, s) -> f acc (Some a) ~k:(loop s next finish f)
    in
    match s with
    | Sequence(s, next) -> loop s next finish f init
  ;;
end

let unfold_step ~init ~f =
  Sequence (init,f)

let unfold ~init ~f =
  unfold_step ~init
    ~f:(fun s ->
      match f s with
      | None -> Step.Done
      | Some(a,s) -> Step.Yield(a,s))

let unfold_with s ~init ~f =
  match s with
  | Sequence(s, next) ->
    Sequence((init, s) ,
             (fun (seed, s) ->
                match next s with
                | Done -> Done
                | Skip s -> Skip (seed, s)
                | Yield(a,s) ->
                  match f seed a with
                  | Done -> Done
                  | Skip seed -> Skip (seed, s)
                  | Yield(a,seed) -> Yield(a,(seed,s))))

let unfold_with_and_finish s ~init ~running_step ~inner_finished ~finishing_step =
  match s with
  | Sequence (s, next) ->
    Sequence (`Inner_running (init, s), (fun state ->
      match state with
      | `Inner_running (state, inner_state) -> begin
          match next inner_state with
          | Done ->
            Skip (`Inner_finished (inner_finished state))
          | Skip inner_state ->
            Skip (`Inner_running (state, inner_state))
          | Yield (x, inner_state) ->
            match running_step state x with
            | Done -> Done
            | Skip state ->
              Skip (`Inner_running (state, inner_state))
            | Yield (y, state) ->
              Yield (y, `Inner_running (state, inner_state))
        end
      | `Inner_finished state -> begin
          match finishing_step state with
          | Done -> Done
          | Skip state ->
            Skip (`Inner_finished state)
          | Yield (y, state) ->
            Yield (y, `Inner_finished state)
        end))

let of_list l =
  unfold_step ~init:l
    ~f:(function
      | [] -> Done
      | x::l -> Yield(x,l))


let fold t ~init ~f =
  let rec loop seed v next f =
    match next seed with
    | Done -> v
    | Skip s -> loop s v next f
    | Yield(a,s) -> loop s (f v a) next f
  in
  match t with
  | Sequence(seed, next) -> loop seed init next f

let to_list_rev t =
  fold t ~init:[] ~f:(fun l x -> x::l)


let to_list (Sequence(s,next)) =
  let safe_to_list t =
    List.rev (to_list_rev t)
  in
  let rec to_list s next i =
    if i = 0 then safe_to_list (Sequence(s,next))
    else
      match next s with
      | Done -> []
      | Skip s -> to_list s next i
      | Yield(a,s) -> a::(to_list s next (i-1))
  in
  to_list s next 500

let sexp_of_t sexp_of_a t = sexp_of_list sexp_of_a (to_list t)

let range ?(stride=1) ?(start=`inclusive) ?(stop=`exclusive) start_v stop_v =
  let step =
    match stop with
    | `inclusive when stride >= 0 ->
      fun i -> if i > stop_v then Done else Yield(i, i + stride)
    | `inclusive ->
      fun i -> if i < stop_v then Done else Yield(i, i + stride)
    | `exclusive when stride >= 0 ->
      fun i -> if i >= stop_v then Done else Yield(i,i + stride)
    | `exclusive ->
      fun i -> if i <= stop_v then Done else Yield(i,i + stride)
  in
  let init =
    match start with
    | `inclusive -> start_v
    | `exclusive -> start_v + stride
  in
  unfold_step ~init ~f:step

let of_lazy t_lazy =
  unfold_step ~init:t_lazy ~f:(fun t_lazy ->
    let Sequence (s, next) = Lazy.force t_lazy in
    match next s with
    | Done         -> Done
    | Skip s       -> Skip (let v = Sequence (s, next) in lazy v)
    | Yield (x, s) -> Yield (x, let v = Sequence (s, next) in lazy v))

let map t ~f =
  match t with
  | Sequence(seed, next) ->
    Sequence(seed,
             fun seed ->
               match next seed with
               | Done -> Done
               | Skip s -> Skip s
               | Yield(a,s) -> Yield(f a,s))


let mapi t ~f =
  match t with
  | Sequence(s, next) ->
    Sequence((0,s),
             fun (i,s) ->
               match next s with
               | Done -> Done
               | Skip s -> Skip (i,s)
               | Yield(a,s) -> Yield(f i a,(i+1,s)))

let folding_map t ~init ~f =
  unfold_with t ~init ~f:(fun acc x ->
    let acc, x = f acc x in
    Yield (x, acc))

let folding_mapi t ~init ~f =
  unfold_with t ~init:(0, init) ~f:(fun (i, acc) x ->
    let acc, x = f i acc x in
    Yield (x, (i+1, acc)))

let filter t ~f =
  match t with
  | Sequence(seed, next) ->
    Sequence(seed,
             fun seed ->
               match next seed with
               | Done -> Done
               | Skip s -> Skip s
               | Yield(a,s) when f a -> Yield(a,s)
               | Yield (_,s) -> Skip s)

let filteri t ~f =
  map ~f:snd (
    filter (mapi t ~f:(fun i s -> (i,s)))
      ~f:(fun (i,s) -> f i s))

let length t =
  let rec loop i s next =
    match next s with
    | Done -> i
    | Skip s -> loop i s next
    | Yield(_,s) -> loop (i+1) s next
  in
  match t with
  | Sequence (seed, next) -> loop 0 seed next

let to_list_rev_with_length t =
  fold t ~init:([],0) ~f:(fun (l,i) x -> (x::l,i+1))

let to_array t =
  let (l,len) = to_list_rev_with_length t in
  match l with
  | [] -> [||]
  | x::l ->
    let a = Array.create ~len x in
    let rec loop i l =
      match l with
      | [] -> assert (i = -1)
      | x::l -> a.(i) <- x; loop (i-1) l
    in
    loop (len - 2) l;
    a

let find t ~f =
  let rec loop s next f =
    match next s with
    | Done -> None
    | Yield(a,_) when f a -> Some a
    | Yield(_,s) | Skip s -> loop s next f
  in
  match t with
  | Sequence (seed, next) -> loop seed next f

let find_map t ~f =
  let rec loop s next f =
    match next s with
    | Done -> None
    | Yield(a,s) ->
      (match f a with
       | None -> loop s next f
       | some_b -> some_b)
    | Skip s -> loop s next f
  in
  match t with
  | Sequence (seed, next) -> loop seed next f


let find_mapi t ~f =
  let rec loop s next f i =
    match next s with
    | Done -> None
    | Yield(a,s) ->
      (match f i a with
       | None -> loop s next f (i+1)
       | some_b -> some_b)
    | Skip s -> loop s next f i
  in
  match t with
  | Sequence (seed, next) -> loop seed next f 0

let for_all t ~f =
  let rec loop s next f =
    match next s with
    | Done -> true
    | Yield(a,_) when not (f a) -> false
    | Yield (_,s) | Skip s -> loop s next f
  in
  match t with
  | Sequence (seed, next) -> loop seed next f

let for_alli t ~f =
  let rec loop s next f i =
    match next s with
    | Done -> true
    | Yield(a,_) when not (f i a) -> false
    | Yield (_,s) -> loop s next f (i+1)
    | Skip s -> loop s next f i
  in
  match t with
  | Sequence (seed, next) -> loop seed next f 0

let exists t ~f =
  let rec loop s next f =
    match next s with
    | Done -> false
    | Yield(a,_) when f a -> true
    | Yield(_,s) | Skip s -> loop s next f
  in
  match t with
  | Sequence (seed, next) -> loop seed next f

let existsi t ~f =
  let rec loop s next f i =
    match next s with
    | Done -> false
    | Yield(a,_) when f i a -> true
    | Yield(_,s) -> loop s next f (i+1)
    | Skip s -> loop s next f i
  in
  match t with
  | Sequence (seed, next) -> loop seed next f 0

let iter t ~f =
  let rec loop seed next f =
    match next seed with
    | Done -> ()
    | Skip s -> loop s next f
    | Yield(a,s) ->
      begin
        f a;
        loop s next f
      end
  in
  match t with
  | Sequence(seed, next) -> loop seed next f

let is_empty t =
  let rec loop s next =
    match next s with
    | Done -> true
    | Skip s -> loop s next
    | Yield _ -> false
  in
  match t with
  | Sequence(seed, next) -> loop seed next

let mem t a ~equal =
  let rec loop s next a =
    match next s with
    | Done -> false
    | Yield(b,_) when equal a b -> true
    | Yield(_,s) | Skip s -> loop s next a
  in
  match t with
  | Sequence(seed, next) -> loop seed next a

let empty =
  Sequence((), fun () -> Done)

let bind t ~f =
  unfold_step
    ~f:(function
      | Sequence(seed,next), rest ->
        match next seed with
        | Done ->
          begin
            match rest with
            | Sequence(seed, next) ->
              match next seed with
              | Done -> Done
              | Skip s -> Skip (empty, Sequence(s, next))
              | Yield(a, s) -> Skip(f a, Sequence(s, next))
          end
        | Skip s -> Skip (Sequence(s,next), rest)
        | Yield(a,s) -> Yield(a, (Sequence(s,next) , rest)))
    ~init:(empty,t)

let return x =
  unfold_step ~init:(Some x)
    ~f:(function
      | None -> Done
      | Some x -> Yield(x,None))

include Monad.Make(struct
    type nonrec 'a t = 'a t
    let map = `Custom map
    let bind = bind
    let return = return
  end)

let nth s n =
  if n < 0 then None
  else
    let rec loop i s next =
      match next s with
      | Done -> None
      | Skip s -> loop i s next
      | Yield(a,s) -> if phys_equal i 0 then Some a else loop (i-1) s next
    in
    match s with
    | Sequence(s,next) ->
      loop n s next

let nth_exn s n =
  if n < 0 then raise (Invalid_argument "Sequence.nth")
  else
    match nth s n with
    | None -> failwith "Sequence.nth"
    | Some x -> x

module Merge_with_duplicates_element = struct
  type ('a, 'b) t =
    | Left of 'a
    | Right of 'b
    | Both of 'a * 'b
  [@@deriving_inline compare, hash, sexp]
  let compare :
    'a 'b .
    ('a -> 'a -> int) -> ('b -> 'b -> int) -> ('a, 'b) t -> ('a, 'b) t -> int
    =
    fun _cmp__a ->
    fun _cmp__b ->
    fun a__001_ ->
    fun b__002_ ->
      if Ppx_compare_lib.phys_equal a__001_ b__002_
      then 0
      else
        (match (a__001_, b__002_) with
         | (Left _a__003_, Left _b__004_) -> _cmp__a _a__003_ _b__004_
         | (Left _, _) -> (-1)
         | (_, Left _) -> 1
         | (Right _a__005_, Right _b__006_) -> _cmp__b _a__005_ _b__006_
         | (Right _, _) -> (-1)
         | (_, Right _) -> 1
         | (Both (_a__007_, _a__009_), Both (_b__008_, _b__010_)) ->
           (match _cmp__a _a__007_ _b__008_ with
            | 0 -> _cmp__b _a__009_ _b__010_
            | n -> n))
  let hash_fold_t : type a b.
    (Ppx_hash_lib.Std.Hash.state -> a -> Ppx_hash_lib.Std.Hash.state) ->
    (Ppx_hash_lib.Std.Hash.state -> b -> Ppx_hash_lib.Std.Hash.state) ->
    Ppx_hash_lib.Std.Hash.state -> (a, b) t -> Ppx_hash_lib.Std.Hash.state
    =
    fun _hash_fold_a ->
    fun _hash_fold_b ->
    fun hsv ->
    fun arg ->
      match arg with
      | Left _a0 ->
        let hsv = Ppx_hash_lib.Std.Hash.fold_int hsv 0 in
        let hsv = hsv in _hash_fold_a hsv _a0
      | Right _a0 ->
        let hsv = Ppx_hash_lib.Std.Hash.fold_int hsv 1 in
        let hsv = hsv in _hash_fold_b hsv _a0
      | Both (_a0, _a1) ->
        let hsv = Ppx_hash_lib.Std.Hash.fold_int hsv 2 in
        let hsv = let hsv = hsv in _hash_fold_a hsv _a0 in
        _hash_fold_b hsv _a1
  let t_of_sexp : type a b.
    (Ppx_sexp_conv_lib.Sexp.t -> a) ->
    (Ppx_sexp_conv_lib.Sexp.t -> b) -> Ppx_sexp_conv_lib.Sexp.t -> (a, b) t
    =
    let _tp_loc = "src/sequence.ml.Merge_with_duplicates_element.t" in
    fun _of_a ->
    fun _of_b ->
      function
      | Ppx_sexp_conv_lib.Sexp.List ((Ppx_sexp_conv_lib.Sexp.Atom
                                        ("left"|"Left" as _tag))::sexp_args) as _sexp ->
        (match sexp_args with
         | v0::[] -> let v0 = _of_a v0 in Left v0
         | _ ->
           Ppx_sexp_conv_lib.Conv_error.stag_incorrect_n_args _tp_loc
             _tag _sexp)
      | Ppx_sexp_conv_lib.Sexp.List ((Ppx_sexp_conv_lib.Sexp.Atom
                                        ("right"|"Right" as _tag))::sexp_args) as _sexp ->
        (match sexp_args with
         | v0::[] -> let v0 = _of_b v0 in Right v0
         | _ ->
           Ppx_sexp_conv_lib.Conv_error.stag_incorrect_n_args _tp_loc
             _tag _sexp)
      | Ppx_sexp_conv_lib.Sexp.List ((Ppx_sexp_conv_lib.Sexp.Atom
                                        ("both"|"Both" as _tag))::sexp_args) as _sexp ->
        (match sexp_args with
         | v0::v1::[] ->
           let v0 = _of_a v0
           and v1 = _of_b v1 in Both (v0, v1)
         | _ ->
           Ppx_sexp_conv_lib.Conv_error.stag_incorrect_n_args _tp_loc
             _tag _sexp)
      | Ppx_sexp_conv_lib.Sexp.Atom ("left"|"Left") as sexp ->
        Ppx_sexp_conv_lib.Conv_error.stag_takes_args _tp_loc sexp
      | Ppx_sexp_conv_lib.Sexp.Atom ("right"|"Right") as sexp ->
        Ppx_sexp_conv_lib.Conv_error.stag_takes_args _tp_loc sexp
      | Ppx_sexp_conv_lib.Sexp.Atom ("both"|"Both") as sexp ->
        Ppx_sexp_conv_lib.Conv_error.stag_takes_args _tp_loc sexp
      | Ppx_sexp_conv_lib.Sexp.List ((Ppx_sexp_conv_lib.Sexp.List _)::_) as
        sexp ->
        Ppx_sexp_conv_lib.Conv_error.nested_list_invalid_sum _tp_loc sexp
      | Ppx_sexp_conv_lib.Sexp.List [] as sexp ->
        Ppx_sexp_conv_lib.Conv_error.empty_list_invalid_sum _tp_loc sexp
      | sexp -> Ppx_sexp_conv_lib.Conv_error.unexpected_stag _tp_loc sexp
  let sexp_of_t : type a b.
    (a -> Ppx_sexp_conv_lib.Sexp.t) ->
    (b -> Ppx_sexp_conv_lib.Sexp.t) -> (a, b) t -> Ppx_sexp_conv_lib.Sexp.t
    =
    fun _of_a ->
    fun _of_b ->
      function
      | Left v0 ->
        let v0 = _of_a v0 in
        Ppx_sexp_conv_lib.Sexp.List
          [Ppx_sexp_conv_lib.Sexp.Atom "Left"; v0]
      | Right v0 ->
        let v0 = _of_b v0 in
        Ppx_sexp_conv_lib.Sexp.List
          [Ppx_sexp_conv_lib.Sexp.Atom "Right"; v0]
      | Both (v0, v1) ->
        let v0 = _of_a v0
        and v1 = _of_b v1 in
        Ppx_sexp_conv_lib.Sexp.List
          [Ppx_sexp_conv_lib.Sexp.Atom "Both"; v0; v1]
  [@@@end]
end

let merge_with_duplicates (Sequence (s1, next1)) (Sequence (s2, next2)) ~compare =
  let unshadowed_compare = compare in
  let open Merge_with_duplicates_element in
  let next  = function
    | Skip s1, s2 -> Skip (next1 s1, s2)
    | s1, Skip s2 -> Skip (s1, next2 s2)
    | (Yield (a, s1') as s1), (Yield (b, s2') as s2) ->
      let comparison = unshadowed_compare a b in
      if comparison < 0
      then Yield (Left a, (Skip s1', s2))
      else if comparison = 0
      then Yield (Both (a, b), (Skip s1', Skip s2'))
      else Yield (Right b, (s1, Skip s2'))
    | Done, Done -> Done
    | Yield (a, s1), Done -> Yield (Left  a, (Skip s1, Done))
    | Done, Yield (b, s2) -> Yield (Right b, (Done, Skip s2))
  in
  Sequence((Skip s1, Skip s2), next)

let merge s1 s2 ~compare =
  map (merge_with_duplicates s1 s2 ~compare)
    ~f:(function Left x | Right x | Both (x, _) -> x)

let hd s =
  let rec loop s next =
    match next s with
    | Done -> None
    | Skip s -> loop s next
    | Yield(a,_) -> Some a
  in
  match s with
  | Sequence (s,next) -> loop s next

let hd_exn s =
  match hd s with
  | None -> failwith "hd_exn"
  | Some a -> a

let tl s =
  let rec loop s next =
    match next s with
    | Done -> None
    | Skip s -> loop s next
    | Yield(_,a) -> Some a
  in
  match s with
  | Sequence (s,next) ->
    match loop s next with
    | None -> None
    | Some s -> Some (Sequence(s,next))

let tl_eagerly_exn s =
  match tl s with
  | None -> failwith "Sequence.tl_exn"
  | Some s -> s

let lift_identity next s =
  match next s with
  | Done -> Done
  | Skip s -> Skip (`Identity s)
  | Yield(a,s) -> Yield(a, `Identity s)

let next s =
  let rec loop s next =
    match next s with
    | Done -> None
    | Skip s -> loop s next
    | Yield(a,s) -> Some (a, Sequence(s, next))
  in
  match s with
  | Sequence(s, next) -> loop s next

let filter_opt s =
  match s with
  | Sequence(s, next) ->
    Sequence(s,
             fun s ->
               match next s with
               | Done -> Done
               | Skip s -> Skip s
               | Yield(None, s) -> Skip s
               | Yield(Some a, s) -> Yield(a, s))

let filter_map s ~f =
  filter_opt (map s ~f)

let filter_mapi s ~f =
  filter_map (mapi s ~f:(fun i s -> (i,s)))
    ~f:(fun (i, s) -> f i s)

let split_n s n =
  let rec loop s i accum next =
    if i <= 0 then
      (List.rev accum, Sequence(s,next))
    else
      match next s with
      | Done -> (List.rev accum, empty)
      | Skip s -> loop s i accum next
      | Yield(a,s) -> loop s (i-1) (a::accum) next
  in
  match s with
  | Sequence(s, next) -> loop s n [] next

let chunks_exn t n =
  if n <= 0
  then raise (Invalid_argument "Sequence.chunks_exn")
  else
    unfold_step ~init:t ~f:(fun t ->
      match split_n t n with
      | [], _empty -> Done
      | _::_ as xs, t -> Yield (xs, t))

let findi s ~f =
  find (mapi s ~f:(fun i s -> (i,s)))
    ~f:(fun (i,s) -> f i s)

let find_exn s ~f =
  match find s ~f with
  | None -> failwith "Sequence.find_exn"
  | Some x -> x

let append s1 s2 =
  match s1, s2 with
  | Sequence(s1, next1), Sequence(s2, next2) ->
    Sequence(`First_list s1,
             function
             | `First_list s1 ->
               begin
                 match next1 s1 with
                 | Done -> Skip (`Second_list s2)
                 | Skip s1 -> Skip (`First_list s1)
                 | Yield(a,s1) -> Yield(a, `First_list s1)
               end
             | `Second_list s2 ->
               begin
                 match next2 s2 with
                 | Done -> Done
                 | Skip s2 -> Skip (`Second_list s2)
                 | Yield(a,s2) -> Yield(a, `Second_list s2)
               end)

let concat_map s ~f = bind s ~f

let concat s = concat_map s ~f:Fn.id

let concat_mapi s ~f =
  concat_map (mapi s  ~f:(fun i s -> (i,s)))
    ~f:(fun (i,s) -> f i s)

let zip (Sequence (s1, next1)) (Sequence (s2, next2)) =
  let next = function
    | Yield (a, s1), Yield (b, s2) -> Yield ((a, b), (Skip s1, Skip s2))
    | Done, _
    | _, Done -> Done
    | Skip s1, s2 -> Skip (next1 s1, s2)
    | s1, Skip s2 -> Skip (s1, next2 s2)
  in
  Sequence ((Skip s1, Skip s2), next)

let zip_full (Sequence(s1,next1)) (Sequence(s2,next2)) =
  let next = function
    | Yield (a, s1), Yield (b, s2) -> Yield (`Both (a, b), (Skip s1, Skip s2))
    | Done, Done -> Done
    | Skip s1, s2 -> Skip (next1 s1, s2)
    | s1, Skip s2 -> Skip (s1, next2 s2)
    | Done, Yield(b, s2) -> Yield((`Right b), (Done, next2 s2))
    | Yield(a, s1), Done -> Yield((`Left a), (next1 s1, Done))
  in
  Sequence ((Skip s1, Skip s2), next)

let bounded_length (Sequence(seed,next)) ~at_most =
  let rec loop i seed next =
    if i > at_most then `Greater
    else
      match next seed with
      | Done -> `Is i
      | Skip seed -> loop i seed next
      | Yield(_, seed) -> loop (i+1) seed next
  in
  loop 0 seed next

let length_is_bounded_by ?(min=(-1)) ?max t =
  let length_is_at_least (Sequence(s,next)) =
    let rec loop s acc =
      if acc >= min then true else
        match next s with
        | Done -> false
        | Skip s -> loop s acc
        | Yield(_,s) -> loop s (acc + 1)
    in loop s 0
  in
  match max with
  | None -> length_is_at_least t
  | Some max ->
    begin
      match bounded_length t ~at_most:max with
      | `Is len when len >= min -> true
      | _ -> false
    end

let iteri s ~f =
  iter (mapi s ~f:(fun i s -> (i, s)))
    ~f:(fun (i, s) -> f i s)

let foldi s ~init ~f =
  fold ~init (mapi s ~f:(fun i s -> (i,s)))
    ~f:(fun acc (i, s) -> f i acc s)

let reduce s ~f =
  match next s with
  | None -> None
  | Some(a, s) -> Some (fold s ~init:a ~f)

let reduce_exn s ~f =
  match reduce s ~f with
  | None -> failwith "Sequence.reduce_exn"
  | Some res -> res

let group (Sequence (s, next)) ~break =
  unfold_step ~init:(Some ([], s)) ~f:(function
    | None          -> Done
    | Some (acc, s) ->
      match acc, next s with
      | _,                Skip s         -> Skip (Some (acc, s))
      | [],               Done           -> Done
      | acc,              Done           -> Yield (List.rev acc, None)
      | [],               Yield (cur, s) -> Skip (Some ([cur], s))
      | prev :: _ as acc, Yield (cur, s) ->
        if break prev cur
        then Yield (List.rev acc, Some ([cur], s))
        else Skip (Some (cur :: acc, s)))
;;

let find_consecutive_duplicate (Sequence(s, next)) ~equal =
  let rec loop last_elt s =
    match next s with
    | Done -> None
    | Skip s -> loop last_elt s
    | Yield(a,s) ->
      match last_elt with
      | Some b when equal a b -> Some (b, a)
      | None | Some _ -> loop (Some a) s
  in
  loop None s

let remove_consecutive_duplicates s ~equal =
  unfold_with s ~init:None
    ~f:(fun prev a ->
      match prev with
      | Some b when equal a b -> Skip(Some a)
      | None | Some _ -> Yield(a, Some a))

let count s ~f =
  length (filter s ~f)

let counti t ~f =
  length (filteri t ~f)

let sum m t ~f = Container.sum ~fold m t ~f
let min_elt t ~compare = Container.min_elt ~fold t ~compare
let max_elt t ~compare = Container.max_elt ~fold t ~compare

let init n ~f =
  unfold_step ~init:0
    ~f:(fun i ->
      if i >= n then Done
      else Yield(f i, i + 1))

let sub s ~pos ~len =
  if pos < 0 || len < 0 then failwith "Sequence.sub";
  match s with
  | Sequence(s, next) ->
    Sequence((0,s),
             (fun (i, s) ->
                if i - pos >= len then Done
                else
                  match next s with
                  | Done -> Done
                  | Skip s -> Skip (i, s)
                  | Yield(a, s) when i >= pos -> Yield (a,(i + 1, s))
                  | Yield(_, s) -> Skip(i + 1, s)))

let take s len =
  if len < 0 then failwith "Sequence.take";
  match s with
  | Sequence(s, next) ->
    Sequence((0,s),
             (fun (i, s) ->
                if i >= len then Done
                else
                  match next s with
                  | Done -> Done
                  | Skip s -> Skip (i, s)
                  | Yield(a, s) -> Yield (a,(i + 1, s))))

let drop s len =
  if len < 0 then failwith "Sequence.drop";
  match s with
  | Sequence(s, next) ->
    Sequence((0,s),
             (fun (i, s) ->
                match next s with
                | Done -> Done
                | Skip s -> Skip (i, s)
                | Yield(a, s) when i >= len -> Yield (a,(i + 1, s))
                | Yield(_, s) -> Skip (i+1, s)))

let take_while s ~f =
  match s with
  | Sequence(s, next) ->
    Sequence(s,
             fun s ->
               match next s with
               | Done -> Done
               | Skip s -> Skip s
               | Yield (a, s) when f a -> Yield(a,s)
               | Yield (_,_) -> Done)

let drop_while s ~f =
  match s with
  | Sequence(s, next) ->
    Sequence(`Dropping s,
             function
             |`Dropping s ->
               begin
                 match next s with
                 | Done -> Done
                 | Skip s -> Skip (`Dropping s)
                 | Yield(a, s) when f a -> Skip (`Dropping s)
                 | Yield(a, s) -> Yield(a, `Identity s)
               end
             | `Identity s -> lift_identity next s)

let shift_right s x =
  match s with
  | Sequence(seed, next) ->
    Sequence(`Consing (seed, x),
             function
             | `Consing (seed, x) -> Yield(x, `Identity seed)
             | `Identity s -> lift_identity next s)

let shift_right_with_list s l =
  append (of_list l) s

let shift_left = drop

module Infix = struct
  let (@) = append
end

let intersperse s ~sep =
  match s with
  | Sequence(s, next) ->
    Sequence(`Init s,
             function
             | `Init s ->
               begin
                 match next s with
                 | Done -> Done
                 | Skip s -> Skip (`Init s)
                 | Yield(a, s) -> Yield(a, `Running s)
               end
             | `Running s ->
               begin
                 match next s with
                 | Done -> Done
                 | Skip s -> Skip (`Running s)
                 | Yield(a, s) -> Yield(sep, `Putting(a,s))
               end
             | `Putting(a,s) -> Yield(a,`Running s))

let repeat x =
  unfold_step ~init:x ~f:(fun x -> Yield(x, x))

let cycle_list_exn xs =
  if List.is_empty xs then raise (Invalid_argument "Sequence.cycle_list_exn");
  let s = of_list xs in
  concat_map ~f:(fun () -> s) (repeat ())

let cartesian_product sa sb =
  concat_map sa
    ~f:(fun a -> zip (repeat a) sb)

let singleton x = return x

let delayed_fold s ~init ~f ~finish =
  Expert.delayed_fold_step s ~init ~finish ~f:(fun acc option ~k ->
    match option with
    | None   -> k acc
    | Some a -> f acc a ~k)

let fold_m ~bind ~return t ~init ~f =
  Expert.delayed_fold_step t
    ~init
    ~f:(fun acc option ~k ->
      match option with
      | None   -> bind (return acc) ~f:k
      | Some a -> bind (f acc a)    ~f:k)
    ~finish:return
;;

let iter_m ~bind ~return t ~f =
  Expert.delayed_fold_step t
    ~init:()
    ~f:(fun () option ~k ->
      match option with
      | None   -> bind (return ()) ~f:k
      | Some a -> bind (f a)       ~f:k)
    ~finish:return
;;

let fold_until s ~init ~f ~finish =
  let rec loop s next f =
    fun acc ->
      match next s with
      | Done   -> finish acc
      | Skip s ->  loop s next f acc
      | Yield(a, s) -> match (f acc a : ('a, 'b) Continue_or_stop.t) with
        | Stop x -> x
        | Continue acc -> loop s next f acc
  in
  match s with
  | Sequence(s, next) -> loop s next f init

let fold_result s ~init ~f =
  let rec loop s next f =
    fun acc ->
      match next s with
      | Done   -> Result.return acc
      | Skip s ->  loop s next f acc
      | Yield(a, s) -> match (f acc a : (_, _) Result.t) with
        | Error _ as e -> e
        | Ok acc -> loop s next f acc
  in
  match s with
  | Sequence(s, next) -> loop s next f init

let force_eagerly t = of_list (to_list t)

let memoize (type a) (Sequence (s, next)) =
  let module M = struct
    type t = T of (a, t) Step.t Lazy.t
  end in
  let rec memoize s = M.T (lazy (find_step s))
  and find_step s =
    match next s with
    | Done -> Done
    | Skip s -> find_step s
    | Yield (a, s) -> Yield (a, memoize s)
  in
  Sequence (memoize s, (fun (M.T l) -> Lazy.force l))

let drop_eagerly s len =
  let rec loop i ~len s next =
    if i >= len then Sequence(s, next)
    else
      match next s with
      | Done -> empty
      | Skip s -> loop i ~len s next
      | Yield(_,s) -> loop (i+1) ~len s  next
  in
  match s with
  | Sequence(s, next) -> loop 0 ~len s next

let drop_while_option (Sequence (s, next)) ~f =
  let rec loop s =
    match next s with
    | Done -> None
    | Skip s -> loop s
    | Yield (x, s) -> if f x then loop s else Some (x, Sequence (s, next))
  in
  loop s

let compare compare_a t1 t2 =
  With_return.with_return (fun r ->
    iter (zip_full t1 t2) ~f:(function
      | `Left _        -> r.return 1
      | `Right _       -> r.return (-1)
      | `Both (v1, v2) ->
        let c = compare_a v1 v2 in
        if c <> 0
        then r.return c);
    0);
;;

let round_robin list =
  let next (todo_stack, done_stack) =
    match todo_stack with
    | Sequence (s, f) :: todo_stack ->
      begin
        match f s with
        | Yield (x, s) -> Yield (x, (todo_stack, Sequence (s, f) :: done_stack))
        | Skip      s  -> Skip      (Sequence (s, f) :: todo_stack, done_stack)
        | Done         -> Skip                         (todo_stack, done_stack)
      end
    | [] ->
      if List.is_empty done_stack
      then Done
      else Skip (List.rev done_stack, [])
  in
  let state = list, [] in
  Sequence (state, next)

let interleave (Sequence (s1, f1)) =
  let next (todo_stack, done_stack, s1) =
    match todo_stack with
    | Sequence (s2, f2) :: todo_stack ->
      begin
        match f2 s2 with
        | Yield (x, s2) -> Yield (x, (todo_stack, Sequence (s2, f2) :: done_stack, s1))
        | Skip s2       -> Skip      (todo_stack, Sequence (s2, f2) :: done_stack, s1)
        | Done          -> Skip      (todo_stack,                      done_stack, s1)
      end
    | [] ->
      begin
        match f1 s1, done_stack with
        | Yield (t, s1), _    -> Skip (List.rev (t :: done_stack), [], s1)
        | Skip      s1 , _    -> Skip (List.rev       done_stack , [], s1)
        | Done         , _::_ -> Skip (List.rev       done_stack , [], s1)
        | Done         , []   -> Done
      end
  in
  let state = [], [], s1 in
  Sequence (state, next)

let interleaved_cartesian_product s1 s2 =
  map s1 ~f:(fun x1 ->
    map s2 ~f:(fun x2 ->
      (x1, x2)))
  |> interleave

module Generator = struct

  type 'elt steps = Wrap of ('elt, unit -> 'elt steps) Step.t

  let unwrap (Wrap step) = step

  module T = struct
    type ('a, 'elt) t = ('a -> 'elt steps) -> 'elt steps
    let return x = (); fun k -> k x
    let bind m ~f = (); fun k -> m (fun a -> let m' = f a in m' k)
    let map m ~f = (); fun k -> m (fun a -> k (f a))
    let map = `Custom map
  end
  include T
  include Monad.Make2 (T)

  let yield e = (); fun k -> Wrap (Yield (e, k))

  let to_steps t = t (fun () -> Wrap Done)

  let of_sequence sequence =
    delayed_fold sequence
      ~init:()
      ~f:(fun () x ~k f -> Wrap (Yield (x, fun () -> k () f)))
      ~finish:return

  let run t =
    let init () = to_steps t in
    let f thunk = unwrap (thunk ()) in
    unfold_step ~init ~f

end