From 62d76ce5557a6a9500d638671c1a5076bcda3e34 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 11 Aug 2017 00:41:26 +0200
Subject: [PATCH] change signature of `CCWBTree.get_rank`

---
 src/data/CCWBTree.ml  | 22 +++++++---------------
 src/data/CCWBTree.mli |  8 ++------
 2 files changed, 9 insertions(+), 21 deletions(-)

diff --git a/src/data/CCWBTree.ml b/src/data/CCWBTree.ml
index 5139d5d3..7277707e 100644
--- a/src/data/CCWBTree.ml
+++ b/src/data/CCWBTree.ml
@@ -82,14 +82,10 @@ 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
+  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
       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}
+      [let (`At n) = get_rank k m in nth_exn n m = get m k] should hold.
       @since NEXT_RELEASE *)
 
   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)
   *)
 
-  let get_rank_exn k m : int =
+  let get_rank k m =
     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, _) ->
         match K.compare k k' with
-          | 0 -> i + weight l
+          | 0 -> `At (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)
+          | `First | `After _ -> true
+          | `At n -> (k,v) = M.nth_exn n m)
         l)
   *)
 
diff --git a/src/data/CCWBTree.mli b/src/data/CCWBTree.mli
index 0f59dc21..40f89b55 100644
--- a/src/data/CCWBTree.mli
+++ b/src/data/CCWBTree.mli
@@ -47,14 +47,10 @@ 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
+  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
       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}
+      [let (`At n) = get_rank k m in nth_exn n m = get m k] should hold.
       @since NEXT_RELEASE *)
 
   val add : key -> 'a -> 'a t -> 'a t