ocaml-containers/src/data/CCFQueue.ml

526 lines
13 KiB
OCaml

(* This file is free software, part of containers. See file "license" for more details. *)
(** {1 Functional queues (fifo)} *)
type 'a sequence = ('a -> unit) -> unit
type 'a klist = unit -> [`Nil | `Cons of 'a * 'a klist]
type 'a equal = 'a -> 'a -> bool
type 'a printer = Format.formatter -> 'a -> unit
(*$inject
let pp_ilist = CCFormat.(to_string (list int))
*)
(** {2 Basics} *)
type 'a digit =
| Zero
| One of 'a
| Two of 'a * 'a
| Three of 'a * 'a * 'a
(* store the size in deep version *)
type 'a t =
| Shallow of 'a digit
| Deep of int * 'a digit * ('a * 'a) t lazy_t * 'a digit
let empty = Shallow Zero
(*$R
let q = empty in
OUnit.assert_bool "is_empty" (is_empty q)
*)
exception Empty
let is_not_zero = function
| Zero -> false
| One _ | Two _ | Three _ -> true
let _single x = Shallow (One x)
let _double x y = Shallow (Two (x,y))
let _deep n hd middle tl =
assert (is_not_zero hd && is_not_zero tl);
Deep (n, hd, middle, tl)
let is_empty = function
| Shallow Zero -> true
| _ -> false
let singleton x = _single x
let doubleton x y = _double x y
let _empty = Lazy.from_val empty
let rec cons : 'a. 'a -> 'a t -> 'a t
= fun x q -> match q with
| Shallow Zero -> _single x
| Shallow (One y) -> Shallow (Two (x,y))
| Shallow (Two (y,z)) -> Shallow (Three (x,y,z))
| Shallow (Three (y,z,z')) ->
_deep 4 (Two (x,y)) _empty (Two (z,z'))
| Deep (_, Zero, _middle, _tl) -> assert false
| Deep (n,One y, middle, tl) -> _deep (n+1) (Two (x,y)) middle tl
| Deep (n,Two (y,z), middle, tl) -> _deep (n+1)(Three (x,y,z)) middle tl
| Deep (n,Three (y,z,z'), lazy q', tail) ->
_deep (n+1) (Two (x,y)) (lazy (cons (z,z') q')) tail
(*$Q
(Q.pair Q.int (Q.list Q.int)) (fun (x,l) -> \
cons x (of_list l) |> to_list = x::l)
*)
let rec snoc : 'a. 'a t -> 'a -> 'a t
= fun q x -> match q with
| Shallow Zero -> _single x
| Shallow (One y) -> Shallow (Two (y,x))
| Shallow (Two (y,z)) -> Shallow (Three (y,z,x))
| Shallow (Three (y,z,z')) ->
_deep 4 (Two (y,z)) _empty (Two (z',x))
| Deep (_,_hd, _middle, Zero) -> assert false
| Deep (n,hd, middle, One y) -> _deep (n+1) hd middle (Two(y,x))
| Deep (n,hd, middle, Two (y,z)) -> _deep (n+1) hd middle (Three(y,z,x))
| Deep (n,hd, lazy q', Three (y,z,z')) ->
_deep (n+1) hd (lazy (snoc q' (y,z))) (Two(z',x))
(*$Q
(Q.pair Q.int (Q.list Q.int)) (fun (x,l) -> \
snoc (of_list l) x |> to_list = l @ [x])
*)
(*$R
let q = List.fold_left snoc empty [1;2;3;4;5] in
let q = tail q in
let q = List.fold_left snoc q [6;7;8] in
let l = Sequence.to_list (to_seq q) in
OUnit.assert_equal ~printer:pp_ilist [2;3;4;5;6;7;8] l
*)
let rec take_front_exn : 'a. 'a t -> ('a *'a t)
= fun q -> match q with
| Shallow Zero -> raise Empty
| Shallow (One x) -> x, empty
| Shallow (Two (x,y)) -> x, Shallow (One y)
| Shallow (Three (x,y,z)) -> x, Shallow (Two (y,z))
| Deep (_,Zero, _, _) -> assert false
| Deep (n,One x, lazy q', tail) ->
if is_empty q'
then x, Shallow tail
else
let (y,z), q' = take_front_exn q' in
x, _deep (n-1)(Two (y,z)) (Lazy.from_val q') tail
| Deep (n,Two (x,y), middle, tail) ->
x, _deep (n-1) (One y) middle tail
| Deep (n,Three (x,y,z), middle, tail) ->
x, _deep (n-1) (Two(y,z)) middle tail
(*$Q
(Q.pair Q.int (Q.list Q.int)) (fun (x,l) -> \
let x', q = cons x (of_list l) |> take_front_exn in \
x'=x && to_list q = l)
*)
(*$R
let q = of_list [1;2;3;4] in
let x, q = take_front_exn q in
OUnit.assert_equal 1 x;
let q = List.fold_left snoc q [5;6;7] in
OUnit.assert_equal 2 (first_exn q);
let x, q = take_front_exn q in
OUnit.assert_equal 2 x;
*)
let take_front q =
try Some (take_front_exn q)
with Empty -> None
(*$T
take_front empty = None
*)
let take_front_l n q =
if n<0 then (
invalid_arg "take_back_l: cannot take negative number of arguments"
);
let rec aux acc q n =
if n=0 || is_empty q then List.rev acc, q
else
let x,q' = take_front_exn q in
aux (x::acc) q' (n-1)
in aux [] q n
(*$T
let l, q = take_front_l 5 (1 -- 10) in \
l = [1;2;3;4;5] && to_list q = [6;7;8;9;10]
*)
let take_front_while p q =
let rec aux acc q =
if is_empty q then List.rev acc, q
else
let x,q' = take_front_exn q in
if p x then aux (x::acc) q' else List.rev acc, q
in aux [] q
(*$T
take_front_while (fun x-> x<5) (1 -- 10) |> fst = [1;2;3;4]
*)
let rec take_back_exn : 'a. 'a t -> 'a t * 'a
= fun q -> match q with
| Shallow Zero -> raise Empty
| Shallow (One x) -> empty, x
| Shallow (Two (x,y)) -> _single x, y
| Shallow (Three (x,y,z)) -> Shallow (Two(x,y)), z
| Deep (_, _hd, _middle, Zero) -> assert false
| Deep (n, hd, lazy q', One x) ->
if is_empty q'
then Shallow hd, x
else
let q'', (y,z) = take_back_exn q' in
_deep (n-1) hd (Lazy.from_val q'') (Two (y,z)), x
| Deep (n, hd, middle, Two(x,y)) -> _deep (n-1) hd middle (One x), y
| Deep (n, hd, middle, Three(x,y,z)) -> _deep (n-1) hd middle (Two (x,y)), z
(*$Q
(Q.pair Q.int (Q.list Q.int)) (fun (x,l) -> \
let q,x' = snoc (of_list l) x |> take_back_exn in \
x'=x && to_list q = l)
*)
let take_back q =
try Some (take_back_exn q)
with Empty -> None
(*$T
take_back empty = None
*)
let take_back_l n q =
if n<0 then (
invalid_arg "take_back_l: cannot take negative number of arguments"
);
let rec aux acc q n =
if n=0 || is_empty q then q, acc
else
let q',x = take_back_exn q in
aux (x::acc) q' (n-1)
in aux [] q n
let take_back_while p q =
let rec aux acc q =
if is_empty q then q, acc
else
let q',x = take_back_exn q in
if p x then aux (x::acc) q' else q, acc
in aux [] q
(** {2 Individual extraction} *)
let first q =
try Some (fst (take_front_exn q))
with Empty -> None
let first_exn q = fst (take_front_exn q)
let last q =
try Some (snd (take_back_exn q))
with Empty -> None
let last_exn q = snd (take_back_exn q)
let _size_digit = function
| Zero -> 0
| One _ -> 1
| Two _ -> 2
| Three _ -> 3
let size : 'a. 'a t -> int
= function
| Shallow d -> _size_digit d
| Deep (n, _, _, _) -> n
(*$Q
(Q.list Q.int) (fun l -> \
size (of_list l) = List.length l)
*)
let _nth_digit i d = match i, d with
| _, Zero -> raise Not_found
| 0, One x -> x
| 0, Two (x,_) -> x
| 1, Two (_,x) -> x
| 0, Three (x,_,_) -> x
| 1, Three (_,x,_) -> x
| 2, Three (_,_,x) -> x
| _, _ -> raise Not_found
let rec nth_exn : 'a. int -> 'a t -> 'a
= fun i q -> match i, q with
| _, Shallow Zero -> raise Not_found
| 0, Shallow (One x) -> x
| 0, Shallow (Two (x,_)) -> x
| 1, Shallow (Two (_,x)) -> x
| 0, Shallow (Three (x,_,_)) -> x
| 1, Shallow (Three (_,x,_)) -> x
| 2, Shallow (Three (_,_,x)) -> x
| _, Shallow _ -> raise Not_found
| _, Deep (_, l, q, r) ->
if i<_size_digit l
then _nth_digit i l
else
let i' = i - _size_digit l in
let q' = Lazy.force q in
if i'<2*size q'
then
let (x,y) = nth_exn (i'/2) q' in
if i' mod 2 = 0 then x else y
else
_nth_digit (i'-2*size q') r
(*$T
let l = CCList.(0--100) in let q = of_list l in \
List.map (fun i->nth_exn i q) l = l
*)
let nth i q =
try Some (nth_exn i q)
with Failure _ -> None
(*$Q & ~count:30
(Q.list Q.int) (fun l -> \
let len = List.length l in let idx = CCList.(0 -- (len - 1)) in \
let q = of_list l in \
l = [] || List.for_all (fun i -> nth i q = Some (List.nth l i)) idx)
*)
let init q =
try fst (take_back_exn q)
with Empty -> q
(*$Q
(Q.list Q.int) (fun l -> \
l = [] || (of_list l |> init |> to_list = List.rev (List.tl (List.rev l))))
*)
let tail q =
try snd (take_front_exn q)
with Empty -> q
(*$Q
(Q.list Q.int) (fun l -> \
l = [] || (of_list l |> tail |> to_list = List.tl l))
*)
let add_seq_front seq q =
let l = ref [] in
(* reversed seq *)
seq (fun x -> l := x :: !l);
List.fold_left (fun q x -> cons x q) q !l
(*$Q
Q.(pair (list int) (list int)) (fun (l1, l2) -> \
add_seq_front (Sequence.of_list l1) (of_list l2) |> to_list = l1 @ l2)
*)
let add_seq_back q seq =
let q = ref q in
seq (fun x -> q := snoc !q x);
!q
let _digit_to_seq d k = match d with
| Zero -> ()
| One x -> k x
| Two (x,y) -> k x; k y
| Three (x,y,z) -> k x; k y; k z
let rec to_seq : 'a. 'a t -> 'a sequence
= fun q k -> match q with
| Shallow d -> _digit_to_seq d k
| Deep (_, hd, lazy q', tail) ->
_digit_to_seq hd k;
to_seq q' (fun (x,y) -> k x; k y);
_digit_to_seq tail k
(*$Q
(Q.list Q.int) (fun l -> \
of_list l |> to_seq |> Sequence.to_list = l)
*)
let append q1 q2 =
match q1, q2 with
| Shallow Zero, _ -> q2
| _, Shallow Zero -> q1
| _ -> add_seq_back q1 (to_seq q2)
(*$Q
(Q.pair (Q.list Q.int)(Q.list Q.int)) (fun (l1,l2) -> \
append (of_list l1) (of_list l2) |> to_list = l1 @ l2)
*)
(*$R
let q1 = of_seq (Sequence.of_list [1;2;3;4]) in
let q2 = of_seq (Sequence.of_list [5;6;7;8]) in
let q = append q1 q2 in
let l = Sequence.to_list (to_seq q) in
OUnit.assert_equal ~printer:pp_ilist [1;2;3;4;5;6;7;8] l
*)
let _map_digit f d = match d with
| Zero -> Zero
| One x -> One (f x)
| Two (x,y) -> Two (f x, f y)
| Three (x,y,z) -> Three (f x, f y, f z)
let rec map : 'a 'b. ('a -> 'b) -> 'a t -> 'b t
= fun f q -> match q with
| Shallow d -> Shallow (_map_digit f d)
| Deep (size, hd, lazy q', tl) ->
let q'' = map (fun (x,y) -> f x, f y) q' in
_deep size (_map_digit f hd) (Lazy.from_val q'') (_map_digit f tl)
(*$Q
(Q.list Q.int) (fun l -> \
of_list l |> map string_of_int |> to_list = List.map string_of_int l)
*)
let (>|=) q f = map f q
let _fold_digit f acc d = match d with
| Zero -> acc
| One x -> f acc x
| Two (x,y) -> f (f acc x) y
| Three (x,y,z) -> f (f (f acc x) y) z
let rec fold : 'a 'b. ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b
= fun f acc q -> match q with
| Shallow d -> _fold_digit f acc d
| Deep (_, hd, lazy q', tl) ->
let acc = _fold_digit f acc hd in
let acc = fold (fun acc (x,y) -> f (f acc x) y) acc q' in
_fold_digit f acc tl
(*$Q
(Q.list Q.int) (fun l -> \
of_list l |> fold (fun acc x->x::acc) [] = List.rev l)
*)
(*$R
let q = of_seq (Sequence.of_list [1;2;3;4]) in
let n = fold (+) 0 q in
OUnit.assert_equal 10 n;
*)
let iter f q = to_seq q f
let of_list l = List.fold_left snoc empty l
let to_list q =
let l = ref [] in
to_seq q (fun x -> l := x :: !l);
List.rev !l
let of_seq seq = add_seq_front seq empty
(*$Q
(Q.list Q.int) (fun l -> \
Sequence.of_list l |> of_seq |> to_list = l)
*)
let rev q =
let q' = ref empty in
iter (fun x -> q' := cons x !q') q;
!q'
(*$Q
(Q.list Q.int) (fun l -> \
of_list l |> rev |> to_list = List.rev l)
*)
let _nil () = `Nil
let _single x cont () = `Cons (x, cont)
let _double x y cont () = `Cons (x, _single y cont)
let _triple x y z cont () = `Cons (x, _double y z cont)
let _digit_to_klist d cont = match d with
| Zero -> _nil
| One x -> _single x cont
| Two (x,y) -> _double x y cont
| Three (x,y,z) -> _triple x y z cont
let rec _flat_klist : 'a. ('a * 'a) klist -> 'a klist -> 'a klist
= fun l cont () -> match l () with
| `Nil -> cont ()
| `Cons ((x,y),l') -> _double x y (_flat_klist l' cont) ()
let to_klist q =
let rec aux : 'a. 'a t -> 'a klist -> 'a klist
= fun q cont () -> match q with
| Shallow d -> _digit_to_klist d cont ()
| Deep (_, hd, lazy q', tl) ->
_digit_to_klist hd
(_flat_klist
(aux q' _nil)
(_digit_to_klist tl cont))
()
in
aux q _nil
let of_klist l =
let rec seq l k = match l() with
| `Nil -> ()
| `Cons(x,l') -> k x; seq l' k
in
add_seq_front (seq l) empty
let rec _equal_klist eq l1 l2 = match l1(), l2() with
| `Nil, `Nil -> true
| `Nil, _
| _, `Nil -> false
| `Cons(x1,l1'), `Cons(x2,l2') ->
eq x1 x2 && _equal_klist eq l1' l2'
let equal eq q1 q2 = _equal_klist eq (to_klist q1) (to_klist q2)
(*$T
let q1 = 1 -- 10 and q2 = append (1 -- 5) (6 -- 10) in \
equal (=) q1 q2
*)
let (--) a b =
let rec up_to q a b = if a = b
then snoc q a
else up_to (snoc q a) (a+1) b
and down_to q a b = if a = b then snoc q a
else down_to (snoc q a) (a-1) b
in
if a <= b then up_to empty a b else down_to empty a b
(*$T
1 -- 5 |> to_list = [1;2;3;4;5]
5 -- 1 |> to_list = [5;4;3;2;1]
0 -- 0 |> to_list = [0]
*)
let (--^) a b =
if a=b then empty
else if a<b then a -- (b-1)
else a -- (b+1)
(*$T
1 --^ 5 |> to_list = [1;2;3;4]
5 --^ 1 |> to_list = [5;4;3;2]
1 --^ 2 |> to_list = [1]
0 --^ 0 |> to_list = []
*)
let print pp_x out d =
let first = ref true in
Format.fprintf out "@[<hov2>queue {";
iter
(fun x ->
if !first then first:= false else Format.fprintf out ";@ ";
pp_x out x
) d;
Format.fprintf out "}@]"