From 2cdbd862a4399cb7d96bf83038791ba613d2ddda Mon Sep 17 00:00:00 2001
From: Daniel <daniel@chainflip.io>
Date: Wed, 27 Sep 2023 13:21:36 +0200
Subject: [PATCH] feat: runtime upgrade utils fix for no-std

---
 state-chain/TROUBLESHOOTING.md                |   5 +-
 .../runtime-upgrade-utilities/src/lib.rs      | 119 +++++++-----------
 2 files changed, 48 insertions(+), 76 deletions(-)

diff --git a/state-chain/TROUBLESHOOTING.md b/state-chain/TROUBLESHOOTING.md
index fded58363ec..995b2c9d590 100644
--- a/state-chain/TROUBLESHOOTING.md
+++ b/state-chain/TROUBLESHOOTING.md
@@ -23,7 +23,8 @@ SKIP_WASM_BUILD=1 cargo build --release --features=try-runtime
 Now you can run your tests like so (using an appropriate public rpc node):
 
 ```sh
-./target/release/chainflip-node try-runtime --execution native \
+./target/release/chainflip-node try-runtime \
+    --runtime ./path/to/state_chain_runtime.wasm \
     on-runtime-upgrade live --uri wss://perseverance.chainflip.xyz:443
 ```
 
@@ -48,8 +49,6 @@ Once the node has synced, in another terminal window, run the checks:
     on-runtime-upgrade live --uri ws://localhost:9944
 ```
 
-> *Note: Using `--execution native` ensures faster execution and also prevents log messages from being scrubbed.
-
 ### General tips and guidelines
 
 - There are some useful storage conversion utilities in `frame_support::storage::migration`.
diff --git a/state-chain/runtime-upgrade-utilities/src/lib.rs b/state-chain/runtime-upgrade-utilities/src/lib.rs
index c5dfa0b6ed7..37cead8be31 100644
--- a/state-chain/runtime-upgrade-utilities/src/lib.rs
+++ b/state-chain/runtime-upgrade-utilities/src/lib.rs
@@ -31,83 +31,56 @@ pub struct VersionedMigration<
 
 #[cfg(feature = "try-runtime")]
 mod try_runtime_helpers {
-	use frame_support::traits::PalletInfoAccess;
-	use sp_std::vec::Vec;
-
-	#[cfg(feature = "std")]
-	pub use with_std::*;
-
-	#[cfg(not(feature = "std"))]
-	pub use without_std::*;
-
-	#[cfg(feature = "std")]
-	mod with_std {
-		use super::*;
-		use core::cell::RefCell;
-		use sp_std::{
-			cmp::{max, min},
-			collections::btree_map::BTreeMap,
-		};
-
-		thread_local! {
-			pub static MIGRATION_BOUNDS: RefCell<BTreeMap<&'static str, (u16, u16)>> = Default::default();
-			#[allow(clippy::type_complexity)]
-			pub static MIGRATION_STATE: RefCell<BTreeMap<&'static str, BTreeMap<(u16, u16), Vec<u8>>>> = Default::default();
-		}
-
-		pub fn update_migration_bounds<T: PalletInfoAccess, const FROM: u16, const TO: u16>() {
-			MIGRATION_BOUNDS.with(|cell| {
-				cell.borrow_mut()
-					.entry(T::name())
-					.and_modify(|(from, to)| {
-						*from = min(*from, FROM);
-						*to = max(*to, TO);
-					})
-					.or_insert((FROM, TO));
-			});
-		}
-
-		pub fn get_migration_bounds<T: PalletInfoAccess>() -> Option<(u16, u16)> {
-			MIGRATION_BOUNDS.with(|cell| cell.borrow().get(T::name()).copied())
-		}
-
-		pub fn save_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>(s: Vec<u8>) {
-			MIGRATION_STATE
-				.with(|cell| cell.borrow_mut().entry(T::name()).or_default().insert((FROM, TO), s));
-		}
-
-		pub fn restore_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>() -> Vec<u8> {
-			MIGRATION_STATE.with(|cell| {
-				cell.borrow()
-					.get(T::name())
-					.cloned()
-					.unwrap_or_default()
-					.get(&(FROM, TO))
-					.cloned()
-					.unwrap_or_default()
-			})
-		}
+	use frame_support::{pallet_prelude::ValueQuery, storage_alias, traits::PalletInfoAccess};
+	use sp_std::{
+		cmp::{max, min},
+		collections::btree_map::BTreeMap,
+		vec::Vec,
+	};
+
+	#[storage_alias]
+	pub type MigrationBounds =
+		StorageValue<CfUpgradeUtilities, BTreeMap<Vec<u8>, (u16, u16)>, ValueQuery>;
+
+	#[storage_alias]
+	pub type MigrationState = StorageValue<
+		CfUpgradeUtilities,
+		BTreeMap<Vec<u8>, BTreeMap<(u16, u16), Vec<u8>>>,
+		ValueQuery,
+	>;
+
+	pub fn update_migration_bounds<T: PalletInfoAccess, const FROM: u16, const TO: u16>() {
+		MigrationBounds::mutate(|bounds| {
+			bounds
+				.entry(T::name().as_bytes().to_vec())
+				.and_modify(|(from, to)| {
+					*from = min(*from, FROM);
+					*to = max(*to, TO);
+				})
+				.or_insert((FROM, TO));
+		});
 	}
 
-	#[cfg(not(feature = "std"))]
-	mod without_std {
-		use super::*;
-
-		pub fn update_migration_bounds<T: PalletInfoAccess, const FROM: u16, const TO: u16>() {
-			log::warn!("❗️ Runtime upgrade utilities are not supported in no-std.");
-		}
-
-		pub fn get_migration_bounds<T: PalletInfoAccess>() -> Option<(u16, u16)> {
-			Default::default()
-		}
+	pub fn get_migration_bounds<T: PalletInfoAccess>() -> Option<(u16, u16)> {
+		MigrationBounds::get().get(T::name().as_bytes()).copied()
+	}
 
-		pub fn save_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>(s: Vec<u8>) {
-			log::warn!("❗️ Runtime upgrade utilities are not supported in no-std.");
-		}
+	pub fn save_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>(s: Vec<u8>) {
+		MigrationState::mutate(|state| {
+			state.entry(T::name().as_bytes().to_vec()).or_default().insert((FROM, TO), s)
+		});
+	}
 
-		pub fn restore_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>() -> Vec<u8> {
-			Default::default()
-		}
+	pub fn restore_state<T: PalletInfoAccess, const FROM: u16, const TO: u16>() -> Vec<u8> {
+		MigrationState::mutate(|state| {
+			state
+				.get(T::name().as_bytes())
+				.cloned()
+				.unwrap_or_default()
+				.get(&(FROM, TO))
+				.cloned()
+				.unwrap_or_default()
+		})
 	}
 }