include Stdlib.Seq

type 'a span = Empty_seq | Span of 'a list * 'a t

let span' p seq =
  let rec go acc = function
    | Nil -> begin
        match acc with
        | []   -> Empty_seq
        | _::_ -> Span (List.rev acc, empty)
      end
    | Cons (x, xs) ->
        if p x then
          go (x :: acc) (xs ())
        else
          Span (List.rev acc, cons x xs) in
  go [] (seq ())

let%test_module _ = (module struct
  let (@=) sp (l2, s2) =
    match sp with
    | Span (l1, s1) -> l1 = l2 && List.of_seq s1 = s2
    | Empty_seq     -> false

  let%test _ =
    span' (fun x -> x < 4) (List.to_seq [1;2;3;4;5;6;7]) @=
      ([1;2;3], [4;5;6;7])

  let%test _ =
    span' (fun _ -> true) (List.to_seq []) = Empty_seq

  let%test _ =
    span' (fun x -> x < 0) (List.to_seq [1;2;3;4;5;6;7]) @=
      ([], [1;2;3;4;5;6;7])
end)

let break' p = span' (Fun.negate p)

let from_span = function
  | Empty_seq       -> [], empty
  | Span (lst, seq) -> lst, seq

let span p seq = from_span (span' p seq)
let break p seq = from_span (break' p seq)


let rec drop_while p seq =
  match seq () with
  | Nil -> empty
  | Cons (x, xs) ->
      if p x then drop_while p xs else cons x xs

let chunks p =
  unfold (fun seq ->
    match break' p seq with
    | Empty_seq -> None
    | Span (lst, seq) -> Some (lst, drop_while p seq))

let line_chunks' = chunks (fun s -> s = "")
let line_chunks ?(join=" ") seq = map (String.concat join) (line_chunks' seq)

let%test_module _ = (module struct
  let lc lst = List.(of_seq (line_chunks (to_seq lst)))

  let%test "ε"         = lc [] = []
  let%test "a b c"     = lc ["a"; "b"; "c"] = ["a b c"]
  let%test "a b | c"   = lc ["a"; "b"; ""; "c"] = ["a b"; "c"]
  let%test "a b || c"  = lc ["a"; "b"; ""; ""; "c"] = ["a b"; "c"]
  let%test "a b | c |" = lc ["a"; "b"; ""; "c"; ""] = ["a b"; "c"]
end)