aoc2020/seq.ml

88 lines
2.4 KiB
OCaml

include Stdlib.Seq
let for_all p = fold_left (fun acc x -> acc && p x) true
let exists p = fold_left (fun acc x -> acc || p x) false
let span' p seq =
let rec go acc = function
| Nil -> begin
match acc with
| [] -> None
| _::_ -> Some (List.rev acc, empty)
end
| Cons (x, xs) ->
if p x then
go (x :: acc) (xs ())
else
Some (List.rev acc, cons x xs) in
go [] (seq ())
let%test_module _ = (module struct
let (@=) sp (l2, s2) =
match sp with
| Some (l1, s1) -> l1 = l2 && List.of_seq s1 = s2
| None -> 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 []) = None
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 = Option.value ~default:([], empty)
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 ~sep =
let drop_snd (lst, seq) = lst, drop_while sep seq in
unfold (fun seq -> break' sep seq |> Option.map drop_snd)
let is_space c = c = ' ' || c = '\t' || c = '\n'
let all_space str = String.to_seq str |> for_all is_space
let line_chunks' = chunks ~sep:all_space
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)
let mapi f seq =
let rec go i seq =
match seq () with
| Nil -> empty
| Cons (x, xs) -> cons (f i x) (go (i + 1) xs) in
go 0 seq
let with_index seq = mapi (fun i x -> i, x) seq
let%test_module _ = (module struct
let mapi' f lst = List.of_seq (mapi f (List.to_seq lst))
let with_index' lst = List.of_seq (with_index (List.to_seq lst))
let%test _ = mapi' (fun i _ -> i) [(); (); ()] = [0; 1; 2]
let%test _ = with_index' ["a"; "b"; "c"] = [0, "a"; 1, "b"; 2, "c"]
end)