diff --git a/Cargo.lock b/Cargo.lock index cd148c698..d36efed44 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "access-control-example" @@ -290,7 +290,7 @@ dependencies = [ "async-stream", "async-trait", "auto_impl", - "dashmap", + "dashmap 5.5.3", "futures", "futures-utils-wasm", "lru", @@ -1320,6 +1320,20 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+/~https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "der" version = "0.7.9" @@ -2387,6 +2401,7 @@ name = "motsu" version = "0.1.0" dependencies = [ "const-hex", + "dashmap 6.1.0", "motsu-proc", "once_cell", "stylus-sdk", @@ -4173,7 +4188,7 @@ dependencies = [ "cc", "cfg-if", "corosensei", - "dashmap", + "dashmap 5.5.3", "derivative", "enum-iterator", "fnv", diff --git a/Cargo.toml b/Cargo.toml index 32510ad44..1043b69e7 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/contracts/src/token/erc721/extensions/enumerable.rs b/contracts/src/token/erc721/extensions/enumerable.rs index 3511489eb..f35bc97c0 100644 --- a/contracts/src/token/erc721/extensions/enumerable.rs +++ b/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}; diff --git a/contracts/src/utils/structs/bitmap.rs b/contracts/src/utils/structs/bitmap.rs index b36c0037e..8c9a8d188 100644 --- a/contracts/src/utils/structs/bitmap.rs +++ b/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; diff --git a/lib/motsu-proc/src/test.rs b/lib/motsu-proc/src/test.rs index fe4af7d00..2e2425d1e 100644 --- a/lib/motsu-proc/src/test.rs +++ b/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() diff --git a/lib/motsu/Cargo.toml b/lib/motsu/Cargo.toml index 5f3e5958a..df8595555 100644 --- a/lib/motsu/Cargo.toml +++ b/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 diff --git a/lib/motsu/src/context.rs b/lib/motsu/src/context.rs index 265790f14..381d504b9 100644 --- a/lib/motsu/src/context.rs +++ b/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> = + 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(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, +} + +/// 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 diff --git a/lib/motsu/src/lib.rs b/lib/motsu/src/lib.rs index 98fcd5dbd..88eb49707 100644 --- a/lib/motsu/src/lib.rs +++ b/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; diff --git a/lib/motsu/src/prelude.rs b/lib/motsu/src/prelude.rs index b4b541a8f..5c8798297 100644 --- a/lib/motsu/src/prelude.rs +++ b/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, }; diff --git a/lib/motsu/src/shims.rs b/lib/motsu/src/shims.rs index 0860a6142..a20ebdfce 100644 --- a/lib/motsu/src/shims.rs +++ b/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(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, diff --git a/lib/motsu/src/storage.rs b/lib/motsu/src/storage.rs deleted file mode 100644 index 4e8d23d4a..000000000 --- a/lib/motsu/src/storage.rs +++ /dev/null @@ -1,32 +0,0 @@ -//! Shims for storage operations. -use std::{collections::HashMap, ptr, sync::Mutex}; - -use once_cell::sync::Lazy; - -use crate::shims::{Bytes32, WORD_BYTES}; - -/// Storage mock: A global mutable key-value store. -pub(crate) static STORAGE: Lazy>> = - Lazy::new(|| Mutex::new(HashMap::new())); - -/// Read the word at address `key`. -pub(crate) unsafe fn read_bytes32(key: *const u8) -> Bytes32 { - let mut res = Bytes32::default(); - ptr::copy(key, res.as_mut_ptr(), WORD_BYTES); - res -} - -/// Write the word `val` to the location pointed by `key`. -pub(crate) unsafe fn write_bytes32(key: *mut u8, val: Bytes32) { - ptr::copy(val.as_ptr(), key, WORD_BYTES); -} - -/// Clears storage, removing all key-value pairs. -/// -/// # Panics -/// -/// May panic if the storage lock is already held by the current thread. -#[allow(clippy::module_name_repetitions)] -pub fn reset_storage() { - STORAGE.lock().unwrap().clear(); -}