feat: bitmap

contracts/Cargo.toml

@@ -17,6 +17,7 @@ mini-alloc.workspace = true
 cfg-if = "1.0"
 
 [dev-dependencies]
+alloy-primitives = { version = "0.3.1", features = ["arbitrary"] }
 motsu = { path = "../lib/motsu" }
 rand = "0.8.5"
 

contracts/src/utils/mod.rs

@@ -1,5 +1,7 @@
 //! Common Smart Contracts utilities.
 
+pub mod structs;
+
 cfg_if::cfg_if! {
     if #[cfg(any(test, feature = "erc20_metadata", feature = "erc721_metadata"))] {
         pub mod metadata;

contracts/src/utils/structs/bitmap.rs

@@ -0,0 +1,140 @@
+//! Contract module for managing `U256` to boolean mapping in a compact and
+//! efficient way, provided the keys are sequential. Largely inspired by
+//! Uniswap's [merkle-distributor].
+//!
+//! `BitMap` packs 256 booleans across each bit of a single 256-bit slot of
+//! `U256` type. Hence, booleans corresponding to 256 _sequential_ indices
+//! would only consume a single slot, unlike the regular boolean which would
+//! consume an entire slot for a single value.
+//!
+//! This results in gas savings in two ways:
+//!
+//! - Setting a zero value to non-zero only once every 256 times
+//! - Accessing the same warm slot for every 256 _sequential_ indices
+//!
+//! [merkle-distributor]: https://github.com/Uniswap/merkle-distributor/blob/master/contracts/MerkleDistributor.sol
+use alloy_primitives::U256;
+use stylus_proc::sol_storage;
+
+const ONE: U256 = U256::from_limbs([1, 0, 0, 0]);
+const HEX_FF: U256 = U256::from_limbs([255, 0, 0, 0]);
+
+sol_storage! {
+    /// State of bit map.
+    #[cfg_attr(all(test, feature = "std"), derive(motsu::DefaultStorageLayout))]
+    pub struct BitMap {
+        /// Inner laying mapping.
+        mapping(uint256 => uint256) _data;
+    }
+}
+
+impl BitMap {
+    /// Returns whether the bit at `index` is set.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - index of the boolean value in the bit map.
+    #[must_use]
+    pub fn get(&self, index: U256) -> bool {
+        let bucket = Self::get_bucket(index);
+        let mask = Self::get_mask(index);
+        let value = self._data.get(bucket);
+        (value & mask) != U256::ZERO
+    }
+
+    /// Sets the bit at `index` to the boolean `value`.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - index of boolean value in the bit map.
+    /// * `value` - boolean value to set in the bit map.
+    pub fn set_to(&mut self, index: U256, value: bool) {
+        if value {
+            self.set(index);
+        } else {
+            self.unset(index);
+        }
+    }
+
+    /// Sets the bit at `index`.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - index of boolean value that should be set `true`.
+    pub fn set(&mut self, index: U256) {
+        let bucket = Self::get_bucket(index);
+        let mask = Self::get_mask(index);
+        let mut value = self._data.setter(bucket);
+        let prev = value.get();
+        value.set(prev | mask);
+    }
+
+    /// Unsets the bit at `index`.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - index of boolean value that should be set `false`.
+    pub fn unset(&mut self, index: U256) {
+        let bucket = Self::get_bucket(index);
+        let mask = Self::get_mask(index);
+        let mut value = self._data.setter(bucket);
+        let prev = value.get();
+        value.set(prev & !mask);
+    }
+
+    /// Get mask of value in the bucket.
+    fn get_mask(index: U256) -> U256 {
+        ONE << (index & HEX_FF)
+    }
+
+    /// Get bucket index.
+    fn get_bucket(index: U256) -> U256 {
+        index >> 8
+    }
+}
+
+#[cfg(all(test, feature = "std"))]
+mod tests {
+    use alloy_primitives::{
+        private::proptest::{
+            prelude::{Arbitrary, ProptestConfig},
+            proptest,
+        },
+        U256,
+    };
+    use stylus_sdk::{prelude::*, storage::StorageMap};
+
+    use crate::utils::structs::bitmap::BitMap;
+
+    #[motsu::test]
+    fn set_value() {
+        proptest!(|(value: U256)| {
+            let mut bit_map = BitMap::default();
+            assert_eq!(bit_map.get(value), false);
+            bit_map.set(value);
+            assert_eq!(bit_map.get(value), true);
+        });
+    }
+
+    #[motsu::test]
+    fn unset_value() {
+        proptest!(|(value: U256)| {
+            let mut bit_map = BitMap::default();
+            bit_map.set(value);
+            assert_eq!(bit_map.get(value), true);
+            bit_map.unset(value);
+            assert_eq!(bit_map.get(value), false);
+        });
+    }
+
+    #[motsu::test]
+    fn set_to_value() {
+        proptest!(|(value: U256)| {
+            let mut bit_map = BitMap::default();
+            bit_map.set_to(value, true);
+            assert_eq!(bit_map.get(value), true);
+            bit_map.set_to(value, false);
+            assert_eq!(bit_map.get(value), false);
+        });
+    }
+}

contracts/src/utils/structs/mod.rs

@@ -0,0 +1,2 @@
+//! Solidity storage types used by other contracts.
+pub mod bitmap;

lib/motsu-proc/src/test.rs

@@ -18,11 +18,15 @@ pub fn test(_attr: TokenStream, input: TokenStream) -> TokenStream {
     // If the test function has no params, then it doesn't need access to the
     // contract, so it is just a regular test.
     if fn_args.is_empty() {
-        let vis = &item_fn.vis;
         return quote! {
             #( #attrs )*
             #[test]
-            #vis #sig #fn_block
+            fn #fn_name() #fn_return_type {
+                let _lock = ::motsu::prelude::acquire_storage();
+                let res = #fn_block;
+                ::motsu::prelude::reset_storage();
+                res
+            }
         }
         .into();
     }

lib/motsu/src/context.rs

@@ -13,7 +13,7 @@ use crate::storage::reset_storage;
 pub(crate) static STORAGE_MUTEX: Mutex<()> = Mutex::new(());
 
 /// Acquires access to storage.
-pub(crate) fn acquire_storage() -> MutexGuard<'static, ()> {
+pub fn acquire_storage() -> MutexGuard<'static, ()> {
     STORAGE_MUTEX.lock().unwrap_or_else(|e| {
         reset_storage();
         e.into_inner()

lib/motsu/src/prelude.rs

@@ -1,2 +1,6 @@
 //! Common imports for `motsu` tests.
-pub use crate::{context::with_context, shims::*, storage::reset_storage};
+pub use crate::{
+    context::{acquire_storage, with_context},
+    shims::*,
+    storage::reset_storage,
+};