add CCWBTree.get_rank{,_exn}

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, _) ->

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