simple bitmap implementation

Signed-off-by: Uncle Stretch <uncle.stretch@ghostchain.io>

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<BucketId, Bucket>,
+    max_length: u32,
+    max_value: u32,
+}
+
+impl BitMap {
+    pub fn new(max_value: u32) -> Self {
+        let mut buckets: BTreeMap<BucketId, Bucket> = 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::<Bucket>() * 8).trailing_zeros()
+    }
+
+    pub fn iter_buckets(&self) -> impl Iterator<Item = (&BucketId, &Bucket)> {
+        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)
+    }
+}