use commonware_utils::bitmap::BitMap;
use std::{collections::BTreeMap, num::NonZeroU64};
pub fn bits_for_indices<const N: usize>(
items_per_blob: NonZeroU64,
indices: impl IntoIterator<Item = u64>,
) -> BTreeMap<u64, Option<BitMap<N>>> {
let items_per_blob = items_per_blob.get();
let mut bits = BTreeMap::new();
for index in indices {
let blob = index / items_per_blob;
let offset = index % items_per_blob;
bits.entry(blob)
.or_insert_with(|| Some(BitMap::zeroes(items_per_blob)))
.as_mut()
.unwrap()
.set(offset, true);
}
bits
}
#[cfg(test)]
mod tests {
use super::*;
use commonware_utils::NZU64;
#[test]
fn test_bits_for_indices() {
let empty = bits_for_indices::<1>(NZU64!(10), core::iter::empty());
assert!(empty.is_empty());
let bits = bits_for_indices::<1>(NZU64!(10), [0, 1, 9, 10, 25]);
assert_eq!(bits.len(), 3);
let blob_0 = bits.get(&0).unwrap().as_ref().unwrap();
assert_eq!(blob_0.len(), 10);
assert_eq!(blob_0.count_ones(), 3);
assert!(blob_0.get(0));
assert!(blob_0.get(1));
assert!(blob_0.get(9));
let blob_1 = bits.get(&1).unwrap().as_ref().unwrap();
assert_eq!(blob_1.len(), 10);
assert_eq!(blob_1.count_ones(), 1);
assert!(blob_1.get(0));
let blob_2 = bits.get(&2).unwrap().as_ref().unwrap();
assert_eq!(blob_2.len(), 10);
assert_eq!(blob_2.count_ones(), 1);
assert!(blob_2.get(5));
}
}