From a7d04757c18ed57461cbdfff7a9df84d572ee17b Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Mon, 3 Feb 2025 20:37:42 +0400
Subject: [PATCH 01/44] add msg sender and balance logic

---
 crates/motsu/src/context.rs | 189 +++++++++++++++++++++++++++++++++---
 crates/motsu/src/shims.rs   |  19 +++-
 2 files changed, 193 insertions(+), 15 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index c742bbc..196b8f2 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -2,7 +2,7 @@
 
 use std::{collections::HashMap, ptr, slice, thread::ThreadId};
 
-use alloy_primitives::Address;
+use alloy_primitives::{Address, U256};
 use dashmap::{mapref::one::RefMut, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
@@ -144,14 +144,40 @@ impl Context {
         calldata_len: usize,
         return_data_len: *mut usize,
     ) -> u8 {
-        let address_bytes = slice::from_raw_parts(address, 20);
-        let address = Address::from_slice(address_bytes);
+        let address = decode_address(address);
+        let (selector, input) = decode_calldata(calldata, calldata_len);
 
-        let input = slice::from_raw_parts(calldata, calldata_len);
-        let selector =
-            u32::from_be_bytes(TryInto::try_into(&input[..4]).unwrap());
+        let result = self.call_contract(address, selector, &input);
+        self.process_arb_result_raw(result, return_data_len)
+    }
 
-        match self.call_contract(address, selector, &input[4..]) {
+    /// Call the contract at raw `address` with the given raw `calldata` and
+    /// `value`.
+    pub(crate) unsafe fn call_contract_with_value_raw(
+        self,
+        address: *const u8,
+        calldata: *const u8,
+        calldata_len: usize,
+        value: *const u8,
+        return_data_len: *mut usize,
+    ) -> u8 {
+        let address = decode_address(address);
+        let value = decode_u256(value);
+        let (selector, input) = decode_calldata(calldata, calldata_len);
+
+        let result =
+            self.call_contract_with_value(address, selector, &input, value);
+        self.process_arb_result_raw(result, return_data_len)
+    }
+
+    /// Based on `result`, set the return data.
+    /// Return 0 if `result` is `Ok`, otherwise return 1.
+    unsafe fn process_arb_result_raw(
+        self,
+        result: ArbResult,
+        return_data_len: *mut usize,
+    ) -> u8 {
+        match result {
             Ok(res) => {
                 return_data_len.write(res.len());
                 self.set_return_data(res);
@@ -165,8 +191,30 @@ impl Context {
         }
     }
 
-    /// Call the function associated with the given `selector` and pass `input`
-    /// to it, at the given `contract_address`.
+    /// Call the function associated with the given `selector` at the given
+    /// `contract_address`. Pass `input` and `value` to it.
+    fn call_contract_with_value(
+        self,
+        contract_address: Address,
+        selector: u32,
+        input: &[u8],
+        value: U256,
+    ) -> ArbResult {
+        // Set new msg_value, and store the previous one.
+        let previous_msg_value = self.set_msg_value(value);
+
+        let result = self.call_contract(contract_address, selector, input);
+
+        // Set the previous msg_value if there is any.
+        if let Some(previous) = previous_msg_value {
+            let _ = self.set_msg_value(previous);
+        }
+
+        result
+    }
+
+    /// Call the function associated with the given `selector` at the given
+    /// `contract_address`. Pass `input` to it.
     fn call_contract(
         self,
         contract_address: Address,
@@ -231,8 +279,7 @@ impl Context {
 
     /// Check if the contract at raw `address` has code.
     pub(crate) unsafe fn has_code_raw(self, address: *const u8) -> bool {
-        let address_bytes = slice::from_raw_parts(address, 20);
-        let address = Address::from_slice(address_bytes);
+        let address = decode_address(address);
         self.has_code(address)
     }
 
@@ -242,6 +289,71 @@ impl Context {
         self.router(address).exists()
     }
 
+    pub(crate) unsafe fn balance_raw(self, address: *const u8) -> U256 {
+        // TODO#q: write to destination here
+        let address = decode_address(address);
+        self.balance(address)
+    }
+
+    fn balance(self, address: Address) -> U256 {
+        self.storage().balances.get(&address).copied().unwrap_or_default()
+    }
+
+    pub(crate) fn set_msg_value(self, value: U256) -> Option<U256> {
+        self.storage().msg_value.replace(value)
+    }
+
+    /// Write the value sent to the contract to `output`.
+    pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
+        let value: U256 = self.msg_value();
+        std::ptr::copy(value.as_le_slice().as_ptr(), output, 32);
+    }
+
+    /// Get the value sent to the contract as [`U256`].
+    pub(crate) fn msg_value(self) -> U256 {
+        self.storage().msg_value.unwrap_or_default()
+    }
+
+    fn transfer_funds_from_caller(self) {
+        let sender = self.msg_sender().expect("msg_sender should be set");
+        let value = self.msg_value();
+        self.add_assign_balance(sender, value);
+    }
+
+    fn checked_transfer(
+        self,
+        from: Address,
+        to: Address,
+        value: U256,
+    ) -> Option<()> {
+        let _ = self.checked_sub_assign_balance(from, value)?;
+        self.add_assign_balance(to, value);
+        Some(())
+    }
+
+    fn checked_sub_assign_balance(
+        self,
+        address: Address,
+        value: U256,
+    ) -> Option<U256> {
+        let mut storage = self.storage();
+        let balance = storage.balances.entry(address).or_default();
+        if *balance < value {
+            return None;
+        }
+        *balance -= value;
+        Some(*balance)
+    }
+
+    fn add_assign_balance(self, address: Address, value: U256) -> U256 {
+        *self
+            .storage()
+            .balances
+            .entry(address)
+            .and_modify(|v| *v += value)
+            .or_insert(value)
+    }
+
     /// Get reference to the storage for the current test thread.
     fn storage(self) -> RefMut<'static, Context, MockStorage> {
         STORAGE.get_mut(&self).expect("contract should be initialised first")
@@ -265,15 +377,44 @@ unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
     ptr::copy(bytes.as_ptr(), ptr, WORD_BYTES);
 }
 
+/// Decode the [`Address`] from the raw pointer.
+unsafe fn decode_address(ptr: *const u8) -> Address {
+    let address_bytes = slice::from_raw_parts(ptr, 20);
+    Address::from_slice(address_bytes)
+}
+
+/// Decode the [`U256`] from the raw pointer.
+unsafe fn decode_u256(ptr: *const u8) -> U256 {
+    let address_bytes = slice::from_raw_parts(ptr, 32);
+    U256::from_le_slice(address_bytes)
+}
+
+/// Decode the selector as [`u32`], and function input as [`Vec<u8>`] from the
+/// raw pointer.
+unsafe fn decode_calldata(
+    calldata: *const u8,
+    calldata_len: usize,
+) -> (u32, Vec<u8>) {
+    let calldata = slice::from_raw_parts(calldata, calldata_len);
+    let selector =
+        u32::from_be_bytes(TryInto::try_into(&calldata[..4]).unwrap());
+    let input = calldata[4..].to_vec();
+    (selector, input)
+}
+
 /// Storage for unit test's mock data.
 #[derive(Default)]
 struct MockStorage {
     /// Address of the message sender.
     msg_sender: Option<Address>,
+    /// The ETH value in wei sent to the program.
+    msg_value: Option<U256>,
     /// Address of the contract that is currently receiving the message.
     contract_address: Option<Address>,
     /// Contract's address to mock data storage mapping.
     contract_data: HashMap<Address, ContractStorage>,
+    /// Account's address to balance mapping.
+    balances: HashMap<Address, U256>,
     // Output of a contract call.
     call_output: Option<Vec<u8>>,
     // Output length of a contract call.
@@ -287,6 +428,7 @@ type ContractStorage = HashMap<Bytes32, Bytes32>;
 pub struct ContractCall<'a, ST: StorageType> {
     storage: ST,
     caller_address: Address,
+    value: Option<U256>,
     /// We need to hold a reference to [`Contract<ST>`], because
     /// `Contract::<ST>::new().sender(alice)` can accidentally drop
     /// [`Contract<ST>`].
@@ -304,6 +446,9 @@ impl<ST: StorageType> ContractCall<'_, ST> {
 
     /// Preset the call parameters.
     fn set_call_params(&self) {
+        if let Some(value) = self.value {
+            let _ = Context::current().set_msg_value(value);
+        }
         let _ = Context::current().set_msg_sender(self.caller_address);
         let _ = Context::current().set_contract_address(self.address());
     }
@@ -388,6 +533,22 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
             caller_address: account.into(),
+            value: None,
+            contract_ref: self,
+        }
+    }
+
+    /// Call contract `self` with `account` as a sender and `value`.
+    #[must_use]
+    pub fn sender_and_value<A: Into<Address>, V: Into<U256>>(
+        &self,
+        account: A,
+        value: V,
+    ) -> ContractCall<ST> {
+        ContractCall {
+            storage: unsafe { ST::new(uint!(0_U256), 0) },
+            caller_address: account.into(),
+            value: Some(value.into()),
             contract_ref: self,
         }
     }
@@ -423,4 +584,10 @@ impl Account {
     pub fn address(&self) -> Address {
         self.address
     }
+
+    /// Get account's balance.
+    #[must_use]
+    pub fn balance(&self) -> U256 {
+        Context::current().balance(self.address)
+    }
 }
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 95c639a..20e08a6 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -155,8 +155,7 @@ pub unsafe extern "C" fn msg_sender(sender: *mut u8) {
 /// Get the ETH value (U256) in wei sent to the program.
 #[no_mangle]
 pub unsafe extern "C" fn msg_value(value: *mut u8) {
-    let dummy_msg_value: Bytes32 = Bytes32::default();
-    std::ptr::copy(dummy_msg_value.as_ptr(), value, 32);
+    Context::current().msg_value_raw(value);
 }
 
 /// Gets the address of the current program. The semantics are equivalent to
@@ -227,6 +226,17 @@ pub unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
     std::ptr::copy(account_codehash.as_ptr(), dest, 32);
 }
 
+/// Gets the ETH balance in wei of the account at the given address.
+/// The semantics are equivalent to that of the EVM's [`BALANCE`] opcode.
+///
+/// [`BALANCE`]: https://www.evm.codes/#31
+pub unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
+    let balance = Context::current().balance_raw(address);
+    let balance_bytes = balance.as_le_slice();
+
+    std::ptr::copy(balance_bytes.as_ptr(), dest, 32);
+}
+
 /// Returns the length of the last EVM call or deployment return result, or `0`
 /// if neither have happened during the program's execution.
 ///
@@ -276,14 +286,15 @@ pub unsafe extern "C" fn call_contract(
     contract: *const u8,
     calldata: *const u8,
     calldata_len: usize,
-    _value: *const u8,
+    value: *const u8,
     _gas: u64,
     return_data_len: *mut usize,
 ) -> u8 {
-    Context::current().call_contract_raw(
+    Context::current().call_contract_with_value_raw(
         contract,
         calldata,
         calldata_len,
+        value,
         return_data_len,
     )
 }

From c145a71afdb7b96a8ede8ddb52e65425f920fa17 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Mon, 3 Feb 2025 21:22:02 +0400
Subject: [PATCH 02/44] ++

---
 crates/motsu/src/context.rs | 73 +++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 39 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 196b8f2..89c644c 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -144,10 +144,10 @@ impl Context {
         calldata_len: usize,
         return_data_len: *mut usize,
     ) -> u8 {
-        let address = decode_address(address);
+        let address = read_address(address);
         let (selector, input) = decode_calldata(calldata, calldata_len);
 
-        let result = self.call_contract(address, selector, &input);
+        let result = self.call_contract(address, selector, &input, None);
         self.process_arb_result_raw(result, return_data_len)
     }
 
@@ -161,12 +161,11 @@ impl Context {
         value: *const u8,
         return_data_len: *mut usize,
     ) -> u8 {
-        let address = decode_address(address);
-        let value = decode_u256(value);
+        let address = read_address(address);
+        let value = read_u256(value);
         let (selector, input) = decode_calldata(calldata, calldata_len);
 
-        let result =
-            self.call_contract_with_value(address, selector, &input, value);
+        let result = self.call_contract(address, selector, &input, Some(value));
         self.process_arb_result_raw(result, return_data_len)
     }
 
@@ -192,34 +191,13 @@ impl Context {
     }
 
     /// Call the function associated with the given `selector` at the given
-    /// `contract_address`. Pass `input` and `value` to it.
-    fn call_contract_with_value(
-        self,
-        contract_address: Address,
-        selector: u32,
-        input: &[u8],
-        value: U256,
-    ) -> ArbResult {
-        // Set new msg_value, and store the previous one.
-        let previous_msg_value = self.set_msg_value(value);
-
-        let result = self.call_contract(contract_address, selector, input);
-
-        // Set the previous msg_value if there is any.
-        if let Some(previous) = previous_msg_value {
-            let _ = self.set_msg_value(previous);
-        }
-
-        result
-    }
-
-    /// Call the function associated with the given `selector` at the given
-    /// `contract_address`. Pass `input` to it.
+    /// `contract_address`. Pass `input` and optional `value` to it.
     fn call_contract(
         self,
         contract_address: Address,
         selector: u32,
         input: &[u8],
+        value: Option<U256>,
     ) -> ArbResult {
         // Set the caller contract as message sender and callee contract as
         // a receiver (`contract_address`).
@@ -230,6 +208,10 @@ impl Context {
             .set_msg_sender(previous_contract_address)
             .expect("msg_sender should be set");
 
+        // Set new msg_value, and store the previous one.
+        let previous_msg_value =
+            value.and_then(|value| self.set_msg_value(value));
+
         // Call external contract.
         let result = self
             .router(contract_address)
@@ -238,10 +220,29 @@ impl Context {
                 panic!("selector not found - selector is {selector}")
             });
 
+        // If the call was successful,
+        if result.is_ok() {
+            // and there is a `value` to pay,
+            if let Some(value) = value {
+                // transfer it to the callee contract.
+                self.checked_transfer(
+                    previous_contract_address,
+                    contract_address,
+                    value,
+                )
+                    .expect("should have enough funds for transfer");
+            }
+        }
+
         // Set the previous message sender and contract address back.
         let _ = self.set_contract_address(previous_contract_address);
         let _ = self.set_msg_sender(previous_msg_sender);
 
+        // Set the previous msg_value if there is any.
+        if let Some(previous) = previous_msg_value {
+            let _ = self.set_msg_value(previous);
+        }
+
         result
     }
 
@@ -279,7 +280,7 @@ impl Context {
 
     /// Check if the contract at raw `address` has code.
     pub(crate) unsafe fn has_code_raw(self, address: *const u8) -> bool {
-        let address = decode_address(address);
+        let address = read_address(address);
         self.has_code(address)
     }
 
@@ -291,7 +292,7 @@ impl Context {
 
     pub(crate) unsafe fn balance_raw(self, address: *const u8) -> U256 {
         // TODO#q: write to destination here
-        let address = decode_address(address);
+        let address = read_address(address);
         self.balance(address)
     }
 
@@ -314,12 +315,6 @@ impl Context {
         self.storage().msg_value.unwrap_or_default()
     }
 
-    fn transfer_funds_from_caller(self) {
-        let sender = self.msg_sender().expect("msg_sender should be set");
-        let value = self.msg_value();
-        self.add_assign_balance(sender, value);
-    }
-
     fn checked_transfer(
         self,
         from: Address,
@@ -378,13 +373,13 @@ unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
 }
 
 /// Decode the [`Address`] from the raw pointer.
-unsafe fn decode_address(ptr: *const u8) -> Address {
+unsafe fn read_address(ptr: *const u8) -> Address {
     let address_bytes = slice::from_raw_parts(ptr, 20);
     Address::from_slice(address_bytes)
 }
 
 /// Decode the [`U256`] from the raw pointer.
-unsafe fn decode_u256(ptr: *const u8) -> U256 {
+unsafe fn read_u256(ptr: *const u8) -> U256 {
     let address_bytes = slice::from_raw_parts(ptr, 32);
     U256::from_le_slice(address_bytes)
 }

From 66096467b5663c1e721f63a3a88e100be13a45f1 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Mon, 3 Feb 2025 21:56:15 +0400
Subject: [PATCH 03/44] add test

---
 crates/motsu/src/context.rs | 124 ++++++++++++++++++++++++++++++------
 crates/motsu/src/lib.rs     |  66 +++++++++++++++++--
 crates/motsu/src/prelude.rs |   2 +-
 3 files changed, 166 insertions(+), 26 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 89c644c..686bf67 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -2,7 +2,7 @@
 
 use std::{collections::HashMap, ptr, slice, thread::ThreadId};
 
-use alloy_primitives::{Address, U256};
+use alloy_primitives::{Address, B256, U256};
 use dashmap::{mapref::one::RefMut, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
@@ -212,6 +212,8 @@ impl Context {
         let previous_msg_value =
             value.and_then(|value| self.set_msg_value(value));
 
+        // TODO#q: add check if sender has enough funds
+
         // Call external contract.
         let result = self
             .router(contract_address)
@@ -222,16 +224,9 @@ impl Context {
 
         // If the call was successful,
         if result.is_ok() {
-            // and there is a `value` to pay,
-            if let Some(value) = value {
-                // transfer it to the callee contract.
-                self.checked_transfer(
-                    previous_contract_address,
-                    contract_address,
-                    value,
-                )
-                    .expect("should have enough funds for transfer");
-            }
+            // transfer value that wasn't transferred yet.
+            self.try_transfer_value()
+                .expect("should have enough funds for transfer");
         }
 
         // Set the previous message sender and contract address back.
@@ -307,7 +302,7 @@ impl Context {
     /// Write the value sent to the contract to `output`.
     pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
         let value: U256 = self.msg_value();
-        std::ptr::copy(value.as_le_slice().as_ptr(), output, 32);
+        write_u256(value, output);
     }
 
     /// Get the value sent to the contract as [`U256`].
@@ -315,6 +310,33 @@ impl Context {
         self.storage().msg_value.unwrap_or_default()
     }
 
+    /// Transfer unsent value from the message sender to the contract.
+    ///
+    /// Returns `None` if there is not enough funds to transfer.
+    fn try_transfer_value(self) -> Option<()> {
+        let mut storage = self.storage();
+        let Some(msg_sender) = storage.msg_sender else {
+            return Some(());
+        };
+        let Some(contract_address) = storage.contract_address else {
+            return Some(());
+        };
+
+        // We should transfer the value only if it is set.
+        // And we should transfer the value only once, despite this function
+        // being called multiple times.
+        if let Some(msg_value) = storage.msg_value.take() {
+            // Drop storage to avoid deadlock.
+            drop(storage);
+            self.checked_transfer(msg_sender, contract_address, msg_value)
+        } else {
+            Some(())
+        }
+    }
+
+    /// Transfer `value` from `from` to `to`.
+    /// 
+    /// Returns `None` if there is not enough funds to transfer.
     fn checked_transfer(
         self,
         from: Address,
@@ -326,6 +348,9 @@ impl Context {
         Some(())
     }
 
+    /// Subtract `value` from the balance of `address`.
+    /// 
+    /// Returns `None` if there is not enough of funds.
     fn checked_sub_assign_balance(
         self,
         address: Address,
@@ -340,6 +365,7 @@ impl Context {
         Some(*balance)
     }
 
+    /// Add `value` to the balance of `address`.
     fn add_assign_balance(self, address: Address, value: U256) -> U256 {
         *self
             .storage()
@@ -372,16 +398,25 @@ unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
     ptr::copy(bytes.as_ptr(), ptr, WORD_BYTES);
 }
 
-/// Decode the [`Address`] from the raw pointer.
+// TODO#q: add write_address
+
+/// Read the [`Address`] from the raw pointer.
 unsafe fn read_address(ptr: *const u8) -> Address {
     let address_bytes = slice::from_raw_parts(ptr, 20);
     Address::from_slice(address_bytes)
 }
 
-/// Decode the [`U256`] from the raw pointer.
+/// Read the [`U256`] from the raw pointer.
 unsafe fn read_u256(ptr: *const u8) -> U256 {
-    let address_bytes = slice::from_raw_parts(ptr, 32);
-    U256::from_le_slice(address_bytes)
+    let mut data = B256::ZERO;
+    ptr::copy(ptr, data.as_mut_ptr(), 32);
+    data.into()
+}
+
+/// Write the [`U256`] `value` to the location pointed by `ptr`.
+unsafe fn write_u256(value: U256, ptr: *mut u8) {
+    let bytes: B256 = value.into();
+    ptr::copy(bytes.as_ptr(), ptr, 32);
 }
 
 /// Decode the selector as [`u32`], and function input as [`Vec<u8>`] from the
@@ -439,6 +474,13 @@ impl<ST: StorageType> ContractCall<'_, ST> {
         self.contract_ref.address
     }
 
+    /// Apply previously not reverted transactions.
+    fn apply_not_reverted_transactions(&self) {
+        Context::current()
+            .try_transfer_value()
+            .expect("should have enough funds for transfer");
+    }
+
     /// Preset the call parameters.
     fn set_call_params(&self) {
         if let Some(value) = self.value {
@@ -454,6 +496,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
+        self.apply_not_reverted_transactions();
         self.set_call_params();
         &self.storage
     }
@@ -462,6 +505,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 impl<ST: StorageType> ::core::ops::DerefMut for ContractCall<'_, ST> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
+        self.apply_not_reverted_transactions();
         self.set_call_params();
         &mut self.storage
     }
@@ -540,6 +584,7 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
         account: A,
         value: V,
     ) -> ContractCall<ST> {
+        // TODO#q: add check if sender has enough funds
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
             caller_address: account.into(),
@@ -579,10 +624,49 @@ impl Account {
     pub fn address(&self) -> Address {
         self.address
     }
+}
 
-    /// Get account's balance.
-    #[must_use]
-    pub fn balance(&self) -> U256 {
-        Context::current().balance(self.address)
+/// Fund the account.
+pub trait Funding {
+    /// Fund the account with the given `value`.
+    fn fund(&self, value: U256);
+
+    /// Get the balance of the account.
+    fn balance(&self) -> U256;
+}
+
+impl Funding for Address {
+    fn fund(&self, value: U256) {
+        Context::current().add_assign_balance(*self, value);
+    }
+
+    fn balance(&self) -> U256 {
+        // Before querying the balance, we should ensure all msg values been
+        // transferred.
+        // TODO#q: move error message inside `try_transfer_value`
+        Context::current()
+            .try_transfer_value()
+            .expect("should have enough funds for transfer");
+        Context::current().balance(*self)
+    }
+}
+
+impl Funding for Account {
+    fn fund(&self, value: U256) {
+        self.address().fund(value);
+    }
+
+    fn balance(&self) -> U256 {
+        self.address().balance()
+    }
+}
+
+impl<ST: StorageType + TestRouter + 'static> Funding for Contract<ST> {
+    fn fund(&self, value: U256) {
+        self.address().fund(value);
+    }
+
+    fn balance(&self) -> U256 {
+        self.address().balance()
     }
 }
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 620f5a6..48c0c0b 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -270,16 +270,19 @@ mod proxies_tests {
     use alloy_primitives::{uint, Address, U256};
     use stylus_sdk::{
         call::Call,
+        msg,
         prelude::{public, storage, TopLevelStorage},
         storage::StorageAddress,
     };
 
-    use crate::context::{Account, Contract};
+    use crate::prelude::*;
 
     stylus_sdk::stylus_proc::sol_interface! {
         interface IProxy {
             #[allow(missing_docs)]
             function callProxy(uint256 value) external returns (uint256);
+            #[allow(missing_docs)]
+            function payProxy() external payable;
         }
     }
 
@@ -306,12 +309,31 @@ mod proxies_tests {
                 proxy.call_proxy(call, value).expect("should call proxy")
             }
         }
+
+        #[payable]
+        fn pay_proxy(&mut self) {
+            let next_proxy = self.next_proxy.get();
+
+            // If there is a next proxy.
+            if !next_proxy.is_zero() {
+                // Add one to the message value.
+                let value = msg::value() + uint!(1_U256);
+
+                // Pay the next proxy.
+                let proxy = IProxy::new(next_proxy);
+                let call = Call::new_in(self).value(value);
+                proxy.pay_proxy(call).expect("should pay proxy");
+            }
+        }
     }
 
     unsafe impl TopLevelStorage for Proxy {}
 
+    const ONE: U256 = uint!(1_U256);
+    const TEN: U256 = uint!(10_U256);
+
     #[motsu_proc::test]
-    fn three_proxies(
+    fn call_three_proxies(
         proxy1: Contract<Proxy>,
         proxy2: Contract<Proxy>,
         proxy3: Contract<Proxy>,
@@ -330,10 +352,44 @@ mod proxies_tests {
         });
 
         // Call the first proxy.
-        let value = uint!(10_U256);
-        let result = proxy1.sender(alice).call_proxy(value);
+        let result = proxy1.sender(alice).call_proxy(TEN);
 
         // The value is incremented by 1 for each proxy.
-        assert_eq!(result, value + uint!(3_U256));
+        assert_eq!(result, TEN + ONE + ONE + ONE);
+    }
+
+    #[motsu_proc::test]
+    fn pay_three_proxies(
+        proxy1: Contract<Proxy>,
+        proxy2: Contract<Proxy>,
+        proxy3: Contract<Proxy>,
+        alice: Account,
+    ) {
+        // Set up a chain of three proxies.
+        // With the given call chain: proxy1 -> proxy2 -> proxy3.
+        proxy1.init(alice, |storage| {
+            storage.next_proxy.set(proxy2.address());
+        });
+        proxy2.init(alice, |storage| {
+            storage.next_proxy.set(proxy3.address());
+        });
+        proxy3.init(alice, |storage| {
+            storage.next_proxy.set(Address::ZERO);
+        });
+
+        // Fund accounts.
+        alice.fund(TEN);
+        proxy1.fund(TEN);
+        proxy2.fund(TEN);
+        proxy3.fund(TEN);
+
+        // Call the first proxy.
+        proxy1.sender_and_value(alice, ONE).pay_proxy();
+
+        // By the end, each actor will lose `ONE`, except last proxy.
+        assert_eq!(alice.balance(), TEN - ONE);
+        assert_eq!(proxy1.balance(), TEN - ONE);
+        assert_eq!(proxy2.balance(), TEN - ONE);
+        assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE);
     }
 }
diff --git a/crates/motsu/src/prelude.rs b/crates/motsu/src/prelude.rs
index a881a05..2534856 100644
--- a/crates/motsu/src/prelude.rs
+++ b/crates/motsu/src/prelude.rs
@@ -1,5 +1,5 @@
 //! Common imports for `motsu` tests.
 pub use crate::{
-    context::{Account, Context, Contract, ContractCall},
+    context::{Account, Context, Contract, ContractCall, Funding},
     shims::*,
 };

From c2099131c8120902c2a8d93c103959c1a05c540e Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 17:18:47 +0400
Subject: [PATCH 04/44] add write_address

---
 crates/motsu/src/context.rs | 11 +++++++----
 crates/motsu/src/shims.rs   | 13 ++++++-------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 686bf67..5690f09 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -302,7 +302,7 @@ impl Context {
     /// Write the value sent to the contract to `output`.
     pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
         let value: U256 = self.msg_value();
-        write_u256(value, output);
+        write_u256(output, value);
     }
 
     /// Get the value sent to the contract as [`U256`].
@@ -398,14 +398,17 @@ unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
     ptr::copy(bytes.as_ptr(), ptr, WORD_BYTES);
 }
 
-// TODO#q: add write_address
-
 /// Read the [`Address`] from the raw pointer.
 unsafe fn read_address(ptr: *const u8) -> Address {
     let address_bytes = slice::from_raw_parts(ptr, 20);
     Address::from_slice(address_bytes)
 }
 
+/// Write the [`Address`] `address` to the location pointed by `ptr`.
+pub(crate) unsafe fn write_address(ptr: *mut u8, address: Address){
+    ptr::copy(address.as_ptr(), ptr, 20);
+}
+
 /// Read the [`U256`] from the raw pointer.
 unsafe fn read_u256(ptr: *const u8) -> U256 {
     let mut data = B256::ZERO;
@@ -414,7 +417,7 @@ unsafe fn read_u256(ptr: *const u8) -> U256 {
 }
 
 /// Write the [`U256`] `value` to the location pointed by `ptr`.
-unsafe fn write_u256(value: U256, ptr: *mut u8) {
+pub(crate) unsafe fn write_u256(ptr: *mut u8, value: U256) {
     let bytes: B256 = value.into();
     ptr::copy(bytes.as_ptr(), ptr, 32);
 }
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 20e08a6..b36ffad 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -39,10 +39,10 @@
 //! ```
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
-
+use alloy_primitives::ruint::aliases::B256;
 use tiny_keccak::{Hasher, Keccak};
 
-use crate::context::Context;
+use crate::context::{write_address, write_u256, Context};
 
 pub(crate) const WORD_BYTES: usize = 32;
 pub(crate) type Bytes32 = [u8; WORD_BYTES];
@@ -149,7 +149,8 @@ pub const CA_CODEHASH: &[u8; 66] =
 pub unsafe extern "C" fn msg_sender(sender: *mut u8) {
     let msg_sender =
         Context::current().msg_sender().expect("msg_sender should be set");
-    std::ptr::copy(msg_sender.as_ptr(), sender, 20);
+    
+    write_address(sender, msg_sender);
 }
 
 /// Get the ETH value (U256) in wei sent to the program.
@@ -171,7 +172,7 @@ pub unsafe extern "C" fn contract_address(address: *mut u8) {
     let contract_address = Context::current()
         .contract_address()
         .expect("contract_address should be set");
-    std::ptr::copy(contract_address.as_ptr(), address, 20);
+    write_address(address, contract_address);
 }
 
 /// Gets the chain ID of the current chain. The semantics are equivalent to
@@ -232,9 +233,7 @@ pub unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
 /// [`BALANCE`]: https://www.evm.codes/#31
 pub unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
     let balance = Context::current().balance_raw(address);
-    let balance_bytes = balance.as_le_slice();
-
-    std::ptr::copy(balance_bytes.as_ptr(), dest, 32);
+    write_u256(dest, balance);
 }
 
 /// Returns the length of the last EVM call or deployment return result, or `0`

From c6b0071d86c7d77ace79b093330863aa6fa2a624 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 17:31:22 +0400
Subject: [PATCH 05/44] make shims private

---
 crates/motsu/src/context.rs | 30 +++++++---------
 crates/motsu/src/prelude.rs |  1 -
 crates/motsu/src/shims.rs   | 68 +++++++++++++++++--------------------
 3 files changed, 44 insertions(+), 55 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 5690f09..2f09791 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -7,10 +7,7 @@ use dashmap::{mapref::one::RefMut, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
 
-use crate::{
-    prelude::{Bytes32, WORD_BYTES},
-    router::{RouterContext, TestRouter},
-};
+use crate::router::{RouterContext, TestRouter};
 
 /// Storage mock.
 ///
@@ -285,13 +282,7 @@ impl Context {
         self.router(address).exists()
     }
 
-    pub(crate) unsafe fn balance_raw(self, address: *const u8) -> U256 {
-        // TODO#q: write to destination here
-        let address = read_address(address);
-        self.balance(address)
-    }
-
-    fn balance(self, address: Address) -> U256 {
+    pub(crate) fn balance(self, address: Address) -> U256 {
         self.storage().balances.get(&address).copied().unwrap_or_default()
     }
 
@@ -335,7 +326,7 @@ impl Context {
     }
 
     /// Transfer `value` from `from` to `to`.
-    /// 
+    ///
     /// Returns `None` if there is not enough funds to transfer.
     fn checked_transfer(
         self,
@@ -349,7 +340,7 @@ impl Context {
     }
 
     /// Subtract `value` from the balance of `address`.
-    /// 
+    ///
     /// Returns `None` if there is not enough of funds.
     fn checked_sub_assign_balance(
         self,
@@ -387,30 +378,30 @@ impl Context {
 }
 
 /// Read the word from location pointed by `ptr`.
-unsafe fn read_bytes32(ptr: *const u8) -> Bytes32 {
+pub(crate) 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) {
+pub(crate) unsafe fn write_bytes32(ptr: *mut u8, bytes: Bytes32) {
     ptr::copy(bytes.as_ptr(), ptr, WORD_BYTES);
 }
 
 /// Read the [`Address`] from the raw pointer.
-unsafe fn read_address(ptr: *const u8) -> Address {
+pub(crate) unsafe fn read_address(ptr: *const u8) -> Address {
     let address_bytes = slice::from_raw_parts(ptr, 20);
     Address::from_slice(address_bytes)
 }
 
 /// Write the [`Address`] `address` to the location pointed by `ptr`.
-pub(crate) unsafe fn write_address(ptr: *mut u8, address: Address){
+pub(crate) unsafe fn write_address(ptr: *mut u8, address: Address) {
     ptr::copy(address.as_ptr(), ptr, 20);
 }
 
 /// Read the [`U256`] from the raw pointer.
-unsafe fn read_u256(ptr: *const u8) -> U256 {
+pub(crate) unsafe fn read_u256(ptr: *const u8) -> U256 {
     let mut data = B256::ZERO;
     ptr::copy(ptr, data.as_mut_ptr(), 32);
     data.into()
@@ -454,7 +445,10 @@ struct MockStorage {
     call_output_len: Option<usize>,
 }
 
+/// Contract's byte storage
 type ContractStorage = HashMap<Bytes32, Bytes32>;
+pub(crate) const WORD_BYTES: usize = 32;
+pub(crate) type Bytes32 = [u8; WORD_BYTES];
 
 /// Contract call entity, related to the contract type `ST` and the caller's
 /// account.
diff --git a/crates/motsu/src/prelude.rs b/crates/motsu/src/prelude.rs
index 2534856..b0c323f 100644
--- a/crates/motsu/src/prelude.rs
+++ b/crates/motsu/src/prelude.rs
@@ -1,5 +1,4 @@
 //! Common imports for `motsu` tests.
 pub use crate::{
     context::{Account, Context, Contract, ContractCall, Funding},
-    shims::*,
 };
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index b36ffad..ae60554 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -39,13 +39,21 @@
 //! ```
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
-use alloy_primitives::ruint::aliases::B256;
 use tiny_keccak::{Hasher, Keccak};
 
-use crate::context::{write_address, write_u256, Context};
+use crate::context::{read_address, write_address, write_bytes32, write_u256, Context, WORD_BYTES};
 
-pub(crate) const WORD_BYTES: usize = 32;
-pub(crate) type Bytes32 = [u8; WORD_BYTES];
+/// Arbitrum's CHAID ID.
+const CHAIN_ID: u64 = 42161;
+
+/// Externally Owned Account (EOA) code hash (wallet account).
+const EOA_CODEHASH: &[u8; 66] =
+    b"0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470";
+
+/// Contract Account (CA) code hash (smart contract code).
+/// NOTE: can be any 256-bit value to pass `has_code` check.
+const CA_CODEHASH: &[u8; 66] =
+    b"0x1111111111111111111111111111111111111111111111111111111111111111";
 
 /// Efficiently computes the [`keccak256`] hash of the given preimage.
 /// The semantics are equivalent to that of the EVM's [`SHA3`] opcode.
@@ -53,7 +61,7 @@ pub(crate) type Bytes32 = [u8; WORD_BYTES];
 /// [`keccak256`]: https://en.wikipedia.org/wiki/SHA-3
 /// [`SHA3`]: https://www.evm.codes/#20
 #[no_mangle]
-pub unsafe extern "C" fn native_keccak256(
+unsafe extern "C" fn native_keccak256(
     bytes: *const u8,
     len: usize,
     output: *mut u8,
@@ -80,7 +88,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) {
+unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
     Context::current().get_bytes_raw(key, out);
 }
 
@@ -100,7 +108,7 @@ pub unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
 ///
 /// May panic if unable to lock `STORAGE`.
 #[no_mangle]
-pub unsafe extern "C" fn storage_cache_bytes32(
+unsafe extern "C" fn storage_cache_bytes32(
     key: *const u8,
     value: *const u8,
 ) {
@@ -113,22 +121,10 @@ pub unsafe extern "C" fn storage_cache_bytes32(
 ///
 /// [`SSTORE`]: https://www.evm.codes/#55
 #[no_mangle]
-pub unsafe extern "C" fn storage_flush_cache(_: bool) {
+unsafe extern "C" fn storage_flush_cache(_: bool) {
     // No-op: we don't use the cache in our unit-tests.
 }
 
-/// Arbitrum's CHAID ID.
-pub const CHAIN_ID: u64 = 42161;
-
-/// Externally Owned Account (EOA) code hash (wallet account).
-pub const EOA_CODEHASH: &[u8; 66] =
-    b"0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470";
-
-/// Contract Account (CA) code hash (smart contract code).
-/// NOTE: can be any 256-bit value to pass `has_code` check.
-pub const CA_CODEHASH: &[u8; 66] =
-    b"0x1111111111111111111111111111111111111111111111111111111111111111";
-
 /// Gets the address of the account that called the program.
 ///
 /// For normal L2-to-L2 transactions the semantics are equivalent to that of the
@@ -146,16 +142,15 @@ pub const CA_CODEHASH: &[u8; 66] =
 ///
 /// May panic if fails to parse `MSG_SENDER` as an address.
 #[no_mangle]
-pub unsafe extern "C" fn msg_sender(sender: *mut u8) {
+unsafe extern "C" fn msg_sender(sender: *mut u8) {
     let msg_sender =
         Context::current().msg_sender().expect("msg_sender should be set");
-    
     write_address(sender, msg_sender);
 }
 
 /// Get the ETH value (U256) in wei sent to the program.
 #[no_mangle]
-pub unsafe extern "C" fn msg_value(value: *mut u8) {
+unsafe extern "C" fn msg_value(value: *mut u8) {
     Context::current().msg_value_raw(value);
 }
 
@@ -168,7 +163,7 @@ pub unsafe extern "C" fn msg_value(value: *mut u8) {
 ///
 /// May panic if fails to parse `CONTRACT_ADDRESS` as an address.
 #[no_mangle]
-pub unsafe extern "C" fn contract_address(address: *mut u8) {
+unsafe extern "C" fn contract_address(address: *mut u8) {
     let contract_address = Context::current()
         .contract_address()
         .expect("contract_address should be set");
@@ -180,7 +175,7 @@ pub unsafe extern "C" fn contract_address(address: *mut u8) {
 ///
 /// [`CHAINID`]: https://www.evm.codes/#46
 #[no_mangle]
-pub unsafe extern "C" fn chainid() -> u64 {
+unsafe extern "C" fn chainid() -> u64 {
     CHAIN_ID
 }
 
@@ -197,7 +192,7 @@ pub unsafe extern "C" fn chainid() -> u64 {
 /// [`LOG3`]: https://www.evm.codes/#a3
 /// [`LOG4`]: https://www.evm.codes/#a4
 #[no_mangle]
-pub unsafe extern "C" fn emit_log(_: *const u8, _: usize, _: usize) {
+unsafe extern "C" fn emit_log(_: *const u8, _: usize, _: usize) {
     // No-op: we don't check for events in our unit-tests.
 }
 
@@ -214,7 +209,7 @@ pub unsafe extern "C" fn emit_log(_: *const u8, _: usize, _: usize) {
 ///
 /// May panic if fails to parse `ACCOUNT_CODEHASH` as a keccack hash.
 #[no_mangle]
-pub unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
+unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
     let code_hash = if Context::current().has_code_raw(address) {
         CA_CODEHASH
     } else {
@@ -224,15 +219,16 @@ pub unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
     let account_codehash =
         const_hex::const_decode_to_array::<32>(code_hash).unwrap();
 
-    std::ptr::copy(account_codehash.as_ptr(), dest, 32);
+    write_bytes32(dest, account_codehash);
 }
 
 /// Gets the ETH balance in wei of the account at the given address.
 /// The semantics are equivalent to that of the EVM's [`BALANCE`] opcode.
 ///
 /// [`BALANCE`]: https://www.evm.codes/#31
-pub unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
-    let balance = Context::current().balance_raw(address);
+unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
+    let address = read_address(address);
+    let balance = Context::current().balance(address);
     write_u256(dest, balance);
 }
 
@@ -244,7 +240,7 @@ pub unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
 ///
 /// [`RETURN_DATA_SIZE`]: https://www.evm.codes/#3d
 #[no_mangle]
-pub unsafe extern "C" fn return_data_size() -> usize {
+unsafe extern "C" fn return_data_size() -> usize {
     Context::current().return_data_size()
 }
 
@@ -258,7 +254,7 @@ pub unsafe extern "C" fn return_data_size() -> usize {
 ///
 /// [`RETURN_DATA_COPY`]: https://www.evm.codes/#3e
 #[no_mangle]
-pub unsafe extern "C" fn read_return_data(
+unsafe extern "C" fn read_return_data(
     dest: *mut u8,
     _offset: usize,
     size: usize,
@@ -281,7 +277,7 @@ pub unsafe extern "C" fn read_return_data(
 ///
 /// [`CALL`]: https://www.evm.codes/#f1
 #[no_mangle]
-pub unsafe extern "C" fn call_contract(
+unsafe extern "C" fn call_contract(
     contract: *const u8,
     calldata: *const u8,
     calldata_len: usize,
@@ -313,7 +309,7 @@ pub unsafe extern "C" fn call_contract(
 ///
 /// [`STATIC_CALL`]: https://www.evm.codes/#FA
 #[no_mangle]
-pub unsafe extern "C" fn static_call_contract(
+unsafe extern "C" fn static_call_contract(
     contract: *const u8,
     calldata: *const u8,
     calldata_len: usize,
@@ -343,7 +339,7 @@ pub unsafe extern "C" fn static_call_contract(
 ///
 /// [`DELEGATE_CALL`]: https://www.evm.codes/#F4
 #[no_mangle]
-pub unsafe extern "C" fn delegate_call_contract(
+unsafe extern "C" fn delegate_call_contract(
     contract: *const u8,
     calldata: *const u8,
     calldata_len: usize,
@@ -364,7 +360,7 @@ pub unsafe extern "C" fn delegate_call_contract(
 ///
 /// [`Block Numbers and Time`]: https://developer.arbitrum.io/time
 #[no_mangle]
-pub unsafe extern "C" fn block_timestamp() -> u64 {
+unsafe extern "C" fn block_timestamp() -> u64 {
     // Epoch timestamp: 1st January 2025 00::00::00
     1_735_689_600
 }

From fb73762aedc1562f6bbfb74b8c7e97c90ebcc1cd Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 17:38:52 +0400
Subject: [PATCH 06/44] ++

---
 crates/motsu/src/context.rs | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 2f09791..2aa8868 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -74,7 +74,7 @@ impl Context {
             .insert(key, value);
     }
 
-    /// Set the message sender address.
+    /// Set the message sender address and return the previous sender if any.
     fn set_msg_sender(self, msg_sender: Address) -> Option<Address> {
         self.storage().msg_sender.replace(msg_sender)
     }
@@ -90,6 +90,22 @@ impl Context {
         self.storage().contract_address.replace(address)
     }
 
+    /// Set message value to `value` and return the previous value if any.
+    pub(crate) fn set_msg_value(self, value: U256) -> Option<U256> {
+        self.storage().msg_value.replace(value)
+    }
+
+    /// Write the value sent to the contract to `output`.
+    pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
+        let value: U256 = self.msg_value();
+        write_u256(output, value);
+    }
+
+    /// Get the value sent to the contract as [`U256`].
+    pub(crate) fn msg_value(self) -> U256 {
+        self.storage().msg_value.unwrap_or_default()
+    }
+
     /// Get the address of the contract, that is called.
     pub(crate) fn contract_address(self) -> Option<Address> {
         self.storage().contract_address
@@ -282,25 +298,11 @@ impl Context {
         self.router(address).exists()
     }
 
+    /// Get the balance of account at `address`.
     pub(crate) fn balance(self, address: Address) -> U256 {
         self.storage().balances.get(&address).copied().unwrap_or_default()
     }
 
-    pub(crate) fn set_msg_value(self, value: U256) -> Option<U256> {
-        self.storage().msg_value.replace(value)
-    }
-
-    /// Write the value sent to the contract to `output`.
-    pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
-        let value: U256 = self.msg_value();
-        write_u256(output, value);
-    }
-
-    /// Get the value sent to the contract as [`U256`].
-    pub(crate) fn msg_value(self) -> U256 {
-        self.storage().msg_value.unwrap_or_default()
-    }
-
     /// Transfer unsent value from the message sender to the contract.
     ///
     /// Returns `None` if there is not enough funds to transfer.

From 67bb7f42348f25438c712eb8cc878fd86a0c2c79 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 17:57:55 +0400
Subject: [PATCH 07/44] ++

---
 crates/motsu/src/context.rs | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 2aa8868..44fa142 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -238,8 +238,7 @@ impl Context {
         // If the call was successful,
         if result.is_ok() {
             // transfer value that wasn't transferred yet.
-            self.try_transfer_value()
-                .expect("should have enough funds for transfer");
+            self.try_transfer_value();
         }
 
         // Set the previous message sender and contract address back.
@@ -305,25 +304,26 @@ impl Context {
 
     /// Transfer unsent value from the message sender to the contract.
     ///
-    /// Returns `None` if there is not enough funds to transfer.
-    fn try_transfer_value(self) -> Option<()> {
+    /// # Panics
+    ///
+    /// If there is not enough funds to transfer.
+    fn try_transfer_value(self) {
         let mut storage = self.storage();
         let Some(msg_sender) = storage.msg_sender else {
-            return Some(());
+            return;
         };
         let Some(contract_address) = storage.contract_address else {
-            return Some(());
+            return;
         };
 
         // We should transfer the value only if it is set.
         // And we should transfer the value only once, despite this function
-        // being called multiple times.
+        // being called multiple times (using `Option::take(..)`).
         if let Some(msg_value) = storage.msg_value.take() {
-            // Drop storage to avoid deadlock.
+            // Drop storage to avoid a deadlock.
             drop(storage);
             self.checked_transfer(msg_sender, contract_address, msg_value)
-        } else {
-            Some(())
+                .unwrap_or_else(|| panic!("should have enough funds for transfer - value is {msg_value}"));
         }
     }
 
@@ -475,9 +475,7 @@ impl<ST: StorageType> ContractCall<'_, ST> {
 
     /// Apply previously not reverted transactions.
     fn apply_not_reverted_transactions(&self) {
-        Context::current()
-            .try_transfer_value()
-            .expect("should have enough funds for transfer");
+        Context::current().try_transfer_value();
     }
 
     /// Preset the call parameters.
@@ -642,10 +640,7 @@ impl Funding for Address {
     fn balance(&self) -> U256 {
         // Before querying the balance, we should ensure all msg values been
         // transferred.
-        // TODO#q: move error message inside `try_transfer_value`
-        Context::current()
-            .try_transfer_value()
-            .expect("should have enough funds for transfer");
+        Context::current().try_transfer_value();
         Context::current().balance(*self)
     }
 }

From dd9e8ee5a51626d1b9d96e502e5662a40de03dc1 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 18:20:44 +0400
Subject: [PATCH 08/44] ++

---
 crates/motsu/src/context.rs | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 44fa142..a895450 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -225,7 +225,10 @@ impl Context {
         let previous_msg_value =
             value.and_then(|value| self.set_msg_value(value));
 
-        // TODO#q: add check if sender has enough funds
+        // If value set, check that sender contract has enough funds.
+        if let Some(value) = value {
+            self.assert_enough_funds(previous_contract_address, value);
+        }
 
         // Call external contract.
         let result = self
@@ -297,6 +300,18 @@ impl Context {
         self.router(address).exists()
     }
 
+    /// Check if `address` has enough funds to transfer `value`
+    ///
+    /// # Panics
+    ///
+    /// Fail if there is not enough funds to transfer.
+    fn assert_enough_funds(self, address: Address, value: U256) {
+        assert!(
+            self.balance(address) >= value,
+            "{address} account should have enough funds to transfer {value} value"
+        );
+    }
+
     /// Get the balance of account at `address`.
     pub(crate) fn balance(self, address: Address) -> U256 {
         self.storage().balances.get(&address).copied().unwrap_or_default()
@@ -306,7 +321,7 @@ impl Context {
     ///
     /// # Panics
     ///
-    /// If there is not enough funds to transfer.
+    /// Fail if there is not enough funds to transfer.
     fn try_transfer_value(self) {
         let mut storage = self.storage();
         let Some(msg_sender) = storage.msg_sender else {
@@ -322,8 +337,10 @@ impl Context {
         if let Some(msg_value) = storage.msg_value.take() {
             // Drop storage to avoid a deadlock.
             drop(storage);
+
+            // Transfer and panic if there is not enough funds.
             self.checked_transfer(msg_sender, contract_address, msg_value)
-                .unwrap_or_else(|| panic!("should have enough funds for transfer - value is {msg_value}"));
+                .unwrap_or_else(|| panic!("{msg_sender} account should have enough funds to transfer {msg_value} value"));
         }
     }
 
@@ -581,11 +598,14 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
         account: A,
         value: V,
     ) -> ContractCall<ST> {
-        // TODO#q: add check if sender has enough funds
+        let caller_address = account.into();
+        let value = value.into();
+        Context::current().assert_enough_funds(caller_address, value);
+
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
-            caller_address: account.into(),
-            value: Some(value.into()),
+            caller_address,
+            value: Some(value),
             contract_ref: self,
         }
     }

From 57554906891af90172840433b0ab7ae257204e45 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 18:24:30 +0400
Subject: [PATCH 09/44] ++

---
 crates/motsu/src/prelude.rs |  4 +---
 crates/motsu/src/shims.rs   | 10 +++++-----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/crates/motsu/src/prelude.rs b/crates/motsu/src/prelude.rs
index b0c323f..e2f7dd4 100644
--- a/crates/motsu/src/prelude.rs
+++ b/crates/motsu/src/prelude.rs
@@ -1,4 +1,2 @@
 //! Common imports for `motsu` tests.
-pub use crate::{
-    context::{Account, Context, Contract, ContractCall, Funding},
-};
+pub use crate::context::{Account, Context, Contract, ContractCall, Funding};
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index ae60554..2c5b19a 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -39,9 +39,12 @@
 //! ```
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
+
 use tiny_keccak::{Hasher, Keccak};
 
-use crate::context::{read_address, write_address, write_bytes32, write_u256, Context, WORD_BYTES};
+use crate::context::{
+    read_address, write_address, write_bytes32, write_u256, Context, WORD_BYTES,
+};
 
 /// Arbitrum's CHAID ID.
 const CHAIN_ID: u64 = 42161;
@@ -108,10 +111,7 @@ unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
 ///
 /// May panic if unable to lock `STORAGE`.
 #[no_mangle]
-unsafe extern "C" fn storage_cache_bytes32(
-    key: *const u8,
-    value: *const u8,
-) {
+unsafe extern "C" fn storage_cache_bytes32(key: *const u8, value: *const u8) {
     Context::current().set_bytes_raw(key, value);
 }
 

From a6f825168dd9b2a671853df4d57c6337ef01c5be Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 18:29:17 +0400
Subject: [PATCH 10/44] ++

---
 crates/motsu/src/context.rs | 12 ++++++------
 crates/motsu/src/shims.rs   |  1 +
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index a895450..8e8d3a2 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -344,7 +344,7 @@ impl Context {
         }
     }
 
-    /// Transfer `value` from `from` to `to`.
+    /// Transfer `value` from `from` account to `to` account.
     ///
     /// Returns `None` if there is not enough funds to transfer.
     fn checked_transfer(
@@ -358,7 +358,7 @@ impl Context {
         Some(())
     }
 
-    /// Subtract `value` from the balance of `address`.
+    /// Subtract `value` from the balance of `address` account.
     ///
     /// Returns `None` if there is not enough of funds.
     fn checked_sub_assign_balance(
@@ -375,7 +375,7 @@ impl Context {
         Some(*balance)
     }
 
-    /// Add `value` to the balance of `address`.
+    /// Add `value` to the balance of `address` account.
     fn add_assign_balance(self, address: Address, value: U256) -> U256 {
         *self
             .storage()
@@ -491,7 +491,7 @@ impl<ST: StorageType> ContractCall<'_, ST> {
     }
 
     /// Apply previously not reverted transactions.
-    fn apply_not_reverted_transactions(&self) {
+    fn apply_not_reverted_transactions() {
         Context::current().try_transfer_value();
     }
 
@@ -510,7 +510,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        self.apply_not_reverted_transactions();
+        Self::apply_not_reverted_transactions();
         self.set_call_params();
         &self.storage
     }
@@ -519,7 +519,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 impl<ST: StorageType> ::core::ops::DerefMut for ContractCall<'_, ST> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.apply_not_reverted_transactions();
+        Self::apply_not_reverted_transactions();
         self.set_call_params();
         &mut self.storage
     }
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 2c5b19a..8f995aa 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -37,6 +37,7 @@
 //!     }
 //! }
 //! ```
+#![allow(dead_code)]
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
 

From 3e8c5ef753919a2bb1d7cf5528eaacd40c7b111d Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 18:58:45 +0400
Subject: [PATCH 11/44] update docs

---
 crates/motsu/src/shims.rs | 29 ++++-------------------------
 1 file changed, 4 insertions(+), 25 deletions(-)

diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 8f995aa..dcfa425 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -8,7 +8,7 @@
 //!
 //! ## Motivation
 //!
-//! Without these shims we can't currently run unit tests for stylus contracts,
+//! Without these shims, we can't currently run unit tests for stylus contracts,
 //! since the symbols the compiled binaries expect to find are not there.
 //!
 //! If you run `cargo test` on a fresh Stylus project, it will error with:
@@ -16,27 +16,6 @@
 //! ```terminal
 //! dyld[97792]: missing symbol called
 //! ```
-//!
-//! This crate is a temporary solution until the Stylus team provides us with a
-//! different and more stable mechanism for unit-testing our contracts.
-//!
-//! ## Usage
-//!
-//! Import these shims in your test modules as `motsu::prelude::*` to populate
-//! the namespace with the appropriate symbols.
-//!
-//! ```rust,ignore
-//! #[cfg(test)]
-//! mod tests {
-//!     use contracts::token::erc20::Erc20;
-//!
-//!     #[motsu::test]
-//!     fn reads_balance(contract: Erc20) {
-//!         let balance = contract.balance_of(Address::ZERO); // Access storage.
-//!         assert_eq!(balance, U256::ZERO);
-//!     }
-//! }
-//! ```
 #![allow(dead_code)]
 #![allow(clippy::missing_safety_doc)]
 use std::slice;
@@ -234,7 +213,7 @@ unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
 }
 
 /// Returns the length of the last EVM call or deployment return result, or `0`
-/// if neither have happened during the program's execution.
+/// if neither has happened during the program's execution.
 ///
 /// The semantics are equivalent to that of the EVM's [`RETURN_DATA_SIZE`]
 /// opcode.
@@ -265,9 +244,9 @@ unsafe extern "C" fn read_return_data(
 
 /// Calls the contract at the given address with options for passing value and
 /// to limit the amount of gas supplied. The return status indicates whether the
-/// call succeeded, and is nonzero on failure.
+/// call succeeded and is nonzero on failure.
 ///
-/// In both cases `return_data_len` will store the length of the result, the
+/// In both cases, `return_data_len` will store the length of the result, the
 /// bytes of which can be read via the `read_return_data` hostio. The bytes are
 /// not returned directly so that the programmer can potentially save gas by
 /// choosing which subset of the return result they'd like to copy.

From e0615f30c0bb206d7e83456b2a1990ff5a196f6c Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:05:51 +0400
Subject: [PATCH 12/44] review fixes

---
 crates/motsu/src/context.rs | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 8e8d3a2..da6a97d 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -143,7 +143,7 @@ impl Context {
             // drop guard to a concurrent hash map to avoid a deadlock,
             drop(storage);
             // and erase the test context.
-            let _ = STORAGE.remove(&self);
+            _ = STORAGE.remove(&self);
         }
 
         self.router(contract_address).reset_storage();
@@ -245,12 +245,12 @@ impl Context {
         }
 
         // Set the previous message sender and contract address back.
-        let _ = self.set_contract_address(previous_contract_address);
-        let _ = self.set_msg_sender(previous_msg_sender);
+        _ = self.set_contract_address(previous_contract_address);
+        _ = self.set_msg_sender(previous_msg_sender);
 
         // Set the previous msg_value if there is any.
         if let Some(previous) = previous_msg_value {
-            let _ = self.set_msg_value(previous);
+            _ = self.set_msg_value(previous);
         }
 
         result
@@ -259,8 +259,8 @@ impl Context {
     /// Set return data as bytes.
     fn set_return_data(self, data: Vec<u8>) {
         let mut call_storage = self.storage();
-        let _ = call_storage.call_output_len.insert(data.len());
-        let _ = call_storage.call_output.insert(data);
+        _ = call_storage.call_output_len.insert(data.len());
+        _ = call_storage.call_output.insert(data);
     }
 
     /// Read the return data (with a given `size`) from the last contract call
@@ -304,7 +304,7 @@ impl Context {
     ///
     /// # Panics
     ///
-    /// Fail if there is not enough funds to transfer.
+    /// * If there is not enough funds to transfer.
     fn assert_enough_funds(self, address: Address, value: U256) {
         assert!(
             self.balance(address) >= value,
@@ -321,7 +321,7 @@ impl Context {
     ///
     /// # Panics
     ///
-    /// Fail if there is not enough funds to transfer.
+    /// * If there is not enough funds to transfer.
     fn try_transfer_value(self) {
         let mut storage = self.storage();
         let Some(msg_sender) = storage.msg_sender else {
@@ -353,7 +353,7 @@ impl Context {
         to: Address,
         value: U256,
     ) -> Option<()> {
-        let _ = self.checked_sub_assign_balance(from, value)?;
+        _ = self.checked_sub_assign_balance(from, value)?;
         self.add_assign_balance(to, value);
         Some(())
     }
@@ -498,10 +498,10 @@ impl<ST: StorageType> ContractCall<'_, ST> {
     /// Preset the call parameters.
     fn set_call_params(&self) {
         if let Some(value) = self.value {
-            let _ = Context::current().set_msg_value(value);
+            _ = Context::current().set_msg_value(value);
         }
-        let _ = Context::current().set_msg_sender(self.caller_address);
-        let _ = Context::current().set_contract_address(self.address());
+        _ = Context::current().set_msg_sender(self.caller_address);
+        _ = Context::current().set_contract_address(self.address());
     }
 }
 

From af670d9bc5f86fbbaab395df57a5f2c3c141ddc7 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:06:38 +0400
Subject: [PATCH 13/44] review fixes

---
 crates/motsu/src/lib.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 48c0c0b..72dd58b 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -215,7 +215,7 @@ mod ping_pong_tests {
         assert_eq!(ping.sender(alice).pinged_from.get(), Address::ZERO);
         assert_eq!(pong.sender(alice).ponged_from.get(), Address::ZERO);
 
-        let _ = ping
+        _ = ping
             .sender(alice)
             .ping(pong.address(), uint!(10_U256))
             .expect("should ping successfully");
@@ -242,7 +242,7 @@ mod ping_pong_tests {
         assert_eq!(ping.sender(alice).contract_address.get(), Address::ZERO);
         assert_eq!(pong.sender(alice).contract_address.get(), Address::ZERO);
 
-        let _ = ping
+        _ = ping
             .sender(alice)
             .ping(pong.address(), uint!(10_U256))
             .expect("should ping successfully");
@@ -259,7 +259,7 @@ mod ping_pong_tests {
         let pong = Contract::<PongContract>::new();
         let pong = pong.sender(alice);
 
-        let _ = ping
+        _ = ping
             .ping(pong.address(), uint!(10_U256))
             .expect("contract ping should not drop");
     }

From 410c775f422c51da5d6b70c88fc92393fec0023e Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:13:56 +0400
Subject: [PATCH 14/44] ++

---
 crates/motsu/src/lib.rs | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 72dd58b..da0fef5 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -87,7 +87,7 @@ mod ping_pong_tests {
             let value = receiver.pong(call, value)?;
 
             let pings_count = self.pings_count.get();
-            self.pings_count.set(pings_count + uint!(1_U256));
+            self.pings_count.set(pings_count + ONE);
 
             self.pinged_from.set(msg::sender());
             self.contract_address.set(contract::address());
@@ -131,6 +131,9 @@ mod ping_pong_tests {
 
     const MAGIC_ERROR_VALUE: U256 = uint!(42_U256);
 
+    const ONE: U256 = uint!(1_U256);
+    const TEN: U256 = uint!(10_U256);
+
     #[storage]
     struct PongContract {
         pongs_count: StorageU256,
@@ -146,12 +149,12 @@ mod ping_pong_tests {
             }
 
             let pongs_count = self.pongs_count.get();
-            self.pongs_count.set(pongs_count + uint!(1_U256));
+            self.pongs_count.set(pongs_count + ONE);
 
             self.ponged_from.set(msg::sender());
             self.contract_address.set(contract::address());
 
-            Ok(value + uint!(1_U256))
+            Ok(value + ONE)
         }
 
         fn can_pong(&self) -> bool {
@@ -167,15 +170,15 @@ mod ping_pong_tests {
         pong: Contract<PongContract>,
         alice: Account,
     ) {
-        let value = uint!(10_U256);
+        let value = TEN;
         let ponged_value = ping
             .sender(alice)
             .ping(pong.address(), value)
             .expect("should ping successfully");
 
-        assert_eq!(ponged_value, value + uint!(1_U256));
-        assert_eq!(ping.sender(alice).pings_count.get(), uint!(1_U256));
-        assert_eq!(pong.sender(alice).pongs_count.get(), uint!(1_U256));
+        assert_eq!(ponged_value, value + ONE);
+        assert_eq!(ping.sender(alice).pings_count.get(), ONE);
+        assert_eq!(pong.sender(alice).pongs_count.get(), ONE);
     }
 
     #[motsu_proc::test]
@@ -217,7 +220,7 @@ mod ping_pong_tests {
 
         _ = ping
             .sender(alice)
-            .ping(pong.address(), uint!(10_U256))
+            .ping(pong.address(), TEN)
             .expect("should ping successfully");
 
         assert_eq!(ping.sender(alice).pinged_from.get(), alice.address());
@@ -244,7 +247,7 @@ mod ping_pong_tests {
 
         _ = ping
             .sender(alice)
-            .ping(pong.address(), uint!(10_U256))
+            .ping(pong.address(), TEN)
             .expect("should ping successfully");
 
         assert_eq!(ping.sender(alice).contract_address.get(), ping.address());
@@ -260,7 +263,7 @@ mod ping_pong_tests {
         let pong = pong.sender(alice);
 
         _ = ping
-            .ping(pong.address(), uint!(10_U256))
+            .ping(pong.address(), TEN)
             .expect("contract ping should not drop");
     }
 }
@@ -297,7 +300,7 @@ mod proxies_tests {
             let next_proxy = self.next_proxy.get();
 
             // Add one to the value.
-            let value = value + uint!(1_U256);
+            let value = value + ONE;
 
             // If there is no next proxy, return the value.
             if next_proxy.is_zero() {
@@ -317,7 +320,7 @@ mod proxies_tests {
             // If there is a next proxy.
             if !next_proxy.is_zero() {
                 // Add one to the message value.
-                let value = msg::value() + uint!(1_U256);
+                let value = msg::value() + ONE;
 
                 // Pay the next proxy.
                 let proxy = IProxy::new(next_proxy);

From 0b48d1c962e07e20d8b4a01890b418ef5677a649 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:37:07 +0400
Subject: [PATCH 15/44] add missing #[no_mangle]

---
 crates/motsu/src/shims.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index dcfa425..2165e2c 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -206,6 +206,7 @@ unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
 /// The semantics are equivalent to that of the EVM's [`BALANCE`] opcode.
 ///
 /// [`BALANCE`]: https://www.evm.codes/#31
+#[no_mangle]
 unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
     let address = read_address(address);
     let balance = Context::current().balance(address);

From ee6eebc719c52eea06c31b2a19db9a2fcce79101 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:58:24 +0400
Subject: [PATCH 16/44] ++

---
 crates/motsu/src/shims.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 2165e2c..02bbaad 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -197,7 +197,7 @@ unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
     };
 
     let account_codehash =
-        const_hex::const_decode_to_array::<32>(code_hash).unwrap();
+        const_hex::const_decode_to_array::<WORD_BYTES>(code_hash).unwrap();
 
     write_bytes32(dest, account_codehash);
 }

From c903db7a7bb6cd7e35ee1936033ce8c970655975 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 19:59:38 +0400
Subject: [PATCH 17/44] ++

---
 crates/motsu/src/shims.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 02bbaad..8cb53ab 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -107,7 +107,7 @@ unsafe extern "C" fn storage_flush_cache(_: bool) {
 
 /// Gets the address of the account that called the program.
 ///
-/// For normal L2-to-L2 transactions the semantics are equivalent to that of the
+/// For normal L2-to-L2 transactions, the semantics are equivalent to that of the
 /// EVM's [`CALLER`] opcode, including in cases arising from [`DELEGATE_CALL`].
 ///
 /// For L1-to-L2 retryable ticket transactions, the top-level sender's address
@@ -253,7 +253,7 @@ unsafe extern "C" fn read_return_data(
 /// choosing which subset of the return result they'd like to copy.
 ///
 /// The semantics are equivalent to that of the EVM's [`CALL`] opcode, including
-/// callvalue stipends and the 63/64 gas rule. This means that supplying the
+/// call value stipends and the 63/64 gas rule. This means that supplying the
 /// `u64::MAX` gas can be used to send as much as possible.
 ///
 /// [`CALL`]: https://www.evm.codes/#f1
@@ -307,9 +307,9 @@ unsafe extern "C" fn static_call_contract(
 
 /// Delegate calls the contract at the given address, with the option to limit
 /// the amount of gas supplied. The return status indicates whether the call
-/// succeeded, and is nonzero on failure.
+/// succeeded and is nonzero on failure.
 ///
-/// In both cases `return_data_len` will store the length of the result, the
+/// In both cases, `return_data_len` will store the length of the result, the
 /// bytes of which can be read via the `read_return_data` hostio. The bytes are
 /// not returned directly so that the programmer can potentially save gas by
 /// choosing which subset of the return result they'd like to copy.

From 155327b54a121a2a93945cbda77f9bfeb55c4a3a Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Tue, 4 Feb 2025 20:03:36 +0400
Subject: [PATCH 18/44] fix fmt

---
 crates/motsu/src/shims.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 8cb53ab..5d8eabf 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -107,8 +107,9 @@ unsafe extern "C" fn storage_flush_cache(_: bool) {
 
 /// Gets the address of the account that called the program.
 ///
-/// For normal L2-to-L2 transactions, the semantics are equivalent to that of the
-/// EVM's [`CALLER`] opcode, including in cases arising from [`DELEGATE_CALL`].
+/// For normal L2-to-L2 transactions, the semantics are equivalent to that of
+/// the EVM's [`CALLER`] opcode, including in cases arising from
+/// [`DELEGATE_CALL`].
 ///
 /// For L1-to-L2 retryable ticket transactions, the top-level sender's address
 /// will be aliased. See [`Retryable Ticket Address Aliasing`][aliasing] for

From 910db33a5781b5f6d9bdae66f01892d8014a778a Mon Sep 17 00:00:00 2001
From: Nenad <xinef.it@gmail.com>
Date: Wed, 5 Feb 2025 19:51:53 +0100
Subject: [PATCH 19/44] new test cases (#33)

---
 crates/motsu/src/lib.rs | 125 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 124 insertions(+), 1 deletion(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index da0fef5..1ecb323 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -273,7 +273,7 @@ mod proxies_tests {
     use alloy_primitives::{uint, Address, U256};
     use stylus_sdk::{
         call::Call,
-        msg,
+        contract, msg,
         prelude::{public, storage, TopLevelStorage},
         storage::StorageAddress,
     };
@@ -286,6 +286,10 @@ mod proxies_tests {
             function callProxy(uint256 value) external returns (uint256);
             #[allow(missing_docs)]
             function payProxy() external payable;
+            #[allow(missing_docs)]
+            function passThisMuchToNextProxy(uint256 pass_value) external payable;
+            #[allow(missing_docs)]
+            function passHalfValueToNextProxy() external payable;
         }
     }
 
@@ -328,6 +332,41 @@ mod proxies_tests {
                 proxy.pay_proxy(call).expect("should pay proxy");
             }
         }
+
+        #[payable]
+        fn pass_this_much_to_next_proxy(&mut self, pass_value: U256) {
+            let next_proxy = self.next_proxy.get();
+
+            // If there is a next proxy.
+            if !next_proxy.is_zero() {
+                // Pay the next proxy.
+                let proxy = IProxy::new(next_proxy);
+                let call = Call::new_in(self).value(pass_value);
+                let value_for_next_next_proxy = pass_value / U256::from(2);
+                proxy
+                    .pass_this_much_to_next_proxy(
+                        call,
+                        value_for_next_next_proxy,
+                    )
+                    .expect("should pass half the value to the next proxy");
+            }
+        }
+
+        #[payable]
+        fn pass_half_value_to_next_proxy(&mut self) {
+            let next_proxy = self.next_proxy.get();
+
+            // If there is a next proxy.
+            if !next_proxy.is_zero() {
+                let half_balance = contract::balance() / U256::from(2);
+                // Pay the next proxy.
+                let proxy = IProxy::new(next_proxy);
+                let call = Call::new_in(self).value(half_balance);
+                proxy
+                    .pass_half_value_to_next_proxy(call)
+                    .expect("should pass half the value to the next proxy");
+            }
+        }
     }
 
     unsafe impl TopLevelStorage for Proxy {}
@@ -395,4 +434,88 @@ mod proxies_tests {
         assert_eq!(proxy2.balance(), TEN - ONE);
         assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE);
     }
+
+    #[motsu_proc::test]
+    fn error_when_transferring_values_for_nested_contract_calls(
+        proxy1: Contract<Proxy>,
+        proxy2: Contract<Proxy>,
+        proxy3: Contract<Proxy>,
+        alice: Account,
+    ) {
+        // Set up a chain of three proxies.
+        // With the given call chain: proxy1 -> proxy2 -> proxy3.
+        proxy1.init(alice, |storage| {
+            storage.next_proxy.set(proxy2.address());
+        });
+        proxy2.init(alice, |storage| {
+            storage.next_proxy.set(proxy3.address());
+        });
+        proxy3.init(alice, |storage| {
+            storage.next_proxy.set(Address::ZERO);
+        });
+
+        let zero = U256::ZERO;
+        let two = U256::from(2);
+        let four = U256::from(4);
+        let eight = U256::from(8);
+
+        // Fund alice, proxies have no funds.
+        alice.fund(eight);
+
+        assert_eq!(alice.balance(), eight);
+        assert_eq!(proxy1.balance(), zero);
+        assert_eq!(proxy2.balance(), zero);
+        assert_eq!(proxy3.balance(), zero);
+
+        // Call the first proxy.
+        proxy1
+            .sender_and_value(alice, eight)
+            .pass_this_much_to_next_proxy(four);
+
+        assert_eq!(alice.balance(), zero);
+        assert_eq!(proxy1.balance(), four);
+        assert_eq!(proxy2.balance(), two);
+        assert_eq!(proxy3.balance(), two);
+    }
+
+    #[motsu_proc::test]
+    fn wrong_contract_balances_for_nested_contract_calls(
+        proxy1: Contract<Proxy>,
+        proxy2: Contract<Proxy>,
+        proxy3: Contract<Proxy>,
+        alice: Account,
+    ) {
+        // Set up a chain of three proxies.
+        // With the given call chain: proxy1 -> proxy2 -> proxy3.
+        proxy1.init(alice, |storage| {
+            storage.next_proxy.set(proxy2.address());
+        });
+        proxy2.init(alice, |storage| {
+            storage.next_proxy.set(proxy3.address());
+        });
+        proxy3.init(alice, |storage| {
+            storage.next_proxy.set(Address::ZERO);
+        });
+
+        let zero = U256::ZERO;
+        let two = U256::from(2);
+        let four = U256::from(4);
+        let eight = U256::from(8);
+
+        // Fund alice, proxies have no funds.
+        alice.fund(eight);
+
+        assert_eq!(alice.balance(), eight);
+        assert_eq!(proxy1.balance(), zero);
+        assert_eq!(proxy2.balance(), zero);
+        assert_eq!(proxy3.balance(), zero);
+
+        // Call the first proxy.
+        proxy1.sender_and_value(alice, eight).pass_half_value_to_next_proxy();
+
+        assert_eq!(alice.balance(), zero);
+        assert_eq!(proxy1.balance(), four);
+        assert_eq!(proxy2.balance(), two);
+        assert_eq!(proxy3.balance(), two);
+    }
 }

From 6457ad4d80057cd87879a8da474679e403342ef7 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Wed, 5 Feb 2025 23:34:37 +0400
Subject: [PATCH 20/44] ++

---
 crates/motsu/src/context.rs | 57 +++++++++++++------------------------
 1 file changed, 19 insertions(+), 38 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index da6a97d..133a46e 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -225,10 +225,8 @@ impl Context {
         let previous_msg_value =
             value.and_then(|value| self.set_msg_value(value));
 
-        // If value set, check that sender contract has enough funds.
-        if let Some(value) = value {
-            self.assert_enough_funds(previous_contract_address, value);
-        }
+        // Transfer value sent by message sender.
+        self.transfer_value();
 
         // Call external contract.
         let result = self
@@ -238,12 +236,6 @@ impl Context {
                 panic!("selector not found - selector is {selector}")
             });
 
-        // If the call was successful,
-        if result.is_ok() {
-            // transfer value that wasn't transferred yet.
-            self.try_transfer_value();
-        }
-
         // Set the previous message sender and contract address back.
         _ = self.set_contract_address(previous_contract_address);
         _ = self.set_msg_sender(previous_msg_sender);
@@ -300,29 +292,18 @@ impl Context {
         self.router(address).exists()
     }
 
-    /// Check if `address` has enough funds to transfer `value`
-    ///
-    /// # Panics
-    ///
-    /// * If there is not enough funds to transfer.
-    fn assert_enough_funds(self, address: Address, value: U256) {
-        assert!(
-            self.balance(address) >= value,
-            "{address} account should have enough funds to transfer {value} value"
-        );
-    }
-
     /// Get the balance of account at `address`.
     pub(crate) fn balance(self, address: Address) -> U256 {
         self.storage().balances.get(&address).copied().unwrap_or_default()
     }
 
-    /// Transfer unsent value from the message sender to the contract.
+    /// Transfer value from the message sender to the contract.
+    /// No-op if `msg_sender` or `contract_address` weren't set.
     ///
     /// # Panics
     ///
     /// * If there is not enough funds to transfer.
-    fn try_transfer_value(self) {
+    fn transfer_value(self) {
         let mut storage = self.storage();
         let Some(msg_sender) = storage.msg_sender else {
             return;
@@ -339,11 +320,21 @@ impl Context {
             drop(storage);
 
             // Transfer and panic if there is not enough funds.
-            self.checked_transfer(msg_sender, contract_address, msg_value)
-                .unwrap_or_else(|| panic!("{msg_sender} account should have enough funds to transfer {msg_value} value"));
+            self.transfer(msg_sender, contract_address, msg_value);
         }
     }
 
+    /// Transfer `value` from `from` account to `to` account.
+    ///
+    /// # Panics
+    ///
+    /// * If there is not enough funds to transfer.
+    fn transfer(self, from: Address, to: Address, value: U256) {
+        // Transfer and panic if there is not enough funds.
+        self.checked_transfer(from, to, value)
+            .unwrap_or_else(|| panic!("{from} account should have enough funds to transfer {value} value"));
+    }
+
     /// Transfer `value` from `from` account to `to` account.
     ///
     /// Returns `None` if there is not enough funds to transfer.
@@ -490,11 +481,6 @@ impl<ST: StorageType> ContractCall<'_, ST> {
         self.contract_ref.address
     }
 
-    /// Apply previously not reverted transactions.
-    fn apply_not_reverted_transactions() {
-        Context::current().try_transfer_value();
-    }
-
     /// Preset the call parameters.
     fn set_call_params(&self) {
         if let Some(value) = self.value {
@@ -510,7 +496,6 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        Self::apply_not_reverted_transactions();
         self.set_call_params();
         &self.storage
     }
@@ -519,7 +504,6 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
 impl<ST: StorageType> ::core::ops::DerefMut for ContractCall<'_, ST> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
-        Self::apply_not_reverted_transactions();
         self.set_call_params();
         &mut self.storage
     }
@@ -600,8 +584,8 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     ) -> ContractCall<ST> {
         let caller_address = account.into();
         let value = value.into();
-        Context::current().assert_enough_funds(caller_address, value);
-
+        
+        Context::current().transfer(caller_address, self.address, value);
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
             caller_address,
@@ -658,9 +642,6 @@ impl Funding for Address {
     }
 
     fn balance(&self) -> U256 {
-        // Before querying the balance, we should ensure all msg values been
-        // transferred.
-        Context::current().try_transfer_value();
         Context::current().balance(*self)
     }
 }

From 74ff90ae8de2695170fa3003595088df04bd58fc Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 10:54:59 +0400
Subject: [PATCH 21/44] reflect value in balance soon

---
 crates/motsu/src/context.rs |  11 ++--
 crates/motsu/src/lib.rs     | 100 ++++++++++++++++--------------------
 2 files changed, 50 insertions(+), 61 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 133a46e..ac184ac 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -304,7 +304,7 @@ impl Context {
     ///
     /// * If there is not enough funds to transfer.
     fn transfer_value(self) {
-        let mut storage = self.storage();
+        let storage = self.storage();
         let Some(msg_sender) = storage.msg_sender else {
             return;
         };
@@ -313,9 +313,7 @@ impl Context {
         };
 
         // We should transfer the value only if it is set.
-        // And we should transfer the value only once, despite this function
-        // being called multiple times (using `Option::take(..)`).
-        if let Some(msg_value) = storage.msg_value.take() {
+        if let Some(msg_value) = storage.msg_value {
             // Drop storage to avoid a deadlock.
             drop(storage);
 
@@ -497,6 +495,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
     #[inline]
     fn deref(&self) -> &Self::Target {
         self.set_call_params();
+        Context::current().transfer_value();
         &self.storage
     }
 }
@@ -505,6 +504,7 @@ impl<ST: StorageType> ::core::ops::DerefMut for ContractCall<'_, ST> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.set_call_params();
+        Context::current().transfer_value();
         &mut self.storage
     }
 }
@@ -584,8 +584,7 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     ) -> ContractCall<ST> {
         let caller_address = account.into();
         let value = value.into();
-        
-        Context::current().transfer(caller_address, self.address, value);
+
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
             caller_address,
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 1ecb323..291100f 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -270,10 +270,12 @@ mod ping_pong_tests {
 
 #[cfg(test)]
 mod proxies_tests {
-    use alloy_primitives::{uint, Address, U256};
+    use alloy_primitives::{address, uint, Address, U256};
     use stylus_sdk::{
         call::Call,
-        contract, msg,
+        contract,
+        contract::address,
+        msg,
         prelude::{public, storage, TopLevelStorage},
         storage::StorageAddress,
     };
@@ -334,15 +336,15 @@ mod proxies_tests {
         }
 
         #[payable]
-        fn pass_this_much_to_next_proxy(&mut self, pass_value: U256) {
+        fn pass_proxy_with_fixed_value(&mut self, this_value: U256) {
             let next_proxy = self.next_proxy.get();
 
             // If there is a next proxy.
             if !next_proxy.is_zero() {
                 // Pay the next proxy.
                 let proxy = IProxy::new(next_proxy);
-                let call = Call::new_in(self).value(pass_value);
-                let value_for_next_next_proxy = pass_value / U256::from(2);
+                let call = Call::new_in(self).value(this_value);
+                let value_for_next_next_proxy = this_value / U256::from(2);
                 proxy
                     .pass_this_much_to_next_proxy(
                         call,
@@ -353,7 +355,7 @@ mod proxies_tests {
         }
 
         #[payable]
-        fn pass_half_value_to_next_proxy(&mut self) {
+        fn pay_proxy_with_half_balance(&mut self) {
             let next_proxy = self.next_proxy.get();
 
             // If there is a next proxy.
@@ -369,10 +371,22 @@ mod proxies_tests {
         }
     }
 
+    impl Proxy {
+        fn init(&mut self, next_proxy: Address){
+            self.next_proxy.set(next_proxy);
+        }
+    }
+    
+
     unsafe impl TopLevelStorage for Proxy {}
 
     const ONE: U256 = uint!(1_U256);
+    const TWO: U256 = uint!(2_U256);
+    const FOUR: U256 = uint!(4_U256);
+    const EIGHT: U256 = uint!(8_U256);
+
     const TEN: U256 = uint!(10_U256);
+    
 
     #[motsu_proc::test]
     fn call_three_proxies(
@@ -383,15 +397,9 @@ mod proxies_tests {
     ) {
         // Set up a chain of three proxies.
         // With the given call chain: proxy1 -> proxy2 -> proxy3.
-        proxy1.init(alice, |storage| {
-            storage.next_proxy.set(proxy2.address());
-        });
-        proxy2.init(alice, |storage| {
-            storage.next_proxy.set(proxy3.address());
-        });
-        proxy3.init(alice, |storage| {
-            storage.next_proxy.set(Address::ZERO);
-        });
+        proxy1.sender(alice).init(proxy2.address());
+        proxy2.sender(alice).init(proxy3.address());
+        proxy3.sender(alice).init(Address::ZERO);
 
         // Call the first proxy.
         let result = proxy1.sender(alice).call_proxy(TEN);
@@ -409,15 +417,9 @@ mod proxies_tests {
     ) {
         // Set up a chain of three proxies.
         // With the given call chain: proxy1 -> proxy2 -> proxy3.
-        proxy1.init(alice, |storage| {
-            storage.next_proxy.set(proxy2.address());
-        });
-        proxy2.init(alice, |storage| {
-            storage.next_proxy.set(proxy3.address());
-        });
-        proxy3.init(alice, |storage| {
-            storage.next_proxy.set(Address::ZERO);
-        });
+        proxy1.sender(alice).init(proxy2.address());
+        proxy2.sender(alice).init(proxy3.address());
+        proxy3.sender(alice).init(Address::ZERO);
 
         // Fund accounts.
         alice.fund(TEN);
@@ -429,14 +431,16 @@ mod proxies_tests {
         proxy1.sender_and_value(alice, ONE).pay_proxy();
 
         // By the end, each actor will lose `ONE`, except last proxy.
-        assert_eq!(alice.balance(), TEN - ONE);
-        assert_eq!(proxy1.balance(), TEN - ONE);
-        assert_eq!(proxy2.balance(), TEN - ONE);
-        assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE);
+        assert_eq!(alice.balance(), TEN - ONE); // 10 - 1
+        assert_eq!(proxy1.balance(), TEN - ONE); // 10 + 1 - 2
+        assert_eq!(proxy2.balance(), TEN - ONE); // 10 + 2 - 3
+        assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE); // 10 + 3
     }
 
+    // TODO#q: if the msg value wasn't set we shouldn't set it
+
     #[motsu_proc::test]
-    fn error_when_transferring_values_for_nested_contract_calls(
+    fn pay_proxy_fixed_value(
         proxy1: Contract<Proxy>,
         proxy2: Contract<Proxy>,
         proxy3: Contract<Proxy>,
@@ -444,18 +448,12 @@ mod proxies_tests {
     ) {
         // Set up a chain of three proxies.
         // With the given call chain: proxy1 -> proxy2 -> proxy3.
-        proxy1.init(alice, |storage| {
-            storage.next_proxy.set(proxy2.address());
-        });
-        proxy2.init(alice, |storage| {
-            storage.next_proxy.set(proxy3.address());
-        });
-        proxy3.init(alice, |storage| {
-            storage.next_proxy.set(Address::ZERO);
-        });
+        proxy1.sender(alice).init(proxy2.address());
+        proxy2.sender(alice).init(proxy3.address());
+        proxy3.sender(alice).init(Address::ZERO);
 
         let zero = U256::ZERO;
-        let two = U256::from(2);
+        let two = uint!(2_U256);
         let four = U256::from(4);
         let eight = U256::from(8);
 
@@ -468,9 +466,7 @@ mod proxies_tests {
         assert_eq!(proxy3.balance(), zero);
 
         // Call the first proxy.
-        proxy1
-            .sender_and_value(alice, eight)
-            .pass_this_much_to_next_proxy(four);
+        proxy1.sender_and_value(alice, eight).pass_proxy_with_fixed_value(four);
 
         assert_eq!(alice.balance(), zero);
         assert_eq!(proxy1.balance(), four);
@@ -487,20 +483,14 @@ mod proxies_tests {
     ) {
         // Set up a chain of three proxies.
         // With the given call chain: proxy1 -> proxy2 -> proxy3.
-        proxy1.init(alice, |storage| {
-            storage.next_proxy.set(proxy2.address());
-        });
-        proxy2.init(alice, |storage| {
-            storage.next_proxy.set(proxy3.address());
-        });
-        proxy3.init(alice, |storage| {
-            storage.next_proxy.set(Address::ZERO);
-        });
+        proxy1.sender(alice).init(proxy2.address());
+        proxy2.sender(alice).init(proxy3.address());
+        proxy3.sender(alice).init(Address::ZERO);
 
         let zero = U256::ZERO;
-        let two = U256::from(2);
-        let four = U256::from(4);
-        let eight = U256::from(8);
+        let two = TWO;
+        let four = FOUR;
+        let eight = EIGHT;
 
         // Fund alice, proxies have no funds.
         alice.fund(eight);
@@ -511,7 +501,7 @@ mod proxies_tests {
         assert_eq!(proxy3.balance(), zero);
 
         // Call the first proxy.
-        proxy1.sender_and_value(alice, eight).pass_half_value_to_next_proxy();
+        proxy1.sender_and_value(alice, eight).pay_proxy_with_half_balance();
 
         assert_eq!(alice.balance(), zero);
         assert_eq!(proxy1.balance(), four);

From ef67e91c138203e86577566ead8919686617e988 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 11:44:27 +0400
Subject: [PATCH 22/44] unset msg value if it wasn't set for current call

---
 crates/motsu/src/context.rs | 35 ++++++++++++++++++++---------------
 crates/motsu/src/lib.rs     | 10 ++++------
 2 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index ac184ac..bc09748 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -90,20 +90,30 @@ impl Context {
         self.storage().contract_address.replace(address)
     }
 
-    /// Set message value to `value` and return the previous value if any.
-    pub(crate) fn set_msg_value(self, value: U256) -> Option<U256> {
-        self.storage().msg_value.replace(value)
+    /// Set optional message value to `value` and return the previous value if
+    /// any.
+    ///
+    /// Setting `value` to `None` will effectively clear the message value, e.g.
+    /// for non "payable" call.
+    pub(crate) fn set_optional_msg_value(
+        self,
+        value: Option<U256>,
+    ) -> Option<U256> {
+        let msg_value = &mut self.storage().msg_value;
+        let previous = msg_value.take();
+        *msg_value = value;
+        previous
     }
 
     /// Write the value sent to the contract to `output`.
     pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
-        let value: U256 = self.msg_value();
+        let value: U256 = self.msg_value().expect("msg_value should be set");
         write_u256(output, value);
     }
 
     /// Get the value sent to the contract as [`U256`].
-    pub(crate) fn msg_value(self) -> U256 {
-        self.storage().msg_value.unwrap_or_default()
+    pub(crate) fn msg_value(self) -> Option<U256> {
+        self.storage().msg_value
     }
 
     /// Get the address of the contract, that is called.
@@ -222,8 +232,7 @@ impl Context {
             .expect("msg_sender should be set");
 
         // Set new msg_value, and store the previous one.
-        let previous_msg_value =
-            value.and_then(|value| self.set_msg_value(value));
+        let previous_msg_value = self.set_optional_msg_value(value);
 
         // Transfer value sent by message sender.
         self.transfer_value();
@@ -240,10 +249,8 @@ impl Context {
         _ = self.set_contract_address(previous_contract_address);
         _ = self.set_msg_sender(previous_msg_sender);
 
-        // Set the previous msg_value if there is any.
-        if let Some(previous) = previous_msg_value {
-            _ = self.set_msg_value(previous);
-        }
+        // Set the previous msg_value.
+        self.set_optional_msg_value(previous_msg_value);
 
         result
     }
@@ -481,9 +488,7 @@ impl<ST: StorageType> ContractCall<'_, ST> {
 
     /// Preset the call parameters.
     fn set_call_params(&self) {
-        if let Some(value) = self.value {
-            _ = Context::current().set_msg_value(value);
-        }
+        _ = Context::current().set_optional_msg_value(self.value);
         _ = Context::current().set_msg_sender(self.caller_address);
         _ = Context::current().set_contract_address(self.address());
     }
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 291100f..4027d52 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -431,14 +431,12 @@ mod proxies_tests {
         proxy1.sender_and_value(alice, ONE).pay_proxy();
 
         // By the end, each actor will lose `ONE`, except last proxy.
-        assert_eq!(alice.balance(), TEN - ONE); // 10 - 1
-        assert_eq!(proxy1.balance(), TEN - ONE); // 10 + 1 - 2
-        assert_eq!(proxy2.balance(), TEN - ONE); // 10 + 2 - 3
-        assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE); // 10 + 3
+        assert_eq!(alice.balance(), TEN - ONE);
+        assert_eq!(proxy1.balance(), TEN - ONE);
+        assert_eq!(proxy2.balance(), TEN - ONE);
+        assert_eq!(proxy3.balance(), TEN + ONE + ONE + ONE);
     }
 
-    // TODO#q: if the msg value wasn't set we shouldn't set it
-
     #[motsu_proc::test]
     fn pay_proxy_fixed_value(
         proxy1: Contract<Proxy>,

From f4297ef3a5534597cef1704274655c6e13e0980f Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 12:04:39 +0400
Subject: [PATCH 23/44] fix fmt

---
 crates/motsu/src/lib.rs | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 4027d52..e86d324 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -372,11 +372,10 @@ mod proxies_tests {
     }
 
     impl Proxy {
-        fn init(&mut self, next_proxy: Address){
+        fn init(&mut self, next_proxy: Address) {
             self.next_proxy.set(next_proxy);
         }
     }
-    
 
     unsafe impl TopLevelStorage for Proxy {}
 
@@ -386,7 +385,6 @@ mod proxies_tests {
     const EIGHT: U256 = uint!(8_U256);
 
     const TEN: U256 = uint!(10_U256);
-    
 
     #[motsu_proc::test]
     fn call_three_proxies(

From 22e03c61564925a27840765084a87474f6f2d5b9 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 12:36:22 +0400
Subject: [PATCH 24/44] review fixes

---
 crates/motsu/src/context.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index bc09748..135212f 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -116,7 +116,7 @@ impl Context {
         self.storage().msg_value
     }
 
-    /// Get the address of the contract, that is called.
+    /// Get the address of the contract that is called.
     pub(crate) fn contract_address(self) -> Option<Address> {
         self.storage().contract_address
     }
@@ -349,7 +349,7 @@ impl Context {
         to: Address,
         value: U256,
     ) -> Option<()> {
-        _ = self.checked_sub_assign_balance(from, value)?;
+        self.checked_sub_assign_balance(from, value)?;
         self.add_assign_balance(to, value);
         Some(())
     }
@@ -551,10 +551,10 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     /// the given `account`.
     pub fn init<A: Into<Address>, Output>(
         &self,
-        account: A,
+        sender: A,
         initializer: impl FnOnce(&mut ST) -> Output,
     ) -> Output {
-        initializer(&mut self.sender(account.into()))
+        initializer(&mut self.sender(sender.into()))
     }
 
     /// Create a new contract with default storage on the random address.
@@ -584,10 +584,10 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     #[must_use]
     pub fn sender_and_value<A: Into<Address>, V: Into<U256>>(
         &self,
-        account: A,
+        sender: A,
         value: V,
     ) -> ContractCall<ST> {
-        let caller_address = account.into();
+        let caller_address = sender.into();
         let value = value.into();
 
         ContractCall {

From f33fe12006a10efd838ea59f389adfa76b49a9a1 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 12:47:27 +0400
Subject: [PATCH 25/44] review fixes

---
 crates/motsu/src/context.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 135212f..285c611 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -90,7 +90,7 @@ impl Context {
         self.storage().contract_address.replace(address)
     }
 
-    /// Set optional message value to `value` and return the previous value if
+    /// Set an optional message value to `value` and return the previous value if
     /// any.
     ///
     /// Setting `value` to `None` will effectively clear the message value, e.g.
@@ -193,7 +193,7 @@ impl Context {
     }
 
     /// Based on `result`, set the return data.
-    /// Return 0 if `result` is `Ok`, otherwise return 1.
+    /// Return 0 if `result` is `Ok`, otherwise 1.
     unsafe fn process_arb_result_raw(
         self,
         result: ArbResult,

From 0bf106bc75932d0abdc2fa6aa57dd85cae034642 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:13:42 +0400
Subject: [PATCH 26/44] fix deadlock + deadlock test

---
 crates/motsu/src/context.rs | 17 ++++++++---
 crates/motsu/src/lib.rs     | 22 +++++++++++++++
 crates/motsu/src/router.rs  | 56 +++++++++++++++++++++++++++----------
 3 files changed, 77 insertions(+), 18 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 285c611..20c1aa3 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -4,6 +4,7 @@ use std::{collections::HashMap, ptr, slice, thread::ThreadId};
 
 use alloy_primitives::{Address, B256, U256};
 use dashmap::{mapref::one::RefMut, DashMap};
+use dashmap::try_result::TryResult;
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
 
@@ -17,7 +18,7 @@ use crate::router::{RouterContext, TestRouter};
 ///
 /// The value is the [`MockStorage`], a storage of the test case.
 ///
-/// NOTE: The [`DashMap`] will deadlock execution, when the same key is
+/// NOTE: The [`Context::storage`] will panic on lock, when the same key is
 /// accessed twice from the same thread.
 static STORAGE: Lazy<DashMap<Context, MockStorage>> = Lazy::new(DashMap::new);
 
@@ -150,7 +151,7 @@ impl Context {
 
         // if no more contracts left,
         if storage.contract_data.is_empty() {
-            // drop guard to a concurrent hash map to avoid a deadlock,
+            // drop guard to a concurrent hash map to avoid a panic on lock,
             drop(storage);
             // and erase the test context.
             _ = STORAGE.remove(&self);
@@ -321,7 +322,7 @@ impl Context {
 
         // We should transfer the value only if it is set.
         if let Some(msg_value) = storage.msg_value {
-            // Drop storage to avoid a deadlock.
+            // Drop storage to avoid a panic on lock.
             drop(storage);
 
             // Transfer and panic if there is not enough funds.
@@ -383,7 +384,15 @@ impl Context {
 
     /// Get reference to the storage for the current test thread.
     fn storage(self) -> RefMut<'static, Context, MockStorage> {
-        STORAGE.get_mut(&self).expect("contract should be initialised first")
+        match STORAGE.try_get_mut(&self) {
+            TryResult::Present(router) => router,
+            TryResult::Absent => {
+                panic!("contract should be initialised first")
+            }
+            TryResult::Locked => {
+                panic!("storage is locked")
+            }
+        }
     }
 
     /// Get router for the contract at `address`.
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index e86d324..b0641e1 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -504,4 +504,26 @@ mod proxies_tests {
         assert_eq!(proxy2.balance(), two);
         assert_eq!(proxy3.balance(), two);
     }
+
+    #[motsu_proc::test]
+    fn no_locks_with_panics() {
+        for _ in 0..1000 {
+            let proxy1 = Contract::<Proxy>::new();
+            let proxy2 = Contract::<Proxy>::new();
+            let proxy3 = Contract::<Proxy>::new();
+            let alice = Account::random();
+
+            // Set up a chain of three proxies.
+            // With the given call chain: proxy1 -> proxy2 -> proxy3.
+            proxy1.sender(alice).init(proxy2.address());
+            proxy2.sender(alice).init(proxy3.address());
+            proxy3.sender(alice).init(Address::ZERO);
+
+            // Call the first proxy.
+            let result = proxy1.sender(alice).call_proxy(TEN);
+
+            // The value is incremented by 1 for each proxy.
+            assert_eq!(result, TEN + ONE + ONE + ONE);
+        }
+    }
 }
diff --git a/crates/motsu/src/router.rs b/crates/motsu/src/router.rs
index f12ea20..3916879 100644
--- a/crates/motsu/src/router.rs
+++ b/crates/motsu/src/router.rs
@@ -1,12 +1,16 @@
 //! Router context for external calls mocks.
 //!
-//! NOTE: [`ROUTER_STORAGE`] should be separated from the main test storage to
-//! avoid  deadlocks.
+//! NOTE: [`ROUTER_STORAGE`] is separated from the main test storage to
+//! avoid a panic on lock.
 
-use std::{borrow::BorrowMut, sync::Mutex, thread::ThreadId};
+use std::{
+    borrow::BorrowMut,
+    sync::{Arc, Mutex, TryLockError},
+    thread::ThreadId,
+};
 
 use alloy_primitives::{uint, Address};
-use dashmap::{mapref::one::RefMut, DashMap};
+use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{
     abi::Router,
@@ -19,13 +23,13 @@ use stylus_sdk::{
 /// A global mutable key-value store that allows concurrent access.
 ///
 /// The key is the [`RouterContext`], a combination of [`ThreadId`] and
-/// [`Address`] to avoid a deadlock, while calling more than two contracts
+/// [`Address`] to avoid a panic on lock, while calling more than two contracts
 /// consecutive.
 ///
 /// The value is the [`RouterStorage`], a router of the contract generated by
 /// `stylus-sdk`.
 ///
-/// NOTE: The [`DashMap`] will deadlock execution, when the same key is
+/// NOTE: The [`RouterContext::storage`] will panic on lock, when the same key is
 /// accessed twice from the same thread.
 static ROUTER_STORAGE: Lazy<DashMap<RouterContext, RouterStorage>> =
     Lazy::new(DashMap::new);
@@ -46,9 +50,15 @@ impl RouterContext {
 
     /// Get reference to the call storage for the current test thread.
     fn storage(self) -> RefMut<'static, RouterContext, RouterStorage> {
-        ROUTER_STORAGE
-            .get_mut(&self)
-            .expect("contract should be initialised first")
+        match ROUTER_STORAGE.try_get_mut(&self) {
+            TryResult::Present(router) => router,
+            TryResult::Absent => {
+                panic!("contract should be initialised first")
+            }
+            TryResult::Locked => {
+                panic!("router storage is locked")
+            }
+        }
     }
 
     /// Check if the router exists for the contract.
@@ -61,9 +71,27 @@ impl RouterContext {
         selector: u32,
         input: &[u8],
     ) -> Option<ArbResult> {
-        let router = &self.storage().router;
-        let mut router = router.lock().expect("should lock test router");
-        router.route(selector, input)
+        let storage = self.storage();
+        let router = Arc::clone(&storage.router);
+
+        // Drop the storage reference to avoid a panic on lock.
+        drop(storage);
+
+        // Try to get lock on the router.
+        let lock_result = router.try_lock();
+        match lock_result {
+            // If lock is acquired, route the message.
+            Ok(mut router) => router.route(selector, input),
+            // Panic instead of locking.
+            Err(TryLockError::WouldBlock) => {
+                panic!("recursive calls are not supported in motsu")
+            }
+            // This branch should not be reached, since we don't catch
+            // panics.
+            Err(TryLockError::Poisoned(_)) => {
+                panic!("should not call contract that panicked")
+            }
+        }
     }
 
     /// Initialise contract router for the current test thread and
@@ -74,7 +102,7 @@ impl RouterContext {
             .insert(
                 self,
                 RouterStorage {
-                    router: Mutex::new(Box::new(unsafe {
+                    router: Arc::new(Mutex::new(unsafe {
                         ST::new(uint!(0_U256), 0)
                     })),
                 },
@@ -95,7 +123,7 @@ impl RouterContext {
 struct RouterStorage {
     // Contract's router.
     // NOTE: Mutex is important since contract type is not `Sync`.
-    router: Mutex<Box<dyn TestRouter>>,
+    router: Arc<Mutex<dyn TestRouter>>,
 }
 
 /// A trait for routing messages to the appropriate selector in tests.

From e059c0c83536a681a370c643c8ac0329c3107095 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:22:56 +0400
Subject: [PATCH 27/44] ++

---
 crates/motsu/src/lib.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index b0641e1..a2e6d1b 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -436,7 +436,7 @@ mod proxies_tests {
     }
 
     #[motsu_proc::test]
-    fn pay_proxy_fixed_value(
+    fn pass_proxy_with_fixed_value(
         proxy1: Contract<Proxy>,
         proxy2: Contract<Proxy>,
         proxy3: Contract<Proxy>,
@@ -471,7 +471,7 @@ mod proxies_tests {
     }
 
     #[motsu_proc::test]
-    fn wrong_contract_balances_for_nested_contract_calls(
+    fn pay_proxy_with_half_balance(
         proxy1: Contract<Proxy>,
         proxy2: Contract<Proxy>,
         proxy3: Contract<Proxy>,

From d74687a84cd9f64623ca009b0652a3841512b92e Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:27:02 +0400
Subject: [PATCH 28/44] ++

---
 crates/motsu/src/lib.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index a2e6d1b..b38df9b 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -289,9 +289,9 @@ mod proxies_tests {
             #[allow(missing_docs)]
             function payProxy() external payable;
             #[allow(missing_docs)]
-            function passThisMuchToNextProxy(uint256 pass_value) external payable;
+            function passProxyWithFixedValue(uint256 pass_value) external payable;
             #[allow(missing_docs)]
-            function passHalfValueToNextProxy() external payable;
+            function payProxyWithHalfBalance() external payable;
         }
     }
 
@@ -346,7 +346,7 @@ mod proxies_tests {
                 let call = Call::new_in(self).value(this_value);
                 let value_for_next_next_proxy = this_value / U256::from(2);
                 proxy
-                    .pass_this_much_to_next_proxy(
+                    .pass_proxy_with_fixed_value(
                         call,
                         value_for_next_next_proxy,
                     )
@@ -365,7 +365,7 @@ mod proxies_tests {
                 let proxy = IProxy::new(next_proxy);
                 let call = Call::new_in(self).value(half_balance);
                 proxy
-                    .pass_half_value_to_next_proxy(call)
+                    .pay_proxy_with_half_balance(call)
                     .expect("should pass half the value to the next proxy");
             }
         }

From 6ebd6047537f0dd1756eefe91adba67f1f226226 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:30:23 +0400
Subject: [PATCH 29/44] ++

---
 crates/motsu/src/router.rs | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/crates/motsu/src/router.rs b/crates/motsu/src/router.rs
index 3916879..4cdce4a 100644
--- a/crates/motsu/src/router.rs
+++ b/crates/motsu/src/router.rs
@@ -1,7 +1,4 @@
 //! Router context for external calls mocks.
-//!
-//! NOTE: [`ROUTER_STORAGE`] is separated from the main test storage to
-//! avoid a panic on lock.
 
 use std::{
     borrow::BorrowMut,

From 621b74a6027707e4fc200657a8ec8e6b503f9829 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:37:54 +0400
Subject: [PATCH 30/44] fix fmt

---
 crates/motsu/src/context.rs | 7 +++----
 crates/motsu/src/router.rs  | 4 ++--
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 20c1aa3..9c7cbc8 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -3,8 +3,7 @@
 use std::{collections::HashMap, ptr, slice, thread::ThreadId};
 
 use alloy_primitives::{Address, B256, U256};
-use dashmap::{mapref::one::RefMut, DashMap};
-use dashmap::try_result::TryResult;
+use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
 
@@ -91,8 +90,8 @@ impl Context {
         self.storage().contract_address.replace(address)
     }
 
-    /// Set an optional message value to `value` and return the previous value if
-    /// any.
+    /// Set an optional message value to `value` and return the previous value
+    /// if any.
     ///
     /// Setting `value` to `None` will effectively clear the message value, e.g.
     /// for non "payable" call.
diff --git a/crates/motsu/src/router.rs b/crates/motsu/src/router.rs
index 4cdce4a..bdc6ca1 100644
--- a/crates/motsu/src/router.rs
+++ b/crates/motsu/src/router.rs
@@ -26,8 +26,8 @@ use stylus_sdk::{
 /// The value is the [`RouterStorage`], a router of the contract generated by
 /// `stylus-sdk`.
 ///
-/// NOTE: The [`RouterContext::storage`] will panic on lock, when the same key is
-/// accessed twice from the same thread.
+/// NOTE: The [`RouterContext::storage`] will panic on lock, when the same key
+/// is accessed twice from the same thread.
 static ROUTER_STORAGE: Lazy<DashMap<RouterContext, RouterStorage>> =
     Lazy::new(DashMap::new);
 

From 9953416b994426c27365dba11c1eec4e383c9963 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:39:55 +0400
Subject: [PATCH 31/44] update CHANGELOG.md

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8b2371e..55ecda7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Mocks for the `msg::sender()` #14
+- Mocks for the `msg::value()` and account balance #31
 - Mocks for the external contract calls. 
   Two and more contracts can be injected into test #14
 - Option to inject `Account` or `Address` in the test #14

From e1f842ab047e479e91c848bfe3ef6d3a4ccd5e6e Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:50:01 +0400
Subject: [PATCH 32/44] update tests

---
 crates/motsu/src/lib.rs | 50 +++++++++++++++++------------------------
 1 file changed, 20 insertions(+), 30 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index b38df9b..8ccb306 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -448,26 +448,21 @@ mod proxies_tests {
         proxy2.sender(alice).init(proxy3.address());
         proxy3.sender(alice).init(Address::ZERO);
 
-        let zero = U256::ZERO;
-        let two = uint!(2_U256);
-        let four = U256::from(4);
-        let eight = U256::from(8);
-
         // Fund alice, proxies have no funds.
-        alice.fund(eight);
+        alice.fund(EIGHT);
 
-        assert_eq!(alice.balance(), eight);
-        assert_eq!(proxy1.balance(), zero);
-        assert_eq!(proxy2.balance(), zero);
-        assert_eq!(proxy3.balance(), zero);
+        assert_eq!(alice.balance(), EIGHT);
+        assert_eq!(proxy1.balance(), U256::ZERO);
+        assert_eq!(proxy2.balance(), U256::ZERO);
+        assert_eq!(proxy3.balance(), U256::ZERO);
 
         // Call the first proxy.
-        proxy1.sender_and_value(alice, eight).pass_proxy_with_fixed_value(four);
+        proxy1.sender_and_value(alice, EIGHT).pass_proxy_with_fixed_value(FOUR);
 
-        assert_eq!(alice.balance(), zero);
-        assert_eq!(proxy1.balance(), four);
-        assert_eq!(proxy2.balance(), two);
-        assert_eq!(proxy3.balance(), two);
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(proxy1.balance(), FOUR);
+        assert_eq!(proxy2.balance(), TWO);
+        assert_eq!(proxy3.balance(), TWO);
     }
 
     #[motsu_proc::test]
@@ -483,26 +478,21 @@ mod proxies_tests {
         proxy2.sender(alice).init(proxy3.address());
         proxy3.sender(alice).init(Address::ZERO);
 
-        let zero = U256::ZERO;
-        let two = TWO;
-        let four = FOUR;
-        let eight = EIGHT;
-
         // Fund alice, proxies have no funds.
-        alice.fund(eight);
+        alice.fund(EIGHT);
 
-        assert_eq!(alice.balance(), eight);
-        assert_eq!(proxy1.balance(), zero);
-        assert_eq!(proxy2.balance(), zero);
-        assert_eq!(proxy3.balance(), zero);
+        assert_eq!(alice.balance(), EIGHT);
+        assert_eq!(proxy1.balance(), U256::ZERO);
+        assert_eq!(proxy2.balance(), U256::ZERO);
+        assert_eq!(proxy3.balance(), U256::ZERO);
 
         // Call the first proxy.
-        proxy1.sender_and_value(alice, eight).pay_proxy_with_half_balance();
+        proxy1.sender_and_value(alice, EIGHT).pay_proxy_with_half_balance();
 
-        assert_eq!(alice.balance(), zero);
-        assert_eq!(proxy1.balance(), four);
-        assert_eq!(proxy2.balance(), two);
-        assert_eq!(proxy3.balance(), two);
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(proxy1.balance(), FOUR);
+        assert_eq!(proxy2.balance(), TWO);
+        assert_eq!(proxy3.balance(), TWO);
     }
 
     #[motsu_proc::test]

From 1e87140f9800183521d82fdcc9eefe1e00cff2e2 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 13:59:00 +0400
Subject: [PATCH 33/44] ++

---
 crates/motsu/src/lib.rs | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 8ccb306..d681ed8 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -270,12 +270,10 @@ mod ping_pong_tests {
 
 #[cfg(test)]
 mod proxies_tests {
-    use alloy_primitives::{address, uint, Address, U256};
+    use alloy_primitives::{uint, Address, U256};
     use stylus_sdk::{
         call::Call,
-        contract,
-        contract::address,
-        msg,
+        contract, msg,
         prelude::{public, storage, TopLevelStorage},
         storage::StorageAddress,
     };
@@ -344,7 +342,7 @@ mod proxies_tests {
                 // Pay the next proxy.
                 let proxy = IProxy::new(next_proxy);
                 let call = Call::new_in(self).value(this_value);
-                let value_for_next_next_proxy = this_value / U256::from(2);
+                let value_for_next_next_proxy = this_value / TWO;
                 proxy
                     .pass_proxy_with_fixed_value(
                         call,
@@ -360,7 +358,7 @@ mod proxies_tests {
 
             // If there is a next proxy.
             if !next_proxy.is_zero() {
-                let half_balance = contract::balance() / U256::from(2);
+                let half_balance = contract::balance() / TWO;
                 // Pay the next proxy.
                 let proxy = IProxy::new(next_proxy);
                 let call = Call::new_in(self).value(half_balance);

From 47049ea3692c7646ff03ef3f8f827f1ad127dd60 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 14:58:23 +0400
Subject: [PATCH 34/44] review fixes

---
 crates/motsu/src/context.rs | 34 ++++++++++++++++------------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 9c7cbc8..69faf64 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -75,7 +75,7 @@ impl Context {
     }
 
     /// Set the message sender address and return the previous sender if any.
-    fn set_msg_sender(self, msg_sender: Address) -> Option<Address> {
+    fn replace_msg_sender(self, msg_sender: Address) -> Option<Address> {
         self.storage().msg_sender.replace(msg_sender)
     }
 
@@ -85,24 +85,22 @@ impl Context {
         self.storage().msg_sender
     }
 
-    /// Set the address of the contract, that is called.
-    fn set_contract_address(self, address: Address) -> Option<Address> {
+    /// Replace the address of the contract, and return the previous address if
+    /// any.
+    fn replace_contract_address(self, address: Address) -> Option<Address> {
         self.storage().contract_address.replace(address)
     }
 
-    /// Set an optional message value to `value` and return the previous value
+    /// Replace an optional message with `value` and return the previous value
     /// if any.
     ///
     /// Setting `value` to `None` will effectively clear the message value, e.g.
     /// for non "payable" call.
-    pub(crate) fn set_optional_msg_value(
+    pub(crate) fn replace_optional_msg_value(
         self,
         value: Option<U256>,
     ) -> Option<U256> {
-        let msg_value = &mut self.storage().msg_value;
-        let previous = msg_value.take();
-        *msg_value = value;
-        previous
+        std::mem::replace(&mut self.storage().msg_value, value)
     }
 
     /// Write the value sent to the contract to `output`.
@@ -225,14 +223,14 @@ impl Context {
         // Set the caller contract as message sender and callee contract as
         // a receiver (`contract_address`).
         let previous_contract_address = self
-            .set_contract_address(contract_address)
+            .replace_contract_address(contract_address)
             .expect("contract_address should be set");
         let previous_msg_sender = self
-            .set_msg_sender(previous_contract_address)
+            .replace_msg_sender(previous_contract_address)
             .expect("msg_sender should be set");
 
         // Set new msg_value, and store the previous one.
-        let previous_msg_value = self.set_optional_msg_value(value);
+        let previous_msg_value = self.replace_optional_msg_value(value);
 
         // Transfer value sent by message sender.
         self.transfer_value();
@@ -246,11 +244,11 @@ impl Context {
             });
 
         // Set the previous message sender and contract address back.
-        _ = self.set_contract_address(previous_contract_address);
-        _ = self.set_msg_sender(previous_msg_sender);
+        _ = self.replace_contract_address(previous_contract_address);
+        _ = self.replace_msg_sender(previous_msg_sender);
 
         // Set the previous msg_value.
-        self.set_optional_msg_value(previous_msg_value);
+        self.replace_optional_msg_value(previous_msg_value);
 
         result
     }
@@ -496,9 +494,9 @@ impl<ST: StorageType> ContractCall<'_, ST> {
 
     /// Preset the call parameters.
     fn set_call_params(&self) {
-        _ = Context::current().set_optional_msg_value(self.value);
-        _ = Context::current().set_msg_sender(self.caller_address);
-        _ = Context::current().set_contract_address(self.address());
+        _ = Context::current().replace_optional_msg_value(self.value);
+        _ = Context::current().replace_msg_sender(self.caller_address);
+        _ = Context::current().replace_contract_address(self.address());
     }
 }
 

From ee264138b4cc00137dc2d11448819dbbcdb9e58e Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 15:00:24 +0400
Subject: [PATCH 35/44] review fixes

---
 crates/motsu/src/context.rs | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 69faf64..cd7275a 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -475,8 +475,8 @@ pub(crate) type Bytes32 = [u8; WORD_BYTES];
 /// account.
 pub struct ContractCall<'a, ST: StorageType> {
     storage: ST,
-    caller_address: Address,
-    value: Option<U256>,
+    msg_sender: Address,
+    msg_value: Option<U256>,
     /// We need to hold a reference to [`Contract<ST>`], because
     /// `Contract::<ST>::new().sender(alice)` can accidentally drop
     /// [`Contract<ST>`].
@@ -494,8 +494,8 @@ impl<ST: StorageType> ContractCall<'_, ST> {
 
     /// Preset the call parameters.
     fn set_call_params(&self) {
-        _ = Context::current().replace_optional_msg_value(self.value);
-        _ = Context::current().replace_msg_sender(self.caller_address);
+        _ = Context::current().replace_optional_msg_value(self.msg_value);
+        _ = Context::current().replace_msg_sender(self.msg_sender);
         _ = Context::current().replace_contract_address(self.address());
     }
 }
@@ -580,8 +580,8 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     pub fn sender<A: Into<Address>>(&self, account: A) -> ContractCall<ST> {
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
-            caller_address: account.into(),
-            value: None,
+            msg_sender: account.into(),
+            msg_value: None,
             contract_ref: self,
         }
     }
@@ -598,8 +598,8 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
 
         ContractCall {
             storage: unsafe { ST::new(uint!(0_U256), 0) },
-            caller_address,
-            value: Some(value),
+            msg_sender: caller_address,
+            msg_value: Some(value),
             contract_ref: self,
         }
     }

From a2258710b5cda8c4cafe25112a47bd946524139b Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 15:05:32 +0400
Subject: [PATCH 36/44] update changelog

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 55ecda7..340e73c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Mocks for the `msg::sender()` #14
-- Mocks for the `msg::value()` and account balance #31
+- Mocks for the `msg::value()` and `contract::balance()` #31
 - Mocks for the external contract calls. 
   Two and more contracts can be injected into test #14
 - Option to inject `Account` or `Address` in the test #14

From 5a50fc0daa4993efd3d05930056da2df69148d07 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 15:19:30 +0400
Subject: [PATCH 37/44] review fix

---
 crates/motsu/src/context.rs | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index cd7275a..6fde18a 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -163,13 +163,13 @@ impl Context {
         address: *const u8,
         calldata: *const u8,
         calldata_len: usize,
-        return_data_len: *mut usize,
+        return_data_size: *mut usize,
     ) -> u8 {
         let address = read_address(address);
         let (selector, input) = decode_calldata(calldata, calldata_len);
 
         let result = self.call_contract(address, selector, &input, None);
-        self.process_arb_result_raw(result, return_data_len)
+        self.process_arb_result_raw(result, return_data_size)
     }
 
     /// Call the contract at raw `address` with the given raw `calldata` and
@@ -180,14 +180,14 @@ impl Context {
         calldata: *const u8,
         calldata_len: usize,
         value: *const u8,
-        return_data_len: *mut usize,
+        return_data_size: *mut usize,
     ) -> u8 {
         let address = read_address(address);
         let value = read_u256(value);
         let (selector, input) = decode_calldata(calldata, calldata_len);
 
         let result = self.call_contract(address, selector, &input, Some(value));
-        self.process_arb_result_raw(result, return_data_len)
+        self.process_arb_result_raw(result, return_data_size)
     }
 
     /// Based on `result`, set the return data.
@@ -195,16 +195,16 @@ impl Context {
     unsafe fn process_arb_result_raw(
         self,
         result: ArbResult,
-        return_data_len: *mut usize,
+        return_data_size: *mut usize,
     ) -> u8 {
         match result {
             Ok(res) => {
-                return_data_len.write(res.len());
+                return_data_size.write(res.len());
                 self.set_return_data(res);
                 0
             }
             Err(err) => {
-                return_data_len.write(err.len());
+                return_data_size.write(err.len());
                 self.set_return_data(err);
                 1
             }
@@ -256,8 +256,8 @@ impl Context {
     /// Set return data as bytes.
     fn set_return_data(self, data: Vec<u8>) {
         let mut call_storage = self.storage();
-        _ = call_storage.call_output_len.insert(data.len());
-        _ = call_storage.call_output.insert(data);
+        _ = call_storage.return_data_size.insert(data.len());
+        _ = call_storage.return_data.insert(data);
     }
 
     /// Read the return data (with a given `size`) from the last contract call
@@ -275,14 +275,14 @@ impl Context {
     /// Return data's size in bytes from the last contract call.
     pub(crate) fn return_data_size(self) -> usize {
         self.storage()
-            .call_output_len
+            .return_data_size
             .take()
             .expect("call_output_len should be set")
     }
 
     /// Return data's bytes from the last contract call.
     fn return_data(self) -> Vec<u8> {
-        self.storage().call_output.take().expect("call_output should be set")
+        self.storage().return_data.take().expect("call_output should be set")
     }
 
     /// Check if the contract at raw `address` has code.
@@ -461,9 +461,9 @@ struct MockStorage {
     /// Account's address to balance mapping.
     balances: HashMap<Address, U256>,
     // Output of a contract call.
-    call_output: Option<Vec<u8>>,
+    return_data: Option<Vec<u8>>,
     // Output length of a contract call.
-    call_output_len: Option<usize>,
+    return_data_size: Option<usize>,
 }
 
 /// Contract's byte storage

From 6d3ce0dd9ab1488e700e2f1237199066fe85c4d2 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 15:37:24 +0400
Subject: [PATCH 38/44] add unit test workflow

---
 .github/workflows/test.yml  | 52 +++++++++++++++++++++++++++++++++++++
 crates/motsu/src/context.rs |  2 +-
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/test.yml

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..3eaa959
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,52 @@
+name: test
+# This is the main CI workflow that runs the test suite on all pushes to main
+# and all pull requests. It runs the following jobs:
+# - required: runs the test suite on ubuntu with stable and beta rust
+# toolchains.
+permissions:
+  contents: read
+on:
+  push:
+    branches: [ main, release/* ]
+  pull_request:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+env:
+  CARGO_TERM_COLOR: always
+jobs:
+  required:
+    runs-on: ubuntu-latest
+    name: ubuntu / ${{ matrix.toolchain }}
+    strategy:
+      matrix:
+        # Run on stable and beta to ensure that tests won't break on the next
+        # version of the rust toolchain.
+        toolchain: [ stable, beta ]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: true
+
+      - name: Install rust ${{ matrix.toolchain }}
+        uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          toolchain: ${{ matrix.toolchain }}
+          rustflags: ""
+
+      - name: Install nextest
+        uses: taiki-e/install-action@v2
+        with:
+          tool: cargo-nextest
+
+      - name: Cargo generate-lockfile
+        # Enable this ci template to run regardless of whether the lockfile is
+        # checked in or not.
+        if: hashFiles('Cargo.lock') == ''
+        run: cargo generate-lockfile
+
+      - name: Run unit tests
+        run: cargo nextest run --locked --all-targets
+
+      - name: Run doc tests
+        run: cargo test --locked --doc
\ No newline at end of file
diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 6fde18a..1d7bc67 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -105,7 +105,7 @@ impl Context {
 
     /// Write the value sent to the contract to `output`.
     pub(crate) unsafe fn msg_value_raw(self, output: *mut u8) {
-        let value: U256 = self.msg_value().expect("msg_value should be set");
+        let value: U256 = self.msg_value().unwrap_or_default();
         write_u256(output, value);
     }
 

From d9c27b2725e66218efd6f3226c5d9ea6d27b5001 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 16:30:56 +0400
Subject: [PATCH 39/44] add backoff storage access without deadlocks

---
 crates/motsu/src/context.rs        | 19 ++++-----
 crates/motsu/src/lib.rs            |  1 +
 crates/motsu/src/router.rs         | 14 ++-----
 crates/motsu/src/storage_access.rs | 63 ++++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 22 deletions(-)
 create mode 100644 crates/motsu/src/storage_access.rs

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 1d7bc67..0dc18b0 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -1,13 +1,16 @@
 //! Unit-testing context for Stylus contracts.
 
-use std::{collections::HashMap, ptr, slice, thread::ThreadId};
+use std::{collections::HashMap, hash::Hash, ptr, slice, thread::ThreadId};
 
 use alloy_primitives::{Address, B256, U256};
-use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
+use dashmap::{mapref::one::RefMut, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
 
-use crate::router::{RouterContext, TestRouter};
+use crate::{
+    router::{RouterContext, TestRouter},
+    storage_access::AccessStorage,
+};
 
 /// Storage mock.
 ///
@@ -381,15 +384,7 @@ impl Context {
 
     /// Get reference to the storage for the current test thread.
     fn storage(self) -> RefMut<'static, Context, MockStorage> {
-        match STORAGE.try_get_mut(&self) {
-            TryResult::Present(router) => router,
-            TryResult::Absent => {
-                panic!("contract should be initialised first")
-            }
-            TryResult::Locked => {
-                panic!("storage is locked")
-            }
-        }
+        STORAGE.access_storage(&self)
     }
 
     /// Get router for the contract at `address`.
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index d681ed8..7a55d12 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -54,6 +54,7 @@ mod context;
 pub mod prelude;
 mod router;
 mod shims;
+mod storage_access;
 
 pub use motsu_proc::test;
 
diff --git a/crates/motsu/src/router.rs b/crates/motsu/src/router.rs
index bdc6ca1..9343c3b 100644
--- a/crates/motsu/src/router.rs
+++ b/crates/motsu/src/router.rs
@@ -7,7 +7,7 @@ use std::{
 };
 
 use alloy_primitives::{uint, Address};
-use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
+use dashmap::{mapref::one::RefMut, DashMap};
 use once_cell::sync::Lazy;
 use stylus_sdk::{
     abi::Router,
@@ -15,6 +15,8 @@ use stylus_sdk::{
     ArbResult,
 };
 
+use crate::storage_access::AccessStorage;
+
 /// Router Storage.
 ///
 /// A global mutable key-value store that allows concurrent access.
@@ -47,15 +49,7 @@ impl RouterContext {
 
     /// Get reference to the call storage for the current test thread.
     fn storage(self) -> RefMut<'static, RouterContext, RouterStorage> {
-        match ROUTER_STORAGE.try_get_mut(&self) {
-            TryResult::Present(router) => router,
-            TryResult::Absent => {
-                panic!("contract should be initialised first")
-            }
-            TryResult::Locked => {
-                panic!("router storage is locked")
-            }
-        }
+        ROUTER_STORAGE.access_storage(&self)
     }
 
     /// Check if the router exists for the contract.
diff --git a/crates/motsu/src/storage_access.rs b/crates/motsu/src/storage_access.rs
new file mode 100644
index 0000000..d2d09a8
--- /dev/null
+++ b/crates/motsu/src/storage_access.rs
@@ -0,0 +1,63 @@
+use std::hash::Hash;
+use dashmap::DashMap;
+use dashmap::mapref::one::RefMut;
+use dashmap::try_result::TryResult;
+
+/// Trait for accessing test storage.
+pub(crate) trait AccessStorage {
+    type Key;
+    type Value;
+    
+    /// Get mutable access to storage with `key`.
+    /// 
+    /// # Panics
+    /// 
+    /// * After 10 attempts to access the storage.
+    /// * If the contract wasn't initialized.
+    fn access_storage(
+        &self,
+        key: &Self::Key,
+    ) -> RefMut<Self::Key, Self::Value> {
+        self.access_storage_with_backoff(key, 10)
+    }
+
+    /// Get mutable access to storage with `key`, with `backoff` number of attempts.
+    /// 
+    /// # Panics
+    /// 
+    /// * After `backoff` attempts to access the storage.
+    /// * If the contract wasn't initialized.
+    fn access_storage_with_backoff(
+        &self,
+        key: &Self::Key,
+        backoff: u32,
+    ) -> RefMut<Self::Key, Self::Value>;
+}
+
+impl<K: Hash + Eq, V> AccessStorage for DashMap<K, V> {
+    type Key = K;
+    type Value = V;
+
+    fn access_storage_with_backoff(
+        &self,
+        key: &Self::Key,
+        backoff: u32,
+    ) -> RefMut<Self::Key, Self::Value> {
+        {
+            match self.try_get_mut(key) {
+                TryResult::Present(router) => router,
+                TryResult::Absent => {
+                    panic!("contract should be initialised first")
+                }
+                TryResult::Locked => {
+                    if backoff == 0 {
+                        panic!("storage is locked")
+                    } else {
+                        std::thread::sleep(std::time::Duration::from_millis(1));
+                        self.access_storage_with_backoff(key, backoff - 1)
+                    }
+                }
+            }
+        }
+    }
+}
\ No newline at end of file

From 3cc244f71636ab93baac3d2f9921d3fb4ece57e3 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 16:37:20 +0400
Subject: [PATCH 40/44] remove `address()` from `ContractCall` to avoid method
 collisions with derefed contract

---
 crates/motsu/src/context.rs | 7 +------
 crates/motsu/src/lib.rs     | 3 ++-
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 0dc18b0..2e9b8a7 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -482,16 +482,11 @@ pub struct ContractCall<'a, ST: StorageType> {
 }
 
 impl<ST: StorageType> ContractCall<'_, ST> {
-    /// Get the contract's address.
-    pub fn address(&self) -> Address {
-        self.contract_ref.address
-    }
-
     /// Preset the call parameters.
     fn set_call_params(&self) {
         _ = Context::current().replace_optional_msg_value(self.msg_value);
         _ = Context::current().replace_msg_sender(self.msg_sender);
-        _ = Context::current().replace_contract_address(self.address());
+        _ = Context::current().replace_contract_address(self.contract_ref.address);
     }
 }
 
diff --git a/crates/motsu/src/lib.rs b/crates/motsu/src/lib.rs
index 7a55d12..f348b84 100644
--- a/crates/motsu/src/lib.rs
+++ b/crates/motsu/src/lib.rs
@@ -261,10 +261,11 @@ mod ping_pong_tests {
         let ping = Contract::<PingContract>::new();
         let mut ping = ping.sender(alice);
         let pong = Contract::<PongContract>::new();
+        let pong_address = pong.address();
         let pong = pong.sender(alice);
 
         _ = ping
-            .ping(pong.address(), TEN)
+            .ping(pong_address, TEN)
             .expect("contract ping should not drop");
     }
 }

From e58663e2696fa504637e8db108df8eacf7e07691 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 16:41:23 +0400
Subject: [PATCH 41/44] fix fmt

---
 crates/motsu/src/context.rs        |  3 ++-
 crates/motsu/src/storage_access.rs | 20 ++++++++++----------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 2e9b8a7..23232a0 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -486,7 +486,8 @@ impl<ST: StorageType> ContractCall<'_, ST> {
     fn set_call_params(&self) {
         _ = Context::current().replace_optional_msg_value(self.msg_value);
         _ = Context::current().replace_msg_sender(self.msg_sender);
-        _ = Context::current().replace_contract_address(self.contract_ref.address);
+        _ = Context::current()
+            .replace_contract_address(self.contract_ref.address);
     }
 }
 
diff --git a/crates/motsu/src/storage_access.rs b/crates/motsu/src/storage_access.rs
index d2d09a8..0a36dae 100644
--- a/crates/motsu/src/storage_access.rs
+++ b/crates/motsu/src/storage_access.rs
@@ -1,17 +1,16 @@
 use std::hash::Hash;
-use dashmap::DashMap;
-use dashmap::mapref::one::RefMut;
-use dashmap::try_result::TryResult;
+
+use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
 
 /// Trait for accessing test storage.
 pub(crate) trait AccessStorage {
     type Key;
     type Value;
-    
+
     /// Get mutable access to storage with `key`.
-    /// 
+    ///
     /// # Panics
-    /// 
+    ///
     /// * After 10 attempts to access the storage.
     /// * If the contract wasn't initialized.
     fn access_storage(
@@ -21,10 +20,11 @@ pub(crate) trait AccessStorage {
         self.access_storage_with_backoff(key, 10)
     }
 
-    /// Get mutable access to storage with `key`, with `backoff` number of attempts.
-    /// 
+    /// Get mutable access to storage with `key`, with `backoff` number of
+    /// attempts.
+    ///
     /// # Panics
-    /// 
+    ///
     /// * After `backoff` attempts to access the storage.
     /// * If the contract wasn't initialized.
     fn access_storage_with_backoff(
@@ -60,4 +60,4 @@ impl<K: Hash + Eq, V> AccessStorage for DashMap<K, V> {
             }
         }
     }
-}
\ No newline at end of file
+}

From 15c928b07e972dcb3e6b355d8419551ed43eddea Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 16:42:27 +0400
Subject: [PATCH 42/44] ++

---
 crates/motsu/src/storage_access.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/crates/motsu/src/storage_access.rs b/crates/motsu/src/storage_access.rs
index 0a36dae..8783e9c 100644
--- a/crates/motsu/src/storage_access.rs
+++ b/crates/motsu/src/storage_access.rs
@@ -1,3 +1,5 @@
+//! Traits for accessing storage.
+
 use std::hash::Hash;
 
 use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};

From ee9bd5a31fd93f27ea55c52395576f3a43664a36 Mon Sep 17 00:00:00 2001
From: Nenad <xinef.it@gmail.com>
Date: Thu, 6 Feb 2025 15:19:37 +0100
Subject: [PATCH 43/44] ci: add missing newline at the end of
 .github/workflows/test.yml

---
 .github/workflows/test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 3eaa959..f1dc7a3 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -49,4 +49,4 @@ jobs:
         run: cargo nextest run --locked --all-targets
 
       - name: Run doc tests
-        run: cargo test --locked --doc
\ No newline at end of file
+        run: cargo test --locked --doc

From d72baf350c72dada26020d9dadfa8447484d7b70 Mon Sep 17 00:00:00 2001
From: Alisander Qoshqosh <qalisander@gmail.com>
Date: Thu, 6 Feb 2025 18:34:48 +0400
Subject: [PATCH 44/44] apply Motsu VM naming

---
 crates/motsu/src/context.rs        | 57 +++++++++++++++---------------
 crates/motsu/src/prelude.rs        |  2 +-
 crates/motsu/src/router.rs         | 32 ++++++++---------
 crates/motsu/src/shims.rs          | 27 +++++++-------
 crates/motsu/src/storage_access.rs |  4 +--
 5 files changed, 62 insertions(+), 60 deletions(-)

diff --git a/crates/motsu/src/context.rs b/crates/motsu/src/context.rs
index 23232a0..45f7d75 100644
--- a/crates/motsu/src/context.rs
+++ b/crates/motsu/src/context.rs
@@ -8,30 +8,31 @@ use once_cell::sync::Lazy;
 use stylus_sdk::{alloy_primitives::uint, prelude::StorageType, ArbResult};
 
 use crate::{
-    router::{RouterContext, TestRouter},
+    router::{TestRouter, VMRouterContext},
     storage_access::AccessStorage,
 };
 
-/// Storage mock.
+/// Motsu VM Storage.
 ///
 /// A global mutable key-value store that allows concurrent access.
 ///
-/// The key is the test [`Context`], an id of the test thread.
+/// The key is the test [`VMContext`], an id of the test thread.
 ///
-/// The value is the [`MockStorage`], a storage of the test case.
+/// The value is the [`VMContextStorage`], a storage of the test case.
 ///
-/// NOTE: The [`Context::storage`] will panic on lock, when the same key is
+/// NOTE: The [`VMContext::storage`] will panic on lock, when the same key is
 /// accessed twice from the same thread.
-static STORAGE: Lazy<DashMap<Context, MockStorage>> = Lazy::new(DashMap::new);
+static MOTSU_VM: Lazy<DashMap<VMContext, VMContextStorage>> =
+    Lazy::new(DashMap::new);
 
-/// Context of stylus unit tests associated with the current test thread.
+/// Context of Motsu test VM associated with the current test thread.
 #[allow(clippy::module_name_repetitions)]
 #[derive(Hash, Eq, PartialEq, Copy, Clone)]
-pub struct Context {
+pub struct VMContext {
     thread_id: ThreadId,
 }
 
-impl Context {
+impl VMContext {
     /// Get test context associated with the current test thread.
     #[must_use]
     pub fn current() -> Self {
@@ -84,7 +85,7 @@ impl Context {
 
     /// Get the message sender address.
     #[must_use]
-    pub fn msg_sender(self) -> Option<Address> {
+    pub(crate) fn msg_sender(self) -> Option<Address> {
         self.storage().msg_sender
     }
 
@@ -128,7 +129,7 @@ impl Context {
         self,
         contract_address: Address,
     ) {
-        if STORAGE
+        if MOTSU_VM
             .entry(self)
             .or_default()
             .contract_data
@@ -141,10 +142,10 @@ impl Context {
         self.router(contract_address).init_storage::<ST>();
     }
 
-    /// Reset storage for the current [`Context`] and `contract_address`.
+    /// Reset storage for the current [`VMContext`] and `contract_address`.
     ///
     /// If all test contracts are removed, flush storage for the current
-    /// test [`Context`].
+    /// test [`VMContext`].
     fn reset_storage(self, contract_address: Address) {
         let mut storage = self.storage();
         storage.contract_data.remove(&contract_address);
@@ -154,7 +155,7 @@ impl Context {
             // drop guard to a concurrent hash map to avoid a panic on lock,
             drop(storage);
             // and erase the test context.
-            _ = STORAGE.remove(&self);
+            _ = MOTSU_VM.remove(&self);
         }
 
         self.router(contract_address).reset_storage();
@@ -383,13 +384,13 @@ impl Context {
     }
 
     /// Get reference to the storage for the current test thread.
-    fn storage(self) -> RefMut<'static, Context, MockStorage> {
-        STORAGE.access_storage(&self)
+    fn storage(self) -> RefMut<'static, VMContext, VMContextStorage> {
+        MOTSU_VM.access_storage(&self)
     }
 
     /// Get router for the contract at `address`.
-    fn router(self, address: Address) -> RouterContext {
-        RouterContext::new(self.thread_id, address)
+    fn router(self, address: Address) -> VMRouterContext {
+        VMRouterContext::new(self.thread_id, address)
     }
 }
 
@@ -444,7 +445,7 @@ unsafe fn decode_calldata(
 
 /// Storage for unit test's mock data.
 #[derive(Default)]
-struct MockStorage {
+struct VMContextStorage {
     /// Address of the message sender.
     msg_sender: Option<Address>,
     /// The ETH value in wei sent to the program.
@@ -484,9 +485,9 @@ pub struct ContractCall<'a, ST: StorageType> {
 impl<ST: StorageType> ContractCall<'_, ST> {
     /// Preset the call parameters.
     fn set_call_params(&self) {
-        _ = Context::current().replace_optional_msg_value(self.msg_value);
-        _ = Context::current().replace_msg_sender(self.msg_sender);
-        _ = Context::current()
+        _ = VMContext::current().replace_optional_msg_value(self.msg_value);
+        _ = VMContext::current().replace_msg_sender(self.msg_sender);
+        _ = VMContext::current()
             .replace_contract_address(self.contract_ref.address);
     }
 }
@@ -497,7 +498,7 @@ impl<ST: StorageType> ::core::ops::Deref for ContractCall<'_, ST> {
     #[inline]
     fn deref(&self) -> &Self::Target {
         self.set_call_params();
-        Context::current().transfer_value();
+        VMContext::current().transfer_value();
         &self.storage
     }
 }
@@ -506,7 +507,7 @@ impl<ST: StorageType> ::core::ops::DerefMut for ContractCall<'_, ST> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.set_call_params();
-        Context::current().transfer_value();
+        VMContext::current().transfer_value();
         &mut self.storage
     }
 }
@@ -519,7 +520,7 @@ pub struct Contract<ST: StorageType> {
 
 impl<ST: StorageType> Drop for Contract<ST> {
     fn drop(&mut self) {
-        Context::current().reset_storage(self.address);
+        VMContext::current().reset_storage(self.address);
     }
 }
 
@@ -539,7 +540,7 @@ impl<ST: StorageType + TestRouter + 'static> Contract<ST> {
     /// Create a new contract with the given `address`.
     #[must_use]
     pub fn new_at(address: Address) -> Self {
-        Context::current().init_storage::<ST>(address);
+        VMContext::current().init_storage::<ST>(address);
 
         Self { phantom: ::core::marker::PhantomData, address }
     }
@@ -639,11 +640,11 @@ pub trait Funding {
 
 impl Funding for Address {
     fn fund(&self, value: U256) {
-        Context::current().add_assign_balance(*self, value);
+        VMContext::current().add_assign_balance(*self, value);
     }
 
     fn balance(&self) -> U256 {
-        Context::current().balance(*self)
+        VMContext::current().balance(*self)
     }
 }
 
diff --git a/crates/motsu/src/prelude.rs b/crates/motsu/src/prelude.rs
index e2f7dd4..e52a05a 100644
--- a/crates/motsu/src/prelude.rs
+++ b/crates/motsu/src/prelude.rs
@@ -1,2 +1,2 @@
 //! Common imports for `motsu` tests.
-pub use crate::context::{Account, Context, Contract, ContractCall, Funding};
+pub use crate::context::{Account, Contract, ContractCall, Funding, VMContext};
diff --git a/crates/motsu/src/router.rs b/crates/motsu/src/router.rs
index 9343c3b..14cc6fd 100644
--- a/crates/motsu/src/router.rs
+++ b/crates/motsu/src/router.rs
@@ -17,44 +17,44 @@ use stylus_sdk::{
 
 use crate::storage_access::AccessStorage;
 
-/// Router Storage.
+/// Motsu VM Router Storage.
 ///
 /// A global mutable key-value store that allows concurrent access.
 ///
-/// The key is the [`RouterContext`], a combination of [`ThreadId`] and
+/// The key is the [`VMRouterContext`], a combination of [`ThreadId`] and
 /// [`Address`] to avoid a panic on lock, while calling more than two contracts
 /// consecutive.
 ///
-/// The value is the [`RouterStorage`], a router of the contract generated by
+/// The value is the [`VMRouterStorage`], a router of the contract generated by
 /// `stylus-sdk`.
 ///
-/// NOTE: The [`RouterContext::storage`] will panic on lock, when the same key
+/// NOTE: The [`VMRouterContext::storage`] will panic on lock, when the same key
 /// is accessed twice from the same thread.
-static ROUTER_STORAGE: Lazy<DashMap<RouterContext, RouterStorage>> =
+static MOTSU_VM_ROUTERS: Lazy<DashMap<VMRouterContext, VMRouterStorage>> =
     Lazy::new(DashMap::new);
 
-/// Context for the router of a test contract for current test thread and
+/// Context of Motsu test VM router associated with the current test thread and
 /// contract's address.
 #[derive(Hash, Eq, PartialEq, Copy, Clone)]
-pub(crate) struct RouterContext {
+pub(crate) struct VMRouterContext {
     thread_id: ThreadId,
     contract_address: Address,
 }
 
-impl RouterContext {
+impl VMRouterContext {
     /// Create a new router context.
     pub(crate) fn new(thread: ThreadId, contract_address: Address) -> Self {
         Self { thread_id: thread, contract_address }
     }
 
     /// Get reference to the call storage for the current test thread.
-    fn storage(self) -> RefMut<'static, RouterContext, RouterStorage> {
-        ROUTER_STORAGE.access_storage(&self)
+    fn storage(self) -> RefMut<'static, VMRouterContext, VMRouterStorage> {
+        MOTSU_VM_ROUTERS.access_storage(&self)
     }
 
     /// Check if the router exists for the contract.
     pub(crate) fn exists(self) -> bool {
-        ROUTER_STORAGE.contains_key(&self)
+        MOTSU_VM_ROUTERS.contains_key(&self)
     }
 
     pub(crate) fn route(
@@ -89,10 +89,10 @@ impl RouterContext {
     /// `contract_address`.
     pub(crate) fn init_storage<ST: StorageType + TestRouter + 'static>(self) {
         let contract_address = self.contract_address;
-        if ROUTER_STORAGE
+        if MOTSU_VM_ROUTERS
             .insert(
                 self,
-                RouterStorage {
+                VMRouterStorage {
                     router: Arc::new(Mutex::new(unsafe {
                         ST::new(uint!(0_U256), 0)
                     })),
@@ -104,14 +104,14 @@ impl RouterContext {
         }
     }
 
-    /// Reset router storage for the current [`RouterContext`].
+    /// Reset router storage for the current [`VMRouterContext`].
     pub(crate) fn reset_storage(self) {
-        ROUTER_STORAGE.remove(&self);
+        MOTSU_VM_ROUTERS.remove(&self);
     }
 }
 
 /// Metadata related to the router of an external contract.
-struct RouterStorage {
+struct VMRouterStorage {
     // Contract's router.
     // NOTE: Mutex is important since contract type is not `Sync`.
     router: Arc<Mutex<dyn TestRouter>>,
diff --git a/crates/motsu/src/shims.rs b/crates/motsu/src/shims.rs
index 5d8eabf..1bded9a 100644
--- a/crates/motsu/src/shims.rs
+++ b/crates/motsu/src/shims.rs
@@ -23,7 +23,8 @@ use std::slice;
 use tiny_keccak::{Hasher, Keccak};
 
 use crate::context::{
-    read_address, write_address, write_bytes32, write_u256, Context, WORD_BYTES,
+    read_address, write_address, write_bytes32, write_u256, VMContext,
+    WORD_BYTES,
 };
 
 /// Arbitrum's CHAID ID.
@@ -72,7 +73,7 @@ unsafe extern "C" fn native_keccak256(
 /// May panic if unable to lock `STORAGE`.
 #[no_mangle]
 unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
-    Context::current().get_bytes_raw(key, out);
+    VMContext::current().get_bytes_raw(key, out);
 }
 
 /// Writes a 32-byte value to the permanent storage cache.
@@ -92,7 +93,7 @@ unsafe extern "C" fn storage_load_bytes32(key: *const u8, out: *mut u8) {
 /// May panic if unable to lock `STORAGE`.
 #[no_mangle]
 unsafe extern "C" fn storage_cache_bytes32(key: *const u8, value: *const u8) {
-    Context::current().set_bytes_raw(key, value);
+    VMContext::current().set_bytes_raw(key, value);
 }
 
 /// Persists any dirty values in the storage cache to the EVM state trie,
@@ -125,14 +126,14 @@ unsafe extern "C" fn storage_flush_cache(_: bool) {
 #[no_mangle]
 unsafe extern "C" fn msg_sender(sender: *mut u8) {
     let msg_sender =
-        Context::current().msg_sender().expect("msg_sender should be set");
+        VMContext::current().msg_sender().expect("msg_sender should be set");
     write_address(sender, msg_sender);
 }
 
 /// Get the ETH value (U256) in wei sent to the program.
 #[no_mangle]
 unsafe extern "C" fn msg_value(value: *mut u8) {
-    Context::current().msg_value_raw(value);
+    VMContext::current().msg_value_raw(value);
 }
 
 /// Gets the address of the current program. The semantics are equivalent to
@@ -145,7 +146,7 @@ unsafe extern "C" fn msg_value(value: *mut u8) {
 /// May panic if fails to parse `CONTRACT_ADDRESS` as an address.
 #[no_mangle]
 unsafe extern "C" fn contract_address(address: *mut u8) {
-    let contract_address = Context::current()
+    let contract_address = VMContext::current()
         .contract_address()
         .expect("contract_address should be set");
     write_address(address, contract_address);
@@ -191,7 +192,7 @@ unsafe extern "C" fn emit_log(_: *const u8, _: usize, _: usize) {
 /// May panic if fails to parse `ACCOUNT_CODEHASH` as a keccack hash.
 #[no_mangle]
 unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
-    let code_hash = if Context::current().has_code_raw(address) {
+    let code_hash = if VMContext::current().has_code_raw(address) {
         CA_CODEHASH
     } else {
         EOA_CODEHASH
@@ -210,7 +211,7 @@ unsafe extern "C" fn account_codehash(address: *const u8, dest: *mut u8) {
 #[no_mangle]
 unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
     let address = read_address(address);
-    let balance = Context::current().balance(address);
+    let balance = VMContext::current().balance(address);
     write_u256(dest, balance);
 }
 
@@ -223,7 +224,7 @@ unsafe extern "C" fn account_balance(address: *const u8, dest: *mut u8) {
 /// [`RETURN_DATA_SIZE`]: https://www.evm.codes/#3d
 #[no_mangle]
 unsafe extern "C" fn return_data_size() -> usize {
-    Context::current().return_data_size()
+    VMContext::current().return_data_size()
 }
 
 /// Copies the bytes of the last EVM call or deployment return result.
@@ -241,7 +242,7 @@ unsafe extern "C" fn read_return_data(
     _offset: usize,
     size: usize,
 ) -> usize {
-    Context::current().read_return_data_raw(dest, size)
+    VMContext::current().read_return_data_raw(dest, size)
 }
 
 /// Calls the contract at the given address with options for passing value and
@@ -267,7 +268,7 @@ unsafe extern "C" fn call_contract(
     _gas: u64,
     return_data_len: *mut usize,
 ) -> u8 {
-    Context::current().call_contract_with_value_raw(
+    VMContext::current().call_contract_with_value_raw(
         contract,
         calldata,
         calldata_len,
@@ -298,7 +299,7 @@ unsafe extern "C" fn static_call_contract(
     _gas: u64,
     return_data_len: *mut usize,
 ) -> u8 {
-    Context::current().call_contract_raw(
+    VMContext::current().call_contract_raw(
         contract,
         calldata,
         calldata_len,
@@ -328,7 +329,7 @@ unsafe extern "C" fn delegate_call_contract(
     _gas: u64,
     return_data_len: *mut usize,
 ) -> u8 {
-    Context::current().call_contract_raw(
+    VMContext::current().call_contract_raw(
         contract,
         calldata,
         calldata_len,
diff --git a/crates/motsu/src/storage_access.rs b/crates/motsu/src/storage_access.rs
index 8783e9c..8a583fe 100644
--- a/crates/motsu/src/storage_access.rs
+++ b/crates/motsu/src/storage_access.rs
@@ -1,10 +1,10 @@
-//! Traits for accessing storage.
+//! Tooling for accessing Motsu test VM storage.
 
 use std::hash::Hash;
 
 use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
 
-/// Trait for accessing test storage.
+/// Trait for Motsu test VM storage access.
 pub(crate) trait AccessStorage {
     type Key;
     type Value;