add storage guard

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

@@ -454,7 +454,7 @@ mod tests {
         assert_eq!(owner, alice);
 
         let res =
-            contract._add_token_to_owner_enumeration(alice, token_id, &erc721);
+            contract._add_token_to_owner_enumeration(alice, token_id, &*erc721);
         assert!(res.is_ok());
 
         let test_token_id = contract
@@ -483,7 +483,7 @@ mod tests {
         assert_eq!(owner, alice);
 
         let res =
-            contract._add_token_to_owner_enumeration(alice, token_id, &erc721);
+            contract._add_token_to_owner_enumeration(alice, token_id, &*erc721);
         assert!(res.is_ok());
 
         let err =
@@ -525,7 +525,7 @@ mod tests {
         assert_eq!(owner, alice);
 
         let res =
-            contract._add_token_to_owner_enumeration(alice, token_id, &erc721);
+            contract._add_token_to_owner_enumeration(alice, token_id, &*erc721);
         assert!(res.is_ok());
 
         // Transfer the token from ALICE to BOB.
@@ -538,11 +538,11 @@ mod tests {
         assert_eq!(owner, BOB);
 
         let res = contract
-            ._remove_token_from_owner_enumeration(alice, token_id, &erc721);
+            ._remove_token_from_owner_enumeration(alice, token_id, &*erc721);
         assert!(res.is_ok());
 
         let res =
-            contract._add_token_to_owner_enumeration(BOB, token_id, &erc721);
+            contract._add_token_to_owner_enumeration(BOB, token_id, &*erc721);
         assert!(res.is_ok());
 
         let test_token_id = contract

lib/motsu-proc/src/test.rs

@@ -12,7 +12,7 @@ pub(crate) fn test(_attr: &TokenStream, input: TokenStream) -> TokenStream {
     let sig = &item_fn.sig;
     let fn_name = &sig.ident;
     let fn_return_type = &sig.output;
-    let fn_block = &item_fn.block;
+    let fn_stmts = &item_fn.block.stmts;
     let fn_args = &sig.inputs;
 
     // Currently, more than one contract per unit test is not supported.
@@ -45,7 +45,9 @@ pub(crate) fn test(_attr: &TokenStream, input: TokenStream) -> TokenStream {
         fn #fn_name() #fn_return_type {
             use ::motsu::prelude::DefaultStorage;
             #( #contract_declarations )*
-            let res = #fn_block;
+            let res = {
+                #( #fn_stmts )*
+            };
             ::motsu::prelude::Context::current().reset_storage();
             res
         }

lib/motsu/Cargo.toml

@@ -15,6 +15,7 @@ tiny-keccak.workspace = true
 stylus-sdk.workspace = true
 motsu-proc.workspace = true
 dashmap.workspace = true
+derive_more = { version = "1.0.0", features = ["full"] }
 
 [lints]
 workspace = true

lib/motsu/src/context.rs

@@ -4,7 +4,13 @@ use std::{collections::HashMap, ptr};
 
 use dashmap::DashMap;
 use once_cell::sync::Lazy;
-use stylus_sdk::{alloy_primitives::uint, prelude::StorageType};
+use stylus_sdk::{
+    alloy_primitives::{
+        private::derive_more::{Deref, DerefMut},
+        uint,
+    },
+    prelude::StorageType,
+};
 
 use crate::prelude::{Bytes32, WORD_BYTES};
 
@@ -101,9 +107,12 @@ pub trait DefaultStorage: StorageType {
     /// Initializes fields of contract storage and child contract storages with
     /// default values.
     #[must_use]
-    fn default() -> Self {
-        unsafe { Self::new(uint!(0_U256), 0) }
+    fn default() -> StorageGuard<Self> {
+        StorageGuard(unsafe { Self::new(uint!(0_U256), 0) })
     }
 }
 
 impl<ST: StorageType> DefaultStorage for ST {}
+
+#[derive(Deref, DerefMut)]
+pub struct StorageGuard<ST: StorageType>(ST);