Addressable for auto stack management

lang/codegen/src/trait_definition.rs

@@ -120,7 +120,7 @@ pub fn generate(_attrs: TokenStream, _input: TokenStream) -> TokenStream {
 
         let pub_mock_env_ident = format_ident!("mock_{}", trait_item.ident.to_string().to_lowercase());
         maybe_use_mock_env = quote! {
-            #[cfg(test)]
+            #[cfg(any(test, feature = "mockable"))]
             pub mod #pub_mock_env_ident {
                 pub use super :: #namespace_ident :: { mock_env as env , using , deploy };
             }
@@ -270,11 +270,14 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
 
             let message_test_impl = match &mock_type {
                 Some(_mock_ty) => quote! {
-                    mock_env :: with(|registry| {
-                        let mut mock_ref = registry.get_mut(self).expect("not an address of mocked contract");
-                        mock_ref.borrow_mut(). #message_ident (
+                    mock_env :: with(|ctx| {
+                        let mut mock_ref = ctx.register.get_mut(self).expect("not an address of mocked contract");
+                        ctx.stack.borrow_mut().push(&self);
+                        let result = mock_ref.borrow_mut(). #message_ident (
                             #( #input_bindings , )*
-                        )
+                        );
+                        ctx.stack.borrow_mut().pop();
+                        result
                     }).expect("mock object not set")
                 },
                 None => quote! { ::core::panic!("cross-contract call is not supported in ink tests; try to set a mock object?") }
@@ -286,13 +289,13 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
                     & self
                     #( , #input_bindings : #input_types )*
                 ) -> #output_ty {
-                    #[cfg(not(test))]
+                    #[cfg(not(any(test, feature = "mockable")))]
                     {
                         Self::#message_builder_ident(self #( , #input_bindings)*)
                             .fire()
                             .unwrap_or_else(|err| ::core::panic!("{}: {:?}", #panic_str, err))
                     }
-                    #[cfg(test)]
+                    #[cfg(any(test, feature = "mockable"))]
                     {
                         #message_test_impl
                     }
@@ -328,40 +331,60 @@ fn generate_wrapper(ink_trait: ItemTrait, mock_type: Option<TokenStream>) -> pro
     let def_messages = def_messages.iter();
 
     let maybe_mock_environmental = match mock_type {
-        Some(ty) =>{
+        Some(ty) => {
             quote! {
-                #[cfg(test)]
-                ::environmental::environmental!(
-                    pub mock_env : std::collections::BTreeMap<
+                #[cfg(any(test, feature = "mockable"))]
+                pub struct Context {
+                    pub stack: std::rc::Rc<std::cell::RefCell<
+                        ::openbrush::traits::mock::ManagedCallStack
+                    >>,
+                    pub register: std::collections::BTreeMap<
                         ::openbrush::traits::AccountId,
                         std::rc::Rc<std::cell::RefCell< #ty >>
-                    > 
+                    >
+                }
+
+                #[cfg(any(test, feature = "mockable"))]
+                ::environmental::environmental!(
+                    pub mock_env : Context
                 );
 
-                #[cfg(test)]
-                pub fn using<F: FnOnce()>(f: F) {
-                    let mut env = Default::default();
+                #[cfg(any(test, feature = "mockable"))]
+                pub fn using<F: FnOnce()>(
+                    stack: std::rc::Rc<std::cell::RefCell<::openbrush::traits::mock::ManagedCallStack>>,
+                    f: F
+                ) {
+                    let mut env = Context {
+                        stack,
+                        register: Default::default()
+                    };
                     mock_env::using(&mut env, f);
                 }
 
-                #[cfg(test)]
-                pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::AccountId, std::rc::Rc<std::cell::RefCell< #ty >>) {
+                #[cfg(any(test, feature = "mockable"))]
+                pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::mock::Addressable< #ty >) {
                     let contract: std::rc::Rc<std::cell::RefCell< #ty >> = std::rc::Rc::new(
                         std::cell::RefCell::< #ty >::new(inner_contract)
                     );
-                    mock_env::with(|register| {
-                        let n: u8 = register.len().try_into()
+                    let (account_id, contract, stack) = mock_env::with(|ctx| {
+                        let n: u8 = ctx.register.len().try_into()
                             .expect("too many contracts to fit into u8");
                         let mut pat = [ #( #mock_address_pattern,  )* ];
                         pat[31] = n;
                         let account_id: ::openbrush::traits::AccountId = pat.into();
 
-                        register.insert(account_id.clone(), contract.clone());
-                        (account_id, contract)
-                    }).expect("must call within `using()`")
+                        ctx.register.insert(account_id.clone(), contract.clone());
+                        (account_id, contract, ctx.stack.clone())
+                    }).expect("must call within `using()`");
+
+                    ::openbrush::traits::mock::Addressable::new(
+                        account_id,
+                        contract,
+                        stack,
+                    )
                 }
             }
-        },
+        }
         None => quote! {},
     };
 
@@ -411,17 +434,13 @@ fn remove_ink_attrs(mut trait_item: ItemTrait) -> ItemTrait {
 }
 
 /// Extracts the mocking related macro args out from the input
-/// 
+///
 /// Return a tuple of an optional mock target and the args without the mock target
 fn extract_mock_config(attr: TokenStream) -> (Option<TokenStream>, TokenStream) {
     let attr_args = syn::parse2::<attr_args::AttributeArgs>(attr).expect("unable to parse trait_definition attribute");
 
-    let (mock_args, ink_args): (Vec<_>, Vec<_>) = attr_args
-        .into_iter()
-        .partition(|arg| {
-            arg.name.is_ident("mock")
-        });
-
+    let (mock_args, ink_args): (Vec<_>, Vec<_>) = attr_args.into_iter().partition(|arg| arg.name.is_ident("mock"));
+
     let mock_type = mock_args.first().map(|mock_attr| {
         let ty = &mock_attr.value;
         quote! { #ty }
@@ -452,20 +471,20 @@ mod tests {
                 mock = MyMockType,
                 namespace = ::name::space
             },
-            quote!{
+            quote! {
                 pub trait SubmittableOracle {
                     #[ink(message)]
                     fn admin(&self) -> AccountId;
-            
+
                     #[ink(message)]
                     fn verifier(&self) -> Verifier;
-            
+
                     #[ink(message)]
                     fn attest(&self, arg: String) -> Result<Attestation, ()>;
                 }
-            }
+            },
         );
 
         println!("OUTPUT:\n\n{:}", r);
     }
-}
+}

lang/src/lib.rs

@@ -21,6 +21,8 @@
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
+extern crate alloc;
+
 pub mod derive;
 mod macros;
 pub mod test_utils;

lang/src/traits.rs

@@ -77,3 +77,213 @@ pub trait Flush: ::ink_storage::traits::SpreadLayout + InkStorage {
 }
 
 impl<T: ::ink_storage::traits::SpreadLayout + InkStorage> Flush for T {}
+
+/// Types for managing mock cross-contract calls in unit tests
+pub mod mock {
+    use super::AccountId;
+
+    use alloc::{
+        rc::Rc,
+        vec::Vec,
+    };
+    use core::{
+        cell::{
+            Ref,
+            RefCell,
+            RefMut,
+        },
+        ops::{
+            Deref,
+            DerefMut,
+        },
+    };
+
+    /// A frame in the call stack
+    #[derive(Clone, Debug)]
+    pub struct MockCallContext {
+        pub level: u32,
+        pub caller: Option<AccountId>,
+        pub callee: AccountId,
+    }
+
+    /// A managed call stack for mocking cross-contract call in test environment
+    pub struct ManagedCallStack {
+        stack: Vec<MockCallContext>,
+    }
+
+    impl ManagedCallStack {
+        /// Crates a call stack with the default `account`
+        pub fn new(account: AccountId) -> Self {
+            ManagedCallStack {
+                stack: alloc::vec![MockCallContext {
+                    level: 0,
+                    caller: None,
+                    callee: account,
+                }],
+            }
+        }
+
+        /// Creates a call stack with the default `account` and returns a shared reference
+        pub fn create_shared(account: AccountId) -> Rc<RefCell<Self>> {
+            Rc::new(RefCell::new(Self::new(account)))
+        }
+
+        /// Changes the caller account
+        ///
+        /// Only allowed outside any contract call (when the stack is empty).
+        pub fn switch_account(&mut self, account: AccountId) -> Result<(), ()> {
+            if self.stack.len() != 1 {
+                return Err(())
+            }
+            let ctx = self.stack.get_mut(0).ok_or(())?;
+            ctx.callee = account;
+            Ok(())
+        }
+
+        /// Pushes a new call frame
+        pub fn push(&mut self, callee: &AccountId) {
+            let parent_ctx = self.peek().clone();
+            self.stack.push(MockCallContext {
+                level: parent_ctx.level + 1,
+                caller: Some(parent_ctx.callee),
+                callee: callee.clone(),
+            });
+            self.sync_to_ink();
+        }
+
+        /// Pops the call frame and returns the frame
+        pub fn pop(&mut self) -> Option<MockCallContext> {
+            if self.stack.len() > 1 {
+                let ctx = self.stack.pop();
+                self.sync_to_ink();
+                ctx
+            } else {
+                None
+            }
+        }
+
+        /// Peeks the current call frame
+        pub fn peek(&self) -> &MockCallContext {
+            self.stack.last().expect("stack is never empty; qed.")
+        }
+
+        /// Syncs the top call frame to ink testing environment
+        pub fn sync_to_ink(&self) {
+            let ctx = self.peek();
+            if let Some(caller) = ctx.caller {
+                ink_env::test::set_caller::<ink_env::DefaultEnvironment>(caller);
+            }
+            ink_env::test::set_callee::<ink_env::DefaultEnvironment>(ctx.callee);
+        }
+    }
+
+    /// A wrapper of a contract with an address for call stake auto-management
+    #[derive(Clone)]
+    pub struct Addressable<T> {
+        inner: Rc<RefCell<T>>,
+        id: AccountId,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<T> Addressable<T> {
+        /// Wraps a contract reference with id and a shared call stack
+        pub fn new(id: AccountId, inner: Rc<RefCell<T>>, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            Addressable { inner, id, stack }
+        }
+
+        /// Wraps a native contract object with a simple id
+        ///
+        /// The account id of the contract will be the `id` with zero-padding.
+        pub fn create_native(id: u8, inner: T, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            Addressable {
+                inner: Rc::new(RefCell::new(inner)),
+                id: naive_id(id),
+                stack,
+            }
+        }
+
+        /// Returns the account id of the inner contract
+        pub fn id(&self) -> AccountId {
+            self.id.clone()
+        }
+
+        /// Borrows the contract for _a_ call with the stack auto-managed
+        ///
+        /// Holding the ref for multiple calls or nested call is considered abuse.
+        pub fn call(&self) -> ScopedRef<'_, T> {
+            ScopedRef::new(self.inner.borrow(), &self.id, self.stack.clone())
+        }
+
+        /// Borrows the contract for _a_ mut call with the stack auto-managed
+        ///
+        /// Holding the mut ref for multiple calls or nested call is considered abuse.
+        pub fn call_mut(&self) -> ScopedRefMut<'_, T> {
+            ScopedRefMut::new(self.inner.borrow_mut(), &self.id, self.stack.clone())
+        }
+    }
+
+    /// Push a call stack when the `Ref` in scope
+    pub struct ScopedRef<'b, T: 'b> {
+        inner: Ref<'b, T>,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<'b, T> ScopedRef<'b, T> {
+        fn new(inner: Ref<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            stack.borrow_mut().push(address);
+            Self { inner, stack }
+        }
+    }
+
+    impl<'b, T> Deref for ScopedRef<'b, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.inner.deref()
+        }
+    }
+
+    impl<'b, T> Drop for ScopedRef<'b, T> {
+        fn drop(&mut self) {
+            self.stack.borrow_mut().pop().expect("pop never fails");
+        }
+    }
+
+    /// Push a call stack when the `RefMut` in scope
+    pub struct ScopedRefMut<'b, T: 'b> {
+        inner: RefMut<'b, T>,
+        stack: Rc<RefCell<ManagedCallStack>>,
+    }
+
+    impl<'b, T> ScopedRefMut<'b, T> {
+        fn new(inner: RefMut<'b, T>, address: &AccountId, stack: Rc<RefCell<ManagedCallStack>>) -> Self {
+            stack.borrow_mut().push(address);
+            Self { inner, stack }
+        }
+    }
+
+    impl<'b, T> Deref for ScopedRefMut<'b, T> {
+        type Target = T;
+        fn deref(&self) -> &T {
+            self.inner.deref()
+        }
+    }
+
+    impl<'b, T> DerefMut for ScopedRefMut<'b, T> {
+        fn deref_mut(&mut self) -> &mut T {
+            self.inner.deref_mut()
+        }
+    }
+
+    impl<'b, T> Drop for ScopedRefMut<'b, T> {
+        fn drop(&mut self) {
+            self.stack.borrow_mut().pop().expect("pop never fails");
+        }
+    }
+
+    /// Generates a naive zero-padding account id with a `u8` number
+    pub fn naive_id(id: u8) -> AccountId {
+        let mut address = [0u8; 32];
+        address[31] = id;
+        address.into()
+    }
+}