diff --git a/pallets/slow-clap/src/bitmap.rs b/pallets/slow-clap/src/bitmap.rs new file mode 100644 index 0000000..9e17c80 --- /dev/null +++ b/pallets/slow-clap/src/bitmap.rs @@ -0,0 +1,112 @@ +use codec::{Decode, Encode}; +use scale_info::TypeInfo; +use sp_core::RuntimeDebug; +use sp_std::collections::btree_map::BTreeMap; + +use crate::AuthIndex; + +pub type BucketId = u32; +pub type Bucket = u64; + +#[derive(Clone, PartialEq, Encode, Decode, RuntimeDebug, TypeInfo)] +pub struct BitMap { + buckets: BTreeMap, + max_length: u32, + max_value: u32, +} + +impl BitMap { + pub fn new(max_value: u32) -> Self { + let mut buckets: BTreeMap = Default::default(); + let max_length = (max_value >> Self::size_of_bucket()) + 1; + + for bucket_id in 0..(max_length) { + buckets.insert(bucket_id, 0); + } + + BitMap { buckets, max_length, max_value } + } + + pub fn size_of_bucket() -> u32 { + (core::mem::size_of::() * 8).trailing_zeros() + } + + pub fn iter_buckets(&self) -> impl Iterator { + self.buckets.iter() + } + + pub fn exists(&self, index: &AuthIndex) -> bool { + let (bucket_id, bit_position) = self.bitmap_positions(index); + self.buckets + .get(&bucket_id) + .map(|bucket| (bucket >> bit_position) & 1 == 1) + .unwrap_or_default() + } + + pub fn set(&mut self, index: AuthIndex) { + let (bucket_id, bit_position) = self.bitmap_positions(&index); + if bucket_id < self.max_length { + let bucket = self.buckets.entry(bucket_id).or_insert(0); + *bucket |= 1 << bit_position; + } + } + + pub fn unset(&mut self, index: AuthIndex) { + let (bucket_id, bit_position) = self.bitmap_positions(&index); + if let Some(bucket) = self.buckets.get_mut(&bucket_id) { + *bucket &= !(1 << bit_position); + if *bucket == 0 { + self.buckets.remove(&bucket_id); + } + } + } + + pub fn get_bucket(&self, bucket_id: &BucketId) -> Bucket { + self.buckets.get(bucket_id).copied().unwrap_or_default() + } + + pub fn count_ones(&self) -> u32 { + let zeros: u32 = self.iter_buckets() + .map(|(_, bucket)| bucket.count_zeros()) + .sum(); + let total_bits = self.total_bits(); + total_bits.saturating_sub(zeros) + } + + pub fn bitor(self, rhs: Self) -> Self { + let (mut base, to_merge) = if self.buckets.len() < rhs.buckets.len() { + (rhs, self) + } else { + (self, rhs) + }; + + for (key, rhs_value) in to_merge.buckets { + base.buckets + .entry(key) + .and_modify(|lhs_value| *lhs_value |= rhs_value) + .or_insert(rhs_value); + } + + base + } + + pub fn max_value(&self) -> u32 { + self.max_value + } + + fn total_bits(&self) -> u32 { + let size_of_bucket = Self::size_of_bucket(); + let bucket_length: u32 = 1 << size_of_bucket; + self.max_length.saturating_mul(bucket_length) + } + + fn bitmap_positions(&self, index: &AuthIndex) -> (u32, u32) { + let size_of_bucket = Self::size_of_bucket(); + let bucket_length = 1 << size_of_bucket; + + let bucket_id = index >> size_of_bucket; + let bit_position = index % bucket_length; + + (bucket_id, bit_position) + } +}