1 files changed, 150 insertions, 23 deletions

diff --git a/stdlib/set.ml b/stdlib/set.ml
index 404056308..84a8a942c 100644
--- a/stdlib/set.ml
+++ b/stdlib/set.ml
@@ -23,59 +23,182 @@ module type S =
     val iter: (elt -> 'a) -> t -> unit
     val fold: (elt -> 'a -> 'a) -> t -> 'a -> 'a
     val elements: t -> elt list
+    val choose: t -> elt
   end
 
 module Make(Ord: OrderedType): (S with elt = Ord.t) =
   struct
-    open Baltree
     type elt = Ord.t
-    type t = elt Baltree.t
+    type t = Empty | Node of t * elt * t * int
+
+    (* Sets are represented by balanced binary trees (the heights of the
+       children differ by at most 2 *)
+
+    let height = function
+        Empty -> 0
+      | Node(_, _, _, h) -> h
+
+    (* Creates a new node with left son l, value x and right son r.
+       l and r must be balanced and | height l - height r | <= 2.
+       Inline expansion of height for better speed. *)
+
+    let new l x r =
+      let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
+      let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
+      Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))
+
+    (* Same as new, but performs one step of rebalancing if necessary.
+       Assumes l and r balanced.
+       Inline expansion of new for better speed in the most frequent case
+       where no rebalancing is required. *)
+
+    let bal l x r =
+      let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
+      let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
+      if hl > hr + 2 then begin
+        match l with
+          Empty -> invalid_arg "Set.bal"
+        | Node(ll, lv, lr, _) ->
+            if height ll >= height lr then
+              new ll lv (new lr x r)
+            else begin
+              match lr with
+                Empty -> invalid_arg "Set.bal"
+              | Node(lrl, lrv, lrr, _)->
+                  new (new ll lv lrl) lrv (new lrr x r)
+            end
+      end else if hr > hl + 2 then begin
+        match r with
+          Empty -> invalid_arg "Set.bal"
+        | Node(rl, rv, rr, _) ->
+            if height rr >= height rl then
+              new (new l x rl) rv rr
+            else begin
+              match rl with
+                Empty -> invalid_arg "Set.bal"
+              | Node(rll, rlv, rlr, _) ->
+                  new (new l x rll) rlv (new rlr rv rr)
+            end
+      end else
+        Node(l, x, r, (if hl >= hr then hl + 1 else hr + 1))
+
+    (* Same as bal, but repeat rebalancing until the final result
+       is balanced. *)
+
+    let rec join l x r =
+      match bal l x r with
+        Empty -> invalid_arg "Set.join"
+      | Node(l', x', r', _) as t' ->
+          let d = height l' - height r' in
+          if d < -2 or d > 2 then join l' x' r' else t'
+
+    (* Merge two trees l and r into one.
+       All elements of l must precede the elements of r.
+       Assumes | height l - height r | <= 2. *)
+
+    let rec merge t1 t2 =
+      match (t1, t2) with
+        (Empty, t) -> t
+      | (t, Empty) -> t
+      | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
+          bal l1 v1 (bal (merge r1 l2) v2 r2)
+
+    (* Same as merge, but does not assume anything about l and r. *)
+
+    let rec concat t1 t2 =
+      match (t1, t2) with
+        (Empty, t) -> t
+      | (t, Empty) -> t
+      | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
+          join l1 v1 (join (concat r1 l2) v2 r2)
+
+    (* Splitting *)
+
+    let rec split x = function
+        Empty ->
+          (Empty, None, Empty)
+      | Node(l, v, r, _) ->
+          let c = Ord.compare x v in
+          if c = 0 then (l, Some v, r)
+          else if c < 0 then
+            let (ll, vl, rl) = split x l in (ll, vl, join rl v r)
+          else
+            let (lr, vr, rr) = split x r in (join l v lr, vr, rr)
+
+    (* Implementation of the set operations *)
 
     let empty = Empty
 
     let is_empty = function Empty -> true | _ -> false
 
-    let mem x s =
-      Baltree.contains (Ord.compare x) s
-
-    let add x s =
-      Baltree.add (Ord.compare x) x s
-
-    let remove x s =
-      Baltree.remove (Ord.compare x) s
+    let rec mem x = function
+        Empty -> false
+      | Node(l, v, r, _) ->
+          let c = Ord.compare x v in
+          if c = 0 then true else
+          if c < 0 then mem x l else mem x r
+
+    let rec add x = function
+        Empty -> Node(Empty, x, Empty, 1)
+      | Node(l, v, r, _) as t ->
+          let c = Ord.compare x v in
+          if c = 0 then t else
+          if c < 0 then bal (add x l) v r else bal l v (add x r)
+
+    let rec remove x = function
+        Empty -> Empty
+      | Node(l, v, r, _) ->
+          let c = Ord.compare x v in
+          if c = 0 then merge l r else
+          if c < 0 then bal (remove x l) v r else bal l v (remove x r)
 
     let rec union s1 s2 =
       match (s1, s2) with
         (Empty, t2) -> t2
       | (t1, Empty) -> t1
       | (Node(l1, v1, r1, _), t2) ->
-          let (l2, _, r2) = Baltree.split (Ord.compare v1) t2 in
-          Baltree.join (union l1 l2) v1 (union r1 r2)
+          let (l2, _, r2) = split v1 t2 in
+          join (union l1 l2) v1 (union r1 r2)
 
     let rec inter s1 s2 =
       match (s1, s2) with
         (Empty, t2) -> Empty
       | (t1, Empty) -> Empty
       | (Node(l1, v1, r1, _), t2) ->
-          match Baltree.split (Ord.compare v1) t2 with
-            (l2, Nothing, r2) ->
-              Baltree.concat (inter l1 l2) (inter r1 r2)
-          | (l2, Something _, r2) ->
-              Baltree.join (inter l1 l2) v1 (inter r1 r2)
+          match split v1 t2 with
+            (l2, None, r2) ->
+              concat (inter l1 l2) (inter r1 r2)
+          | (l2, Some _, r2) ->
+              join (inter l1 l2) v1 (inter r1 r2)
 
     let rec diff s1 s2 =
       match (s1, s2) with
         (Empty, t2) -> Empty
       | (t1, Empty) -> t1
       | (Node(l1, v1, r1, _), t2) ->
-          match Baltree.split (Ord.compare v1) t2 with
-            (l2, Nothing, r2) ->
-              Baltree.join (diff l1 l2) v1 (diff r1 r2)
-          | (l2, Something _, r2) ->
-              Baltree.concat (diff l1 l2) (diff r1 r2)
+          match split v1 t2 with
+            (l2, None, r2) ->
+              join (diff l1 l2) v1 (diff r1 r2)
+          | (l2, Some _, r2) ->
+              concat (diff l1 l2) (diff r1 r2)
+
+    let rec compare_aux l1 l2 =
+        match (l1, l2) with
+        ([], []) -> 0
+      | ([], _)  -> -1
+      | (_, []) -> 1
+      | (Empty :: t1, Empty :: t2) ->
+          compare_aux t1 t2
+      | (Node(Empty, v1, r1, _) :: t1, Node(Empty, v2, r2, _) :: t2) ->
+          let c = Ord.compare v1 v2 in
+          if c <> 0 then c else compare_aux (r1::t1) (r2::t2)
+      | (Node(l1, v1, r1, _) :: t1, t2) ->
+          compare_aux (l1 :: Node(Empty, v1, r1, 0) :: t1) t2
+      | (t1, Node(l2, v2, r2, _) :: t2) ->
+          compare_aux t1 (l2 :: Node(Empty, v2, r2, 0) :: t2)
 
     let compare s1 s2 =
-      Baltree.compare Ord.compare s1 s2
+      compare_aux [s1] [s2]
 
     let equal s1 s2 =
       compare s1 s2 = 0
@@ -96,4 +219,8 @@ module Make(Ord: OrderedType): (S with elt = Ord.t) =
     let elements s =
       elements_aux [] s
 
+    let rec choose = function
+        Empty -> raise Not_found
+      | Node(Empty, v, r, _) -> v
+      | Node(l, v, r, _) -> choose l
   end