diff --git a/src/data/CCRAL.ml b/src/data/CCRAL.ml index 7d59b04c..b5cf196f 100644 --- a/src/data/CCRAL.ml +++ b/src/data/CCRAL.ml @@ -47,10 +47,12 @@ let is_empty = function | Nil -> true | Cons _ -> false -let rec get l i = match l with +let rec get_exn l i = match l with | Nil -> raise (Invalid_argument "RAL.get: wrong index") | Cons (size,t, _) when i < size -> tree_lookup size t i - | Cons (size,_, l') -> get l' (i - size) + | Cons (size,_, l') -> get_exn l' (i - size) + +let get l i = try Some (get_exn l i) with Invalid_argument _ -> None let rec set l i v = match l with | Nil -> raise (Invalid_argument "RAL.set: wrong index") @@ -62,14 +64,14 @@ let rec set l i v = match l with l=[] || \ (let i = (abs i) mod (List.length l) in \ let ral = of_list l in let ral = set ral i v in \ - get ral i = v)) + get_exn ral i = v)) *) (*$Q & ~small:List.length Q.(list small_int) (fun l -> \ let l1 = of_list l in \ CCList.Idx.mapi (fun i x -> i,x) l \ - |> List.for_all (fun (i,x) -> get l1 i = x)) + |> List.for_all (fun (i,x) -> get_exn l1 i = x)) *) let cons x l = match l with @@ -79,6 +81,8 @@ let cons x l = match l with else Cons (1, Leaf x, l) | _ -> Cons (1, Leaf x, l) +let cons' l x = cons x l + let hd l = match l with | Nil -> raise (Invalid_argument "RAL.hd: empty list") | Cons (_, Leaf x, _) -> x @@ -126,6 +130,27 @@ let rec map f l = match l with | Nil -> Nil | Cons (i, t, tl) -> Cons (i, _map_tree f t, map f tl) +let mapi f l = + let rec aux f i l = match l with + | Nil -> Nil + | Cons (size, t, tl) -> Cons (size, aux_t f ~size i t, aux f (i+size) tl) + and aux_t f ~size i t = match t with + | Leaf x -> Leaf (f i x) + | Node (x, l, r) -> + let x = f i x in + let l = aux_t f ~size:(size/2) (i+1) l in + Node (x, l, aux_t f ~size:(size/2) (i+1+size/2) r) + in + aux f 0 l + +(*$QR + Q.small_int (fun n -> + let l = CCList.(0 -- n) in + let l' = of_list l |> mapi (fun i x ->i,x) in + List.mapi (fun i x->i,x) l = to_list l' + ) +*) + let rec length l = match l with | Nil -> 0 | Cons (size,_, l') -> size + length l' @@ -164,7 +189,15 @@ and fold_tree_rev t acc f = match t with let acc = fold_tree_rev t1 acc f in f acc x -let rev l = fold (fun acc x -> cons x acc) empty l +let rev_map f l = fold (fun acc x -> cons (f x) acc) empty l + +(*$Q + Q.(list int) (fun l -> \ + let f x = x+1 in \ + of_list l |> rev_map f |> to_list = List.rev_map f l) +*) + +let rev l = fold cons' empty l (*$Q Q.(list small_int) (fun l -> \ @@ -180,6 +213,8 @@ let append l1 l2 = fold_rev (fun l2 x -> cons x l2) l2 l1 append (of_list l1) (of_list l2) = of_list (l1 @ l2)) *) +let append_tree_ t l = fold_tree_rev t l cons' + let filter p l = fold_rev (fun acc x -> if p x then cons x acc else acc) empty l let filter_map f l = @@ -220,6 +255,122 @@ let app funs l = [3; 12; 10; 100] *) +type 'a stack = + | St_nil + | St_list of 'a t * 'a stack + | St_tree of 'a tree * 'a stack + +let rec stack_to_list = function + | St_nil -> Nil + | St_list (l, st') -> append l (stack_to_list st') + | St_tree (t, st') -> append_tree_ t (stack_to_list st') + +let rec take n l = match l with + | Nil -> Nil + | Cons (size, t, tl) -> + if size <= n + then append_tree_ t (take (n-size) tl) + else take_tree_ ~size n t +and take_tree_ ~size n t = match t with + | _ when n=0 -> Nil + | Leaf x -> cons x Nil + | Node (x, l, r) -> + let size' = size/2 in + if size' <= n-1 + then cons x (append_tree_ l (take_tree_ ~size:size' (n-size'-1) r)) + else cons x (take_tree_ ~size:size' (n-1) l) + +(*$T + take 3 (of_list CCList.(1--10)) |> to_list = [1;2;3] + take 5 (of_list CCList.(1--10)) |> to_list = [1;2;3;4;5] + take 0 (of_list CCList.(1--10)) |> to_list = [] +*) + +let take_while p l = + (* st: stack of subtrees *) + let rec aux p st = match st with + | St_nil -> Nil + | St_list (Nil, st') -> aux p st' + | St_list (Cons (_, t, tl), st') -> aux p (St_tree (t, St_list (tl, st'))) + | St_tree (Leaf x, st') -> + if p x then cons x (aux p st') else Nil + | St_tree (Node (x,l,r), st') -> + if p x then cons x (aux p (St_tree (l, St_tree (r, st')))) else Nil + in aux p (St_list (l, St_nil)) + +(*$Q + Q.(list int) (fun l -> \ + let f x = x mod 7 <> 0 in \ + of_list l |> take_while f |> to_list = CCList.take_while f l) +*) + +let rec drop n l = match l with + | _ when n=0 -> l + | Nil -> Nil + | Cons (size, t, tl) -> + if n >= size then drop (n-size) tl + else drop_tree_ ~size n t tl +and drop_tree_ ~size n t tail = match t with + | _ when n=0 -> tail + | Leaf _ -> tail + | Node (_,l,r) -> + if n=1 then append_tree_ l (append_tree_ r tail) + else + let size' = size/2 in + if n-1 < size' + then drop_tree_ ~size:size' (n-1) l (append_tree_ r tail) + else drop_tree_ ~size:size' (n-1-size') r tail + +let drop_while p l = + let rec aux p st = match st with + | St_nil -> Nil + | St_list (Nil, st') -> aux p st' + | St_list (Cons (_, t, tail), st') -> + aux p (St_tree (t, St_list (tail, st'))) + | St_tree (Leaf x, st') -> + if p x then aux p st' else cons x (stack_to_list st') + | St_tree (Node (x,l,r) as tree, st') -> + if p x + then aux p (St_tree (l, St_tree (r, st'))) + else append_tree_ tree (stack_to_list st') + in aux p (St_list (l, St_nil)) + +(*$T + drop 3 (of_list CCList.(1--10)) |> to_list = CCList.(4--10) + drop 5 (of_list CCList.(1--10)) |> to_list = [6;7;8;9;10] + drop 0 (of_list CCList.(1--10)) |> to_list = CCList.(1--10) + drop 15 (of_list CCList.(1--10)) |> to_list = [] +*) + +(*$Q + Q.(list_of_size Gen.(0 -- 200) int) (fun l -> \ + let f x = x mod 10 <> 0 in \ + of_list l |> drop_while f |> to_list = CCList.drop_while f l) +*) + +let take_drop n l = take n l, drop n l + +let equal ?(eq=(=)) l1 l2 = + let rec aux ~eq l1 l2 = match l1, l2 with + | Nil, Nil -> true + | Cons (size1, t1, l1'), Cons (size2, t2, l2') -> + size1 = size2 && aux_t ~eq t1 t2 && aux ~eq l1' l2' + | Nil, Cons _ + | Cons _, Nil -> false + and aux_t ~eq t1 t2 = match t1, t2 with + | Leaf x, Leaf y -> eq x y + | Node (x1, l1, r1), Node (x2, l2, r2) -> + eq x1 x2 && aux_t ~eq l1 l2 && aux_t ~eq r1 r2 + | Leaf _, Node _ + | Node _, Leaf _ -> false + in + aux ~eq l1 l2 + +(*$Q + Q.(pair (list int)(list int)) (fun (l1,l2) -> \ + equal (of_list l1) (of_list l2) = (l1=l2)) +*) + (** {2 Conversions} *) type 'a sequence = ('a -> unit) -> unit @@ -306,9 +457,27 @@ let rec of_list_map f l = match l with let y = f x in cons y (of_list_map f l') +let compare ?(cmp=Pervasives.compare) l1 l2 = + let rec cmp_gen ~cmp g1 g2 = match g1(), g2() with + | None, None -> 0 + | Some _, None -> 1 + | None, Some _ -> -1 + | Some x, Some y -> + let c = cmp x y in + if c<> 0 then c else cmp_gen ~cmp g1 g2 + in + cmp_gen ~cmp (to_gen l1)(to_gen l2) + +(*$Q + Q.(pair (list int)(list int)) (fun (l1,l2) -> \ + compare (of_list l1) (of_list l2) = (Pervasives.compare l1 l2)) +*) + (** {2 Infix} *) module Infix = struct + let (@+) = cons + let (>>=) l f = flat_map f l let (>|=) l f = map f l let (<*>) = app diff --git a/src/data/CCRAL.mli b/src/data/CCRAL.mli index 5c52422e..cb7c4cdb 100644 --- a/src/data/CCRAL.mli +++ b/src/data/CCRAL.mli @@ -32,6 +32,9 @@ val return : 'a -> 'a t val map : ('a -> 'b) -> 'a t -> 'b t (** Map on elements *) +val mapi : (int -> 'a -> 'b) -> 'a t -> 'b t +(** Map with index *) + val hd : 'a t -> 'a (** First element of the list, or @raise Invalid_argument if the list is empty *) @@ -50,8 +53,11 @@ val front_exn : 'a t -> 'a * 'a t val length : 'a t -> int (** Number of elements *) -val get : 'a t -> int -> 'a -(** [get l i] accesses the [i]-th element of the list. O(log(n)). +val get : 'a t -> int -> 'a option +(** [get l i] accesses the [i]-th element of the list. O(log(n)). *) + +val get_exn : 'a t -> int -> 'a +(** Unsafe version of {!get} @raise Invalid_argument if the list has less than [i+1] elements. *) val set : 'a t -> int -> 'a -> 'a t @@ -74,6 +80,18 @@ val flatten : 'a t t -> 'a t val app : ('a -> 'b) t -> 'a t -> 'b t +val take : int -> 'a t -> 'a t + +val take_while : ('a -> bool) -> 'a t -> 'a t + +val drop : int -> 'a t -> 'a t + +val drop_while : ('a -> bool) -> 'a t -> 'a t + +val take_drop : int -> 'a t -> 'a t * 'a t +(** [take_drop n l] splits [l] into [a, b] such that [length a = n] + if [length l >= n], and such that [append a b = l] *) + val iter : ('a -> unit) -> 'a t -> unit (** Iterate on the list's elements *) @@ -83,9 +101,16 @@ val fold : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b val fold_rev : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b (** Fold on the list's elements, in reverse order (starting from the tail) *) +val rev_map : ('a -> 'b) -> 'a t -> 'b t +(** [rev_map f l] is the same as [map f (rev l)] *) + val rev : 'a t -> 'a t (** Reverse the list *) +val equal : ?eq:('a -> 'a -> bool) -> 'a t -> 'a t -> bool + +val compare : ?cmp:('a -> 'a -> int) -> 'a t -> 'a t -> int + (** {2 Conversions} *) type 'a sequence = ('a -> unit) -> unit @@ -116,6 +141,10 @@ val to_gen : 'a t -> 'a gen (** {2 Infix} *) module Infix : sig + val (@+) : 'a -> 'a t -> 'a t + (** Cons (alias to {!cons}) + @since NEXT_RELEASE *) + val (>>=) : 'a t -> ('a -> 'b t) -> 'b t val (>|=) : 'a t -> ('a -> 'b) -> 'b t val (<*>) : ('a -> 'b) t -> 'a t -> 'b t