From 26c1f873117eaa9d9ef782bb113d62c9843832d3 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 21 Mar 2015 22:47:14 +0100
Subject: [PATCH] wip: inter/union for CCIntMap

---
 src/data/CCIntMap.ml  | 43 +++++++++++++++++++++++++++++++++++++++++--
 src/data/CCIntMap.mli |  4 ++--
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/data/CCIntMap.ml b/src/data/CCIntMap.ml
index 200d0ec6..7a0fed15 100644
--- a/src/data/CCIntMap.ml
+++ b/src/data/CCIntMap.ml
@@ -187,9 +187,48 @@ let choose t =
   try Some (choose_exn t)
   with Not_found -> None
 
-let union _ _ _ = assert false
+let rec union f a b = match a, b with
+  | E, o | o, E -> o
+  | L (k, v), o
+  | o, L (k, v) ->
+    (* insert k, v into o *)
+    insert_ (fun ~old v -> f k old v) k v o
+  | N (p1, m1, l1, r1), N (p2, m2, l2, r2) ->
+    if p1 = p2 && m1 = m2
+    then mk_node_ p1 m1 (union f l1 l2) (union f r1 r2)
+    else if m1 < m2 && is_prefix_ ~prefix:p2 p1 ~bit:m1
+      then if bit_is_0_ p2 ~bit:m1
+        then N (p1, m1, union f l1 b, r1)
+        else N (p1, m1, l1, union f r1 b)
+    else if m1 > m2 && is_prefix_ ~prefix:p1 p2 ~bit:m2
+      then if bit_is_0_ p1 ~bit:m2
+        then N (p2, m2, union f l2 a, r2)
+        else N (p2, m2, l2, union f r2 a)
+    else join_ a p1 b p2
 
-let inter _ _ _ = assert false
+let rec inter f a b = match a, b with
+  | E, _ | _, E -> E
+  | L (k, v), o
+  | o, L (k, v) ->
+    begin try
+      let v' = find_exn k o in
+      L (k, f k v v')
+    with Not_found -> E
+    end
+  | N (p1, m1, l1, r1), N (p2, m2, l2, r2) ->
+    if p1 = p2 && m1 = m2
+    then mk_node_ p1 m1 (inter f l1 l2) (inter f r1 r2)
+    else if m1 < m2 && is_prefix_ ~prefix:p2 p1 ~bit:m1
+      then if bit_is_0_ p2 ~bit:m1
+        then inter f l1 b
+        else inter f r1 b
+    else if m1 > m2 && is_prefix_ ~prefix:p1 p2 ~bit:m2
+      then if bit_is_0_ p1 ~bit:m2
+        then inter f l2 a
+        else inter f r2 a
+    else E
+
+(* TODO: write tests *)
 
 (** {2 Whole-collection operations} *)
 
diff --git a/src/data/CCIntMap.mli b/src/data/CCIntMap.mli
index 68269317..4caf0d14 100644
--- a/src/data/CCIntMap.mli
+++ b/src/data/CCIntMap.mli
@@ -58,9 +58,9 @@ val choose : 'a t -> (int * 'a) option
 
 val choose_exn : 'a t -> int * 'a
 
-val union : (int -> 'a -> 'a -> 'a option) -> 'a t -> 'a t -> 'a t
+val union : (int -> 'a -> 'a -> 'a) -> 'a t -> 'a t -> 'a t
 
-val inter : (int -> 'a -> 'a -> 'a option) -> 'a t -> 'a t -> 'a t
+val inter : (int -> 'a -> 'a -> 'a) -> 'a t -> 'a t -> 'a t
 
 (** {2 Whole-collection operations} *)