contracts: Allow ChainExtension::call() to access &mut self (#11874)

* Give chain extensions the ability to store some temporary values

* Update frame/contracts/src/wasm/runtime.rs

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* Rename func_id -> id

* Replace `id` param by two functions on `env`

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

substrate/frame/contracts/src/chain_extension.rs

@@ -33,7 +33,7 @@
 //!
 //! Often there is a need for having multiple chain extensions. This is often the case when
 //! some generally useful off-the-shelf extensions should be included. To have multiple chain
-//! extensions they can be put into a tuple which is then passed to `[Config::ChainExtension]` like
+//! extensions they can be put into a tuple which is then passed to [`Config::ChainExtension`] like
 //! this `type Extensions = (ExtensionA, ExtensionB)`.
 //!
 //! However, only extensions implementing [`RegisteredChainExtension`] can be put into a tuple.
@@ -94,6 +94,12 @@ pub type Result<T> = sp_std::result::Result<T, DispatchError>;
 /// In order to create a custom chain extension this trait must be implemented and supplied
 /// to the pallet contracts configuration trait as the associated type of the same name.
 /// Consult the [module documentation](self) for a general explanation of chain extensions.
+///
+/// # Lifetime
+///
+/// The extension will be [`Default`] initialized at the beginning of each call
+/// (**not** per call stack) and dropped afterwards. Hence any value held inside the extension
+/// can be used as a per-call scratch buffer.
 pub trait ChainExtension<C: Config> {
 	/// Call the chain extension logic.
 	///
@@ -102,8 +108,6 @@ pub trait ChainExtension<C: Config> {
 	/// imported wasm function.
 	///
 	/// # Parameters
-	/// - `func_id`: The first argument to `seal_call_chain_extension`. Usually used to determine
-	///   which function to realize.
 	/// - `env`: Access to the remaining arguments and the execution environment.
 	///
 	/// # Return
@@ -111,7 +115,7 @@ pub trait ChainExtension<C: Config> {
 	/// In case of `Err` the contract execution is immediately suspended and the passed error
 	/// is returned to the caller. Otherwise the value of [`RetVal`] determines the exit
 	/// behaviour.
-	fn call<E>(func_id: u32, env: Environment<E, InitState>) -> Result<RetVal>
+	fn call<E>(&mut self, env: Environment<E, InitState>) -> Result<RetVal>
 	where
 		E: Ext<T = C>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>;
@@ -132,7 +136,7 @@ pub trait ChainExtension<C: Config> {
 ///
 /// An extension that implements this trait can be put in a tuple in order to have multiple
 /// extensions available. The tuple implementation routes requests based on the first two
-/// most significant bytes of the `func_id` passed to `call`.
+/// most significant bytes of the `id` passed to `call`.
 ///
 /// If this extensions is to be used by multiple runtimes consider
 /// [registering it](https://github.com/paritytech/chainextension-registry) to ensure that there
@@ -150,15 +154,15 @@ pub trait RegisteredChainExtension<C: Config>: ChainExtension<C> {
 #[impl_trait_for_tuples::impl_for_tuples(10)]
 #[tuple_types_custom_trait_bound(RegisteredChainExtension<C>)]
 impl<C: Config> ChainExtension<C> for Tuple {
-	fn call<E>(func_id: u32, mut env: Environment<E, InitState>) -> Result<RetVal>
+	fn call<E>(&mut self, mut env: Environment<E, InitState>) -> Result<RetVal>
 	where
 		E: Ext<T = C>,
 		<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 	{
 		for_tuples!(
 			#(
-				if (Tuple::ID == (func_id >> 16) as u16) && Tuple::enabled() {
-					return Tuple::call(func_id, env);
+				if (Tuple::ID == env.ext_id()) && Tuple::enabled() {
+					return Tuple.call(env);
 				}
 			)*
 		);
@@ -206,6 +210,22 @@ impl<'a, 'b, E: Ext, S: state::State> Environment<'a, 'b, E, S>
 where
 	<E::T as SysConfig>::AccountId: UncheckedFrom<<E::T as SysConfig>::Hash> + AsRef<[u8]>,
 {
+	/// The function id within the `id` passed by a contract.
+	///
+	/// It returns the two least significant bytes of the `id` passed by a contract as the other
+	/// two bytes represent the chain extension itself (the code which is calling this function).
+	pub fn func_id(&self) -> u16 {
+		(self.inner.id & 0x00FF) as u16
+	}
+
+	/// The chain extension id within the `id` passed by a contract.
+	///
+	/// It returns the two most significant bytes of the `id` passed by a contract which represent
+	/// the chain extension itself (the code which is calling this function).
+	pub fn ext_id(&self) -> u16 {
+		(self.inner.id >> 16) as u16
+	}
+
 	/// Charge the passed `amount` of weight from the overall limit.
 	///
 	/// It returns `Ok` when there the remaining weight budget is larger than the passed
@@ -251,13 +271,14 @@ impl<'a, 'b, E: Ext> Environment<'a, 'b, E, state::Init> {
 	/// ever create this type. Chain extensions merely consume it.
 	pub(crate) fn new(
 		runtime: &'a mut Runtime<'b, E>,
+		id: u32,
 		input_ptr: u32,
 		input_len: u32,
 		output_ptr: u32,
 		output_len_ptr: u32,
 	) -> Self {
 		Environment {
-			inner: Inner { runtime, input_ptr, input_len, output_ptr, output_len_ptr },
+			inner: Inner { runtime, id, input_ptr, input_len, output_ptr, output_len_ptr },
 			phantom: PhantomData,
 		}
 	}
@@ -406,6 +427,8 @@ struct Inner<'a, 'b, E: Ext> {
 	/// The runtime contains all necessary functions to interact with the running contract.
 	runtime: &'a mut Runtime<'b, E>,
 	/// Verbatim argument passed to `seal_call_chain_extension`.
+	id: u32,
+	/// Verbatim argument passed to `seal_call_chain_extension`.
 	input_ptr: u32,
 	/// Verbatim argument passed to `seal_call_chain_extension`.
 	input_len: u32,