From f95bd729fb4b66ef7d940f9c048649b61cd18cb6 Mon Sep 17 00:00:00 2001
From: Tarek Mohamed Abdalla <tarekkma@gmail.com>
Date: Tue, 1 Oct 2024 08:26:32 +0000
Subject: [PATCH] upgrades

---
 pallets/parachain-staking/src/migrations.rs | 73 +++++++++++++++++++++
 runtime/common/src/migrations.rs            | 29 ++++++++
 2 files changed, 102 insertions(+)

diff --git a/pallets/parachain-staking/src/migrations.rs b/pallets/parachain-staking/src/migrations.rs
index 2e08649a70..422f855947 100644
--- a/pallets/parachain-staking/src/migrations.rs
+++ b/pallets/parachain-staking/src/migrations.rs
@@ -13,3 +13,76 @@
 
 // You should have received a copy of the GNU General Public License
 // along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
+
+use frame_support::{traits::OnRuntimeUpgrade, weights::Weight};
+
+use crate::*;
+
+#[derive(Clone, PartialEq, Eq, parity_scale_codec::Decode, sp_runtime::RuntimeDebug)]
+/// Reserve information { account, percent_of_inflation }
+pub struct OldParachainBondConfig<AccountId> {
+	/// Account which receives funds intended for parachain bond
+	pub account: AccountId,
+	/// Percent of inflation set aside for parachain bond account
+	pub percent: sp_runtime::Percent,
+}
+
+pub struct MigrateParachainBondConfig<T>(sp_std::marker::PhantomData<T>);
+impl<T: Config> OnRuntimeUpgrade for MigrateParachainBondConfig<T> {
+	fn on_runtime_upgrade() -> Weight {
+		let (account, percent) = if let Some(config) =
+			frame_support::storage::migration::get_storage_value::<
+				OldParachainBondConfig<T::AccountId>,
+			>(b"ParachainStaking", b"ParachainBondInfo", &[])
+		{
+			(config.account, config.percent)
+		} else {
+			return Weight::default();
+		};
+
+		let pbr = InflationDistributionAccount { account, percent };
+		let treasury = InflationDistributionAccount::<T::AccountId>::default();
+		let configs: InflationDistributionConfig<T::AccountId> = [pbr, treasury].into();
+
+		//***** Start mutate storage *****//
+
+		InflationDistributionInfo::<T>::put(configs);
+
+		// Remove storage value AssetManager::SupportedFeePaymentAssets
+		frame_support::storage::unhashed::kill(&frame_support::storage::storage_prefix(
+			b"ParachainStaking",
+			b"ParachainBondInfo",
+		));
+
+		Weight::default()
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade(&self) -> Result<Vec<u8>, sp_runtime::DispatchError> {
+		let state = frame_support::storage::migration::get_storage_value::<
+			OldParachainBondConfig<T::AccountId>,
+		>(b"ParachainStaking", b"ParachainBondInfo", &[]);
+
+		ensure!(state.is_some(), "State not found");
+
+		Ok(state.unwrap().encode())
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(&self, state: Vec<u8>) -> Result<(), sp_runtime::DispatchError> {
+		let old_state: OldParachainBondConfig<T::AccountId> =
+			parity_scale_codec::Decode::decode(&mut &state[..])
+				.map_err(|_| sp_runtime::DispatchError::Other("Failed to decode old state"))?;
+
+		let new_state = InflationDistributionInfo::<T>::get();
+
+		let pbr = InflationDistributionAccount {
+			account: old_state.account,
+			percent: old_state.percent,
+		};
+		let treasury = InflationDistributionAccount::<T::AccountId>::default();
+		let expected_new_state: InflationDistributionConfig<T::AccountId> = [pbr, treasury].into();
+
+		ensure!(new_state == expected_new_state, "State migration failed");
+	}
+}
diff --git a/runtime/common/src/migrations.rs b/runtime/common/src/migrations.rs
index ef94f52b87..44f88b6f49 100644
--- a/runtime/common/src/migrations.rs
+++ b/runtime/common/src/migrations.rs
@@ -154,6 +154,32 @@ where
 	}
 }
 
+pub struct MigrateStakingParachainBondConfig<Runtime>(PhantomData<Runtime>);
+impl<Runtime> Migration for MigrateStakingParachainBondConfig<Runtime>
+where
+	Runtime: pallet_parachain_staking::Config,
+{
+	fn friendly_name(&self) -> &str {
+		"MM_MigrateStakingParachainBondConfig"
+	}
+
+	fn migrate(&self, _available_weight: Weight) -> Weight {
+		pallet_parachain_staking::migrations::MigrateParachainBondConfig::<Runtime>::on_runtime_upgrade()
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade(&self) -> Result<Vec<u8>, sp_runtime::DispatchError> {
+		pallet_parachain_staking::migrations::MigrateParachainBondConfig::<Runtime>::pre_upgrade()
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(&self, state: Vec<u8>) -> Result<(), sp_runtime::DispatchError> {
+		pallet_parachain_staking::migrations::MigrateParachainBondConfig::<Runtime>::post_upgrade(
+			state,
+		)
+	}
+}
+
 pub struct CommonMigrations<Runtime>(PhantomData<Runtime>);
 
 impl<Runtime> GetMigrations for CommonMigrations<Runtime>
@@ -303,6 +329,9 @@ where
 			Box::new(MigrateXcmFeesAssetsMeatdata::<Runtime>(Default::default())),
 			// permanent migrations
 			Box::new(MigrateToLatestXcmVersion::<Runtime>(Default::default())),
+			Box::new(MigrateStakingParachainBondConfig::<Runtime>(
+				Default::default(),
+			)),
 		]
 	}
 }