change signature of CCWBTree.get_rank

This commit is contained in:
Simon Cruanes 2017-08-11 00:41:26 +02:00
parent cf09112f79
commit 62d76ce555
2 changed files with 9 additions and 21 deletions

View file

@ -82,14 +82,10 @@ module type S = sig
val nth_exn : int -> 'a t -> key * 'a val nth_exn : int -> 'a t -> key * 'a
(** @raise Not_found if the index is invalid *) (** @raise Not_found if the index is invalid *)
val get_rank : key -> 'a t -> int option val get_rank : key -> 'a t -> [`At of int | `After of int | `First]
(** [get_rank k m] looks for the rank of [k] in [m], i.e. the index (** [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]. 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. [let (`At n) = get_rank k m in nth_exn n 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 *) @since NEXT_RELEASE *)
val add : key -> 'a -> 'a t -> 'a t val add : key -> 'a -> 'a t -> 'a t
@ -381,29 +377,25 @@ 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) List.for_all (fun i -> M.nth_exn i m = (i,i)) CCList.(0--1000)
*) *)
let get_rank_exn k m : int = let get_rank k m =
let rec aux i k m = match m with let rec aux i k m = match m with
| E -> raise Not_found | E -> if i=0 then `First else `After i
| N (k', _, l, r, _) -> | N (k', _, l, r, _) ->
match K.compare k k' with match K.compare k k' with
| 0 -> i + weight l | 0 -> `At (i + weight l)
| n when n<0 -> aux i k l | n when n<0 -> aux i k l
| _ -> aux (1 + weight l + i) k r | _ -> aux (1 + weight l + i) k r
in in
aux 0 k m 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 (*$QR & ~count:1_000
Q.(list_of_size Gen.(0 -- 30) (pair small_int small_int)) (fun l -> 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 l = CCList.sort_uniq ~cmp:(CCFun.compose_binop fst compare) l in
let m = M.of_list l in let m = M.of_list l in
List.for_all List.for_all
(fun (k,v) -> match M.get_rank k m with (fun (k,v) -> match M.get_rank k m with
| None -> true | `First | `After _ -> true
| Some n -> (k,v) = M.nth_exn n m) | `At n -> (k,v) = M.nth_exn n m)
l) l)
*) *)

View file

@ -47,14 +47,10 @@ module type S = sig
val nth_exn : int -> 'a t -> key * 'a val nth_exn : int -> 'a t -> key * 'a
(** @raise Not_found if the index is invalid *) (** @raise Not_found if the index is invalid *)
val get_rank : key -> 'a t -> int option val get_rank : key -> 'a t -> [`At of int | `After of int | `First]
(** [get_rank k m] looks for the rank of [k] in [m], i.e. the index (** [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]. 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. [let (`At n) = get_rank k m in nth_exn n 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 *) @since NEXT_RELEASE *)
val add : key -> 'a -> 'a t -> 'a t val add : key -> 'a -> 'a t -> 'a t