diff --git a/src/data/CCWBTree.ml b/src/data/CCWBTree.ml index 27d47a4e..5139d5d3 100644 --- a/src/data/CCWBTree.ml +++ b/src/data/CCWBTree.ml @@ -82,6 +82,16 @@ module type S = sig val nth_exn : int -> 'a t -> key * 'a (** @raise Not_found if the index is invalid *) + val get_rank : key -> 'a t -> int option + (** [get_rank k m] looks for the rank of [k] in [m], i.e. the index + of [k] in the sorted list of bindings of [m]. + [nth_exn (get_rank k m |> Opt.get_exn) m = get m k] should hold. + @since NEXT_RELEASE *) + + val get_rank_exn : key -> 'a t -> int + (** Unsafe version of {!get_rank} + @since NEXT_RELEASE *) + val add : key -> 'a -> 'a t -> 'a t val remove : key -> 'a t -> 'a t @@ -98,8 +108,14 @@ module type S = sig val fold : f:('b -> key -> 'a -> 'b) -> x:'b -> 'a t -> 'b val mapi : f:(key -> 'a -> 'b) -> 'a t -> 'b t + (** Map values, giving both key and value. Will use {!WORD.of_list} to rebuild keys. + @since 0.17 + *) val map : f:('a -> 'b) -> 'a t -> 'b t + (** Map values, giving only the value. + @since 0.17 + *) val iter : f:(key -> 'a -> unit) -> 'a t -> unit @@ -365,6 +381,32 @@ module MakeFull(K : KEY) : S with type key = K.t = struct List.for_all (fun i -> M.nth_exn i m = (i,i)) CCList.(0--1000) *) + let get_rank_exn k m : int = + let rec aux i k m = match m with + | E -> raise Not_found + | N (k', _, l, r, _) -> + match K.compare k k' with + | 0 -> i + weight l + | n when n<0 -> aux i k l + | _ -> aux (1 + weight l + i) k r + in + aux 0 k m + + let get_rank k m : int option = + try Some (get_rank_exn k m) + with Not_found -> None + + (*$QR & ~count:1_000 + Q.(list_of_size Gen.(0 -- 30) (pair small_int small_int)) (fun l -> + let l = CCList.sort_uniq ~cmp:(CCFun.compose_binop fst compare) l in + let m = M.of_list l in + List.for_all + (fun (k,v) -> match M.get_rank k m with + | None -> true + | Some n -> (k,v) = M.nth_exn n m) + l) + *) + let rec fold ~f ~x:acc m = match m with | E -> acc | N (k, v, l, r, _) -> diff --git a/src/data/CCWBTree.mli b/src/data/CCWBTree.mli index c3c85ef8..0f59dc21 100644 --- a/src/data/CCWBTree.mli +++ b/src/data/CCWBTree.mli @@ -47,6 +47,16 @@ module type S = sig val nth_exn : int -> 'a t -> key * 'a (** @raise Not_found if the index is invalid *) + val get_rank : key -> 'a t -> int option + (** [get_rank k m] looks for the rank of [k] in [m], i.e. the index + of [k] in the sorted list of bindings of [m]. + [nth_exn (get_rank k m |> Opt.get_exn) m = get m k] should hold. + @since NEXT_RELEASE *) + + val get_rank_exn : key -> 'a t -> int + (** Unsafe version of {!get_rank} + @since NEXT_RELEASE *) + val add : key -> 'a -> 'a t -> 'a t val remove : key -> 'a t -> 'a t