From 9f88c08dd30909164efa9aef826a2d89f190c304 Mon Sep 17 00:00:00 2001
From: Zachary Dremann <dremann@gmail.com>
Date: Sat, 13 Jan 2024 19:37:12 -0500
Subject: [PATCH] fix(treemap): xor can no longer leave empty bitmaps in the
 tree

---
 croaring/src/treemap/imp.rs | 74 ++++++++++++++++++++++++++++++-------
 1 file changed, 60 insertions(+), 14 deletions(-)

diff --git a/croaring/src/treemap/imp.rs b/croaring/src/treemap/imp.rs
index 05e9201..5abb300 100644
--- a/croaring/src/treemap/imp.rs
+++ b/croaring/src/treemap/imp.rs
@@ -842,20 +842,10 @@ impl Treemap {
     /// ```
     #[must_use]
     pub fn xor(&self, other: &Self) -> Self {
-        let mut treemap = self.clone();
-
-        for (key, other_bitmap) in &other.map {
-            match treemap.map.entry(*key) {
-                Entry::Vacant(current_map) => {
-                    current_map.insert(other_bitmap.clone());
-                }
-                Entry::Occupied(mut bitmap) => {
-                    bitmap.get_mut().xor_inplace(other_bitmap);
-                }
-            };
-        }
-
-        treemap
+        binop(self, other, |args| match args {
+            BinopArgs::Both(lhs, rhs) => Some(lhs.xor(rhs)),
+            BinopArgs::Lhs(bitmap) | BinopArgs::Rhs(bitmap) => Some(bitmap.clone()),
+        })
     }
 
     /// Inplace version of xor, stores result in the current treemap.
@@ -1228,6 +1218,62 @@ impl Treemap {
     }
 }
 
+enum BinopArgs<'a> {
+    Both(&'a Bitmap, &'a Bitmap),
+    Lhs(&'a Bitmap),
+    Rhs(&'a Bitmap),
+}
+
+#[must_use]
+fn binop<F>(lhs: &Treemap, rhs: &Treemap, f: F) -> Treemap
+where
+    F: Fn(BinopArgs<'_>) -> Option<Bitmap>,
+{
+    let mut treemap = Treemap::new();
+    let mut lhs_iter = lhs.map.iter();
+    let mut rhs_iter = rhs.map.iter();
+
+    let mut lhs_next = lhs_iter.next();
+    let mut rhs_next = rhs_iter.next();
+
+    loop {
+        let (key, bitmap) = match (lhs_next, rhs_next) {
+            (Some((&lhs_key, lhs_bitmap)), Some((&rhs_key, rhs_bitmap))) => {
+                if lhs_key == rhs_key {
+                    let bitmap = f(BinopArgs::Both(lhs_bitmap, rhs_bitmap));
+                    lhs_next = lhs_iter.next();
+                    rhs_next = rhs_iter.next();
+                    (lhs_key, bitmap)
+                } else if lhs_key < rhs_key {
+                    let bitmap = f(BinopArgs::Lhs(lhs_bitmap));
+                    lhs_next = lhs_iter.next();
+                    (lhs_key, bitmap)
+                } else {
+                    let bitmap = f(BinopArgs::Rhs(rhs_bitmap));
+                    rhs_next = rhs_iter.next();
+                    (rhs_key, bitmap)
+                }
+            }
+            (Some((&lhs_key, lhs_bitmap)), None) => {
+                let bitmap = f(BinopArgs::Lhs(lhs_bitmap));
+                lhs_next = lhs_iter.next();
+                (lhs_key, bitmap)
+            }
+            (None, Some((&rhs_key, rhs_bitmap))) => {
+                let bitmap = f(BinopArgs::Rhs(rhs_bitmap));
+                rhs_next = rhs_iter.next();
+                (rhs_key, bitmap)
+            }
+            (None, None) => break,
+        };
+        if let Some(bitmap) = bitmap.filter(|b| !b.is_empty()) {
+            treemap.map.insert(key, bitmap);
+        }
+    }
+
+    treemap
+}
+
 fn range_to_inclusive<R: RangeBounds<u64>>(range: R) -> (u64, u64) {
     let start = match range.start_bound() {
         Bound::Included(&i) => i,