Make extension trait.

Author: kianenigma

frame/support/src/traits.rs

@@ -1554,26 +1554,12 @@ pub trait OnGenesis {
 }
 
 /// Prefix to be used (optionally) for implementing [`OnRuntimeUpgrade::storage_key`].
+#[cfg(feature = "try-runtime")]
 pub const ON_RUNTIME_UPGRADE_PREFIX: &[u8] = b"__ON_RUNTIME_UPGRADE__";
 
-/// The runtime upgrade trait.
-///
-/// Implementing this lets you express what should happen when the runtime upgrades,
-/// and changes may need to occur to your module.
-pub trait OnRuntimeUpgrade {
-	/// Perform a module upgrade.
-	///
-	/// # Warning
-	///
-	/// This function will be called before we initialized any runtime state, aka `on_initialize`
-	/// wasn't called yet. So, information like the block number and any other
-	/// block local data are not accessible.
-	///
-	/// Return the non-negotiable weight consumed for runtime upgrade.
-	fn on_runtime_upgrade() -> crate::weights::Weight {
-		0
-	}
-
+/// Some helper functions for [`OnRuntimeUpgrade`] during `try-runtime` testing.
+#[cfg(feature = "try-runtime")]
+pub trait OnRuntimeUpgradeHelpersExt {
 	/// Generate a storage key unique to this runtime upgrade.
 	///
 	/// This can be used to communicate data from pre-upgrade to post-upgrade state and check
@@ -1590,6 +1576,17 @@ pub trait OnRuntimeUpgrade {
 		final_key
 	}
 
+	/// Get temporary storage data written by [`set_temp_storage`].
+	///
+	/// Returns `None` if either the data is unavailable or un-decodable.
+	///
+	/// A `at` storage identifier must be provided to indicate where the storage is being read from.
+	#[cfg(feature = "try-runtime")]
+	fn get_temp_storage<T: Decode>(at: &str) -> Option<T> {
+		sp_io::storage::get(&Self::storage_key(at))
+			.and_then(|bytes| Decode::decode(&mut &*bytes).ok())
+	}
+
 	/// Write some temporary data to a specific storage that can be read (potentially in
 	/// post-upgrade hook) via [`get_temp_storage`].
 	///
@@ -1599,16 +1596,27 @@ pub trait OnRuntimeUpgrade {
 	fn set_temp_storage<T: Encode>(data: T, at: &str) {
 		sp_io::storage::set(&Self::storage_key(at), &data.encode());
 	}
+}
 
-	/// Get temporary storage data written by [`set_temp_storage`].
+#[cfg(feature = "try-runtime")]
+impl<U: OnRuntimeUpgrade> OnRuntimeUpgradeHelpersExt for U {}
+
+/// The runtime upgrade trait.
+///
+/// Implementing this lets you express what should happen when the runtime upgrades,
+/// and changes may need to occur to your module.
+pub trait OnRuntimeUpgrade {
+	/// Perform a module upgrade.
 	///
-	/// Returns `None` if either the data is unavailable or un-decodable.
+	/// # Warning
 	///
-	/// A `at` storage identifier must be provided to indicate where the storage is being read from.
-	#[cfg(feature = "try-runtime")]
-	fn get_temp_storage<T: Decode>(at: &str) -> Option<T> {
-		sp_io::storage::get(&Self::storage_key(at))
-			.and_then(|bytes| Decode::decode(&mut &*bytes).ok())
+	/// This function will be called before we initialized any runtime state, aka `on_initialize`
+	/// wasn't called yet. So, information like the block number and any other
+	/// block local data are not accessible.
+	///
+	/// Return the non-negotiable weight consumed for runtime upgrade.
+	fn on_runtime_upgrade() -> crate::weights::Weight {
+		0
 	}
 
 	/// Execute some pre-checks prior to a runtime upgrade.
@@ -1649,12 +1657,6 @@ impl OnRuntimeUpgrade for Tuple {
 		for_tuples!( #( result = result.and(Tuple::post_upgrade()); )* );
 		result
 	}
-
-	#[cfg(feature = "try-runtime")]
-	fn storage_key(ident: &str) -> [u8; 32] {
-		// We will never use the key on a tuple, effectively _don't care_.
-		unreachable!()
-	}
 }
 
 /// Off-chain computation trait.