From 1520623b906033cbd89a58136b488cdfd159ce8e Mon Sep 17 00:00:00 2001
From: Gavin Wood <gavin@parity.io>
Date: Sat, 3 Dec 2022 16:36:25 +0000
Subject: [PATCH] Tweak to active total migrations (#12832)

* Tweak to active total migrations

* Formatting

* Expose trait

* Remove empty pre_ and post_upgrade hooks.

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 substrate/frame/balances/src/migration.rs | 35 ++++++++++++++++-------
 substrate/frame/support/src/traits.rs     |  4 +--
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/substrate/frame/balances/src/migration.rs b/substrate/frame/balances/src/migration.rs
index e27efc2174..1dd3e6f51f 100644
--- a/substrate/frame/balances/src/migration.rs
+++ b/substrate/frame/balances/src/migration.rs
@@ -19,9 +19,7 @@ use frame_support::{pallet_prelude::*, traits::OnRuntimeUpgrade, weights::Weight
 
 // NOTE: This must be used alongside the account whose balance is expected to be inactive.
 // Generally this will be used for the XCM teleport checking account.
-pub struct MigrateToTrackInactive<T: Config, A: Get<T::AccountId>>(
-	sp_std::marker::PhantomData<(T, A)>,
-);
+pub struct MigrateToTrackInactive<T, A>(PhantomData<(T, A)>);
 impl<T: Config, A: Get<T::AccountId>> OnRuntimeUpgrade for MigrateToTrackInactive<T, A> {
 	fn on_runtime_upgrade() -> Weight {
 		let current_version = Pallet::<T>::current_storage_version();
@@ -32,20 +30,35 @@ impl<T: Config, A: Get<T::AccountId>> OnRuntimeUpgrade for MigrateToTrackInactiv
 			Pallet::<T>::deactivate(b);
 			current_version.put::<Pallet<T>>();
 			log::info!(target: "runtime::balances", "Storage to version {:?}", current_version);
-			T::DbWeight::get().reads_writes(3, 3)
+			T::DbWeight::get().reads_writes(4, 3)
 		} else {
 			log::info!(target: "runtime::balances",  "Migration did not execute. This probably should be removed");
 			T::DbWeight::get().reads(2)
 		}
 	}
+}
 
-	#[cfg(feature = "try-runtime")]
-	fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
-		Ok(vec![])
-	}
+// NOTE: This must be used alongside the account whose balance is expected to be inactive.
+// Generally this will be used for the XCM teleport checking account.
+pub struct MigrateManyToTrackInactive<T, A>(PhantomData<(T, A)>);
+impl<T: Config, A: Get<Vec<T::AccountId>>> OnRuntimeUpgrade for MigrateManyToTrackInactive<T, A> {
+	fn on_runtime_upgrade() -> Weight {
+		let current_version = Pallet::<T>::current_storage_version();
+		let onchain_version = Pallet::<T>::on_chain_storage_version();
 
-	#[cfg(feature = "try-runtime")]
-	fn post_upgrade(total: Vec<u8>) -> Result<(), &'static str> {
-		Ok(())
+		if onchain_version == 0 && current_version == 1 {
+			let accounts = A::get();
+			let total = accounts
+				.iter()
+				.map(|a| Pallet::<T>::total_balance(a))
+				.fold(T::Balance::zero(), |a, e| a.saturating_add(e));
+			Pallet::<T>::deactivate(total);
+			current_version.put::<Pallet<T>>();
+			log::info!(target: "runtime::balances", "Storage to version {:?}", current_version);
+			T::DbWeight::get().reads_writes(3 + accounts.len() as u64, 3)
+		} else {
+			log::info!(target: "runtime::balances",  "Migration did not execute. This probably should be removed");
+			T::DbWeight::get().reads(2)
+		}
 	}
 }
diff --git a/substrate/frame/support/src/traits.rs b/substrate/frame/support/src/traits.rs
index f09b715a97..9b5300dfc5 100644
--- a/substrate/frame/support/src/traits.rs
+++ b/substrate/frame/support/src/traits.rs
@@ -22,8 +22,8 @@
 pub mod tokens;
 pub use tokens::{
 	currency::{
-		Currency, LockIdentifier, LockableCurrency, NamedReservableCurrency, ReservableCurrency,
-		TotalIssuanceOf, VestingSchedule,
+		ActiveIssuanceOf, Currency, LockIdentifier, LockableCurrency, NamedReservableCurrency,
+		ReservableCurrency, TotalIssuanceOf, VestingSchedule,
 	},
 	fungible, fungibles,
 	imbalance::{Imbalance, OnUnbalanced, SignedImbalance},