From 9ebc2cda0ca6db392cde19e463b09c035492c542 Mon Sep 17 00:00:00 2001
From: Zachary Dremann <dremann@gmail.com>
Date: Sun, 14 Jan 2024 02:29:52 -0500
Subject: [PATCH] fix(treemap): Don't leave empty items after binary ops

---
 croaring/src/treemap/imp.rs | 52 +++++++++++++------------------------
 1 file changed, 18 insertions(+), 34 deletions(-)

diff --git a/croaring/src/treemap/imp.rs b/croaring/src/treemap/imp.rs
index 5abb300..36b37a9 100644
--- a/croaring/src/treemap/imp.rs
+++ b/croaring/src/treemap/imp.rs
@@ -658,16 +658,10 @@ impl Treemap {
     /// ```
     #[must_use]
     pub fn and(&self, other: &Self) -> Self {
-        let mut treemap = Treemap::new();
-
-        for (key, bitmap) in &self.map {
-            other
-                .map
-                .get(key)
-                .map(|other_bitmap| treemap.map.insert(*key, bitmap.and(other_bitmap)));
-        }
-
-        treemap
+        binop(self, other, |args| match args {
+            BinopArgs::Both(lhs, rhs) => Some(lhs.and(rhs)),
+            BinopArgs::Lhs(_) | BinopArgs::Rhs(_) => None,
+        })
     }
 
     /// Computes the intersection between two treemaps and stores the result
@@ -754,20 +748,10 @@ impl Treemap {
     /// ```
     #[must_use]
     pub fn or(&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().or_inplace(other_bitmap);
-                }
-            };
-        }
-
-        treemap
+        binop(self, other, |args| match args {
+            BinopArgs::Both(lhs, rhs) => Some(lhs.or(rhs)),
+            BinopArgs::Lhs(bitmap) | BinopArgs::Rhs(bitmap) => Some(bitmap.clone()),
+        })
     }
 
     /// Computes the intersection between two bitmaps and stores the result
@@ -877,8 +861,6 @@ impl Treemap {
     /// assert!(treemap3.contains(15));
     /// ```
     pub fn xor_inplace(&mut self, other: &Self) {
-        let mut keys_to_remove: Vec<u32> = Vec::new();
-
         for (key, other_bitmap) in &other.map {
             match self.map.entry(*key) {
                 Entry::Vacant(bitmap) => {
@@ -887,15 +869,11 @@ impl Treemap {
                 Entry::Occupied(mut bitmap) => {
                     bitmap.get_mut().xor_inplace(other_bitmap);
                     if bitmap.get().is_empty() {
-                        keys_to_remove.push(*key);
+                        bitmap.remove();
                     }
                 }
             };
         }
-
-        for key in keys_to_remove {
-            self.map.remove(&key);
-        }
     }
 
     /// Computes the difference between two bitmaps and returns the result.
@@ -982,9 +960,15 @@ impl Treemap {
     /// assert!(treemap3.contains(15));
     /// ```
     pub fn andnot_inplace(&mut self, other: &Self) {
-        for (key, bitmap) in &mut self.map {
-            if let Some(other_bitmap) = other.map.get(key) {
-                bitmap.andnot_inplace(other_bitmap);
+        for (&key, to_sub) in &other.map {
+            match self.map.entry(key) {
+                Entry::Vacant(_) => (),
+                Entry::Occupied(mut entry) => {
+                    entry.get_mut().andnot_inplace(to_sub);
+                    if entry.get().is_empty() {
+                        entry.remove();
+                    }
+                }
             }
         }
     }