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 = 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(); + } + } } } }