feat: add hooks

Cargo.toml

@@ -12,15 +12,9 @@ description = "Library for performing storage migrations in Substrate pallets"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
-syn = "2.0.32"
-proc-macro2 = "1.0.64"
-quote = "1.0.33"
-proc-macro-crate = "1.3.1"
-log = { version = "0.4", default-features = false }
-codec = { package = "parity-scale-codec", version = "3.6.5", default-features = false, features = ["derive", "max-encoded-len"] }
+parity-scale-codec = { version = "3.6", default-features = false, features = ["derive", "max-encoded-len"] }
 migratable-procedural = { path = "./procedural", default-features = false }
 frame-support = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }
-frame-support-procedural = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }
 frame-system = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }
 sp-std = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }
 sp-runtime = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }
@@ -29,7 +23,7 @@ impl-trait-for-tuples = "0.2"
 [features]
 default = ["std"]
 std = [
-	"codec/std",
+	"parity-scale-codec/std",
 	"frame-support/std",
 	"frame-system/std",
 	"sp-runtime/std",

README.md

@@ -0,0 +1,3 @@
+#Migratable
+
+This project is a library used to have `pallet-contracts`-like migration schema for Substrate pallets.

procedural/Cargo.toml

@@ -15,6 +15,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 proc-macro = true
 
 [dependencies]
-syn = { version = "2.0.32", features = ["full"] }
+syn = { version = "2.0", features = ["full"] }
 proc-macro2 = "1.0.64"
 quote = "1.0.33"
+sp-std = { default-features = false, git = "https://github.com/paritytech/substrate.git", branch = "polkadot-v1.0.0" }

procedural/src/lib.rs

@@ -1,22 +1,113 @@
-use quote::quote;
-use syn::{parse_quote, parse_macro_input};
+use quote::{quote, ToTokens};
+use syn::{parse_macro_input, parse_quote};
+
+#[proc_macro_attribute]
+pub fn hooks(
+    _attr: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
+) -> proc_macro::TokenStream {
+    let mut input = parse_macro_input!(item as syn::ItemImpl);
+    let on_idle_impl = quote!(
+		{
+			let mut weight_sum = remaining_weight.clone();
+			loop {
+				let (result, weight) = Migration::<T>::migrate(remaining_weight);
+				remaining_weight.saturating_reduce(weight);
+
+				match result {
+					// There is not enough weight to perform a migration, or make any progress, we
+					// just return the remaining weight.
+					migratable::MigrateResult::NoMigrationPerformed | migratable::MigrateResult::InProgress { steps_done: 0 } => return remaining_weight,
+					// Migration is still in progress, we can start the next step.
+					migratable::MigrateResult::InProgress { .. } => continue,
+					// Either no migration is in progress, or we are done with all migrations, we
+					// can do some more other work with the remaining weight.
+					migratable::MigrateResult::Completed | migratable::MigrateResult::NoMigrationInProgress => break,
+				}
+			};
+			weight_sum.saturating_reduce(remaining_weight);
+			weight_sum
+		}
+	);
+    let integrity_test_impl = quote!(
+		{
+			Migration::<T>::integrity_test();
+		}
+	);
+    let custom_code_map: sp_std::collections::btree_map::BTreeMap<&str, String> = [
+        ("on_idle", on_idle_impl.to_string()),
+        ("integrity_test", integrity_test_impl.to_string()),
+    ]
+        .iter()
+        .cloned()
+        .collect();
+    // check the required functions are present
+    let implemented_fn = input.items.iter().filter(|item| {
+        if let syn::ImplItem::Fn(method) = item {
+            if custom_code_map.contains_key(method.sig.ident.to_string().as_str()) {
+                true
+            } else { false }
+        } else { false }
+    }).collect::<Vec<_>>();
+    if implemented_fn.len() != 2 {
+        let required_hooks = custom_code_map.keys().map(|k| *k).collect::<Vec<_>>().join(", ");
+        panic!("You must define the following hooks: {}", required_hooks);
+    };
+    // modify the actual hooks
+    for item in &mut input.items {
+        if let syn::ImplItem::Fn(method) = item {
+            let fn_name = method.sig.ident.to_string();
+            if let Some(custom_code) = custom_code_map.get(&fn_name.as_str()) {
+                let new_code = syn::parse_str::<proc_macro2::TokenStream>(custom_code).expect("Invalid custom code");
+                let curr_impl = method.block.to_token_stream();
+                match fn_name.as_str() {
+                    "on_idle" => {
+                        method.block = parse_quote! {
+							{
+								let migration_weight = #new_code;
+								let mut weight = #curr_impl;
+								weight.saturating_add(migration_weight);
+								weight
+							}
+						};
+                    },
+					"integrity_test" => {
+						method.block = parse_quote! {
+							{
+								#new_code
+								#curr_impl
+							}
+						};
+					},
+					_ => {}
+                }
+            } else {
+                panic!("\"{}\" not found in pallet hooks", fn_name);
+            }
+        }
+    }
+    let modified_impl = quote! {
+        #input
+    };
+    modified_impl.into()
+}
 
 /// Adds the `Migrations` type to `Config`
 #[proc_macro_attribute]
 pub fn config(
-	_attr: proc_macro::TokenStream,
-	item: proc_macro::TokenStream,
+    _attr: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
-	let mut input = parse_macro_input!(item as syn::ItemTrait);
-	let migrations = quote!(
+    let mut input = parse_macro_input!(item as syn::ItemTrait);
+    let migrations = quote!(
 		/// The sequence of migration steps that will be applied during a migration.
 		type Migrations: migratable::MigrateSequence;
 	);
-	input.items.push(parse_quote! { #migrations });
-	let output = quote! {
+    input.items.push(parse_quote! { #migrations });
+    let output = quote! {
         #input
     };
-	output.into()
+    output.into()
 }
 
 /// Adds the following to the pallet module:
@@ -27,36 +118,36 @@ pub fn pallet(
     _attr: proc_macro::TokenStream,
     item: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
-	let mut input = parse_macro_input!(item as syn::ItemMod);
-	let content = &mut input.content.as_mut().unwrap().1;
+    let mut input = parse_macro_input!(item as syn::ItemMod);
+    let content = &mut input.content.as_mut().unwrap().1;
 
-	// add storage
-	let storage = quote!(
+    // add storage
+    let storage = quote!(
 		/// A migration can span across multiple blocks. This storage defines a cursor to track the
-		/// progress of the migration, enabling us to resume from the last completed position.
+        /// progress of the migration, enabling us to resume from the last completed position.
 		#[pallet::storage]
 		pub type MigrationInProgress<T: Config> = StorageValue<
 			_,
 			migratable::Cursor,
 			frame_support::storage::types::OptionQuery,
 		>;
 	);
-	content.push(parse_quote! { #storage });
+    content.push(parse_quote! { #storage });
 
-	// add migration struct
-	let migration = quote!(
+    // add migration struct
+    let migration = quote!(
 		/// Performs all necessary migrations based on `StorageVersion`.
-		///
-		/// If `TEST_ALL_STEPS == true` and `try-runtime` is enabled, this will run all the migrations
-		/// inside `on_runtime_upgrade`. This should be set to false in tests that want to ensure the step
-		/// by step migratable works.
+        ///
+        /// If `TEST_ALL_STEPS == true` and `try-runtime` is enabled, this will run all the migrations
+        /// inside `on_runtime_upgrade`. This should be set to false in tests that want to ensure the step
+        /// by step migratable works.
 		struct Migration<T: Config, const TEST_ALL_STEPS: bool = true>(core::marker::PhantomData<T>);
 	);
     content.push(parse_quote! { #migration });
 
-	// add migration logic
-	let expand = generate_mod_expand();
-	content.push(parse_quote! { #expand });
+    // add migration logic
+    let expand = generate_mod_expand();
+    content.push(parse_quote! { #expand });
 
     let output = quote! {
         #input
@@ -178,14 +269,14 @@ fn generate_mod_expand() -> proc_macro2::TokenStream {
 				use migratable::weights::WeightInfo;
 				impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
 					/// Verify that each migratable's step of the [`Config::Migrations`] sequence fits into
-					/// `Cursor`.
+                    /// `Cursor`.
 					pub(crate) fn integrity_test() {
 						let max_weight = <T as frame_system::Config>::BlockWeights::get().max_block;
 						T::Migrations::integrity_test(max_weight)
 					}
 
 					/// Migrate
-					/// Return the weight used and whether or not a migratable is in progress
+                    /// Return the weight used and whether or not a migratable is in progress
 					pub(crate) fn migrate(weight_limit: frame_support::weights::Weight) -> (migratable::MigrateResult, frame_support::weights::Weight) {
 						let name = <Pallet<T>>::name();
 						let mut weight_left = weight_limit;

src/lib.rs

@@ -1,3 +1,5 @@
+#![cfg_attr(not(feature = "std"), no_std)]
+
 //! Multi-block Migration framework.
 //!
 //! This module allows us to define a migratable as a sequence of [`MigrationStep`]s that can be
@@ -40,14 +42,12 @@
 //! While the migratable is in progress, all dispatchables except `migrate`, are blocked, and returns
 //! a `MigrationInProgress` error.
 
-pub use migratable_procedural::{config, pallet};
+pub use migratable_procedural::{config, hooks, pallet};
 pub mod weights;
 
-pub use log;
-
 extern crate alloc;
 
-use codec::{Codec, Decode};
+use parity_scale_codec::{Codec, Decode};
 use frame_support::{
     pallet_prelude::{BoundedVec, Encode, MaxEncodedLen, StorageVersion, Weight},
     traits::ConstU32,
@@ -61,7 +61,6 @@ use sp_std::prelude::*;
 const PROOF_ENCODE: &str = "Tuple::max_encoded_len() < Cursor::max_encoded_len()` is verified in `Self::integrity_test()`; qed";
 const PROOF_DECODE: &str =
     "We encode to the same type in this trait only. No other code touches this item; qed";
-const LOG_TARGET: &str = "migratable";
 
 fn invalid_version(version: StorageVersion) -> ! {
     panic!("Required migratable {version:?} not supported by this runtime. This is a bug.");
@@ -140,7 +139,6 @@ impl<const N: u16> MigrationStep for NoopMigration<N> {
         Weight::zero()
     }
     fn step(&mut self) -> (IsFinished, Weight) {
-        log::debug!(target: LOG_TARGET, "Noop migratable for version {}", N);
         (IsFinished::Yes, Weight::zero())
     }
 }