From cc898077b4f0f48839ce69933fdf8336f1ef37bc Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Wed, 13 Sep 2023 10:31:54 +0200
Subject: [PATCH] runtime-sdk: Implement automatic Module derivation

---
 runtime-sdk-macros/src/module_derive/mod.rs   | 17 +++++-
 .../src/module_derive/module.rs               | 54 +++++++++++++++++++
 runtime-sdk/modules/contracts/src/lib.rs      | 11 ++--
 runtime-sdk/modules/evm/src/lib.rs            | 11 ++--
 runtime-sdk/src/dispatcher.rs                 |  7 +--
 runtime-sdk/src/modules/accounts/mod.rs       | 11 ++--
 runtime-sdk/src/modules/accounts/test.rs      |  7 +--
 runtime-sdk/src/modules/consensus/mod.rs      | 13 ++---
 .../src/modules/consensus_accounts/mod.rs     | 15 ++----
 runtime-sdk/src/modules/core/mod.rs           | 11 ++--
 runtime-sdk/src/modules/core/test.rs          |  7 +--
 runtime-sdk/src/modules/rewards/mod.rs        |  7 +--
 12 files changed, 104 insertions(+), 67 deletions(-)
 create mode 100644 runtime-sdk-macros/src/module_derive/module.rs

diff --git a/runtime-sdk-macros/src/module_derive/mod.rs b/runtime-sdk-macros/src/module_derive/mod.rs
index 084ea07839..7e4467cbd4 100644
--- a/runtime-sdk-macros/src/module_derive/mod.rs
+++ b/runtime-sdk-macros/src/module_derive/mod.rs
@@ -1,5 +1,6 @@
 mod method_handler;
 mod migration_handler;
+mod module;
 
 use proc_macro2::TokenStream;
 use quote::quote;
