Merge pull request #3 from nbraud/avl

AVL: Change API (why not, it makes for a usable polymorphic map)

misc/AVL.ml

@@ -28,13 +28,17 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 See https://en.wikipedia.org/wiki/AVL_tree *)
 
-type ('a,'b) t =
-  | Empty
-  | Node of ('a,'b) t * 'a * 'b * ('a,'b) t * int
-
 type 'a comparator = 'a -> 'a -> int
 
-let empty = Empty
+type ('a,'b) tree =
+  | Empty
+  | Node of ('a,'b) tree * 'a * 'b * ('a,'b) tree * int
+
+type ('a,'b) t =
+	{ cmp: 'a comparator;
+	    t: ('a,'b) tree }
+
+let empty ~cmp = { cmp; t=Empty }
 
 let _height = function
   | Empty -> 0
@@ -46,7 +50,8 @@ let _balance l r = _height l - _height r
 let _make l x y r =
   Node (l, x, y, r, 1 + max (_height l) (_height r))
 
-let singleton k v = _make empty k v empty
+let _singleton k v = _make Empty k v Empty
+let singleton ~cmp k v = { cmp; t = _singleton k v }
 
 (* balance tree [t] *)
 let _rebalance t = match t with
@@ -93,32 +98,40 @@ let _rebalance t = match t with
 let _make_balance l k v r =
   _rebalance (_make l k v r)
 
-let rec fold f acc t = match t with
+let rec _fold f acc t = match t with
   | Empty -> acc
   | Node (l, x, y, r, _) ->
-      let acc = fold f acc l in
+      let acc = _fold f acc l in
       let acc = f acc x y in
-      fold f acc r
+      _fold f acc r
 
-let rec for_all p t = match t with
+let fold f acc {t; _} = _fold f acc t
+
+let rec _for_all p t = match t with
   | Empty -> true
   | Node (l, x, y, r, _) ->
-      p x y && for_all p l && for_all p r
+      p x y && _for_all p l && _for_all p r
 
-let rec exists p t = match t with
+let for_all p {t; _} = _for_all p t
+
+let rec _exists p t = match t with
   | Empty -> false
   | Node (l, x, y, r, _) ->
-      p x y || exists p l || exists p r
+      p x y || _exists p l || _exists p r
 
-let rec insert ~cmp t k v = match t with
+let exists p {t; _} = _exists p t
-  | Empty -> _make empty k v empty
+
+let rec _insert ~cmp t k v = match t with
+  | Empty -> _make Empty k v Empty
   | Node (l, k1, v1, r, _) ->
       let c = cmp k k1 in
       if c < 0
-        then _make_balance (insert ~cmp l k v) k1 v1 r
+        then _make_balance (_insert ~cmp l k v) k1 v1 r
       else if c = 0
         then _make l k v r
-      else _make_balance l k1 v1 (insert ~cmp r k v)
+      else _make_balance l k1 v1 (_insert ~cmp r k v)
+
+let insert {cmp; t} k v = {cmp; t=_insert ~cmp t k v}
 
 (* remove the maximal value in the given tree (the only which only has a left
   child), and return its key/value pair *)
@@ -132,7 +145,7 @@ let rec _remove_max t = match t with
 
 exception NoSuchElement
 
-let remove ~cmp t key =
+let _remove ~cmp t key =
   let rec _remove t = match t with
     | Empty -> raise NoSuchElement
     | Node (l, k, v, r, _) ->
@@ -153,30 +166,34 @@ let remove ~cmp t key =
   try _remove t
   with NoSuchElement -> t  (* element not found *)
 
-let update ~cmp t key f = failwith "update: not implemented"
+let remove {cmp; t} k = {cmp; t=_remove ~cmp t k}
 
-let rec find_exn ~cmp t key = match t with
+let _update ~cmp t key f = failwith "update: not implemented"
+let update {cmp; t} = _update ~cmp t
+
+let rec _find_exn ~cmp t key = match t with
   | Empty -> raise Not_found
   | Node (l, k, v, r, _) ->
       let c = cmp key k in
-      if c < 0 then find_exn ~cmp l key
+      if c < 0 then _find_exn ~cmp l key
-      else if c > 0 then find_exn ~cmp r key
+      else if c > 0 then _find_exn ~cmp r key
       else v
+let find_exn {cmp; t} = _find_exn ~cmp t
 
-let find ~cmp t key =
+let find t key =
-  try Some (find_exn ~cmp t key)
+  try Some (find_exn t key)
   with Not_found -> None
 
 (* add k,v as strictly maximal element to t. [t] must not contain
   any key >= k *)
 let rec _add_max k v t = match t with
-  | Empty -> singleton k v
+  | Empty -> _singleton k v
   | Node (l, k', v', r, _) ->
       _make_balance l k' v' (_add_max k v r)
-and
+
 (* same for minimal value *)
-_add_min k v t = match t with
+let rec _add_min k v t = match t with
-  | Empty -> singleton k v
+  | Empty -> _singleton k v
   | Node (l, k', v', r, _) ->
       _make_balance (_add_min k v l) k' v' r
 
