From c0a3b2ef6a80ae5d8f5b85247d354a9d3348cccb Mon Sep 17 00:00:00 2001 From: h4x3rotab Date: Tue, 21 Jun 2022 11:04:58 +0000 Subject: [PATCH] Addressable for auto stack management --- lang/codegen/src/trait_definition.rs | 89 +++++++----- lang/src/lib.rs | 2 + lang/src/traits.rs | 210 +++++++++++++++++++++++++++ 3 files changed, 266 insertions(+), 35 deletions(-) diff --git a/lang/codegen/src/trait_definition.rs b/lang/codegen/src/trait_definition.rs index e8dfbe9510..4dc1078f04 100644 --- a/lang/codegen/src/trait_definition.rs +++ b/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) -> 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) -> 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) -> 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>, + pub register: std::collections::BTreeMap< ::openbrush::traits::AccountId, std::rc::Rc> - > + > + } + + #[cfg(any(test, feature = "mockable"))] + ::environmental::environmental!( + pub mock_env : Context ); - #[cfg(test)] - pub fn using(f: F) { - let mut env = Default::default(); + #[cfg(any(test, feature = "mockable"))] + pub fn using( + stack: std::rc::Rc>, + 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>) { + #[cfg(any(test, feature = "mockable"))] + pub fn deploy(inner_contract : #ty) -> (::openbrush::traits::mock::Addressable< #ty >) { let contract: std::rc::Rc> = 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) { let attr_args = syn::parse2::(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; } - } + }, ); println!("OUTPUT:\n\n{:}", r); } -} \ No newline at end of file +} diff --git a/lang/src/lib.rs b/lang/src/lib.rs index 0f44d3e573..de5e94f9a8 100644 --- a/lang/src/lib.rs +++ b/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; diff --git a/lang/src/traits.rs b/lang/src/traits.rs index 7aba40000f..861ce2c09c 100644 --- a/lang/src/traits.rs +++ b/lang/src/traits.rs @@ -77,3 +77,213 @@ pub trait Flush: ::ink_storage::traits::SpreadLayout + InkStorage { } impl 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, + pub callee: AccountId, + } + + /// A managed call stack for mocking cross-contract call in test environment + pub struct ManagedCallStack { + stack: Vec, + } + + 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> { + 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 { + 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::(caller); + } + ink_env::test::set_callee::(ctx.callee); + } + } + + /// A wrapper of a contract with an address for call stake auto-management + #[derive(Clone)] + pub struct Addressable { + inner: Rc>, + id: AccountId, + stack: Rc>, + } + + impl Addressable { + /// Wraps a contract reference with id and a shared call stack + pub fn new(id: AccountId, inner: Rc>, stack: Rc>) -> 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>) -> 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>, + } + + impl<'b, T> ScopedRef<'b, T> { + fn new(inner: Ref<'b, T>, address: &AccountId, stack: Rc>) -> 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>, + } + + impl<'b, T> ScopedRefMut<'b, T> { + fn new(inner: RefMut<'b, T>, address: &AccountId, stack: Rc>) -> 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() + } +}