diff --git a/benchmarks/benches/lib.rs b/benchmarks/benches/lib.rs index 23f6a6c3..a2518aa5 100644 --- a/benchmarks/benches/lib.rs +++ b/benchmarks/benches/lib.rs @@ -1,5 +1,6 @@ use itertools::Itertools; use std::cmp::Reverse; +use std::io::Cursor; use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Sub, SubAssign}; use criterion::measurement::Measurement; @@ -117,6 +118,114 @@ fn pairwise_binary_op_matrix( group.finish(); } +// Pairwise compare frame for the *_with_serialized evaluation: natural pair +// order (first bitmap of each pair is the materialized lhs, second is the +// serialized rhs — no small/large swap), serialization outside the timed +// setup, one group per op with a "ref_ref" (borrowing) and an "assign_ref" +// (in-place) bench per group. This commit sits BEFORE the *_with_serialized +// implementations land, so both slots run the deserialize baselines; a later +// commit switches the slots to the ws / ws_assign implementations, and +// criterion's saved-baseline comparison reports the change per slot. +fn pairwise_ops_with_serialized( + c: &mut Criterion, + op_name: &str, + op_ref_ref: fn(&RoaringBitmap, &[u8]) -> RoaringBitmap, + op_assign_ref: fn(&mut RoaringBitmap, &[u8]) -> (), +) { + let mut group = c.benchmark_group(format!("pairwise_{op_name}")); + + for dataset in Datasets { + let pairs = dataset + .bitmaps + .iter() + .tuple_windows::<(_, _)>() + .map(|(a, b)| { + let mut buf = Vec::new(); + b.serialize_into(&mut buf).unwrap(); + (a.clone(), Box::from(buf)) + }) + .collect::>(); + + group.bench_function(BenchmarkId::new("ref_ref", &dataset.name), |b| { + b.iter_batched_ref( + || pairs.clone(), + |bitmaps| { + for (a, b) in bitmaps { + black_box(op_ref_ref(a, b)); + } + }, + BatchSize::SmallInput, + ); + }); + + group.bench_function(BenchmarkId::new("assign_ref", &dataset.name), |b| { + b.iter_batched_ref( + || pairs.clone(), + |bitmaps| { + for (a, b) in bitmaps { + black_box(op_assign_ref(a, b)); + } + }, + BatchSize::SmallInput, + ); + }); + } + + group.finish(); +} + +// Successive chains for the *_with_serialized evaluation, same shape as +// pairwise_ops_with_serialized: fold every bitmap but the first into an +// accumulator seeded from the first bitmap (all four ops need a non-empty +// seed for a uniform shape), one group per op, "ref_ref" (borrowing chain, +// per-step result assigned back) and "assign_ref" (in-place chain) benches +// per group. Before the ws implementations land, both chains run the +// deserialize baselines; a later commit switches them to ws / ws_assign. +fn successive_ops_with_serialized( + c: &mut Criterion, + group_name: &str, + op_ref_ref: fn(&mut RoaringBitmap, &[u8]), + op_assign_ref: fn(&mut RoaringBitmap, &[u8]), +) { + let mut group = c.benchmark_group(group_name); + + for dataset in Datasets { + let first = dataset.bitmaps[0].clone(); + let serialized: Vec> = dataset + .bitmaps + .iter() + .skip(1) + .map(|b| { + let mut buf = Vec::new(); + b.serialize_into(&mut buf).unwrap(); + Box::from(buf) + }) + .collect(); + + group.bench_function(BenchmarkId::new("ref_ref", &dataset.name), |b| { + b.iter(|| { + let mut acc = first.clone(); + for bytes in &serialized { + op_ref_ref(&mut acc, bytes.as_ref()); + } + black_box(&acc); + }); + }); + + group.bench_function(BenchmarkId::new("assign_ref", &dataset.name), |b| { + b.iter(|| { + let mut acc = first.clone(); + for bytes in &serialized { + op_assign_ref(&mut acc, bytes.as_ref()); + } + black_box(&acc); + }); + }); + } + + group.finish(); +} + fn pairwise_binary_op( group: &mut BenchmarkGroup, op_name: &str, @@ -557,6 +666,78 @@ fn successive_or(c: &mut Criterion) { group.finish(); } +fn intersection_with_serialized(c: &mut Criterion) { + pairwise_ops_with_serialized( + c, + "intersection_with_serialized", + |a, b| a.intersection_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |a, b| a.intersection_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn union_with_serialized(c: &mut Criterion) { + pairwise_ops_with_serialized( + c, + "union_with_serialized", + |a, b| a.union_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |a, b| a.union_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn difference_with_serialized(c: &mut Criterion) { + pairwise_ops_with_serialized( + c, + "difference_with_serialized", + |a, b| a.difference_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |a, b| a.difference_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn symmetric_difference_with_serialized(c: &mut Criterion) { + pairwise_ops_with_serialized( + c, + "symmetric_difference_with_serialized", + |a, b| a.symmetric_difference_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |a, b| a.symmetric_difference_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn successive_and_with_serialized(c: &mut Criterion) { + successive_ops_with_serialized( + c, + "Successive And With Serialized", + |acc, b| *acc = acc.intersection_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |acc, b| acc.intersection_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn successive_or_with_serialized(c: &mut Criterion) { + successive_ops_with_serialized( + c, + "Successive Or With Serialized", + |acc, b| *acc = acc.union_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |acc, b| acc.union_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn successive_sub_with_serialized(c: &mut Criterion) { + successive_ops_with_serialized( + c, + "Successive Sub With Serialized", + |acc, b| *acc = acc.difference_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |acc, b| acc.difference_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + +fn successive_xor_with_serialized(c: &mut Criterion) { + successive_ops_with_serialized( + c, + "Successive Xor With Serialized", + |acc, b| *acc = acc.symmetric_difference_with_serialized_unchecked(Cursor::new(b)).unwrap(), + |acc, b| acc.symmetric_difference_assign_with_serialized_unchecked(Cursor::new(b)).unwrap(), + ) +} + // LEGACY BENCHMARKS // ================= @@ -740,6 +921,14 @@ criterion_group!( serialization, deserialization, successive_and, - successive_or + successive_or, + intersection_with_serialized, + union_with_serialized, + difference_with_serialized, + symmetric_difference_with_serialized, + successive_and_with_serialized, + successive_or_with_serialized, + successive_sub_with_serialized, + successive_xor_with_serialized ); criterion_main!(benches); diff --git a/roaring/src/bitmap/ops_with_serialized.rs b/roaring/src/bitmap/ops_with_serialized.rs index 3bae76c0..92e52571 100644 --- a/roaring/src/bitmap/ops_with_serialized.rs +++ b/roaring/src/bitmap/ops_with_serialized.rs @@ -235,9 +235,1200 @@ impl RoaringBitmap { run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); let store = if is_run_container { - let runs = reader.read_u16::().unwrap(); + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let mut other_container = Container { key, store }; + other_container &= container; + if !other_container.is_empty() { + containers.push(other_container); + } + } + + Ok(RoaringBitmap { containers }) + } + + /// Computes the intersection between this [`RoaringBitmap`] and a serialized one, + /// in place. + /// + /// This is the in-place counterpart of + /// [`RoaringBitmap::intersection_with_serialized_unchecked`]: containers of `self` are merged + /// in place (reusing their allocations) and containers absent from `other` are removed, + /// so no result bitmap is allocated. Prefer it when the intersection result should replace + /// `self`. + /// + /// # Errors + /// + /// If error happens, the operation stops early and `self` is left in a partially-updated state. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let mut rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// + /// rb1.intersection_assign_with_serialized_unchecked(Cursor::new(&bytes)).unwrap(); + /// assert_eq!(rb1, (3..4).collect::()); + /// ``` + pub fn intersection_assign_with_serialized_unchecked(&mut self, other: R) -> io::Result<()> + where + R: io::Read + io::Seek, + { + RoaringBitmap::intersection_assign_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn intersection_assign_with_serialized_impl( + &mut self, + mut reader: R, + a: A, + b: B, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[key, len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + let mut offsets = vec![0; size]; + reader.read_exact(cast_slice_mut(&mut offsets))?; + offsets.iter_mut().for_each(|offset| *offset = u32::from_le(*offset)); + return self.intersection_assign_with_serialized_impl_with_offsets( + reader, + a, + b, + &descriptions, + &offsets, + run_container_bitmap.as_deref(), + ); + } + + // No offsets table + let mut write = 0usize; + let mut read = 0usize; + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while read < self.containers.len() && self.containers[read].key < key { + read += 1; // lhs-only + } + + let lhs_matches = read < self.containers.len() && self.containers[read].key == key; + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + if !lhs_matches { + if is_run_container { + let runs = reader.read_u16::()?; + let runs_size = mem::size_of::() * 2 * runs as usize; + reader.seek(SeekFrom::Current(runs_size as i64))?; + } else if cardinality <= ARRAY_LIMIT { + let array_size = mem::size_of::() * cardinality as usize; + reader.seek(SeekFrom::Current(array_size as i64))?; + } else { + let bitmap_size = mem::size_of::() * BITMAP_LENGTH; + reader.seek(SeekFrom::Current(bitmap_size as i64))?; + } + continue; + } + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let rhs_container = Container { key, store }; + let k = self.containers[read].key; + let mut lhs_container = mem::replace(&mut self.containers[read], Container::new(k)); + lhs_container &= rhs_container; + if !lhs_container.is_empty() { + self.containers[write] = lhs_container; + write += 1; + } + read += 1; + } + + self.containers.truncate(write); + Ok(()) + } + + fn intersection_assign_with_serialized_impl_with_offsets( + &mut self, + mut reader: R, + a: A, + b: B, + descriptions: &[[u16; 2]], + offsets: &[u32], + run_container_bitmap: Option<&[u8]>, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let mut write = 0usize; + for read in 0..self.containers.len() { + let i = match descriptions.binary_search_by_key(&self.containers[read].key, |[k, _]| *k) + { + Ok(index) => index, + Err(_) => continue, // not part of the intersection, dropped + }; + + // Skip rhs using offsets table. + reader.seek(SeekFrom::Start(offsets[i] as u64))?; + + let [key, len_minus_one] = descriptions[i]; + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let rhs_container = Container { key, store }; + let placeholder = Container::new(self.containers[read].key); + let mut lhs_container = mem::replace(&mut self.containers[read], placeholder); + lhs_container &= rhs_container; + if !lhs_container.is_empty() { + self.containers[write] = lhs_container; + write += 1; + } + } + + self.containers.truncate(write); + Ok(()) + } + + /// Computes the union between a materialized [`RoaringBitmap`] and a serialized one. + /// + /// This is faster and more space efficient when you only need the union result. + /// It reduces the number of deserialized internal container and therefore + /// the number of allocations and copies of bytes. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// // Let's say the rb2 bitmap is serialized + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// let rb2_bytes = Cursor::new(bytes); + /// + /// assert_eq!( + /// rb1.union_with_serialized_unchecked(rb2_bytes).unwrap(), + /// rb1 | rb2, + /// ); + /// ``` + pub fn union_with_serialized_unchecked(&self, other: R) -> io::Result + where + R: io::Read + io::Seek, + { + RoaringBitmap::union_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn union_with_serialized_impl( + &self, + mut reader: R, + a: A, + b: B, + ) -> io::Result + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[key, len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + reader.seek(SeekFrom::Current(size as i64 * 4))?; + } + + let mut containers = Vec::with_capacity(self.containers.len() + descriptions.len()); + let mut left_containers = self.containers.iter().peekable(); + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while left_containers.peek().is_some_and(|container| container.key < key) { + containers.push(left_containers.next().unwrap().clone()); + } + + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let mut right_container = Container { key, store }; + if left_containers.peek().is_some_and(|container| container.key == key) { + right_container |= left_containers.next().unwrap(); + } + if !right_container.is_empty() { + containers.push(right_container); + } + } + + containers.extend(left_containers.cloned()); + Ok(RoaringBitmap { containers }) + } + + /// Computes the union between this [`RoaringBitmap`] and a serialized one, + /// in place. + /// + /// This is the in-place counterpart of + /// [`RoaringBitmap::union_with_serialized_unchecked`]: containers already + /// present in `self` are merged in place (reusing their allocations) and + /// containers only present in `other` are moved into `self`, so no result + /// bitmap is allocated. Prefer it when the union result should replace + /// `self`. + /// + /// # Errors + /// + /// If error happens, the operation stops early and `self` is left in a partially-updated state. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let mut rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// + /// rb1.union_assign_with_serialized_unchecked(Cursor::new(&bytes)).unwrap(); + /// assert_eq!(rb1, (1..5).collect::()); + /// ``` + pub fn union_assign_with_serialized_unchecked(&mut self, other: R) -> io::Result<()> + where + R: io::Read + io::Seek, + { + RoaringBitmap::union_assign_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn union_assign_with_serialized_impl( + &mut self, + mut reader: R, + a: A, + b: B, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[key, len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + reader.seek(SeekFrom::Current(size as i64 * 4))?; + } + + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let rhs_container = Container { key, store }; + match self.containers.binary_search_by_key(&key, |c| c.key) { + Ok(loc) => { + self.containers[loc] |= rhs_container; + } + Err(loc) => { + self.containers.insert(loc, rhs_container); + } + } + } + + Ok(()) + } + + /// Computes the difference between a materialized [`RoaringBitmap`] and a serialized one. + /// + /// This is faster and more space efficient when you only need the difference result. + /// It reduces the number of deserialized internal container and therefore + /// the number of allocations and copies of bytes. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// // Let's say the rb2 bitmap is serialized + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// let rb2_bytes = Cursor::new(bytes); + /// + /// assert_eq!( + /// rb1.difference_with_serialized_unchecked(rb2_bytes).unwrap(), + /// &rb1 - &rb2, + /// ); + /// ``` + pub fn difference_with_serialized_unchecked(&self, other: R) -> io::Result + where + R: io::Read + io::Seek, + { + RoaringBitmap::difference_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn difference_with_serialized_impl( + &self, + mut reader: R, + a: A, + b: B, + ) -> io::Result + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + // First read the cookie to determine which version of the format we are reading + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + // Read the run container bitmap if necessary + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + // Read the container descriptions + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[ref mut key, ref mut len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + let mut offsets = vec![0; size]; + reader.read_exact(cast_slice_mut(&mut offsets))?; + offsets.iter_mut().for_each(|offset| *offset = u32::from_le(*offset)); + return self.difference_with_serialized_impl_with_offsets( + reader, + a, + b, + &descriptions, + &offsets, + run_container_bitmap.as_deref(), + ); + } + + // Walk rhs in order. Drain smaller lhs containers into result unchanged, + // subtract when keys match, skip rhs bytes when lhs has no matching key. + let mut containers = Vec::with_capacity(self.containers.len()); + let lhs = self.containers.as_slice(); + let mut lhs_idx = 0usize; + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while lhs_idx < lhs.len() && lhs[lhs_idx].key < key { + containers.push(lhs[lhs_idx].clone()); + lhs_idx += 1; + } + + let lhs_matches = lhs_idx < lhs.len() && lhs[lhs_idx].key == key; + let cardinality = u64::from(len_minus_one) + 1; + + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + if !lhs_matches { + if is_run_container { + let runs = reader.read_u16::()?; + let runs_size = mem::size_of::() * 2 * runs as usize; + reader.seek(SeekFrom::Current(runs_size as i64))?; + } else if cardinality <= ARRAY_LIMIT { + let array_size = mem::size_of::() * cardinality as usize; + reader.seek(SeekFrom::Current(array_size as i64))?; + } else { + let bitmap_size = mem::size_of::() * BITMAP_LENGTH; + reader.seek(SeekFrom::Current(bitmap_size as i64))?; + } + continue; + } + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let other_container = Container { key, store }; + let mut lhs_container = lhs[lhs_idx].clone(); + lhs_container -= &other_container; + if !lhs_container.is_empty() { + containers.push(lhs_container); + } + lhs_idx += 1; + } + + containers.extend_from_slice(&lhs[lhs_idx..]); + Ok(RoaringBitmap { containers }) + } + + fn difference_with_serialized_impl_with_offsets( + &self, + mut reader: R, + a: A, + b: B, + descriptions: &[[u16; 2]], + offsets: &[u32], + run_container_bitmap: Option<&[u8]>, + ) -> io::Result + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let mut containers = Vec::with_capacity(self.containers.len()); + for container in &self.containers { + let i = match descriptions.binary_search_by_key(&container.key, |[k, _]| *k) { + Ok(index) => index, + Err(_) => { + containers.push(container.clone()); + continue; + } + }; + + // Seek to the bytes of the container we want. + reader.seek(SeekFrom::Start(offsets[i] as u64))?; + + let [key, len_minus_one] = descriptions[i]; + let cardinality = u64::from(len_minus_one) + 1; + + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let other_container = Container { key, store }; + let mut lhs_container = container.clone(); + lhs_container -= &other_container; + if !lhs_container.is_empty() { + containers.push(lhs_container); + } + } + + Ok(RoaringBitmap { containers }) + } + + /// Computes the difference between a serialized [`RoaringBitmap`] and this one, + /// in place: `self` keeps the values absent from `other`. + /// + /// This is the in-place counterpart of + /// [`RoaringBitmap::difference_with_serialized_unchecked`]: containers of + /// `self` are subtracted in place (reusing their allocations) and moved + /// within the same Vec — no result bitmap is allocated and nothing is + /// cloned. Prefer it when the difference result should replace `self`. + /// + /// # Errors + /// + /// If error happens, the operation stops early and `self` is left in a partially-updated state. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let mut rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// + /// rb1.difference_assign_with_serialized_unchecked(Cursor::new(&bytes)).unwrap(); + /// assert_eq!(rb1, (1..3).collect::()); + /// ``` + pub fn difference_assign_with_serialized_unchecked(&mut self, other: R) -> io::Result<()> + where + R: io::Read + io::Seek, + { + RoaringBitmap::difference_assign_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn difference_assign_with_serialized_impl( + &mut self, + mut reader: R, + a: A, + b: B, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[key, len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + let mut offsets = vec![0; size]; + reader.read_exact(cast_slice_mut(&mut offsets))?; + offsets.iter_mut().for_each(|offset| *offset = u32::from_le(*offset)); + return self.difference_assign_with_serialized_impl_with_offsets( + reader, + a, + b, + &descriptions, + &offsets, + run_container_bitmap.as_deref(), + ); + } + + // No offsets table + let mut write = 0usize; + let mut read = 0usize; + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while read < self.containers.len() && self.containers[read].key < key { + self.containers.swap(write, read); // lhs-only: kept + write += 1; + read += 1; + } + + let lhs_matches = read < self.containers.len() && self.containers[read].key == key; + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + if !lhs_matches { + if is_run_container { + let runs = reader.read_u16::()?; + let runs_size = mem::size_of::() * 2 * runs as usize; + reader.seek(SeekFrom::Current(runs_size as i64))?; + } else if cardinality <= ARRAY_LIMIT { + let array_size = mem::size_of::() * cardinality as usize; + reader.seek(SeekFrom::Current(array_size as i64))?; + } else { + let bitmap_size = mem::size_of::() * BITMAP_LENGTH; + reader.seek(SeekFrom::Current(bitmap_size as i64))?; + } + continue; + } + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let other_container = Container { key, store }; + self.containers.swap(write, read); // lhs to the write slot + self.containers[write] -= &other_container; + if !self.containers[write].is_empty() { + write += 1; + } + read += 1; + } + + for idx in read..self.containers.len() { + self.containers.swap(write, idx); + write += 1; + } + self.containers.truncate(write); + Ok(()) + } + + fn difference_assign_with_serialized_impl_with_offsets( + &mut self, + mut reader: R, + a: A, + b: B, + descriptions: &[[u16; 2]], + offsets: &[u32], + run_container_bitmap: Option<&[u8]>, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let mut write = 0usize; + for read in 0..self.containers.len() { + let i = match descriptions.binary_search_by_key(&self.containers[read].key, |[k, _]| *k) + { + Ok(index) => index, + Err(_) => { + self.containers.swap(write, read); // lhs-only: kept + write += 1; + continue; + } + }; + + // Skip rhs using offsets table. + reader.seek(SeekFrom::Start(offsets[i] as u64))?; + + let [key, len_minus_one] = descriptions[i]; + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let other_container = Container { key, store }; + self.containers.swap(write, read); // lhs to the write slot + self.containers[write] -= &other_container; + if !self.containers[write].is_empty() { + write += 1; + } + } + + self.containers.truncate(write); + Ok(()) + } + + /// Computes the symmetric difference between a materialized [`RoaringBitmap`] and a serialized one. + /// + /// This is faster and more space efficient when you only need the symmetric difference result. + /// It reduces the number of deserialized internal container and therefore + /// the number of allocations and copies of bytes. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// // Let's say the rb2 bitmap is serialized + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// let rb2_bytes = Cursor::new(bytes); + /// + /// assert_eq!( + /// rb1.symmetric_difference_with_serialized_unchecked(rb2_bytes).unwrap(), + /// &rb1 ^ &rb2, + /// ); + /// ``` + pub fn symmetric_difference_with_serialized_unchecked( + &self, + other: R, + ) -> io::Result + where + R: io::Read + io::Seek, + { + RoaringBitmap::symmetric_difference_with_serialized_impl::( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn symmetric_difference_with_serialized_impl( + &self, + mut reader: R, + a: A, + b: B, + ) -> io::Result + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + // First read the cookie to determine which version of the format we are reading + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + // Read the run container bitmap if necessary + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + // Read the container descriptions + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[ref mut key, ref mut len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + // Skip the offsets if present, we don't need them for symmetric difference + if has_offsets { + reader.seek(SeekFrom::Current(size as i64 * 4))?; + } + + // Walk rhs sequentially. Drain smaller lhs containers into result unchanged, + // xor when keys match, push rhs-only containers directly. + let mut containers = Vec::with_capacity(self.containers.len() + descriptions.len()); + let lhs = self.containers.as_slice(); + let mut lhs_idx = 0usize; + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while lhs_idx < lhs.len() && lhs[lhs_idx].key < key { + containers.push(lhs[lhs_idx].clone()); + lhs_idx += 1; + } + + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; let mut intervals = vec![[0, 0]; runs as usize]; - reader.read_exact(cast_slice_mut(&mut intervals)).unwrap(); + reader.read_exact(cast_slice_mut(&mut intervals))?; intervals.iter_mut().for_each(|[s, len]| { *s = u16::from_le(*s); *len = u16::from_le(*len); @@ -253,28 +1444,188 @@ impl RoaringBitmap { store } else if cardinality <= ARRAY_LIMIT { let mut values = vec![0; cardinality as usize]; - reader.read_exact(cast_slice_mut(&mut values)).unwrap(); + reader.read_exact(cast_slice_mut(&mut values))?; values.iter_mut().for_each(|n| *n = u16::from_le(*n)); let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; Store::Array(array) } else { let mut values = Box::new([0; BITMAP_LENGTH]); - reader.read_exact(cast_slice_mut(&mut values[..])).unwrap(); + reader.read_exact(cast_slice_mut(&mut values[..]))?; values.iter_mut().for_each(|n| *n = u64::from_le(*n)); let bitmap = b(cardinality, values) .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; Store::Bitmap(bitmap) }; - let mut other_container = Container { key, store }; - other_container &= container; - if !other_container.is_empty() { - containers.push(other_container); + let mut container = Container { key, store }; + if lhs_idx < lhs.len() && lhs[lhs_idx].key == key { + container ^= &lhs[lhs_idx]; + lhs_idx += 1; + } + if !container.is_empty() { + containers.push(container); } } + containers.extend_from_slice(&lhs[lhs_idx..]); Ok(RoaringBitmap { containers }) } + + /// Computes the symmetric difference between this [`RoaringBitmap`] and a + /// serialized one, in place. + /// + /// This is the in-place counterpart of + /// [`RoaringBitmap::symmetric_difference_with_serialized_unchecked`]: + /// containers of `self` are moved through the merge (lhs-only ones kept, + /// matching ones xor-ed in place reusing their allocations) and + /// rhs-only containers are moved in as they are read — nothing is cloned. + /// Prefer it when the symmetric difference result should replace `self`. + /// + /// # Errors + /// + /// If error happens, the operation stops early and `self` is left in a partially-updated state. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringBitmap; + /// use std::io::Cursor; + /// + /// let mut rb1: RoaringBitmap = (1..4).collect(); + /// let rb2: RoaringBitmap = (3..5).collect(); + /// + /// let mut bytes = Vec::new(); + /// rb2.serialize_into(&mut bytes).unwrap(); + /// + /// rb1.symmetric_difference_assign_with_serialized_unchecked(Cursor::new(&bytes)).unwrap(); + /// assert_eq!(rb1, (1..3).chain(4..5).collect::()); + /// ``` + pub fn symmetric_difference_assign_with_serialized_unchecked( + &mut self, + other: R, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + { + RoaringBitmap::symmetric_difference_assign_with_serialized_impl::< + R, + _, + Infallible, + _, + Infallible, + >( + self, + other, + |values| Ok(ArrayStore::from_vec_unchecked(values)), + |len, values| Ok(BitmapStore::from_unchecked(len, values)), + ) + } + + fn symmetric_difference_assign_with_serialized_impl( + &mut self, + mut reader: R, + a: A, + b: B, + ) -> io::Result<()> + where + R: io::Read + io::Seek, + A: Fn(Vec) -> Result, + AErr: Error + Send + Sync + 'static, + B: Fn(u64, Box<[u64; 1024]>) -> Result, + BErr: Error + Send + Sync + 'static, + { + let (size, has_offsets, has_run_containers) = { + let cookie = reader.read_u32::()?; + if cookie == SERIAL_COOKIE_NO_RUNCONTAINER { + (reader.read_u32::()? as usize, true, false) + } else if (cookie as u16) == SERIAL_COOKIE { + let size = ((cookie >> 16) + 1) as usize; + (size, size >= NO_OFFSET_THRESHOLD, true) + } else { + return Err(io::Error::other("unknown cookie value")); + } + }; + + let run_container_bitmap = if has_run_containers { + let mut bitmap = vec![0u8; size.div_ceil(8)]; + reader.read_exact(&mut bitmap)?; + Some(bitmap) + } else { + None + }; + + if size > u16::MAX as usize + 1 { + return Err(io::Error::other("size is greater than supported")); + } + + let mut descriptions = vec![[0; 2]; size]; + reader.read_exact(cast_slice_mut(&mut descriptions))?; + descriptions.iter_mut().for_each(|[ref mut key, ref mut len]| { + *key = u16::from_le(*key); + *len = u16::from_le(*len); + }); + + if has_offsets { + reader.seek(SeekFrom::Current(size as i64 * 4))?; + } + + let mut lhs_iter = mem::take(&mut self.containers).into_iter().peekable(); + for (i, &[key, len_minus_one]) in descriptions.iter().enumerate() { + while lhs_iter.peek().is_some_and(|container| container.key < key) { + self.containers.push(lhs_iter.next().unwrap()); // lhs-only: kept + } + + let cardinality = u64::from(len_minus_one) + 1; + let is_run_container = + run_container_bitmap.as_ref().is_some_and(|bm| bm[i / 8] & (1 << (i % 8)) != 0); + + let store = if is_run_container { + let runs = reader.read_u16::()?; + let mut intervals = vec![[0, 0]; runs as usize]; + reader.read_exact(cast_slice_mut(&mut intervals))?; + intervals.iter_mut().for_each(|[s, len]| { + *s = u16::from_le(*s); + *len = u16::from_le(*len); + }); + + let cardinality = intervals.iter().map(|[_, len]| *len as usize).sum(); + let mut store = Store::with_capacity(cardinality); + intervals.into_iter().try_for_each(|[s, len]| -> Result<(), io::ErrorKind> { + let end = s.checked_add(len).ok_or(io::ErrorKind::InvalidData)?; + store.insert_range(RangeInclusive::new(s, end)); + Ok(()) + })?; + store + } else if cardinality <= ARRAY_LIMIT { + let mut values = vec![0; cardinality as usize]; + reader.read_exact(cast_slice_mut(&mut values))?; + values.iter_mut().for_each(|n| *n = u16::from_le(*n)); + let array = a(values).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Array(array) + } else { + let mut values = Box::new([0; BITMAP_LENGTH]); + reader.read_exact(cast_slice_mut(&mut values[..]))?; + values.iter_mut().for_each(|n| *n = u64::from_le(*n)); + let bitmap = b(cardinality, values) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Store::Bitmap(bitmap) + }; + + let rhs_container = Container { key, store }; + if lhs_iter.peek().is_some_and(|container| container.key == key) { + let mut lhs_container = lhs_iter.next().unwrap(); + lhs_container ^= &rhs_container; + if !lhs_container.is_empty() { + self.containers.push(lhs_container); + } + } else if !rhs_container.is_empty() { + self.containers.push(rhs_container); + } + } + + self.containers.extend(lhs_iter); + Ok(()) + } } #[cfg(test)] @@ -296,5 +1647,103 @@ mod test { prop_assert_eq!(a.intersection_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(), a & b); } + + #[test] + fn intersection_assign_with_serialized_eq_materialized_intersection( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + let mut assigned = a.clone(); + assigned.intersection_assign_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(); + prop_assert_eq!(assigned, a & b); + } + } + + proptest! { + #[test] + fn union_with_serialized_eq_materialized_union( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + prop_assert_eq!(a.union_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(), a | b); + } + + #[test] + fn union_assign_with_serialized_eq_materialized_union( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + let mut assigned = a.clone(); + assigned.union_assign_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(); + prop_assert_eq!(assigned, a | b); + } + } + + proptest! { + #[test] + fn difference_with_serialized_eq_materialized_difference( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + prop_assert_eq!(a.difference_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(), &a - &b); + } + + #[test] + fn difference_assign_with_serialized_eq_materialized_difference( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + let mut assigned = a.clone(); + assigned.difference_assign_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(); + prop_assert_eq!(assigned, &a - &b); + } + } + + proptest! { + #[test] + fn symmetric_difference_with_serialized_eq_materialized_symmetric_difference( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + prop_assert_eq!(a.symmetric_difference_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(), &a ^ &b); + } + + #[test] + fn symmetric_difference_assign_with_serialized_eq_materialized_symmetric_difference( + a in RoaringBitmap::arbitrary(), + b in RoaringBitmap::arbitrary() + ) { + let mut serialized_bytes_b = Vec::new(); + b.serialize_into(&mut serialized_bytes_b).unwrap(); + let serialized_bytes_b = &serialized_bytes_b[..]; + + let mut assigned = a.clone(); + assigned.symmetric_difference_assign_with_serialized_unchecked(Cursor::new(serialized_bytes_b)).unwrap(); + prop_assert_eq!(assigned, &a ^ &b); + } } } diff --git a/roaring/src/treemap/entry.rs b/roaring/src/treemap/entry.rs new file mode 100644 index 00000000..2580f434 --- /dev/null +++ b/roaring/src/treemap/entry.rs @@ -0,0 +1,202 @@ +use alloc::collections::btree_map; + +use crate::{RoaringBitmap, RoaringTreemap}; + +/// A view into a single partition of a [`RoaringTreemap`], obtained from +/// [`RoaringTreemap::bitmap_entry`]. +/// +/// This is the entry-point style access for per-partition in-place +/// mutation: it lets callers locate, mutate, insert, or take out the +/// [`RoaringBitmap`] of one partition without rebuilding the treemap +/// around it. +pub enum BitmapEntry<'a> { + /// The partition is absent. + Vacant(VacantBitmapEntry<'a>), + /// The partition exists. + Occupied(OccupiedBitmapEntry<'a>), +} + +/// A vacant partition view, obtained from [`RoaringTreemap::bitmap_entry`]. +/// +/// Use [`VacantBitmapEntry::insert`] to create that partition. +pub struct VacantBitmapEntry<'a> { + inner: btree_map::VacantEntry<'a, u32, RoaringBitmap>, +} + +/// An occupied partition view, obtained from [`RoaringTreemap::bitmap_entry`]. +/// +/// Use [`OccupiedBitmapEntry::get_mut`] to borrow its bitmap mutably and +/// [`OccupiedBitmapEntry::remove`] to take it out. +pub struct OccupiedBitmapEntry<'a> { + inner: btree_map::OccupiedEntry<'a, u32, RoaringBitmap>, +} + +impl<'a> VacantBitmapEntry<'a> { + /// Inserts a bitmap into the vacant partition and returns a mutable reference to it. + /// + /// # Examples + /// + /// ```rust + /// use roaring::{RoaringBitmap, RoaringTreemap}; + /// use roaring::treemap::BitmapEntry; + /// + /// let mut treemap = RoaringTreemap::new(); + /// match treemap.bitmap_entry(3) { + /// BitmapEntry::Vacant(entry) => { + /// let bitmap = entry.insert(RoaringBitmap::from([1])); + /// bitmap.insert(2); + /// } + /// BitmapEntry::Occupied(_) => unreachable!(), + /// } + /// assert_eq!(treemap.len(), 2); + /// ``` + pub fn insert(self, bitmap: RoaringBitmap) -> &'a mut RoaringBitmap { + self.inner.insert(bitmap) + } +} + +impl OccupiedBitmapEntry<'_> { + /// Returns a mutable reference to the partition's bitmap. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringTreemap; + /// use roaring::treemap::BitmapEntry; + /// + /// let mut treemap = RoaringTreemap::new(); + /// treemap.insert(1); + /// if let BitmapEntry::Occupied(mut entry) = treemap.bitmap_entry(0) { + /// entry.get_mut().insert(2); + /// } + /// assert_eq!(treemap.len(), 2); + /// ``` + pub fn get_mut(&mut self) -> &mut RoaringBitmap { + self.inner.get_mut() + } + + /// Takes the partition's bitmap out of the treemap, removing the partition. + /// + /// # Examples + /// + /// ```rust + /// use roaring::RoaringTreemap; + /// use roaring::treemap::BitmapEntry; + /// + /// let mut treemap = RoaringTreemap::new(); + /// treemap.insert(1); + /// if let BitmapEntry::Occupied(entry) = treemap.bitmap_entry(0) { + /// assert_eq!(entry.remove().len(), 1); + /// } + /// assert!(treemap.is_empty()); + /// ``` + pub fn remove(self) -> RoaringBitmap { + self.inner.remove() + } +} + +impl RoaringTreemap { + /// Returns a view into the partition with the given prefix (the 32 most significant bits). + /// + /// # Examples + /// + /// ```rust + /// use roaring::{RoaringBitmap, RoaringTreemap}; + /// use roaring::treemap::BitmapEntry; + /// + /// let mut treemap = RoaringTreemap::new(); + /// + /// // Vacant: insert a new partition. + /// match treemap.bitmap_entry(0) { + /// BitmapEntry::Vacant(entry) => { entry.insert(RoaringBitmap::from([1])); } + /// BitmapEntry::Occupied(_) => unreachable!(), + /// } + /// + /// // Occupied: mutate the partition in place. + /// match treemap.bitmap_entry(0) { + /// BitmapEntry::Occupied(mut entry) => { entry.get_mut().insert(2); } + /// BitmapEntry::Vacant(_) => unreachable!(), + /// } + /// assert_eq!(treemap.len(), 2); + /// ``` + pub fn bitmap_entry(&mut self, prefix: u32) -> BitmapEntry<'_> { + match self.map.entry(prefix) { + btree_map::Entry::Vacant(inner) => BitmapEntry::Vacant(VacantBitmapEntry { inner }), + btree_map::Entry::Occupied(inner) => { + BitmapEntry::Occupied(OccupiedBitmapEntry { inner }) + } + } + } +} + +#[cfg(test)] +mod test { + use crate::{RoaringBitmap, RoaringTreemap}; + use proptest::prelude::*; + + use super::BitmapEntry; + + proptest! { + #[test] + fn occupied_entry_remove_and_insert( + treemap in RoaringTreemap::arbitrary(), + ) { + let mut treemap = treemap; + let Some(prefix) = treemap.iter().next().map(|v| (v >> 32) as u32) else { + // empty treemap: a lookup is vacant + return Ok(()); + }; + + // Occupied: take the bitmap out, then put it back. + let expected = treemap.clone(); + let bitmap = match treemap.bitmap_entry(prefix) { + BitmapEntry::Occupied(entry) => entry.remove(), + BitmapEntry::Vacant(_) => unreachable!("prefix exists"), + }; + prop_assert!(matches!(treemap.bitmap_entry(prefix), BitmapEntry::Vacant(_))); + + match treemap.bitmap_entry(prefix) { + BitmapEntry::Vacant(entry) => { + entry.insert(bitmap); + } + BitmapEntry::Occupied(_) => unreachable!("prefix was removed"), + } + prop_assert_eq!(treemap, expected); + } + + #[test] + fn vacant_entry_insert( + treemap in RoaringTreemap::arbitrary(), + prefix in any::(), + bitmap in RoaringBitmap::arbitrary(), + ) { + let mut treemap = treemap; + // Start from a guaranteed-vacant prefix (drop it if present). + if let BitmapEntry::Occupied(entry) = treemap.bitmap_entry(prefix) { + entry.remove(); + } + + // A vacant lookup inserts nothing. + let before = treemap.clone(); + match treemap.bitmap_entry(prefix) { + BitmapEntry::Vacant(_) => {} + BitmapEntry::Occupied(_) => unreachable!("prefix was removed"), + } + prop_assert_eq!(&treemap, &before); + + // Inserting through the vacant entry makes it occupied with that bitmap. + match treemap.bitmap_entry(prefix) { + BitmapEntry::Vacant(entry) => { + entry.insert(bitmap.clone()); + } + BitmapEntry::Occupied(_) => unreachable!("prefix was removed"), + } + match treemap.bitmap_entry(prefix) { + BitmapEntry::Occupied(mut entry) => { + prop_assert_eq!(&*entry.get_mut(), &bitmap); + } + BitmapEntry::Vacant(_) => unreachable!("prefix was inserted"), + } + } + } +} diff --git a/roaring/src/treemap/mod.rs b/roaring/src/treemap/mod.rs index 93ca126d..1d630255 100644 --- a/roaring/src/treemap/mod.rs +++ b/roaring/src/treemap/mod.rs @@ -9,6 +9,7 @@ mod util; // the docs mod arbitrary; mod cmp; +mod entry; mod inherent; mod iter; mod ops; @@ -17,6 +18,7 @@ mod serde; #[cfg(feature = "std")] mod serialization; +pub use self::entry::{BitmapEntry, OccupiedBitmapEntry, VacantBitmapEntry}; pub use self::iter::{BitmapIter, IntoIter, Iter}; /// A compressed bitmap with u64 values.