snowbridge_core/
sparse_bitmap.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: 2023 Snowfork <hello@snowfork.com>
3use frame_support::storage::StorageMap;
4use sp_std::marker::PhantomData;
5
6/// Sparse bitmap interface.
7pub trait SparseBitmap<BitMap>
8where
9	BitMap: StorageMap<u128, u128, Query = u128>,
10{
11	fn get(index: u128) -> bool;
12	fn set(index: u128);
13}
14
15/// Sparse bitmap implementation.
16pub struct SparseBitmapImpl<BitMap>(PhantomData<BitMap>);
17
18impl<BitMap> SparseBitmapImpl<BitMap>
19where
20	BitMap: StorageMap<u128, u128, Query = u128>,
21{
22	/// Computes the bucket index and the bit mask for a given bit index.
23	/// Each bucket contains 128 bits.
24	fn compute_bucket_and_mask(index: u128) -> (u128, u128) {
25		(index >> 7, 1u128 << (index & 127))
26	}
27}
28
29impl<BitMap> SparseBitmap<BitMap> for SparseBitmapImpl<BitMap>
30where
31	BitMap: StorageMap<u128, u128, Query = u128>,
32{
33	fn get(index: u128) -> bool {
34		// Calculate bucket and mask
35		let (bucket, mask) = Self::compute_bucket_and_mask(index);
36
37		// Retrieve bucket and check bit
38		let bucket_value = BitMap::get(bucket);
39		bucket_value & mask != 0
40	}
41
42	fn set(index: u128) {
43		// Calculate bucket and mask
44		let (bucket, mask) = Self::compute_bucket_and_mask(index);
45
46		// Mutate the storage to set the bit
47		BitMap::mutate(bucket, |value| {
48			*value |= mask; // Set the bit in the bucket
49		});
50	}
51}
52
53#[cfg(test)]
54mod tests {
55	use super::*;
56	use frame_support::{
57		storage::{generator::StorageMap as StorageMapHelper, storage_prefix},
58		Twox64Concat,
59	};
60	use sp_io::TestExternalities;
61	pub struct MockStorageMap;
62
63	impl StorageMapHelper<u128, u128> for MockStorageMap {
64		type Query = u128;
65		type Hasher = Twox64Concat;
66		fn pallet_prefix() -> &'static [u8] {
67			b"MyModule"
68		}
69
70		fn storage_prefix() -> &'static [u8] {
71			b"MyStorageMap"
72		}
73
74		fn prefix_hash() -> [u8; 32] {
75			storage_prefix(Self::pallet_prefix(), Self::storage_prefix())
76		}
77
78		fn from_optional_value_to_query(v: Option<u128>) -> Self::Query {
79			v.unwrap_or_default()
80		}
81
82		fn from_query_to_optional_value(v: Self::Query) -> Option<u128> {
83			Some(v)
84		}
85	}
86
87	type TestSparseBitmap = SparseBitmapImpl<MockStorageMap>;
88
89	#[test]
90	fn test_sparse_bitmap_set_and_get() {
91		TestExternalities::default().execute_with(|| {
92			let index = 300;
93			let (bucket, mask) = TestSparseBitmap::compute_bucket_and_mask(index);
94
95			// Test initial state
96			assert_eq!(MockStorageMap::get(bucket), 0);
97			assert!(!TestSparseBitmap::get(index));
98
99			// Set the bit
100			TestSparseBitmap::set(index);
101
102			// Test after setting
103			assert_eq!(MockStorageMap::get(bucket), mask);
104			assert!(TestSparseBitmap::get(index));
105		});
106	}
107
108	#[test]
109	fn test_sparse_bitmap_multiple_sets() {
110		TestExternalities::default().execute_with(|| {
111			let index1 = 300;
112			let index2 = 305; // Same bucket, different bit
113			let (bucket, _) = TestSparseBitmap::compute_bucket_and_mask(index1);
114
115			let (_, mask1) = TestSparseBitmap::compute_bucket_and_mask(index1);
116			let (_, mask2) = TestSparseBitmap::compute_bucket_and_mask(index2);
117
118			// Test initial state
119			assert_eq!(MockStorageMap::get(bucket), 0);
120			assert!(!TestSparseBitmap::get(index1));
121			assert!(!TestSparseBitmap::get(index2));
122
123			// Set the first bit
124			TestSparseBitmap::set(index1);
125
126			// Test after first set
127			assert_eq!(MockStorageMap::get(bucket), mask1);
128			assert!(TestSparseBitmap::get(index1));
129			assert!(!TestSparseBitmap::get(index2));
130
131			// Set the second bit
132			TestSparseBitmap::set(index2);
133
134			// Test after second set
135			assert_eq!(MockStorageMap::get(bucket), mask1 | mask2); // Bucket should contain both masks
136			assert!(TestSparseBitmap::get(index1));
137			assert!(TestSparseBitmap::get(index2));
138		})
139	}
140
141	#[test]
142	fn test_sparse_bitmap_different_buckets() {
143		TestExternalities::default().execute_with(|| {
144			let index1 = 300; // Bucket 1
145			let index2 = 300 + (1 << 7); // Bucket 2 (128 bits apart)
146
147			let (bucket1, _) = TestSparseBitmap::compute_bucket_and_mask(index1);
148			let (bucket2, _) = TestSparseBitmap::compute_bucket_and_mask(index2);
149
150			let (_, mask1) = TestSparseBitmap::compute_bucket_and_mask(index1);
151			let (_, mask2) = TestSparseBitmap::compute_bucket_and_mask(index2);
152
153			// Test initial state
154			assert_eq!(MockStorageMap::get(bucket1), 0);
155			assert_eq!(MockStorageMap::get(bucket2), 0);
156
157			// Set bits in different buckets
158			TestSparseBitmap::set(index1);
159			TestSparseBitmap::set(index2);
160
161			// Test after setting
162			assert_eq!(MockStorageMap::get(bucket1), mask1); // Bucket 1 should contain mask1
163			assert_eq!(MockStorageMap::get(bucket2), mask2); // Bucket 2 should contain mask2
164
165			assert!(TestSparseBitmap::get(index1));
166			assert!(TestSparseBitmap::get(index2));
167		})
168	}
169}
snowbridge_core/sparse_bitmap.rs

snowbridge_core/
sparse_bitmap.rs