@@ -17,6 +18,7 @@ pub fn derive_module(impl_block: syn::ItemImpl) -> TokenStream {
     let mut derivations: Vec<TokenStream> = Vec::new();
 
     let mut derivers: Vec<Box<dyn Deriver>> = vec![
+        module::DeriveModule::new(),
         migration_handler::DeriveMigrationHandler::new(),
         method_handler::DeriveMethodHandler::new(),
     ];
@@ -376,9 +378,14 @@ mod tests {
     }
 
     #[test]
-    fn generate_migration_handler_impl() {
+    fn generate_module_impl() {
         let input = syn::parse_quote!(
             impl<C: Cfg> MyModule<C> {
+                const NAME: &'static str = MODULE_NAME;
+                const VERSION: u32 = 2;
+                type Error = Error;
+                type Event = ();
+                type Parameters = Parameters;
                 type Genesis = Genesis;
 
                 #[migration(init)]
@@ -404,6 +411,14 @@ mod tests {
                     const _: () = {
                         #uses
                         #[automatically_derived]
+                        impl<C: Cfg> sdk::module::Module for MyModule<C> {
+                            const NAME: &'static str = MODULE_NAME;
+                            const VERSION: u32 = 2;
+                            type Error = Error;
+                            type Event = ();
+                            type Parameters = Parameters;
+                        }
+                        #[automatically_derived]
                         impl<C: Cfg> sdk::module::MigrationHandler for MyModule<C> {
                             type Genesis = Genesis;
                             fn init_or_migrate<C: Context>(
diff --git a/runtime-sdk-macros/src/module_derive/module.rs b/runtime-sdk-macros/src/module_derive/module.rs
new file mode 100644
index 0000000000..0eac9d3625
--- /dev/null
+++ b/runtime-sdk-macros/src/module_derive/module.rs
@@ -0,0 +1,54 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+
+/// Deriver for the `Module` trait.
+pub struct DeriveModule {
+    /// Items specifying the module configuration.
+    module_cfg: Vec<syn::ImplItem>,
+}
+
+impl DeriveModule {
+    pub fn new() -> Box<Self> {
+        Box::new(Self { module_cfg: vec![] })
+    }
+}
+
+impl super::Deriver for DeriveModule {
+    fn preprocess(&mut self, item: syn::ImplItem) -> Option<syn::ImplItem> {
+        match item {
+            syn::ImplItem::Type(ref ty) => {
+                match ty.ident.to_string().as_str() {
+                    "Error" | "Event" | "Parameters" => {
+                        self.module_cfg.push(item);
+                        None // Take the item.
+                    }
+                    _ => Some(item), // Return the item.
+                }
+            }
+            syn::ImplItem::Const(ref cnst) => {
+                match cnst.ident.to_string().as_str() {
+                    "NAME" | "VERSION" => {
+                        self.module_cfg.push(item);
+                        None // Take the item.
+                    }
+                    _ => Some(item), // Return the item.
+                }
+            }
+            _ => Some(item), // Return the item.
+        }
+    }
+
+    fn derive(&mut self, generics: &syn::Generics, ty: &Box<syn::Type>) -> TokenStream {
+        if self.module_cfg.is_empty() {
+            return quote! {};
+        }
+        let module_cfg = &self.module_cfg;
+
+        quote! {
+            #[automatically_derived]
+            impl #generics sdk::module::Module for #ty {
+                #(#module_cfg)*
+            }
+        }
+    }
+}
diff --git a/runtime-sdk/modules/contracts/src/lib.rs b/runtime-sdk/modules/contracts/src/lib.rs
index 094b77bb89..29023184f4 100644
--- a/runtime-sdk/modules/contracts/src/lib.rs
+++ b/runtime-sdk/modules/contracts/src/lib.rs
@@ -421,6 +421,10 @@ impl<Cfg: Config> Module<Cfg> {
 
 #[sdk_derive(Module)]
 impl<Cfg: Config> Module<Cfg> {
+    const NAME: &'static str = MODULE_NAME;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
@@ -864,13 +868,6 @@ impl<Cfg: Config> Module<Cfg> {
     }
 }
 
-impl<Cfg: Config> module::Module for Module<Cfg> {
-    const NAME: &'static str = MODULE_NAME;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 impl<Cfg: Config> module::TransactionHandler for Module<Cfg> {}
 impl<Cfg: Config> module::BlockHandler for Module<Cfg> {}
 impl<Cfg: Config> module::InvariantHandler for Module<Cfg> {}
diff --git a/runtime-sdk/modules/evm/src/lib.rs b/runtime-sdk/modules/evm/src/lib.rs
index 7430e50025..bd5c1f52d9 100644
--- a/runtime-sdk/modules/evm/src/lib.rs
+++ b/runtime-sdk/modules/evm/src/lib.rs
@@ -260,13 +260,6 @@ pub enum Event {
     },
 }
 
-impl<Cfg: Config> module::Module for Module<Cfg> {
-    const NAME: &'static str = MODULE_NAME;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 /// Interface that can be called from other modules.
 pub trait API {
     /// Perform an Ethereum CREATE transaction.
@@ -660,6 +653,10 @@ impl<Cfg: Config> Module<Cfg> {
 
 #[sdk_derive(Module)]
 impl<Cfg: Config> Module<Cfg> {
+    const NAME: &'static str = MODULE_NAME;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
diff --git a/runtime-sdk/src/dispatcher.rs b/runtime-sdk/src/dispatcher.rs
index f064c6cf54..2fc3e0b1f9 100644
--- a/runtime-sdk/src/dispatcher.rs
+++ b/runtime-sdk/src/dispatcher.rs
@@ -959,16 +959,13 @@ mod test {
     /// A module with multiple no-op methods; intended for testing routing.
     struct AlphabetModule;
 
-    impl module::Module for AlphabetModule {
+    #[sdk_derive(Module)]
+    impl AlphabetModule {
         const NAME: &'static str = "alphabet";
         const VERSION: u32 = 42;
         type Error = AlphabetError;
         type Event = ();
         type Parameters = ();
-    }
-
-    #[sdk_derive(Module)]
-    impl AlphabetModule {
         type Genesis = ();
 
         #[handler(call = "alphabet.ReadOnly")]
diff --git a/runtime-sdk/src/modules/accounts/mod.rs b/runtime-sdk/src/modules/accounts/mod.rs
index 17cb72aae4..7e7cea6301 100644
--- a/runtime-sdk/src/modules/accounts/mod.rs
+++ b/runtime-sdk/src/modules/accounts/mod.rs
@@ -732,6 +732,10 @@ impl API for Module {
 
 #[sdk_derive(Module)]
 impl Module {
+    const NAME: &'static str = MODULE_NAME;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
@@ -866,13 +870,6 @@ impl Module {
     }
 }
 
-impl module::Module for Module {
-    const NAME: &'static str = MODULE_NAME;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 impl module::TransactionHandler for Module {
     fn authenticate_tx<C: Context>(
         ctx: &mut C,
diff --git a/runtime-sdk/src/modules/accounts/test.rs b/runtime-sdk/src/modules/accounts/test.rs
index cf8187aa92..6e6b82384a 100644
--- a/runtime-sdk/src/modules/accounts/test.rs
+++ b/runtime-sdk/src/modules/accounts/test.rs
@@ -72,15 +72,12 @@ impl Runtime for TestRuntime {
 /// A module with multiple no-op methods; intended for testing routing.
 struct TestModule;
 
-impl module::Module for TestModule {
+#[sdk_derive(Module)]
+impl TestModule {
     const NAME: &'static str = "test";
     type Error = CoreError;
     type Event = ();
     type Parameters = ();
-}
-
-#[sdk_derive(Module)]
-impl TestModule {
     type Genesis = ();
 
     #[handler(call = "test.RefundFee")]
diff --git a/runtime-sdk/src/modules/consensus/mod.rs b/runtime-sdk/src/modules/consensus/mod.rs
index 4862be42c8..9159529101 100644
--- a/runtime-sdk/src/modules/consensus/mod.rs
+++ b/runtime-sdk/src/modules/consensus/mod.rs
@@ -202,6 +202,11 @@ impl Module {
 
 #[sdk_derive(Module)]
 impl Module {
+    const NAME: &'static str = MODULE_NAME;
+    const VERSION: u32 = 1;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
@@ -400,14 +405,6 @@ impl API for Module {
     }
 }
 
-impl module::Module for Module {
-    const NAME: &'static str = MODULE_NAME;
-    const VERSION: u32 = 1;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 impl module::TransactionHandler for Module {}
 
 impl module::BlockHandler for Module {}
diff --git a/runtime-sdk/src/modules/consensus_accounts/mod.rs b/runtime-sdk/src/modules/consensus_accounts/mod.rs
index 11d129ea92..1d034431ff 100644
--- a/runtime-sdk/src/modules/consensus_accounts/mod.rs
+++ b/runtime-sdk/src/modules/consensus_accounts/mod.rs
@@ -370,6 +370,11 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> API
 impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
     Module<Accounts, Consensus>
 {
+    const NAME: &'static str = MODULE_NAME;
+    const VERSION: u32 = 1;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
@@ -694,16 +699,6 @@ impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
     }
 }
 
-impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API> module::Module
-    for Module<Accounts, Consensus>
-{
-    const NAME: &'static str = MODULE_NAME;
-    const VERSION: u32 = 1;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 impl<Accounts: modules::accounts::API, Consensus: modules::consensus::API>
     module::TransactionHandler for Module<Accounts, Consensus>
 {
diff --git a/runtime-sdk/src/modules/core/mod.rs b/runtime-sdk/src/modules/core/mod.rs
index d2261d67b2..8dff159a99 100644
--- a/runtime-sdk/src/modules/core/mod.rs
+++ b/runtime-sdk/src/modules/core/mod.rs
@@ -483,6 +483,10 @@ impl<Cfg: Config> API for Module<Cfg> {
 
 #[sdk_derive(Module)]
 impl<Cfg: Config> Module<Cfg> {
+    const NAME: &'static str = MODULE_NAME;
+    type Error = Error;
+    type Event = Event;
+    type Parameters = Parameters;
     type Genesis = Genesis;
 
     #[migration(init)]
@@ -910,13 +914,6 @@ impl<Cfg: Config> Module<Cfg> {
     }
 }
 
-impl<Cfg: Config> module::Module for Module<Cfg> {
-    const NAME: &'static str = MODULE_NAME;
-    type Error = Error;
-    type Event = Event;
-    type Parameters = Parameters;
-}
-
 impl<Cfg: Config> module::TransactionHandler for Module<Cfg> {
     fn approve_raw_tx<C: Context>(_ctx: &mut C, tx: &[u8]) -> Result<(), Error> {
         let params = Self::params();
diff --git a/runtime-sdk/src/modules/core/test.rs b/runtime-sdk/src/modules/core/test.rs
index 70b5bb12ee..292e8f0724 100644
--- a/runtime-sdk/src/modules/core/test.rs
+++ b/runtime-sdk/src/modules/core/test.rs
@@ -207,16 +207,13 @@ impl GasWasterModule {
     const METHOD_SPECIFIC_GAS_REQUIRED_HUGE: &'static str = "test.SpecificGasRequiredHuge";
 }
 
-impl module::Module for GasWasterModule {
+#[sdk_derive(Module)]
+impl GasWasterModule {
     const NAME: &'static str = "gaswaster";
     const VERSION: u32 = 42;
     type Error = crate::modules::core::Error;
     type Event = ();
     type Parameters = ();
-}
-
-#[sdk_derive(Module)]
-impl GasWasterModule {
     type Genesis = ();
 
     #[handler(call = Self::METHOD_WASTE_GAS)]
diff --git a/runtime-sdk/src/modules/rewards/mod.rs b/runtime-sdk/src/modules/rewards/mod.rs
index 2af655b365..bbcae5ed5f 100644
--- a/runtime-sdk/src/modules/rewards/mod.rs
+++ b/runtime-sdk/src/modules/rewards/mod.rs
@@ -90,16 +90,13 @@ pub struct Module<Accounts: modules::accounts::API> {
 pub static ADDRESS_REWARD_POOL: Lazy<Address> =
     Lazy::new(|| Address::from_module(MODULE_NAME, "reward-pool"));
 
-impl<Accounts: modules::accounts::API> module::Module for Module<Accounts> {
+#[sdk_derive(Module)]
+impl<Accounts: modules::accounts::API> Module<Accounts> {
     const NAME: &'static str = MODULE_NAME;
     const VERSION: u32 = 2;
     type Error = Error;
     type Event = ();
     type Parameters = Parameters;
-}
-
-#[sdk_derive(Module)]
-impl<Accounts: modules::accounts::API> Module<Accounts> {
     type Genesis = Genesis;
 
     #[migration(init)]