diff --git a/Cargo.toml b/Cargo.toml index 6e43a8ac..53d39404 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ itertools = "0.14" once_cell = "1.20" postcard = { version = "1.1", features = [ "alloc" ] } proptest = "1.6.0" +rkyv = { version = "0.8.16", features = [ "alloc" ], default-features = false } serde = "1.0.217" serde_json = "1.0.138" zip = { version = "0.6", default-features = false } diff --git a/roaring/Cargo.toml b/roaring/Cargo.toml index f4f9310d..5b48fb80 100644 --- a/roaring/Cargo.toml +++ b/roaring/Cargo.toml @@ -20,6 +20,7 @@ license = "MIT OR Apache-2.0" bytemuck = { workspace = true, optional = true } byteorder = { workspace = true, optional = true } serde = { workspace = true, optional = true } +rkyv = { workspace = true, optional = true } [features] default = ["std"] @@ -31,3 +32,4 @@ std = ["dep:bytemuck", "dep:byteorder"] proptest = { workspace = true } serde_json = { workspace = true } postcard = { workspace = true } +rkyv = { workspace = true, features = ["bytecheck"] } diff --git a/roaring/src/bitmap/container.rs b/roaring/src/bitmap/container.rs index b11a6446..bcb5a29a 100644 --- a/roaring/src/bitmap/container.rs +++ b/roaring/src/bitmap/container.rs @@ -14,6 +14,7 @@ pub const RUN_MAX_SIZE: u64 = 2048; use alloc::vec::Vec; #[derive(PartialEq, Eq, Clone)] +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] pub(crate) struct Container { pub key: u16, pub store: Store, diff --git a/roaring/src/bitmap/mod.rs b/roaring/src/bitmap/mod.rs index add18dfd..0140fd15 100644 --- a/roaring/src/bitmap/mod.rs +++ b/roaring/src/bitmap/mod.rs @@ -15,8 +15,12 @@ mod iter; mod ops; #[cfg(feature = "std")] mod ops_with_serialized; + +#[cfg(feature = "rkyv")] +mod rkyv; #[cfg(feature = "serde")] mod serde; + #[cfg(feature = "std")] mod serialization; @@ -45,6 +49,7 @@ use alloc::vec::Vec; /// println!("total bits set to true: {}", rb.len()); /// ``` #[derive(PartialEq, Eq)] +#[cfg_attr(feature = "rkyv", derive(::rkyv::Archive, ::rkyv::Serialize, ::rkyv::Deserialize))] pub struct RoaringBitmap { containers: Vec, } diff --git a/roaring/src/bitmap/rkyv.rs b/roaring/src/bitmap/rkyv.rs new file mode 100644 index 00000000..2451fdb9 --- /dev/null +++ b/roaring/src/bitmap/rkyv.rs @@ -0,0 +1,379 @@ +use super::{ + container::ArchivedContainer, + store::{ArchivedArrayStore, ArchivedBitmapStore, ArchivedIntervalStore, ArchivedStore}, +}; + +impl super::ArchivedRoaringBitmap { + /// Returns the number of elements in the archived bitmap. + #[must_use] + pub fn len(&self) -> u64 { + self.containers.iter().map(ArchivedContainer::len).sum() + } + + /// Checks if the archived bitmap is empty. + #[must_use] + pub fn is_empty(&self) -> bool { + self.containers.is_empty() + } + + /// Returns the minimum value in the archived bitmap. + #[must_use] + pub fn min(&self) -> Option { + self.containers + .first() + .and_then(|c| c.min().map(|min| ((c.key.to_native() as u32) << 16) | (min as u32))) + } + + /// Returns the maximum value in the archived bitmap. + #[must_use] + pub fn max(&self) -> Option { + self.containers + .last() + .and_then(|c| c.max().map(|max| ((c.key.to_native() as u32) << 16) | (max as u32))) + } + + /// Checks if the archived bitmap contains the given value. + #[must_use] + pub fn contains(&self, value: u32) -> bool { + let key = (value >> 16) as u16; + let index = value as u16; + + self.containers + .binary_search_by_key(&key, |c| c.key.to_native()) + .is_ok_and(|loc| self.containers[loc].contains(index)) + } + + /// Returns the number of integers that are <= value. + #[must_use] + pub fn rank(&self, value: u32) -> u64 { + let key = (value >> 16) as u16; + let index = value as u16; + + match self.containers.binary_search_by_key(&key, |c| c.key.to_native()) { + Ok(loc) => { + self.containers[loc].rank(index) + + self.containers[..loc].iter().map(ArchivedContainer::len).sum::() + } + Err(loc) => self.containers[..loc].iter().map(ArchivedContainer::len).sum::(), + } + } + + /// Returns the `n`th integer in the set or `None` if `n >= len()` + #[must_use] + pub fn select(&self, mut n: u32) -> Option { + for c in self.containers.iter() { + let len = c.len(); + if (len as u32) > n { + return c + .select(n) + .map(|index| ((c.key.to_native() as u32) << 16) | (index as u32)); + } + n -= len as u32; + } + None + } +} + +impl ArchivedContainer { + pub fn contains(&self, index: u16) -> bool { + self.store.contains(index) + } + pub fn len(&self) -> u64 { + self.store.len() + } + pub fn min(&self) -> Option { + self.store.min() + } + pub fn max(&self) -> Option { + self.store.max() + } + pub fn rank(&self, index: u16) -> u64 { + self.store.rank(index) + } + pub fn select(&self, n: u32) -> Option { + self.store.select(n) + } +} + +impl ArchivedStore { + pub fn contains(&self, index: u16) -> bool { + match self { + ArchivedStore::Array(arr) => arr.contains(index), + ArchivedStore::Bitmap(bit) => bit.contains(index), + ArchivedStore::Run(run) => run.contains(index), + } + } + pub fn len(&self) -> u64 { + match self { + ArchivedStore::Array(arr) => arr.len(), + ArchivedStore::Bitmap(bit) => bit.len(), + ArchivedStore::Run(run) => run.len(), + } + } + pub fn min(&self) -> Option { + match self { + ArchivedStore::Array(arr) => arr.min(), + ArchivedStore::Bitmap(bit) => bit.min(), + ArchivedStore::Run(run) => run.min(), + } + } + pub fn max(&self) -> Option { + match self { + ArchivedStore::Array(arr) => arr.max(), + ArchivedStore::Bitmap(bit) => bit.max(), + ArchivedStore::Run(run) => run.max(), + } + } + pub fn rank(&self, index: u16) -> u64 { + match self { + ArchivedStore::Array(arr) => arr.rank(index), + ArchivedStore::Bitmap(bit) => bit.rank(index), + ArchivedStore::Run(run) => run.rank(index), + } + } + pub fn select(&self, n: u32) -> Option { + match self { + ArchivedStore::Array(arr) => arr.select(n), + ArchivedStore::Bitmap(bit) => bit.select(n), + ArchivedStore::Run(run) => run.select(n), + } + } +} + +impl ArchivedArrayStore { + pub fn contains(&self, index: u16) -> bool { + self.vec.as_slice().binary_search_by_key(&index, |x| x.to_native()).is_ok() + } + pub fn len(&self) -> u64 { + self.vec.len() as u64 + } + pub fn min(&self) -> Option { + self.vec.first().map(|x| x.to_native()) + } + pub fn max(&self) -> Option { + self.vec.last().map(|x| x.to_native()) + } + pub fn rank(&self, index: u16) -> u64 { + match self.vec.as_slice().binary_search_by_key(&index, |x| x.to_native()) { + Ok(loc) => loc as u64 + 1, + Err(loc) => loc as u64, + } + } + pub fn select(&self, n: u32) -> Option { + self.vec.get(n as usize).map(|x| x.to_native()) + } +} + +impl ArchivedBitmapStore { + pub fn contains(&self, index: u16) -> bool { + let key = (index / 64) as usize; + let bit = index % 64; + (self.bits[key].to_native() & (1 << bit)) != 0 + } + pub const fn len(&self) -> u64 { + self.len.to_native() + } + pub fn min(&self) -> Option { + self.bits + .iter() + .enumerate() + .find(|&(_, bit)| bit.to_native() != 0) + .map(|(index, bit)| (index * 64 + bit.to_native().trailing_zeros() as usize) as u16) + } + pub fn max(&self) -> Option { + self.bits.iter().enumerate().rev().find(|&(_, bit)| bit.to_native() != 0).map( + |(index, bit)| (index * 64 + (63 - bit.to_native().leading_zeros() as usize)) as u16, + ) + } + pub fn rank(&self, index: u16) -> u64 { + let key = (index / 64) as usize; + let bit = index % 64; + + self.bits[..key].iter().map(|v| v.to_native().count_ones() as u64).sum::() + + (self.bits[key].to_native() << (63 - bit)).count_ones() as u64 + } + pub fn select(&self, mut n: u32) -> Option { + for (key, word) in self.bits.iter().enumerate() { + let word = word.to_native(); + let count = word.count_ones(); + if n < count { + let mut w = word; + for _ in 0..n { + w &= w - 1; + } + return Some((key * 64 + w.trailing_zeros() as usize) as u16); + } + n -= count; + } + None + } +} + +impl ArchivedIntervalStore { + pub fn contains(&self, index: u16) -> bool { + use core::cmp::Ordering; + + self.0 + .as_slice() + .binary_search_by(|iv| { + let start = iv.start.to_native(); + let end = iv.end.to_native(); + if index < start { + Ordering::Greater + } else if index > end { + Ordering::Less + } else { + Ordering::Equal + } + }) + .is_ok() + } + pub fn len(&self) -> u64 { + self.0.iter().map(|iv| (iv.end.to_native() - iv.start.to_native()) as u64 + 1).sum() + } + pub fn min(&self) -> Option { + self.0.first().map(|iv| iv.start.to_native()) + } + pub fn max(&self) -> Option { + self.0.last().map(|iv| iv.end.to_native()) + } + pub fn rank(&self, index: u16) -> u64 { + let mut rank = 0; + for iv in self.0.iter() { + let start = iv.start.to_native(); + let end = iv.end.to_native(); + if end <= index { + rank += (end - start) as u64 + 1; + } else if start <= index { + rank += (index - start) as u64 + 1; + break; + } else { + break; + } + } + rank + } + pub fn select(&self, mut n: u32) -> Option { + for iv in self.0.iter() { + let start = iv.start.to_native(); + let end = iv.end.to_native(); + let len = (end - start) as u32 + 1; + if n < len { + return Some(start + n as u16); + } + n -= len; + } + None + } +} + +#[cfg(test)] +mod tests { + use rkyv::rancor::Error; + + use crate::{ArchivedRoaringBitmap, RoaringBitmap}; + + /// Helper to serialize a `RoaringBitmap` and check that all methods on the + /// `ArchivedRoaringBitmap` yield the exact same results as the native counterpart. + fn check_archived_methods(rb: &RoaringBitmap) { + let bytes = rkyv::to_bytes::(rb).unwrap(); + let archived = rkyv::access::(&bytes[..]).unwrap(); + + assert_eq!(archived.len(), rb.len(), "Length mismatch"); + assert_eq!(archived.is_empty(), rb.is_empty(), "is_empty mismatch"); + assert_eq!(archived.min(), rb.min(), "Min mismatch"); + assert_eq!(archived.max(), rb.max(), "Max mismatch"); + + // Dynamically check boundaries + if let Some(min) = rb.min() { + assert!(archived.contains(min)); + if min > 0 { + assert_eq!(archived.contains(min - 1), rb.contains(min - 1)); + } + } + + if let Some(max) = rb.max() { + assert!(archived.contains(max)); + if max < u32::MAX { + assert_eq!(archived.contains(max + 1), rb.contains(max + 1)); + } + } + + // Check a varied assortment of presence/absence across potential container boundaries + let check_vals = [0, 1, 100, 4095, 4096, 5000, 10000, 65535, 65536, 100_000, u32::MAX]; + for &v in &check_vals { + assert_eq!(archived.contains(v), rb.contains(v), "Mismatch at {v}"); + assert_eq!(archived.rank(v), rb.rank(v), "Rank mismatch at {v}"); + } + + // Check select + if !rb.is_empty() { + let select_vals = [0, (rb.len() / 2) as u32, (rb.len() - 1) as u32]; + for &n in &select_vals { + assert_eq!(archived.select(n), rb.select(n), "Select mismatch at {n}"); + } + } + assert_eq!(archived.select(rb.len() as u32), None); + } + + #[test] + fn test_empty() { + let rb = RoaringBitmap::new(); + check_archived_methods(&rb); + } + + #[test] + fn test_array_store() { + // Less than 4096 elements creates an ArrayStore + let mut rb = RoaringBitmap::new(); + rb.insert(1); + rb.insert(2); + rb.insert(100); + rb.insert(1000); + + check_archived_methods(&rb); + } + + #[test] + fn test_bitmap_store() { + let mut rb = RoaringBitmap::new(); + // Inserting > 4096 scattered elements forces a BitmapStore + for i in (0..10_000).step_by(2) { + rb.insert(i); + } + + check_archived_methods(&rb); + } + + #[test] + fn test_interval_store() { + let mut rb = RoaringBitmap::new(); + rb.insert_range(100..=5000); + // Calling optimize on a contiguous block converts it to an IntervalStore (Run container) + rb.optimize(); + + check_archived_methods(&rb); + } + + #[test] + fn test_mixed_stores() { + let mut rb = RoaringBitmap::new(); + + // Container 0: ArrayStore + rb.insert(1); + rb.insert(10); + + // Container 1: BitmapStore + let offset = 1 << 16; + for i in (0..10_000).step_by(2) { + rb.insert(offset + i); + } + + // Container 2: IntervalStore + let offset2 = 2 << 16; + rb.insert_range(offset2 + 100..=offset2 + 5000); + rb.optimize(); + + check_archived_methods(&rb); + } +} diff --git a/roaring/src/bitmap/store/array_store/mod.rs b/roaring/src/bitmap/store/array_store/mod.rs index 61967704..7719e9cb 100644 --- a/roaring/src/bitmap/store/array_store/mod.rs +++ b/roaring/src/bitmap/store/array_store/mod.rs @@ -21,8 +21,9 @@ use super::Interval; pub(crate) const ARRAY_ELEMENT_BYTES: usize = 2; #[derive(Clone, Eq, PartialEq)] +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] pub(crate) struct ArrayStore { - vec: Vec, + pub(crate) vec: Vec, } /// Return the first contiguous range of elements in a sorted slice. diff --git a/roaring/src/bitmap/store/bitmap_store.rs b/roaring/src/bitmap/store/bitmap_store.rs index ca32e2b5..c3d83ca5 100644 --- a/roaring/src/bitmap/store/bitmap_store.rs +++ b/roaring/src/bitmap/store/bitmap_store.rs @@ -16,9 +16,10 @@ pub const BITMAP_LENGTH: usize = 1024; pub const BITMAP_BYTES: usize = BITMAP_LENGTH * 8; #[derive(Clone, Eq, PartialEq)] -pub struct BitmapStore { - len: u64, - bits: Box<[u64; BITMAP_LENGTH]>, +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] +pub(crate) struct BitmapStore { + pub(crate) len: u64, + pub(crate) bits: Box<[u64; BITMAP_LENGTH]>, } impl BitmapStore { diff --git a/roaring/src/bitmap/store/interval_store.rs b/roaring/src/bitmap/store/interval_store.rs index 86b20223..7e2a08cb 100644 --- a/roaring/src/bitmap/store/interval_store.rs +++ b/roaring/src/bitmap/store/interval_store.rs @@ -8,7 +8,8 @@ use core::{cmp::Ordering, ops::ControlFlow}; use super::{ArrayStore, BitmapStore}; #[derive(PartialEq, Eq, Clone, Debug)] -pub(crate) struct IntervalStore(Vec); +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] +pub(crate) struct IntervalStore(pub(crate) Vec); pub(crate) const RUN_NUM_BYTES: usize = 2; pub(crate) const RUN_ELEMENT_BYTES: usize = 4; @@ -897,9 +898,10 @@ impl> ExactSizeIterator for RunIter {} /// This interval is inclusive to end. #[derive(PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Debug)] +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] pub(crate) struct Interval { - start: u16, - end: u16, + pub(crate) start: u16, + pub(crate) end: u16, } impl From> for Interval { diff --git a/roaring/src/bitmap/store/mod.rs b/roaring/src/bitmap/store/mod.rs index 11a3de0c..134f30a3 100644 --- a/roaring/src/bitmap/store/mod.rs +++ b/roaring/src/bitmap/store/mod.rs @@ -13,18 +13,26 @@ pub use self::bitmap_store::{BITMAP_BYTES, BITMAP_LENGTH}; use self::Store::{Array, Bitmap, Run}; pub(crate) use self::array_store::ArrayStore; -pub use self::bitmap_store::{BitmapIter, BitmapStore}; +pub(crate) use self::bitmap_store::{BitmapIter, BitmapStore}; pub(crate) use self::interval_store::Interval; pub(crate) use interval_store::{IntervalStore, RunIterBorrowed, RunIterOwned}; #[cfg(feature = "std")] pub(crate) use interval_store::{RUN_ELEMENT_BYTES, RUN_NUM_BYTES}; +#[cfg(feature = "rkyv")] +pub(crate) use self::array_store::ArchivedArrayStore; +#[cfg(feature = "rkyv")] +pub(crate) use self::bitmap_store::ArchivedBitmapStore; +#[cfg(feature = "rkyv")] +pub(crate) use self::interval_store::ArchivedIntervalStore; + use crate::bitmap::container::ARRAY_LIMIT; #[cfg(not(feature = "std"))] use alloc::boxed::Box; #[derive(Clone)] +#[cfg_attr(feature = "rkyv", derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize))] pub(crate) enum Store { Array(ArrayStore), Bitmap(BitmapStore), diff --git a/roaring/src/lib.rs b/roaring/src/lib.rs index f6c4453b..054bfa84 100644 --- a/roaring/src/lib.rs +++ b/roaring/src/lib.rs @@ -33,6 +33,11 @@ pub mod treemap; pub use bitmap::RoaringBitmap; pub use treemap::RoaringTreemap; +#[cfg(feature = "rkyv")] +pub use bitmap::ArchivedRoaringBitmap; +#[cfg(feature = "rkyv")] +pub use treemap::ArchivedRoaringTreemap; + /// An error type that is returned when a `try_push` in a bitmap did not succeed. #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct IntegerTooSmall; diff --git a/roaring/src/treemap/mod.rs b/roaring/src/treemap/mod.rs index 93ca126d..aeed2430 100644 --- a/roaring/src/treemap/mod.rs +++ b/roaring/src/treemap/mod.rs @@ -12,8 +12,12 @@ mod cmp; mod inherent; mod iter; mod ops; + +#[cfg(feature = "rkyv")] +mod rkyv; #[cfg(feature = "serde")] mod serde; + #[cfg(feature = "std")] mod serialization; @@ -37,6 +41,7 @@ pub use self::iter::{BitmapIter, IntoIter, Iter}; /// println!("total bits set to true: {}", rb.len()); /// ``` #[derive(PartialEq, Eq)] +#[cfg_attr(feature = "rkyv", derive(::rkyv::Archive, ::rkyv::Serialize, ::rkyv::Deserialize))] pub struct RoaringTreemap { map: BTreeMap, } diff --git a/roaring/src/treemap/rkyv.rs b/roaring/src/treemap/rkyv.rs new file mode 100644 index 00000000..ed8e8c3c --- /dev/null +++ b/roaring/src/treemap/rkyv.rs @@ -0,0 +1,143 @@ +impl super::ArchivedRoaringTreemap { + /// Returns the number of elements in the archived treemap. + #[must_use] + pub fn len(&self) -> u64 { + self.map.iter().map(|(_, rb)| rb.len()).sum() + } + + /// Checks if the archived treemap is empty. + #[must_use] + pub fn is_empty(&self) -> bool { + self.map.iter().all(|(_, rb)| rb.is_empty()) + } + + /// Returns the minimum value in the archived treemap. + #[must_use] + pub fn min(&self) -> Option { + self.map + .iter() + .find_map(|(k, rb)| rb.min().map(|min| ((k.to_native() as u64) << 32) | (min as u64))) + } + + /// Returns the maximum value in the archived treemap. + #[must_use] + pub fn max(&self) -> Option { + self.map + .iter() + .filter(|(_, rb)| !rb.is_empty()) + .last() + .and_then(|(k, rb)| rb.max().map(|max| ((k.to_native() as u64) << 32) | (max as u64))) + } + + /// Checks if the archived treemap contains the given value. + #[must_use] + pub fn contains(&self, value: u64) -> bool { + let hi = (value >> 32) as u32; + let lo = value as u32; + + let hi_archived = rkyv::Archived::::from_native(hi); + self.map.get(&hi_archived).is_some_and(|rb| rb.contains(lo)) + } + + /// Returns the number of integers that are <= value. + #[must_use] + pub fn rank(&self, value: u64) -> u64 { + let hi = (value >> 32) as u32; + let lo = value as u32; + + let mut rank = 0; + for (k, rb) in self.map.iter() { + let k = k.to_native(); + match k.cmp(&hi) { + std::cmp::Ordering::Less => rank += rb.len(), + std::cmp::Ordering::Equal => { + rank += rb.rank(lo); + break; + } + std::cmp::Ordering::Greater => break, + } + } + rank + } + + /// Returns the `n`th integer in the set or `None` if `n >= len()` + #[must_use] + pub fn select(&self, mut n: u64) -> Option { + for (k, rb) in self.map.iter() { + let len = rb.len(); + if n < len { + return rb.select(n as u32).map(|lo| ((k.to_native() as u64) << 32) | (lo as u64)); + } + n -= len; + } + None + } +} + +#[cfg(test)] +mod tests { + use rkyv::rancor::Error; + + use crate::{ArchivedRoaringTreemap, RoaringTreemap}; + + /// Helper to serialize a `RoaringTreemap` and check that all methods on the + /// `ArchivedRoaringTreemap` yield the exact same results as the native counterpart. + fn check_archived_methods(rt: &RoaringTreemap) { + let bytes = rkyv::to_bytes::(rt).unwrap(); + let archived = rkyv::access::(&bytes[..]).unwrap(); + + assert_eq!(archived.len(), rt.len(), "Length mismatch"); + assert_eq!(archived.is_empty(), rt.is_empty(), "is_empty mismatch"); + assert_eq!(archived.min(), rt.min(), "Min mismatch"); + assert_eq!(archived.max(), rt.max(), "Max mismatch"); + + // Dynamically check boundaries + if let Some(min) = rt.min() { + assert!(archived.contains(min)); + if min > 0 { + assert_eq!(archived.contains(min - 1), rt.contains(min - 1)); + } + } + + if let Some(max) = rt.max() { + assert!(archived.contains(max)); + if max < u64::MAX { + assert_eq!(archived.contains(max + 1), rt.contains(max + 1)); + } + } + + // Check a varied assortment of presence/absence across potential container boundaries + let check_vals = [0, 1, 100, 4095, 4096, 5000, 10000, 65535, 65536, 100_000, u64::MAX]; + for &v in &check_vals { + assert_eq!(archived.contains(v), rt.contains(v), "Mismatch at {v}"); + assert_eq!(archived.rank(v), rt.rank(v), "Rank mismatch at {v}"); + } + + // Check select + if !rt.is_empty() { + let select_vals = [0, rt.len() / 2, rt.len() - 1]; + for &n in &select_vals { + assert_eq!(archived.select(n), rt.select(n), "Select mismatch at {n}"); + } + } + assert_eq!(archived.select(rt.len()), None); + } + + #[test] + fn test_empty() { + let rt = RoaringTreemap::new(); + check_archived_methods(&rt); + } + + #[test] + fn test_basic() { + let mut rt = RoaringTreemap::new(); + rt.insert(1); + rt.insert(2); + rt.insert(100); + rt.insert(1000); + rt.insert(u32::MAX as u64 + 10); + + check_archived_methods(&rt); + } +}