use super::{array, run};
use bytes::{Buf, BufMut};
use commonware_codec::{EncodeSize, Error as CodecError, Read, ReadExt, Write};
use core::ops::Range;
pub const WORDS: usize = 1024;
const ENCODED_BYTES: usize = WORDS * core::mem::size_of::<u64>();
pub const BITS: u32 = WORDS as u32 * 64;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Bitmap {
words: [u64; WORDS],
cardinality: u32,
run_count: u16,
}
impl Default for Bitmap {
fn default() -> Self {
Self::new()
}
}
impl From<&array::Array> for Bitmap {
fn from(array: &array::Array) -> Self {
let mut words = [0u64; WORDS];
for value in array.iter() {
let word_idx = (value >> 6) as usize;
let bit_idx = value & 63;
words[word_idx] |= 1u64 << bit_idx;
}
Self {
words,
cardinality: array.len() as u32,
run_count: count_runs(&words),
}
}
}
impl From<[u64; WORDS]> for Bitmap {
fn from(words: [u64; WORDS]) -> Self {
let cardinality = words.iter().map(|w| w.count_ones()).sum();
let run_count = count_runs(&words);
Self {
words,
cardinality,
run_count,
}
}
}
impl From<&run::Run> for Bitmap {
fn from(run: &run::Run) -> Self {
let mut bitmap = Self::new();
let mut cardinality = 0u32;
for (start, end_inclusive) in run.runs() {
bitmap.fill_range_unchecked(start, end_inclusive);
cardinality += (end_inclusive - start) as u32 + 1;
}
bitmap.cardinality = cardinality;
bitmap.run_count = run.run_count() as u16;
bitmap
}
}
impl Bitmap {
pub const fn new() -> Self {
Self {
words: [0; WORDS],
cardinality: 0,
run_count: 0,
}
}
fn fill_range_unchecked(&mut self, start: u16, end_inclusive: u16) {
let start_word = (start >> 6) as usize;
let end_word = (end_inclusive >> 6) as usize;
let start_bit = start & 63;
let end_bit = end_inclusive & 63;
if start_word == end_word {
let num_bits = end_bit - start_bit + 1;
let mask = if num_bits == 64 {
!0u64
} else {
((1u64 << num_bits) - 1) << start_bit
};
self.words[start_word] |= mask;
} else {
self.words[start_word] |= !0u64 << start_bit;
for word in &mut self.words[start_word + 1..end_word] {
*word = !0u64;
}
let last_mask = if end_bit == 63 {
!0u64
} else {
(1u64 << (end_bit + 1)) - 1
};
self.words[end_word] |= last_mask;
}
}
pub const fn run_count(&self) -> u16 {
self.run_count
}
pub const fn len(&self) -> u32 {
self.cardinality
}
pub const fn is_empty(&self) -> bool {
self.cardinality == 0
}
pub const fn is_full(&self) -> bool {
self.cardinality == BITS
}
pub const fn contains(&self, value: u16) -> bool {
let word_idx = (value >> 6) as usize;
let bit_idx = value & 63;
(self.words[word_idx] & (1u64 << bit_idx)) != 0
}
pub fn is_subset(&self, other: &Self) -> bool {
self.words
.iter()
.zip(other.words.iter())
.all(|(&a, &b)| (a & !b) == 0)
}
pub fn intersects(&self, other: &Self) -> bool {
self.words
.iter()
.zip(other.words.iter())
.any(|(&a, &b)| (a & b) != 0)
}
pub const fn insert(&mut self, value: u16) -> bool {
let word_idx = (value >> 6) as usize;
let bit_idx = value & 63;
let mask = 1u64 << bit_idx;
if (self.words[word_idx] & mask) != 0 {
return false;
}
let left_set = if value == 0 {
false
} else {
let lv = value - 1;
(self.words[(lv >> 6) as usize] & (1u64 << (lv & 63))) != 0
};
let right_set = if value == u16::MAX {
false
} else {
let rv = value + 1;
(self.words[(rv >> 6) as usize] & (1u64 << (rv & 63))) != 0
};
self.words[word_idx] |= mask;
self.cardinality += 1;
match (left_set, right_set) {
(false, false) => self.run_count += 1,
(true, false) | (false, true) => {}
(true, true) => self.run_count -= 1,
}
true
}
pub fn insert_range(&mut self, range: Range<u16>) -> u32 {
if range.is_empty() {
return 0;
}
let start = range.start;
let end = range.end;
let start_word = (start >> 6) as usize;
let end_word = ((end - 1) >> 6) as usize;
let start_bit = start & 63;
let end_bit = (end - 1) & 63;
let mut inserted = 0u32;
if start_word == end_word {
let num_bits = end_bit - start_bit + 1;
let mask = if num_bits == 64 {
!0u64
} else {
((1u64 << num_bits) - 1) << start_bit
};
let old_count = self.words[start_word].count_ones();
self.words[start_word] |= mask;
inserted = self.words[start_word].count_ones() - old_count;
} else {
let first_mask = !0u64 << start_bit;
let old_count = self.words[start_word].count_ones();
self.words[start_word] |= first_mask;
inserted += self.words[start_word].count_ones() - old_count;
for word in &mut self.words[start_word + 1..end_word] {
let old_count = word.count_ones();
*word = !0u64;
inserted += 64 - old_count;
}
let last_mask = if end_bit == 63 {
!0u64
} else {
(1u64 << (end_bit + 1)) - 1
};
let old_count = self.words[end_word].count_ones();
self.words[end_word] |= last_mask;
inserted += self.words[end_word].count_ones() - old_count;
}
self.cardinality += inserted;
self.run_count = count_runs(&self.words);
inserted
}
pub const fn iter(&self) -> Iter<'_> {
Iter {
words: &self.words,
word_idx: 0,
end_word: WORDS - 1,
end_mask: !0u64,
current_word: self.words[0],
}
}
pub fn iter_range(&self, range: Range<u32>) -> Iter<'_> {
let start = range.start.min(BITS);
let end = range.end.min(BITS);
if start >= end {
return Iter::empty(&self.words);
}
let start_word = (start >> 6) as usize;
let end_word = ((end - 1) >> 6) as usize;
let start_mask = !0u64 << (start & 63);
let end_bit = (end - 1) & 63;
let end_mask = if end_bit == 63 {
!0u64
} else {
(1u64 << (end_bit + 1)) - 1
};
let mut current_word = self.words[start_word] & start_mask;
if start_word == end_word {
current_word &= end_mask;
}
Iter {
words: &self.words,
word_idx: start_word,
end_word,
end_mask,
current_word,
}
}
pub const fn words(&self) -> &[u64; WORDS] {
&self.words
}
pub fn min(&self) -> Option<u16> {
for (word_idx, &word) in self.words.iter().enumerate() {
if word != 0 {
let bit_idx = word.trailing_zeros();
return Some((word_idx as u16) << 6 | bit_idx as u16);
}
}
None
}
pub fn max(&self) -> Option<u16> {
for (word_idx, &word) in self.words.iter().enumerate().rev() {
if word != 0 {
let bit_idx = 63 - word.leading_zeros();
return Some((word_idx as u16) << 6 | bit_idx as u16);
}
}
None
}
pub fn or_new(a: &Self, b: &Self) -> Self {
let mut words = [0u64; WORDS];
let mut cardinality = 0u32;
for i in (0..WORDS).step_by(4) {
let w0 = a.words[i] | b.words[i];
let w1 = a.words[i + 1] | b.words[i + 1];
let w2 = a.words[i + 2] | b.words[i + 2];
let w3 = a.words[i + 3] | b.words[i + 3];
words[i] = w0;
words[i + 1] = w1;
words[i + 2] = w2;
words[i + 3] = w3;
cardinality += w0.count_ones() + w1.count_ones() + w2.count_ones() + w3.count_ones();
}
let run_count = count_runs(&words);
Self {
words,
cardinality,
run_count,
}
}
pub fn and_new(a: &Self, b: &Self) -> Self {
let mut words = [0u64; WORDS];
let mut cardinality = 0u32;
for i in (0..WORDS).step_by(4) {
let w0 = a.words[i] & b.words[i];
let w1 = a.words[i + 1] & b.words[i + 1];
let w2 = a.words[i + 2] & b.words[i + 2];
let w3 = a.words[i + 3] & b.words[i + 3];
words[i] = w0;
words[i + 1] = w1;
words[i + 2] = w2;
words[i + 3] = w3;
cardinality += w0.count_ones() + w1.count_ones() + w2.count_ones() + w3.count_ones();
}
let run_count = count_runs(&words);
Self {
words,
cardinality,
run_count,
}
}
pub fn and_not_new(a: &Self, b: &Self) -> Self {
let mut words = [0u64; WORDS];
let mut cardinality = 0u32;
for i in (0..WORDS).step_by(4) {
let w0 = a.words[i] & !b.words[i];
let w1 = a.words[i + 1] & !b.words[i + 1];
let w2 = a.words[i + 2] & !b.words[i + 2];
let w3 = a.words[i + 3] & !b.words[i + 3];
words[i] = w0;
words[i + 1] = w1;
words[i + 2] = w2;
words[i + 3] = w3;
cardinality += w0.count_ones() + w1.count_ones() + w2.count_ones() + w3.count_ones();
}
let run_count = count_runs(&words);
Self {
words,
cardinality,
run_count,
}
}
}
fn count_runs(words: &[u64; WORDS]) -> u16 {
let mut runs: u32 = 0;
let mut prev_bit: u64 = 0;
for &w in words.iter() {
let extended = (w << 1) | prev_bit;
runs += (w & !extended).count_ones();
prev_bit = w >> 63;
}
runs as u16
}
impl Write for Bitmap {
fn write(&self, buf: &mut impl BufMut) {
let mut bytes = [0u8; ENCODED_BYTES];
for (dst, &word) in bytes.chunks_exact_mut(8).zip(self.words.iter()) {
dst.copy_from_slice(&word.to_be_bytes());
}
buf.put_slice(&bytes);
}
}
impl EncodeSize for Bitmap {
fn encode_size(&self) -> usize {
ENCODED_BYTES
}
}
impl Read for Bitmap {
type Cfg = ();
fn read_cfg(buf: &mut impl Buf, _cfg: &Self::Cfg) -> Result<Self, CodecError> {
let bytes = <[u8; ENCODED_BYTES]>::read(buf)?;
let mut words = [0u64; WORDS];
for (word, chunk) in words.iter_mut().zip(bytes.chunks_exact(8)) {
let mut word_bytes = [0u8; 8];
word_bytes.copy_from_slice(chunk);
*word = u64::from_be_bytes(word_bytes);
}
Ok(Self::from(words))
}
}
pub struct Iter<'a> {
words: &'a [u64; WORDS],
word_idx: usize,
end_word: usize,
end_mask: u64,
current_word: u64,
}
impl<'a> Iter<'a> {
const fn empty(words: &'a [u64; WORDS]) -> Self {
Self {
words,
word_idx: WORDS,
end_word: 0,
end_mask: 0,
current_word: 0,
}
}
}
impl Iterator for Iter<'_> {
type Item = u16;
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.current_word != 0 {
let bit_idx = self.current_word.trailing_zeros();
self.current_word &= self.current_word - 1; return Some((self.word_idx as u16) << 6 | bit_idx as u16);
}
if self.word_idx >= self.end_word {
return None;
}
self.word_idx += 1;
self.current_word = self.words[self.word_idx];
if self.word_idx == self.end_word {
self.current_word &= self.end_mask;
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
if self.word_idx >= WORDS || self.word_idx > self.end_word {
return (0, Some(0));
}
let mut remaining = self.current_word.count_ones();
if self.word_idx < self.end_word {
remaining += self.words[self.word_idx + 1..self.end_word]
.iter()
.map(|w| w.count_ones())
.sum::<u32>();
remaining += (self.words[self.end_word] & self.end_mask).count_ones();
}
(remaining as usize, Some(remaining as usize))
}
}
impl ExactSizeIterator for Iter<'_> {}
#[cfg(feature = "arbitrary")]
impl arbitrary::Arbitrary<'_> for Bitmap {
fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
let mut words = [0u64; WORDS];
for word in &mut words {
*word = u.arbitrary()?;
}
Ok(Self::from(words))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_and_empty() {
let container = Bitmap::new();
assert!(container.is_empty());
assert_eq!(container.len(), 0);
assert!(!container.is_full());
}
#[test]
fn test_insert_and_contains() {
let mut container = Bitmap::new();
assert!(container.insert(5));
assert!(container.insert(3));
assert!(container.insert(7));
assert!(!container.insert(5));
assert_eq!(container.len(), 3);
assert!(container.contains(3));
assert!(container.contains(5));
assert!(container.contains(7));
assert!(!container.contains(4));
}
#[test]
fn test_insert_range_single_word() {
let mut container = Bitmap::new();
let inserted = container.insert_range(5..10);
assert_eq!(inserted, 5);
assert_eq!(container.len(), 5);
let values: Vec<_> = container.iter().collect();
assert_eq!(values, vec![5, 6, 7, 8, 9]);
}
#[test]
fn test_insert_range_multiple_words() {
let mut container = Bitmap::new();
let inserted = container.insert_range(60..130);
assert_eq!(inserted, 70);
assert_eq!(container.len(), 70);
for i in 60..130 {
assert!(container.contains(i), "missing value {}", i);
}
}
#[test]
fn test_iterator() {
let mut container = Bitmap::new();
container.insert(100);
container.insert(10);
container.insert(1000);
container.insert(5);
let values: Vec<_> = container.iter().collect();
assert_eq!(values, vec![5, 10, 100, 1000]);
}
#[test]
fn test_iterator_size_hint_after_exhaustion_regression() {
let mut container = Bitmap::new();
container.insert(5);
container.insert(100);
let mut iter = container.iter();
assert_eq!(iter.size_hint(), (2, Some(2)));
assert_eq!(iter.next(), Some(5));
assert_eq!(iter.size_hint(), (1, Some(1)));
assert_eq!(iter.next(), Some(100));
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
let empty = Bitmap::new();
let mut iter = empty.iter();
assert_eq!(iter.size_hint(), (0, Some(0)));
assert_eq!(iter.next(), None);
assert_eq!(iter.size_hint(), (0, Some(0)));
}
#[test]
fn test_min_max() {
let mut container = Bitmap::new();
assert_eq!(container.min(), None);
assert_eq!(container.max(), None);
container.insert(50);
container.insert(10);
container.insert(100);
assert_eq!(container.min(), Some(10));
assert_eq!(container.max(), Some(100));
}
#[test]
fn test_is_full() {
let words = [!0u64; WORDS];
let container = Bitmap::from(words);
assert!(container.is_full());
assert_eq!(container.len(), BITS);
}
#[test]
fn test_bitwise_operations() {
let mut a = Bitmap::new();
let mut b = Bitmap::new();
a.insert(1);
a.insert(2);
a.insert(3);
b.insert(2);
b.insert(3);
b.insert(4);
let result = Bitmap::and_new(&a, &b);
let values: Vec<_> = result.iter().collect();
assert_eq!(values, vec![2, 3]);
let result = Bitmap::or_new(&a, &b);
let values: Vec<_> = result.iter().collect();
assert_eq!(values, vec![1, 2, 3, 4]);
let result = Bitmap::and_not_new(&a, &b);
let values: Vec<_> = result.iter().collect();
assert_eq!(values, vec![1]);
}
#[test]
fn test_is_subset() {
let mut a = Bitmap::new();
a.insert(1);
a.insert(100);
let mut b = Bitmap::new();
b.insert(1);
b.insert(50);
b.insert(100);
let mut c = Bitmap::new();
c.insert(1);
c.insert(99);
assert!(a.is_subset(&b));
assert!(!a.is_subset(&c));
}
#[test]
fn test_intersects() {
let mut a = Bitmap::new();
a.insert(1);
a.insert(1000);
let mut b = Bitmap::new();
b.insert(2);
b.insert(2000);
let mut c = Bitmap::new();
c.insert(1000);
c.insert(3000);
assert!(!a.intersects(&b));
assert!(a.intersects(&c));
}
#[test]
fn test_from_array() {
let mut array = array::Array::new();
array.insert(5);
array.insert(10);
array.insert(15);
let bitmap = Bitmap::from(&array);
assert_eq!(bitmap.len(), 3);
assert!(bitmap.contains(5));
assert!(bitmap.contains(10));
assert!(bitmap.contains(15));
}
#[test]
fn test_insert_range_shift_overflow_regression() {
let mut container = Bitmap::new();
let inserted = container.insert_range(0..64);
assert_eq!(inserted, 64);
assert_eq!(container.len(), 64);
for i in 0..64 {
assert!(container.contains(i), "missing value {}", i);
}
let mut container = Bitmap::new();
let inserted = container.insert_range(32..64);
assert_eq!(inserted, 32);
for i in 32..64 {
assert!(container.contains(i), "missing value {}", i);
}
let mut container = Bitmap::new();
let inserted = container.insert_range(60..128);
assert_eq!(inserted, 68);
for i in 60..128 {
assert!(container.contains(i), "missing value {}", i);
}
let mut container = Bitmap::new();
let inserted = container.insert_range(0..u16::MAX);
assert_eq!(inserted, 65535);
container.insert(u16::MAX);
assert!(container.is_full());
assert_eq!(container.len(), BITS);
let mut container = Bitmap::new();
let inserted = container.insert_range(65500..u16::MAX);
assert_eq!(inserted, 35);
for i in 65500..u16::MAX {
assert!(container.contains(i), "missing value {}", i);
}
let mut container = Bitmap::new();
let inserted = container.insert_range(64..128);
assert_eq!(inserted, 64);
for i in 64..128 {
assert!(container.contains(i), "missing value {}", i);
}
let mut container = Bitmap::new();
let inserted = container.insert_range(0..64);
assert_eq!(inserted, 64);
assert!(container.contains(63));
let mut container = Bitmap::new();
let inserted = container.insert_range(48..64);
assert_eq!(inserted, 16);
for i in 48..64 {
assert!(container.contains(i), "missing value {}", i);
}
}
#[test]
fn test_encode_size() {
let b = Bitmap::new();
assert_eq!(b.encode_size(), ENCODED_BYTES);
}
#[test]
fn test_encode_size_independent_of_cardinality() {
let empty_size = Bitmap::new().encode_size();
let mut full = Bitmap::new();
for i in 0..=u16::MAX {
full.insert(i);
}
assert_eq!(full.encode_size(), empty_size);
}
#[test]
fn test_run_count_empty() {
let b = Bitmap::new();
assert_eq!(b.run_count(), 0);
}
#[test]
fn test_run_count_single_insert() {
let mut b = Bitmap::new();
b.insert(42);
assert_eq!(b.run_count(), 1);
}
#[test]
fn test_run_count_extends_left() {
let mut b = Bitmap::new();
b.insert(10);
b.insert(11);
assert_eq!(b.run_count(), 1);
}
#[test]
fn test_run_count_extends_right() {
let mut b = Bitmap::new();
b.insert(10);
b.insert(9);
assert_eq!(b.run_count(), 1);
}
#[test]
fn test_run_count_isolated_inserts() {
let mut b = Bitmap::new();
b.insert(0);
b.insert(100);
b.insert(1000);
assert_eq!(b.run_count(), 3);
}
#[test]
fn test_run_count_bridge_decrements() {
let mut b = Bitmap::new();
b.insert(10);
b.insert(12);
assert_eq!(b.run_count(), 2);
b.insert(11); assert_eq!(b.run_count(), 1);
}
#[test]
fn test_run_count_idempotent_on_duplicate() {
let mut b = Bitmap::new();
b.insert(42);
let before = b.run_count();
b.insert(42); assert_eq!(b.run_count(), before);
}
#[test]
fn test_run_count_at_boundaries() {
let mut b = Bitmap::new();
b.insert(0);
assert_eq!(b.run_count(), 1);
b.insert(u16::MAX);
assert_eq!(b.run_count(), 2);
b.insert(u16::MAX - 1);
assert_eq!(b.run_count(), 2);
b.insert(1);
assert_eq!(b.run_count(), 2);
}
#[test]
fn test_run_count_after_insert_range() {
let mut b = Bitmap::new();
for i in 0u16..100 {
b.insert(i * 2);
}
assert_eq!(b.run_count(), 100);
b.insert_range(0..200);
assert_eq!(b.run_count(), 1);
}
#[test]
fn test_run_count_matches_scan_after_random_inserts() {
use crate::test_rng;
use rand::Rng;
let mut rng = test_rng();
let mut b = Bitmap::new();
for _ in 0..2000 {
let v: u16 = rng.gen();
b.insert(v);
}
assert_eq!(b.run_count(), super::count_runs(b.words()));
}
#[test]
fn test_from_words_run_count_matches_scan() {
let mut words = [0u64; WORDS];
words[0] = 0xAAAA_AAAA_AAAA_AAAA;
let b = Bitmap::from(words);
assert_eq!(b.run_count(), super::count_runs(b.words()));
assert_eq!(b.run_count(), 32);
}
#[cfg(feature = "arbitrary")]
mod conformance {
use commonware_codec::conformance::CodecConformance;
commonware_conformance::conformance_tests! {
CodecConformance<super::Bitmap>,
}
}
}