feat(test): parallel unit test execution

Cargo.toml

@@ -93,6 +93,7 @@ regex = "1.10.4"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tokio = { version = "1.12.0", features = ["full"] }
 futures = "0.3.30"
+dashmap = "6.1.0"
 
 # procedural macros
 syn = { version = "2.0.58", features = ["full"] }

contracts/src/token/erc721/extensions/enumerable.rs

@@ -327,7 +327,6 @@ impl Erc721Enumerable {
 #[cfg(all(test, feature = "std"))]
 mod tests {
     use alloy_primitives::{address, uint, Address, U256};
-    use motsu::prelude::*;
     use stylus_sdk::msg;
 
     use super::{Erc721Enumerable, Error, IErc721Enumerable};

contracts/src/utils/structs/bitmap.rs

@@ -95,7 +95,6 @@ impl BitMap {
 #[cfg(all(test, feature = "std"))]
 mod tests {
     use alloy_primitives::{private::proptest::proptest, U256};
-    use motsu::prelude::*;
 
     use crate::utils::structs::bitmap::BitMap;
 

lib/motsu-proc/src/test.rs

@@ -15,37 +15,39 @@ pub(crate) fn test(_attr: &TokenStream, input: TokenStream) -> TokenStream {
     let fn_block = &item_fn.block;
     let fn_args = &sig.inputs;
 
-    // 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() {
-        return quote! {
-            #( #attrs )*
-            #[test]
-            fn #fn_name() #fn_return_type {
-                let _lock = ::motsu::prelude::acquire_storage();
-                let res = #fn_block;
-                ::motsu::prelude::reset_storage();
-                res
-            }
-        }
-        .into();
+    // Currently, more than one contract per unit test is not supported.
+    if fn_args.len() > 1 {
+        error!(fn_args, "expected at most one contract in test signature");
     }
 
-    // We can unwrap because we handle the empty case above. We don't support
-    // more than one parameter for now, so we skip them.
-    let arg = fn_args.first().unwrap();
-    let FnArg::Typed(arg) = arg else {
-        error!(arg, "unexpected receiver argument in test signature");
-    };
-    let contract_arg_binding = &arg.pat;
-    let contract_ty = &arg.ty;
+    // Whether 1 or none contracts will be declared.
+    let contract_declarations = fn_args.into_iter().map(|arg| {
+        let FnArg::Typed(arg) = arg else {
+            error!(arg, "unexpected receiver argument in test signature");
+        };
+        let contract_arg_binding = &arg.pat;
+        let contract_ty = &arg.ty;
+
+        // Test case assumes, that contract's variable has `&mut` reference
+        // to contract's type.
+        quote! {
+            let mut #contract_arg_binding = <#contract_ty>::default();
+            let #contract_arg_binding = &mut #contract_arg_binding;
+        }
+    });
+
+    // Output full testcase function.
+    // Declare contract.
+    // And in the end, reset storage for test context.
     quote! {
         #( #attrs )*
         #[test]
         fn #fn_name() #fn_return_type {
-            ::motsu::prelude::with_context::<#contract_ty>(| #contract_arg_binding |
-                #fn_block
-            )
+            use ::motsu::prelude::DefaultStorage;
+            #( #contract_declarations )*
+            let res = #fn_block;
+            ::motsu::prelude::Context::current().reset_storage();
+            res
         }
     }
     .into()

lib/motsu/Cargo.toml

@@ -14,6 +14,7 @@ once_cell.workspace = true
 tiny-keccak.workspace = true
 stylus-sdk.workspace = true
 motsu-proc.workspace = true
+dashmap.workspace = true
 
 [lints]
 workspace = true

lib/motsu/src/context.rs

@@ -1,34 +1,98 @@
 //! Unit-testing context for Stylus contracts.
-use std::sync::{Mutex, MutexGuard};
 
+use std::{collections::HashMap, ptr};
+
+use dashmap::DashMap;
+use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType};
 
-use crate::storage::reset_storage;
+use crate::prelude::{Bytes32, WORD_BYTES};
 
-/// A global static mutex.
-///
-/// We use this for scenarios where concurrent mutation of storage is wanted.
-/// For example, when a test harness is running, this ensures each test
-/// accesses storage in an non-overlapping manner.
+/// Context of stylus unit tests associated with the current test thread.
+#[allow(clippy::module_name_repetitions)]
+pub struct Context {
+    thread_name: ThreadName,
+}
+
+impl Context {
+    /// Get test context associated with the current test thread.
+    #[must_use]
+    pub fn current() -> Self {
+        Self { thread_name: ThreadName::current() }
+    }
+
+    /// Get the value at `key` in storage.
+    pub(crate) fn get_bytes(self, key: &Bytes32) -> Bytes32 {
+        let storage = STORAGE.entry(self.thread_name).or_default();
+        storage.contract_data.get(key).copied().unwrap_or_default()
+    }
+
+    /// Get the raw value at `key` in storage and write it to `value`.
+    pub(crate) unsafe fn get_bytes_raw(self, key: *const u8, value: *mut u8) {
+        let key = read_bytes32(key);
+
+        write_bytes32(value, self.get_bytes(&key));
+    }
+
+    /// Set the value at `key` in storage to `value`.
+    pub(crate) fn set_bytes(self, key: Bytes32, value: Bytes32) {
+        let mut storage = STORAGE.entry(self.thread_name).or_default();
+        storage.contract_data.insert(key, value);
+    }
+
+    /// Set the raw value at `key` in storage to `value`.
+    pub(crate) unsafe fn set_bytes_raw(self, key: *const u8, value: *const u8) {
+        let (key, value) = (read_bytes32(key), read_bytes32(value));
+        self.set_bytes(key, value);
+    }
+
+    /// Clears storage, removing all key-value pairs associated with the current
+    /// test thread.
+    pub fn reset_storage(self) {
+        STORAGE.remove(&self.thread_name);
+    }
+}
+
+/// Storage mock: A global mutable key-value store.
+/// Allows concurrent access.
 ///
-/// See [`with_context`].
-pub(crate) static STORAGE_MUTEX: Mutex<()> = Mutex::new(());
-
-/// Acquires access to storage.
-pub fn acquire_storage() -> MutexGuard<'static, ()> {
-    STORAGE_MUTEX.lock().unwrap_or_else(|e| {
-        reset_storage();
-        e.into_inner()
-    })
+/// The key is the name of the test thread, and the value is the storage of the
+/// test case.
+static STORAGE: Lazy<DashMap<ThreadName, MockStorage>> =
+    Lazy::new(DashMap::new);
+
+/// Test thread name metadata.
+#[derive(Clone, Eq, PartialEq, Hash)]
+struct ThreadName(String);
+
+impl ThreadName {
+    /// Get the name of the current test thread.
+    fn current() -> Self {
+        let current_thread_name = std::thread::current()
+            .name()
+            .expect("should retrieve current thread name")
+            .to_string();
+        Self(current_thread_name)
+    }
 }
 
-/// Decorates a closure by running it with exclusive access to storage.
-#[allow(clippy::module_name_repetitions)]
-pub fn with_context<C: StorageType>(closure: impl FnOnce(&mut C)) {
-    let _lock = acquire_storage();
-    let mut contract = C::default();
-    closure(&mut contract);
-    reset_storage();
+/// Storage for unit test's mock data.
+#[derive(Default)]
+struct MockStorage {
+    /// Contract's mock data storage.
+    contract_data: HashMap<Bytes32, Bytes32>,
+}
+
+/// Read the word from location pointed by `ptr`.
+unsafe fn read_bytes32(ptr: *const u8) -> Bytes32 {
+    let mut res = Bytes32::default();
+    ptr::copy(ptr, res.as_mut_ptr(), WORD_BYTES);
+    res
+}
+
+/// Write the word `bytes` to the location pointed by `ptr`.
+unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
+    ptr::copy(bytes.as_ptr(), ptr, WORD_BYTES);
 }
 
 /// Initializes fields of contract storage and child contract storages with

lib/motsu/src/lib.rs

@@ -42,20 +42,9 @@
 //! }
 //! ```
 //!
-//! Note that currently, test suites using [`motsu::test`][test_attribute] will
-//! run serially because of global access to storage.
-//!
-//! ### Notice
-//!
-//! We maintain this crate on a best-effort basis. We use it extensively on our
-//! own tests, so we will add here any symbols we may need. However, since we
-//! expect this to be a temporary solution, don't expect us to address all
-//! requests.
-//!
 //! [test_attribute]: crate::test
 mod context;
 pub mod prelude;
 mod shims;
-mod storage;
 
 pub use motsu_proc::test;

lib/motsu/src/prelude.rs

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

lib/motsu/src/shims.rs

@@ -37,42 +37,12 @@
 //!     }
 //! }
 //! ```
-//!
-//! Note that for proper usage, tests should have exclusive access to storage,
-//! since they run in parallel, which may cause undesired results.
-//!
-//! One solution is to wrap tests with a function that acquires a global mutex:
-//!
-//! ```rust,no_run
-//! use std::sync::{Mutex, MutexGuard};
-//!
-//! use motsu::prelude::reset_storage;
-//!
-//! pub static STORAGE_MUTEX: Mutex<()> = Mutex::new(());
-//!
-//! pub fn acquire_storage() -> MutexGuard<'static, ()> {
-//!     STORAGE_MUTEX.lock().unwrap()
-//! }
-//!
-//! pub fn with_context<C: Default>(closure: impl FnOnce(&mut C)) {
-//!     let _lock = acquire_storage();
-//!     let mut contract = C::default();
-//!     closure(&mut contract);
-//!     reset_storage();
-//! }
-//!
-//! #[motsu::test]
-//! fn reads_balance() {
-//!     let balance = token.balance_of(Address::ZERO);
-//!     assert_eq!(balance, U256::ZERO);
-//! }
-//! ```
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
 
 use tiny_keccak::{Hasher, Keccak};
 
-use crate::storage::{read_bytes32, write_bytes32, STORAGE};
+use crate::context::Context;
 
 pub(crate) const WORD_BYTES: usize = 32;
 pub(crate) type Bytes32 = [u8; WORD_BYTES];
@@ -90,10 +60,10 @@ pub unsafe extern "C" fn native_keccak256(
 ) {
     let mut hasher = Keccak::v256();
 
-    let data = unsafe { slice::from_raw_parts(bytes, len) };
+    let data = slice::from_raw_parts(bytes, len);
     hasher.update(data);
 
-    let output = unsafe { slice::from_raw_parts_mut(output, WORD_BYTES) };
+    let output = slice::from_raw_parts_mut(output, WORD_BYTES);
     hasher.finalize(output);
 }
 
@@ -110,16 +80,7 @@ pub unsafe extern "C" fn native_keccak256(
 /// May panic if unable to lock `STORAGE`.
 #[no_mangle]
 pub unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
-    let key = unsafe { read_bytes32(key) };
-
-    let value = STORAGE
-        .lock()
-        .unwrap()
-        .get(&key)
-        .map(Bytes32::to_owned)
-        .unwrap_or_default();
-
-    unsafe { write_bytes32(out, value) };
+    Context::current().get_bytes_raw(key, out);
 }
 
 /// Writes a 32-byte value to the permanent storage cache. Stylus's storage
@@ -141,8 +102,7 @@ pub unsafe extern "C" fn storage_cache_bytes32(
     key: *const u8,
     value: *const u8,
 ) {
-    let (key, value) = unsafe { (read_bytes32(key), read_bytes32(value)) };
-    STORAGE.lock().unwrap().insert(key, value);
+    Context::current().set_bytes_raw(key, value);
 }
 
 /// Persists any dirty values in the storage cache to the EVM state trie,