feat: modify dispatch

poc/extensions/Cargo.toml

@@ -4,9 +4,10 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+parity-scale-codec = { version = "3.6.12", default-features = false }
 scale-info = { version = "2.6.0", default-features = false }
 poc-executor = { path = "../executor", default-features = false }
 
 [features]
 default = ["std"]
-std = ["scale-info/std", "poc-executor/std"]
+std = ["parity-scale-codec/std", "scale-info/std", "poc-executor/std"]

poc/extensions/src/dispatchable.rs

@@ -0,0 +1,8 @@
+pub trait Dispatchable {
+    type MethodName;
+    type MethodIndex;
+    fn query_method(method_name: Self::MethodName) -> Self::MethodIndex;
+    fn dispatch(self) -> Vec<u8>;
+}
+
+pub type ExtensionTypeId = u64;

poc/extensions/src/extension_core.rs

@@ -1,28 +1,74 @@
-use super::{Extension, ExtensionTypeId};
-pub trait ExtensionCore: Extension {
+use super::Dispatchable;
+use super::ExtensionError;
+use super::ExtensionTypeId;
+use parity_scale_codec::{Decode, Encode};
+// SDK codes
+pub trait ExtensionCore: Dispatchable + TryFrom<(u32, Vec<u8>)> {
     // ExtensionId should be generated by the macro
     // It should normalize the order of methods and parameter names
     const TYPE_ID: ExtensionTypeId = 0u64;
-    // TODO: Actual args and return values are complex types
-    // and we adapt them to polkavm ABI in `impl XcqExecutorContext for HostFunctions`
-    fn some_host_function() -> u32;
-    fn another_host_function() -> u32;
+    type ArgsOfSomeHostFunction;
+    type ResultOfSomeHostFunction;
+    fn some_host_function(args: Self::ArgsOfSomeHostFunction) -> Self::ResultOfSomeHostFunction;
 }
 
-#[derive(Clone)]
-struct ExtensionCoreImpl;
+// User impls: handwritten
 
-impl Extension for ExtensionCoreImpl {
-    fn methods(&self) -> Vec<String> {
-        vec!["some_host_function".to_string(), "another_host_function".to_string()]
-    }
+#[derive(Encode, Decode)]
+pub struct ArgsImpl {
+    pub a: u32,
+    pub b: u32,
 }
 
 impl ExtensionCore for ExtensionCoreImpl {
-    fn some_host_function() -> u32 {
-        100
+    type ArgsOfSomeHostFunction = ArgsImpl;
+    type ResultOfSomeHostFunction = u32;
+    fn some_host_function(args: Self::ArgsOfSomeHostFunction) -> Self::ResultOfSomeHostFunction {
+        args.a + args.b
+    }
+}
+
+// User impls: generated
+impl Dispatchable for ExtensionCoreImpl {
+    type MethodName = Vec<u8>;
+    type MethodIndex = u32;
+    fn query_method(method_name: Self::MethodName) -> Self::MethodIndex {
+        match method_name.as_slice() {
+            b"query_method" => 0,
+            b"some_host_function" => 1,
+            _ => Self::MethodIndex::MAX,
+        }
     }
-    fn another_host_function() -> u32 {
-        42
+    fn dispatch(self) -> Vec<u8> {
+        match self {
+            Self::QueryMethod { method_name } => Self::query_method(method_name).encode(),
+            Self::SomeHostFunction { args } => Self::some_host_function(args).encode(),
+        }
+    }
+}
+pub enum ExtensionCoreImpl {
+    // args can be decoded later or here
+    QueryMethod {
+        method_name: <Self as Dispatchable>::MethodName,
+    },
+    SomeHostFunction {
+        args: ArgsImpl,
+    },
+}
+
+impl TryFrom<(u32, Vec<u8>)> for ExtensionCoreImpl {
+    type Error = ExtensionError;
+    fn try_from((method_idx, args): (u32, Vec<u8>)) -> Result<Self, Self::Error> {
+        match method_idx {
+            0 => Ok(Self::QueryMethod {
+                method_name: <Self as Dispatchable>::MethodName::decode(&mut &args[..])
+                    .map_err(Self::Error::DecodeError)?,
+            }),
+            1 => {
+                let args = ArgsImpl::decode(&mut &args[..]).map_err(Self::Error::DecodeError)?;
+                Ok(Self::SomeHostFunction { args })
+            }
+            _ => Err(Self::Error::UnsupportedMethod),
+        }
     }
 }

poc/extensions/src/extension_fungibles.rs

@@ -1,28 +1,96 @@
-use super::{Extension, ExtensionTypeId};
-pub trait ExtensionFungibles: Extension {
+use super::Dispatchable;
+use super::ExtensionError;
+use super::ExtensionTypeId;
+use parity_scale_codec::{Decode, Encode};
+pub trait ExtensionFungibles: Dispatchable + TryFrom<(u32, Vec<u8>)> {
     // ExtensionId should be generated by the macro
     // It should normalize the order of methods and parameter names
     const TYPE_ID: ExtensionTypeId = 1u64;
-    // TODO: Actual args and return values are complex types
-    // and we adapt them to polkavm ABI in `impl XcqExecutorContext for HostFunctions`
-    fn transfer(from: u32, to: u32, amount: u64) -> u64;
-    fn balance(account: u32) -> u64;
-}
-
-#[derive(Clone)]
-struct ExtensionFungiblesImpl;
-
-impl Extension for ExtensionFungiblesImpl {
-    fn methods(&self) -> Vec<String> {
-        vec!["transfer".to_string(), "balance".to_string()]
-    }
+    // TODO: In practice, these type are defined together in Config
+    type AccountId;
+    type Balance;
+    fn transfer(from: Self::AccountId, to: Self::AccountId, amount: Self::Balance) -> Self::Balance;
+    fn balance(account: Self::AccountId) -> Self::Balance;
 }
 
 impl ExtensionFungibles for ExtensionFungiblesImpl {
-    fn transfer(_from: u32, _to: u32, _amount: u64) -> u64 {
+    type AccountId = u32;
+    type Balance = u64;
+    fn transfer(_from: Self::AccountId, _to: Self::AccountId, _amount: Self::Balance) -> Self::Balance {
         unimplemented!()
     }
-    fn balance(_account: u32) -> u64 {
+    fn balance(_account: u32) -> Self::Balance {
         unimplemented!()
     }
 }
+
+// User impls: generated
+impl Dispatchable for ExtensionFungiblesImpl {
+    type MethodName = Vec<u8>;
+    type MethodIndex = u32;
+    fn query_method(method_name: Self::MethodName) -> Self::MethodIndex {
+        match method_name.as_slice() {
+            b"query_method" => 0,
+            b"transfer" => 1,
+            b"balance" => 2,
+            _ => Self::MethodIndex::MAX,
+        }
+    }
+    fn dispatch(self) -> Vec<u8> {
+        match self {
+            Self::QueryMethod { method_name } => Self::query_method(method_name).encode(),
+            Self::Transfer { from, to, amount } => Self::transfer(from, to, amount).encode(),
+            Self::Balance_ { account } => Self::balance(account).encode(),
+        }
+    }
+}
+pub enum ExtensionFungiblesImpl {
+    QueryMethod {
+        method_name: <Self as Dispatchable>::MethodName,
+    },
+    Transfer {
+        from: u32,
+        to: u32,
+        amount: u64,
+    },
+    // Differentiate with associated type Balance
+    Balance_ {
+        account: u32,
+    },
+}
+
+impl TryFrom<(u32, Vec<u8>)> for ExtensionFungiblesImpl {
+    type Error = ExtensionError;
+    fn try_from((method_idx, args): (u32, Vec<u8>)) -> Result<Self, Self::Error> {
+        match method_idx {
+            0 => Ok(Self::QueryMethod {
+                method_name: <Self as Dispatchable>::MethodName::decode(&mut &args[..])
+                    .map_err(Self::Error::DecodeError)?,
+            }),
+            1 => {
+                // This struct encoding must be compatible with the guest side
+                #[derive(Encode, Decode)]
+                struct TransferArgs {
+                    from: u32,
+                    to: u32,
+                    amount: u64,
+                }
+                let args = TransferArgs::decode(&mut &args[..]).map_err(Self::Error::DecodeError)?;
+                Ok(Self::Transfer {
+                    from: args.from,
+                    to: args.to,
+                    amount: args.amount,
+                })
+            }
+            2 => {
+                #[derive(Encode, Decode)]
+                struct BalanceArgs {
+                    account: u32,
+                }
+                let args = BalanceArgs::decode(&mut &args[..]).map_err(Self::Error::DecodeError)?;
+                Ok(Self::Balance_ { account: args.account })
+            }
+            _ => Err(Self::Error::UnsupportedMethod),
+        }
+    }
+}

poc/extensions/src/guest.rs

@@ -4,9 +4,9 @@ pub trait Guest {
     fn program(&self) -> &[u8];
 }
 
-struct GuestImpl {
-    extension_type: ExtensionTypeId,
-    program: Vec<u8>,
+pub struct GuestImpl {
+    pub extension_type: ExtensionTypeId,
+    pub program: Vec<u8>,
 }
 
 impl Guest for GuestImpl {
@@ -25,9 +25,9 @@ pub trait Input {
     fn args(&self) -> &[u8];
 }
 
-struct InputImpl {
-    method: Method,
-    args: Vec<u8>,
+pub struct InputImpl {
+    pub method: Method,
+    pub args: Vec<u8>,
 }
 
 impl Input for InputImpl {