From 0363f3611027ef18111353f8366ff0d366c955dd Mon Sep 17 00:00:00 2001 From: Cormac Relf Date: Thu, 27 Nov 2025 11:47:09 +1100 Subject: [PATCH] Support allocative with `features = ["allocative"]` Allocative is a trait that lets you trace allocations. This gives you insight similar to `RoaringBitmap::statistics()`, but built-in to larger allocative traces of arbitrary data structures. --- Cargo.toml | 1 + roaring/Cargo.toml | 2 + roaring/src/bitmap/container.rs | 1 + roaring/src/bitmap/mod.rs | 1 + roaring/src/bitmap/store/array_store/mod.rs | 1 + roaring/src/bitmap/store/bitmap_store.rs | 1 + roaring/src/bitmap/store/interval_store.rs | 2 + roaring/src/bitmap/store/mod.rs | 1 + roaring/src/treemap/mod.rs | 1 + roaring/tests/allocative.rs | 50 +++++++++++++++++++++ 10 files changed, 61 insertions(+) create mode 100644 roaring/tests/allocative.rs diff --git a/Cargo.toml b/Cargo.toml index 6e43a8ac0..4a8b9c410 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ proptest = "1.6.0" serde = "1.0.217" serde_json = "1.0.138" zip = { version = "0.6", default-features = false } +allocative = { version = "0.3" } [profile.test] opt-level = 2 diff --git a/roaring/Cargo.toml b/roaring/Cargo.toml index 255579997..7b441dcd0 100644 --- a/roaring/Cargo.toml +++ b/roaring/Cargo.toml @@ -20,12 +20,14 @@ license = "MIT OR Apache-2.0" bytemuck = { workspace = true, optional = true } byteorder = { workspace = true, optional = true } serde = { workspace = true, optional = true } +allocative = { workspace = true, optional = true } [features] default = ["std"] serde = ["dep:serde", "std"] simd = [] std = ["dep:bytemuck", "dep:byteorder"] +allocative = ["dep:allocative"] [dev-dependencies] proptest = { workspace = true } diff --git a/roaring/src/bitmap/container.rs b/roaring/src/bitmap/container.rs index 49a6564b2..9e0d65f4a 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, Clone)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] 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 a63ebf241..53e7b86f7 100644 --- a/roaring/src/bitmap/mod.rs +++ b/roaring/src/bitmap/mod.rs @@ -45,6 +45,7 @@ use alloc::vec::Vec; /// println!("total bits set to true: {}", rb.len()); /// ``` #[derive(PartialEq)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub struct RoaringBitmap { containers: Vec, } diff --git a/roaring/src/bitmap/store/array_store/mod.rs b/roaring/src/bitmap/store/array_store/mod.rs index 173b7def3..a2f7c6d35 100644 --- a/roaring/src/bitmap/store/array_store/mod.rs +++ b/roaring/src/bitmap/store/array_store/mod.rs @@ -21,6 +21,7 @@ use super::Interval; pub(crate) const ARRAY_ELEMENT_BYTES: usize = 2; #[derive(Clone, Eq, PartialEq)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub(crate) struct ArrayStore { vec: Vec, } diff --git a/roaring/src/bitmap/store/bitmap_store.rs b/roaring/src/bitmap/store/bitmap_store.rs index 98eed58fa..5a81d97b7 100644 --- a/roaring/src/bitmap/store/bitmap_store.rs +++ b/roaring/src/bitmap/store/bitmap_store.rs @@ -16,6 +16,7 @@ pub const BITMAP_LENGTH: usize = 1024; pub const BITMAP_BYTES: usize = BITMAP_LENGTH * 8; #[derive(Clone, Eq, PartialEq)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub struct BitmapStore { len: u64, bits: Box<[u64; BITMAP_LENGTH]>, diff --git a/roaring/src/bitmap/store/interval_store.rs b/roaring/src/bitmap/store/interval_store.rs index 86b20223a..b75d5df1c 100644 --- a/roaring/src/bitmap/store/interval_store.rs +++ b/roaring/src/bitmap/store/interval_store.rs @@ -8,6 +8,7 @@ use core::{cmp::Ordering, ops::ControlFlow}; use super::{ArrayStore, BitmapStore}; #[derive(PartialEq, Eq, Clone, Debug)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub(crate) struct IntervalStore(Vec); pub(crate) const RUN_NUM_BYTES: usize = 2; @@ -897,6 +898,7 @@ impl> ExactSizeIterator for RunIter {} /// This interval is inclusive to end. #[derive(PartialEq, Eq, PartialOrd, Ord, Copy, Clone, Debug)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub(crate) struct Interval { start: u16, end: u16, diff --git a/roaring/src/bitmap/store/mod.rs b/roaring/src/bitmap/store/mod.rs index c57d3afe7..b8eb9a238 100644 --- a/roaring/src/bitmap/store/mod.rs +++ b/roaring/src/bitmap/store/mod.rs @@ -25,6 +25,7 @@ use crate::bitmap::container::ARRAY_LIMIT; use alloc::boxed::Box; #[derive(Clone)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub(crate) enum Store { Array(ArrayStore), Bitmap(BitmapStore), diff --git a/roaring/src/treemap/mod.rs b/roaring/src/treemap/mod.rs index 765f8482f..cfa1248e6 100644 --- a/roaring/src/treemap/mod.rs +++ b/roaring/src/treemap/mod.rs @@ -37,6 +37,7 @@ pub use self::iter::{BitmapIter, IntoIter, Iter}; /// println!("total bits set to true: {}", rb.len()); /// ``` #[derive(PartialEq)] +#[cfg_attr(feature = "allocative", derive(allocative::Allocative))] pub struct RoaringTreemap { map: BTreeMap, } diff --git a/roaring/tests/allocative.rs b/roaring/tests/allocative.rs new file mode 100644 index 000000000..6ee1d1b5f --- /dev/null +++ b/roaring/tests/allocative.rs @@ -0,0 +1,50 @@ +#![cfg(feature = "allocative")] + +use std::fs::File; +use std::io::Write; + +use allocative::FlameGraphBuilder; +use roaring::RoaringBitmap; +use roaring::RoaringTreemap; + +#[test] +fn flamegraph_bitmap() { + let mut foo1 = RoaringBitmap::new(); + foo1.insert_range(0..1_000_000); + foo1.insert(2_000_000); + foo1.insert(9_000_000); + + let mut flamegraph = FlameGraphBuilder::default(); + flamegraph.visit_root(&foo1); + let flamegraph_src = flamegraph.finish().flamegraph().write(); + let mut f = File::create("../target/bitmap.folded").unwrap(); + write!(f, "{}", flamegraph_src).unwrap(); + + /* + cargo test -p roaring --features allocative --test allocative + inferno-flamegraph target/bitmap.folded > target/bitmap.flamegraph.svg + open target/bitmap.flamegraph.svg + */ +} + +#[test] +fn flamegraph_treemap() { + let mut foo1 = RoaringTreemap::new(); + foo1.insert_range(0..1_000_000); + foo1.insert(2_000_000); + foo1.insert(9_000_000); + foo1.insert((1 << 32) + 9000); + foo1.insert((2 << 32) + 9000); + + let mut flamegraph = FlameGraphBuilder::default(); + flamegraph.visit_root(&foo1); + let flamegraph_src = flamegraph.finish().flamegraph().write(); + let mut f = File::create("../target/treemap.folded").unwrap(); + write!(f, "{}", flamegraph_src).unwrap(); + + /* + cargo test -p roaring --features allocative --test allocative + inferno-flamegraph target/treemap.folded > target/treemap.flamegraph.svg + open target/treemap.flamegraph.svg + */ +}