Repository: paritytech/substrate
File: frame/democracy/src/migrations/unlock_and_unreserve_all_funds.rs (Rust)

Code before:
	/// Steps:
	/// 1. Retrieves the deposit and accounts with locks for the pallet.
	/// 2. Unreserves the deposited funds for each account.
	/// 3. Unlocks the staked funds for each account.
	fn on_runtime_upgrade() -> frame_support::weights::Weight {
		// Get staked and deposited balances as reported by this pallet.
		let (account_deposits, account_stakes, initial_reads) =
			Self::get_account_deposits_and_locks();

		// Deposited funds need to be unreserved.
		for (account, unreserve_amount) in account_deposits.iter() {
			if unreserve_amount.is_zero() {
				continue
			}
			T::Currency::unreserve(&account, *unreserve_amount);
		}

		// Staked funds need to be unlocked.
		for account in account_stakes.keys() {
			T::Currency::remove_lock(DEMOCRACY_ID, account);
		}

		T::DbWeight::get().reads_writes(
			account_stakes.len().saturating_add(account_deposits.len()) as u64,
			account_stakes.len().saturating_add(account_deposits.len()) as u64,
		) + initial_reads
	}

	/// Performs post-upgrade sanity checks:
	///
	/// 1. No locks remain for this pallet in Balances.
	/// 2. The reserved balance for each account has been reduced by the expected amount.
	#[cfg(feature = "try-runtime")]
	fn post_upgrade(
		account_reserved_before_bytes: Vec<u8>,
	) -> Result<(), sp_runtime::TryRuntimeError> {
		use codec::Decode;

		let account_reserved_before =
			BTreeMap::<T::AccountId, BalanceOf<T>>::decode(&mut &account_reserved_before_bytes[..])
				.map_err(|_| "Failed to decode account_reserved_before_bytes")?;

		// Get staked and deposited balances as reported by this pallet.
		let (account_deposits, _, _) = Self::get_account_deposits_and_locks();

		// Check that the reserved balance is reduced by the expected deposited amount.
		for (account, actual_reserved_before) in account_reserved_before {
			let actual_reserved_after = T::Currency::reserved_balance(&account);
			let expected_amount_deducted = *account_deposits
				.get(&account)
				.expect("account deposit must exist to be in pre_migration_data, qed");

Review comment (suggestion):
```suggestion
		).saturating_add(initial_reads)
```

Code after:
	/// 1. Retrieves the deposit and accounts with locks for the pallet.
	/// 2. Unreserves the deposited funds for each account.
	/// 3. Unlocks the staked funds for each account.
	fn on_runtime_upgrade() -> frame_support::weights::Weight {
		// Get staked and deposited balances as reported by this pallet.
		let (account_deposits, account_stakes, initial_reads) =
			Self::get_account_deposits_and_locks();

		// Deposited funds need to be unreserved.
		for (account, unreserve_amount) in account_deposits.iter() {
			if unreserve_amount.is_zero() {
				log::warn!(target: LOG_TARGET, "Unexpected zero amount to unreserve!");
				continue
			}
			T::Currency::unreserve(&account, *unreserve_amount);
		}

		// Staked funds need to be unlocked.
		for account in account_stakes.keys() {
			T::Currency::remove_lock(DEMOCRACY_ID, account);
		}

		T::DbWeight::get()
			.reads_writes(
				account_stakes.len().saturating_add(account_deposits.len()) as u64,
				account_stakes.len().saturating_add(account_deposits.len()) as u64,
			)
			.saturating_add(initial_reads)
	}

	/// Performs post-upgrade sanity checks:
	///
	/// 1. No locks remain for this pallet in Balances.
	/// 2. The reserved balance for each account has been reduced by the expected amount.
	#[cfg(feature = "try-runtime")]
	fn post_upgrade(
		account_reserved_before_bytes: Vec<u8>,
	) -> Result<(), sp_runtime::TryRuntimeError> {
		use codec::Decode;

		let account_reserved_before =
			BTreeMap::<T::AccountId, BalanceOf<T>>::decode(&mut &account_reserved_before_bytes[..])
				.map_err(|_| "Failed to decode account_reserved_before_bytes")?;

		// Get staked and deposited balances as reported by this pallet.
		let (account_deposits, _, _) = Self::get_account_deposits_and_locks();

		// Check that the reserved balance is reduced by the expected deposited amount.
		for (account, actual_reserved_before) in account_reserved_before {
			let actual_reserved_after = T::Currency::reserved_balance(&account);
			let expected_amount_deducted = *account_deposits