Fast path for set operations.

We consider an approximation of the size of sets before choosing the most
appropriate algorithm. This drastically affects some universe-polymorphic
code which was doing a lot of set operations on disimilar sizes.

3 files changed, 89 insertions, 21 deletions

diff --git a/lib/cMap.ml b/lib/cMap.ml
index 4b058380c..ba0873ffa 100644
--- a/lib/cMap.ml
+++ b/lib/cMap.ml
@@ -34,6 +34,7 @@ sig
   val fold_right : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
   val smartmap : ('a -> 'a) -> 'a t -> 'a t
   val smartmapi : (key -> 'a -> 'a) -> 'a t -> 'a t
+  val height : 'a t -> int
   module Unsafe :
   sig
     val map : (key -> 'a -> key * 'b) -> 'a t -> 'b t
@@ -57,6 +58,7 @@ sig
   val fold_right : (M.t -> 'a -> 'b -> 'b) -> 'a map -> 'b -> 'b
   val smartmap : ('a -> 'a) -> 'a map -> 'a map
   val smartmapi : (M.t -> 'a -> 'a) -> 'a map -> 'a map
+  val height : 'a map -> int
   module Unsafe :
   sig
     val map : (M.t -> 'a -> M.t * 'b) -> 'a map -> 'b map
@@ -168,6 +170,10 @@ struct
     if l == l' && r == r' && v == v' then s
     else map_inj (MNode (l', k, v', r', h))
 
+  let height s = match map_prj s with
+  | MEmpty -> 0
+  | MNode (_, _, _, _, h) -> h
+
   module Unsafe =
   struct
 
diff --git a/lib/cMap.mli b/lib/cMap.mli
index 3ef7fa2c8..2838b374e 100644
--- a/lib/cMap.mli
+++ b/lib/cMap.mli
@@ -61,6 +61,9 @@ sig
   val smartmapi : (key -> 'a -> 'a) -> 'a t -> 'a t
   (** As [mapi] but tries to preserve sharing. *)
 
+  val height : 'a t -> int
+  (** An indication of the logarithmic size of a map *)
+
   module Unsafe :
   sig
     val map : (key -> 'a -> key * 'b) -> 'a t -> 'b t
diff --git a/lib/hMap.ml b/lib/hMap.ml
index 778c366fd..ea76e7427 100644
--- a/lib/hMap.ml
+++ b/lib/hMap.ml
@@ -68,34 +68,91 @@ struct
         Int.Map.update h m s
     with Not_found -> s
 
+  let height s = Int.Map.height s
+
+  let is_smaller s1 s2 = height s1 <= height s2 + 3
+
+  (** Assumes s1 << s2 *)
+  let fast_union s1 s2 =
+    let fold h s accu =
+      try Int.Map.modify h (fun _ s' -> Set.fold Set.add s s') accu
+      with Not_found -> Int.Map.add h s accu
+    in
+    Int.Map.fold fold s1 s2
+
   let union s1 s2 =
-    let fu _ m1 m2 = match m1, m2 with
-    | None, None -> None
-    | (Some _ as m), None | None, (Some _ as m) -> m
-    | Some m1, Some m2 -> Some (Set.union m1 m2)
+    if is_smaller s1 s2 then fast_union s1 s2
+    else if is_smaller s2 s1 then fast_union s2 s1
+    else
+      let fu _ m1 m2 = match m1, m2 with
+      | None, None -> None
+      | (Some _ as m), None | None, (Some _ as m) -> m
+      | Some m1, Some m2 -> Some (Set.union m1 m2)
+      in
+      Int.Map.merge fu s1 s2
+
+  (** Assumes s1 << s2 *)
+  let fast_inter s1 s2 =
+    let fold h s accu =
+      try
+        let s' = Int.Map.find h s2 in
+        let si = Set.filter (fun e -> Set.mem e s') s in
+        if Set.is_empty si then accu
+        else Int.Map.add h si accu
+     with Not_found -> accu
     in
-    Int.Map.merge fu s1 s2
+    Int.Map.fold fold s1 Int.Map.empty
 
   let inter s1 s2 =
-    let fu _ m1 m2 = match m1, m2 with
-    | None, None -> None
-    | Some _, None | None, Some _ -> None
-    | Some m1, Some m2 ->
-      let m = Set.inter m1 m2 in
-      if Set.is_empty m then None else Some m
+    if is_smaller s1 s2 then fast_inter s1 s2
+    else if is_smaller s2 s1 then fast_inter s2 s1
+    else
+      let fu _ m1 m2 = match m1, m2 with
+      | None, None -> None
+      | Some _, None | None, Some _ -> None
+      | Some m1, Some m2 ->
+        let m = Set.inter m1 m2 in
+        if Set.is_empty m then None else Some m
+      in
+      Int.Map.merge fu s1 s2
+
+  (** Assumes s1 << s2 *)
+  let fast_diff_l s1 s2 =
+    let fold h s accu =
+      try
+        let s' = Int.Map.find h s2 in
+        let si = Set.filter (fun e -> not (Set.mem e s')) s in
+        if Set.is_empty si then accu
+        else Int.Map.add h si accu
+     with Not_found -> Int.Map.add h s accu
+    in
+    Int.Map.fold fold s1 Int.Map.empty
+
+  (** Assumes s2 << s1 *)
+  let fast_diff_r s1 s2 =
+    let fold h s accu =
+      try
+        let s' = Int.Map.find h accu in
+        let si = Set.filter (fun e -> not (Set.mem e s)) s' in
+        if Set.is_empty si then Int.Map.remove h accu
+        else Int.Map.update h si accu
+     with Not_found -> accu
     in
-    Int.Map.merge fu s1 s2
+    Int.Map.fold fold s2 s1
 
   let diff s1 s2 =
-    let fu _ m1 m2 = match m1, m2 with
-    | None, None -> None
-    | (Some _ as m), None -> m
-    | None, Some _ -> None
-    | Some m1, Some m2 ->
-      let m = Set.diff m1 m2 in
-      if Set.is_empty m then None else Some m
-    in
-    Int.Map.merge fu s1 s2
+    if is_smaller s1 s2 then fast_diff_l s1 s2
+    else if is_smaller s2 s2 then fast_diff_r s1 s2
+    else
+      let fu _ m1 m2 = match m1, m2 with
+      | None, None -> None
+      | (Some _ as m), None -> m
+      | None, Some _ -> None
+      | Some m1, Some m2 ->
+        let m = Set.diff m1 m2 in
+        if Set.is_empty m then None else Some m
+      in
+      Int.Map.merge fu s1 s2
 
   let compare s1 s2 = Int.Map.compare Set.compare s1 s2
 
@@ -324,6 +381,8 @@ struct
     let fs m = Map.smartmapi f m in
     Int.Map.smartmap fs s
 
+  let height s = Int.Map.height s
+
   module Unsafe =
   struct
     let map f s =