Support multiple chain extensions (#1958)

* Combined chain extension example

* Make CI happy

* Make CI happy

* Update CHANGELOG.md

* Added `ink::combine_extensions!` helper macro to simplify the definition of the combined chain extension

* Make CI happy

* Make CI happy

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fail when decoding from storage and not all bytes consumed - [#1897](https://github.com/paritytech/ink/pull/1897)
 - [E2E] resolve DispatchError error details for dry-runs - [#1944](https://github.com/paritytech/ink/pull/1994)
 - [E2E] update to new `drink` API - [#2005](https://github.com/paritytech/ink/pull/2005)
+- Support multiple chain extensions - [#1958](https://github.com/paritytech/ink/pull/1958)
+  - New example of how to use multiple chain extensions in one contract.
+  - Affects the usage of the `#[ink::chain_extension]` macro and the definition of the chain extension.
 
 
 ## Version 5.0.0-alpha

crates/engine/src/chain_extension.rs

@@ -29,29 +29,29 @@ pub struct ChainExtensionHandler {
     output: Vec<u8>,
 }
 
-/// The unique ID of the registered chain extension method.
+/// The unique ID of the registered chain extension.
 #[derive(
     Debug, From, scale::Encode, scale::Decode, PartialEq, Eq, PartialOrd, Ord, Hash,
 )]
-pub struct ExtensionId(u32);
+pub struct ExtensionId(u16);
 
 /// Types implementing this trait can be used as chain extensions.
 ///
 /// This trait is only useful for testing contract via the off-chain environment.
 pub trait ChainExtension {
-    /// The static function ID of the chain extension.
+    /// The static ID of the chain extension.
     ///
     /// # Note
     ///
     /// This is expected to return a constant value.
-    fn func_id(&self) -> u32;
+    fn ext_id(&self) -> u16;
 
     /// Calls the chain extension with the given input.
     ///
     /// Returns an error code and may fill the `output` buffer with a SCALE encoded
     /// result.
     #[allow(clippy::ptr_arg)]
-    fn call(&mut self, input: &[u8], output: &mut Vec<u8>) -> u32;
+    fn call(&mut self, func_id: u16, input: &[u8], output: &mut Vec<u8>) -> u32;
 }
 
 impl Default for ChainExtensionHandler {
@@ -79,20 +79,24 @@ impl ChainExtensionHandler {
 
     /// Register a new chain extension.
     pub fn register(&mut self, extension: Box<dyn ChainExtension>) {
-        let func_id = extension.func_id();
-        self.registered
-            .insert(ExtensionId::from(func_id), extension);
+        let ext_id = extension.ext_id();
+        self.registered.insert(ExtensionId::from(ext_id), extension);
     }
 
     /// Evaluates the chain extension with the given parameters.
     ///
     /// Upon success returns the values returned by the evaluated chain extension.
-    pub fn eval(&mut self, func_id: u32, input: &[u8]) -> Result<(u32, &[u8]), Error> {
+    pub fn eval(&mut self, id: u32, input: &[u8]) -> Result<(u32, &[u8]), Error> {
         self.output.clear();
-        let extension_id = ExtensionId::from(func_id);
+
+        let func_id = (id & 0x0000FFFF) as u16;
+        let ext_id = (id >> 16) as u16;
+
+        let extension_id = ExtensionId::from(ext_id);
         match self.registered.entry(extension_id) {
             Entry::Occupied(occupied) => {
-                let status_code = occupied.into_mut().call(input, &mut self.output);
+                let status_code =
+                    occupied.into_mut().call(func_id, input, &mut self.output);
                 Ok((status_code, &mut self.output))
             }
             Entry::Vacant(_vacant) => Err(Error::UnregisteredChainExtension),

crates/engine/src/ext.rs

@@ -442,14 +442,14 @@ impl Engine {
     /// Calls the chain extension method registered at `func_id` with `input`.
     pub fn call_chain_extension(
         &mut self,
-        func_id: u32,
+        id: u32,
         input: &[u8],
         output: &mut &mut [u8],
     ) {
         let encoded_input = input.encode();
         let (status_code, out) = self
             .chain_extension_handler
-            .eval(func_id, &encoded_input)
+            .eval(id, &encoded_input)
             .unwrap_or_else(|error| {
                 panic!(
                     "Encountered unexpected missing chain extension method: {error:?}"

crates/env/src/backend.rs

@@ -327,7 +327,7 @@ pub trait EnvBackend {
     /// drive the decoding and error management process from the outside.
     fn call_chain_extension<I, T, E, ErrorCode, F, D>(
         &mut self,
-        func_id: u32,
+        id: u32,
         input: &I,
         status_to_result: F,
         decode_to_result: D,

crates/env/src/chain_extension.rs

@@ -80,17 +80,17 @@ pub trait FromStatusCode: Sized {
 ///       type. The method just returns `O`.
 #[derive(Debug)]
 pub struct ChainExtensionMethod<I, O, ErrorCode, const IS_RESULT: bool> {
-    func_id: u32,
+    id: u32,
     #[allow(clippy::type_complexity)]
     state: PhantomData<fn() -> (I, O, ErrorCode)>,
 }
 
 impl ChainExtensionMethod<(), (), (), false> {
     /// Creates a new chain extension method instance.
     #[inline]
-    pub fn build(func_id: u32) -> Self {
+    pub fn build(id: u32) -> Self {
         Self {
-            func_id,
+            id,
             state: Default::default(),
         }
     }
@@ -112,7 +112,7 @@ impl<O, ErrorCode, const IS_RESULT: bool>
         I: scale::Encode,
     {
         ChainExtensionMethod {
-            func_id: self.func_id,
+            id: self.id,
             state: Default::default(),
         }
     }
@@ -136,7 +136,7 @@ impl<I, ErrorCode> ChainExtensionMethod<I, (), ErrorCode, false> {
         O: scale::Decode,
     {
         ChainExtensionMethod {
-            func_id: self.func_id,
+            id: self.id,
             state: Default::default(),
         }
     }
@@ -159,7 +159,7 @@ impl<I, O, const IS_RESULT: bool> ChainExtensionMethod<I, O, (), IS_RESULT> {
         self,
     ) -> ChainExtensionMethod<I, O, state::IgnoreErrorCode, IS_RESULT> {
         ChainExtensionMethod {
-            func_id: self.func_id,
+            id: self.id,
             state: Default::default(),
         }
     }
@@ -178,7 +178,7 @@ impl<I, O, const IS_RESULT: bool> ChainExtensionMethod<I, O, (), IS_RESULT> {
         ErrorCode: FromStatusCode,
     {
         ChainExtensionMethod {
-            func_id: self.func_id,
+            id: self.id,
             state: Default::default(),
         }
     }
@@ -269,7 +269,7 @@ where
                 _,
             >(
                 instance,
-                self.func_id,
+                self.id,
                 input,
                 ErrorCode::from_status_code,
                 |mut output| scale::Decode::decode(&mut output).map_err(Into::into),
@@ -338,7 +338,7 @@ where
                 _,
             >(
                 instance,
-                self.func_id,
+                self.id,
                 input,
                 |_status_code| Ok(()),
                 |mut output| scale::Decode::decode(&mut output).map_err(Into::into),
@@ -399,7 +399,7 @@ where
         <EnvInstance as OnInstance>::on_instance(|instance| {
             EnvBackend::call_chain_extension::<I, O, ErrorCode, ErrorCode, _, _>(
                 instance,
-                self.func_id,
+                self.id,
                 input,
                 ErrorCode::from_status_code,
                 |mut output| {
@@ -449,7 +449,7 @@ where
         <EnvInstance as OnInstance>::on_instance(|instance| {
             EnvBackend::call_chain_extension::<I, O, (), (), _, _>(
                 instance,
-                self.func_id,
+                self.id,
                 input,
                 |_status_code| Ok(()),
                 |mut output| {

crates/env/src/engine/off_chain/impls.rs

@@ -358,7 +358,7 @@ impl EnvBackend for EnvInstance {
 
     fn call_chain_extension<I, T, E, ErrorCode, F, D>(
         &mut self,
-        func_id: u32,
+        id: u32,
         input: &I,
         status_to_result: F,
         decode_to_result: D,
@@ -374,7 +374,7 @@ impl EnvBackend for EnvInstance {
         let mut output: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
 
         self.engine
-            .call_chain_extension(func_id, enc_input, &mut &mut output[..]);
+            .call_chain_extension(id, enc_input, &mut &mut output[..]);
         let (status, out): (u32, Vec<u8>) = scale::Decode::decode(&mut &output[..])
             .unwrap_or_else(|error| {
                 panic!("could not decode `call_chain_extension` output: {error:?}")

crates/env/src/engine/off_chain/mod.rs

@@ -20,8 +20,6 @@ mod types;
 #[cfg(test)]
 mod tests;
 
-pub use call_data::CallData;
-
 use super::OnInstance;
 use crate::Error;
 

crates/env/src/engine/on_chain/impls.rs

@@ -344,7 +344,7 @@ impl EnvBackend for EnvInstance {
 
     fn call_chain_extension<I, T, E, ErrorCode, F, D>(
         &mut self,
-        func_id: u32,
+        id: u32,
         input: &I,
         status_to_result: F,
         decode_to_result: D,
@@ -359,7 +359,7 @@ impl EnvBackend for EnvInstance {
         let mut scope = self.scoped_buffer();
         let enc_input = scope.take_encoded(input);
         let output = &mut scope.take_rest();
-        status_to_result(ext::call_chain_extension(func_id, enc_input, output))?;
+        status_to_result(ext::call_chain_extension(id, enc_input, output))?;
         let decoded = decode_to_result(&output[..])?;
         Ok(decoded)
     }

crates/env/src/lib.rs

@@ -35,7 +35,6 @@
     overflowing_literals,
     path_statements,
     patterns_in_fns_without_body,
-    private_in_public,
     unconditional_recursion,
     unused_allocation,
     unused_comparisons,

crates/ink/codegen/src/generator/chain_extension.rs

@@ -36,7 +36,7 @@ impl ChainExtension<'_> {
         let span = method.span();
         let attrs = method.attrs();
         let ident = method.ident();
-        let func_id = method.id().into_u32();
+        let id = method.id().into_u32();
         let sig = method.sig();
         let inputs = &sig.inputs;
         let input_bindings = method.inputs().map(|pat_type| &pat_type.pat);
@@ -110,7 +110,7 @@ impl ChainExtension<'_> {
                 where
                     #where_output_impls_from_error_code
                 {
-                    ::ink::env::chain_extension::ChainExtensionMethod::build(#func_id)
+                    ::ink::env::chain_extension::ChainExtensionMethod::build(#id)
                     .input::<#compound_input_type>()
                     .output::<#output_type, {::ink::is_result_type!(#output_type)}>()
                     #error_code_handling

crates/ink/ir/src/ast/attr_args.rs

@@ -26,11 +26,17 @@ use syn::{
 ///
 /// For example, the segment `env = ::my::env::Environment`
 /// in `#[ink::contract(env = ::my::env::Environment)]`.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct AttributeArgs {
     args: Punctuated<MetaNameValue, Token![,]>,
 }
 
+impl quote::ToTokens for AttributeArgs {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        self.args.to_tokens(tokens)
+    }
+}
+
 impl IntoIterator for AttributeArgs {
     type Item = MetaNameValue;
     type IntoIter = syn::punctuated::IntoIter<MetaNameValue>;