@@ -204,42 +221,51 @@ let _concat t1 t2 = match t1, t2 with
       let t1', k, v = _remove_max t1 in
       _join t1' k v t2
 
-let rec split ~cmp t key = match t with
+let rec _split ~cmp t key = match t with
-  | Empty -> empty, None, empty
+  | Empty -> Empty, None, Empty
   | Node (l, k, v, r, _) ->
       let c = cmp key k in
       if c < 0
         then
-          let ll, result, lr = split ~cmp l key in
+          let ll, result, lr = _split ~cmp l key in
           ll, result, _join lr k v r
       else if c > 0
         then
-          let rl, result, rr = split ~cmp r key in
+          let rl, result, rr = _split ~cmp r key in
           _join l k v rl, result, rr
       else
         l, Some v, r
 
+let split {cmp; t} k =
+  let (t,b,t') = _split ~cmp t k in
+  {cmp; t}, b, {cmp; t=t'}
+
 (* if k = Some v, join l k v r, else concat l v *)
 let _concat_or_join l k result r = match result with
   | None -> _concat l r
   | Some v -> _join l k v r
 
-let rec merge ~cmp f t1 t2 = match t1, t2 with
+let rec _merge ~cmp f t1 t2 = match t1, t2 with
-  | Empty, Empty -> empty
+  | Empty, Empty -> Empty
   | Node (l1, k1, v1, r1, h1), _ when h1 >= _height t2 ->
-      let l2, result2, r2 = split ~cmp t2 k1 in
+      let l2, result2, r2 = _split ~cmp t2 k1 in
       let result = f k1 (Some v1) result2 in
-      let l = merge ~cmp f l1 l2 in
+      let l = _merge ~cmp f l1 l2 in
-      let r = merge ~cmp f r1 r2 in
+      let r = _merge ~cmp f r1 r2 in
       _concat_or_join l k1 result r
   | _, Node (l2, k2, v2, r2, _) ->
-      let l1, result1, r1 = split ~cmp t1 k2 in
+      let l1, result1, r1 = _split ~cmp t1 k2 in
       let result = f k2 result1 (Some v2) in
-      let l = merge ~cmp f l1 l2 in
+      let l = _merge ~cmp f l1 l2 in
-      let r = merge ~cmp f r1 r2 in
+      let r = _merge ~cmp f r1 r2 in
       _concat_or_join l k2 result r
   | _, Empty -> assert false  (* h1 < heigth h2?? *)
 
+let merge f {cmp; t} {cmp=cmp'; t=t'} =
+  if(cmp != cmp') then invalid_arg "AVL.merge: trees wit different
+                                               comparison function";
+  {cmp; t = _merge ~cmp f t t'}
+
 (* invariant: balanced *)
 let rec invariant_balance t = match t with
   | Empty -> true
@@ -253,13 +279,13 @@ let rec invariant_search ~cmp t = match t with
   | Node (l, x, _, r, _) ->
       invariant_search ~cmp l &&
       invariant_search ~cmp r &&
-      for_all (fun x' _ -> cmp x' x < 0) l &&
+      _for_all (fun x' _ -> cmp x' x < 0) l &&
-      for_all (fun x' _ -> cmp x' x > 0) r
+      _for_all (fun x' _ -> cmp x' x > 0) r
 
 let of_list ~cmp l =
-  List.fold_left (fun acc (x,y) -> insert ~cmp acc x y) empty l
+  {cmp; t = List.fold_left (fun acc (x,y) -> _insert ~cmp acc x y) Empty l}
 
-let to_list t =
+let to_list {t; _} =
   let rec aux acc t = match t with
   | Empty -> acc
   | Node (l, k, v, r, _) ->
@@ -273,15 +299,15 @@ let to_list t =
 module type ITERATOR = sig
   type 'a iter
 
-  val after : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
+  val after  : ('a,'b) t -> 'a -> ('a * 'b) iter
-  val before : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
+  val before : ('a,'b) t -> 'a -> ('a * 'b) iter
   val iter   : ('a,'b) t -> ('a * 'b) iter
-  val add : cmp:'a comparator -> ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
+  val add    : ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
 end
 
 type ('a,'b) explore =
   | Yield of 'a * 'b
-  | Explore of ('a, 'b) t
+  | Explore of ('a, 'b) tree
 
 exception EndOfIter
 
@@ -338,16 +364,18 @@ module KList = struct
         let k, v, l' = _pop l in
         `Cons ((k,v), _next l')
 
-  let iter t = _next (_push t []) ()
+  let iter {t; _} = _next (_push t []) ()
 
-  let rec add ~cmp t (l:'a t) = match l with
+  let rec _add ~cmp t (l:'a t) = match l with
     | `Nil -> t
     | `Cons ((k,v), l') ->
-        add ~cmp (insert ~cmp t k v) (l' ())
+        _add ~cmp (_insert ~cmp t k v) (l' ())
 
-  let after ~cmp t key = _next (_after ~cmp [] t key) ()
+  let add {cmp; t} l = {cmp; t=_add ~cmp t l}
 
-  let before ~cmp t key = _next (_before ~cmp [] t key) ()
+  let after {cmp; t} key = _next (_after ~cmp [] t key) ()
+
+  let before {cmp; t} key = _next (_before ~cmp [] t key) ()
 end
 
 module Gen = struct
@@ -355,7 +383,7 @@ module Gen = struct
 
   let _gen stack =
     let stack = ref stack in
-    let rec next () =
+    let next () =
       match !stack with
       | [] -> None
       | l ->
@@ -364,13 +392,15 @@ module Gen = struct
           Some (k, v)
     in next
 
-  let iter t = _gen (_push t [])
+  let iter {t; _} = _gen (_push t [])
 
-  let rec add ~cmp t gen =
+  let rec _add ~cmp t gen =
     match gen() with
     | None -> t
-    | Some (k,v) -> add ~cmp (insert ~cmp t k v) gen
+    | Some (k,v) -> _add ~cmp (_insert ~cmp t k v) gen
 
-  let after ~cmp t key = _gen (_after ~cmp [] t key)
+  let add {cmp; t} l = {cmp; t=_add ~cmp t l}
-  let before ~cmp t key = _gen (_before ~cmp [] t key)
+
+  let after {cmp; t} key = _gen (_after ~cmp [] t key)
+  let before {cmp; t} key = _gen (_before ~cmp [] t key)
 end

misc/AVL.mli

@@ -26,16 +26,20 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 (** {1 AVL trees} *)
 
-type ('a,'b) t = private
-  | Empty
-  | Node of ('a,'b) t * 'a * 'b * ('a,'b) t * int
-
 type 'a comparator = 'a -> 'a -> int
 
-val empty : ('a,'b) t
+type ('a,'b) tree = private
+  | Empty
+  | Node of ('a,'b) tree * 'a * 'b * ('a,'b) tree * int
+
+type ('a,'b) t = private
+	{ cmp: 'a comparator;
+	    t: ('a,'b) tree }
+
+val empty : cmp:'a comparator -> ('a,'b) t
 (** Empty tree *)
 
-val singleton : 'a -> 'b -> ('a,'b) t
+val singleton : cmp:'a comparator -> 'a -> 'b -> ('a,'b) t
 (** Tree with a single node *)
 
 val fold : ('c -> 'a -> 'b -> 'c) -> 'c -> ('a,'b) t -> 'c
@@ -44,29 +48,29 @@ val fold : ('c -> 'a -> 'b -> 'c) -> 'c -> ('a,'b) t -> 'c
 val for_all : ('a -> 'b -> bool) -> ('a,'b) t -> bool
 val exists : ('a -> 'b -> bool) -> ('a,'b) t -> bool
 
-val find : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b option
+val find : ('a,'b) t -> 'a -> 'b option
 (** Find the value associated to the key, if any *)
 
-val find_exn : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b
+val find_exn : ('a,'b) t -> 'a -> 'b
 (** @raise Not_found if the key is not present *)
 
-val insert : cmp:'a comparator -> ('a,'b) t -> 'a -> 'b -> ('a,'b) t
+val insert : ('a,'b) t -> 'a -> 'b -> ('a,'b) t
 (** Insertion in the tree *)
 
-val remove : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a,'b) t
+val remove : ('a,'b) t -> 'a -> ('a,'b) t
 (** Removal from the tree *)
 
-val update : cmp:'a comparator -> ('a,'b) t -> 'a ->
+val update : ('a,'b) t -> 'a ->
             ('b option -> ('a * 'b) option) -> ('a,'b) t
 (** Update of the given key binding (subsumes [insert] and [remove]) *)
 
-val split : cmp:'a comparator -> ('a,'b) t -> 'a ->
+val split : ('a,'b) t -> 'a ->
             ('a,'b) t * 'b option * ('a,'b) t
 (** [split ~cmp t k] splits [t] into a left part that
     is smaller than [k], the possible binding of [k],
     and a part bigger than [k]. *)
 
-val merge : cmp:'a comparator ->
+val merge :
             ('a -> 'b option -> 'c option -> 'd option) ->
             ('a,'b) t -> ('a,'c) t -> ('a,'d) t
 (** Merge two trees together, with the given function *)
@@ -82,10 +86,10 @@ val to_list : ('a,'b) t -> ('a * 'b) list
 module type ITERATOR = sig
   type 'a iter
 
-  val after : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
+  val after : ('a,'b) t -> 'a -> ('a * 'b) iter
-  val before : cmp:'a comparator -> ('a,'b) t -> 'a -> ('a * 'b) iter
+  val before : ('a,'b) t -> 'a -> ('a * 'b) iter
   val iter : ('a,'b) t -> ('a * 'b) iter
-  val add : cmp:'a comparator -> ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
+  val add : ('a,'b) t -> ('a * 'b) iter -> ('a,'b) t
 end
 
 module KList : sig