From c927735d8f43f05ba67cc1e1c4cbb3cb5eb386b5 Mon Sep 17 00:00:00 2001
From: pgherveou <pgherveou@gmail.com>
Date: Mon, 19 Jun 2023 20:55:17 +0200
Subject: [PATCH] Contracts: Better migration types

---
 frame/contracts/src/migration/v10.rs | 18 ++++++++++--------
 frame/contracts/src/migration/v9.rs  | 15 ++++++---------
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/frame/contracts/src/migration/v10.rs b/frame/contracts/src/migration/v10.rs
index d58a9479f2fb9..75d794a6d8167 100644
--- a/frame/contracts/src/migration/v10.rs
+++ b/frame/contracts/src/migration/v10.rs
@@ -42,7 +42,7 @@ use sp_core::hexdisplay::HexDisplay;
 #[cfg(feature = "try-runtime")]
 use sp_runtime::TryRuntimeError;
 use sp_runtime::{traits::Zero, Perbill, Saturating};
-use sp_std::{marker::PhantomData, ops::Deref, prelude::*};
+use sp_std::{ops::Deref, prelude::*};
 
 mod old {
 	use super::*;
@@ -109,8 +109,7 @@ pub struct ContractInfo<T: Config> {
 
 #[derive(Encode, Decode, MaxEncodedLen, DefaultNoBound)]
 pub struct Migration<T: Config> {
-	last_key: Option<BoundedVec<u8, ConstU32<256>>>,
-	_phantom: PhantomData<T>,
+	last_account: Option<T::AccountId>,
 }
 
 #[storage_alias]
@@ -125,8 +124,10 @@ impl<T: Config> MigrationStep for Migration<T> {
 	}
 
 	fn step(&mut self) -> (IsFinished, Weight) {
-		let mut iter = if let Some(last_key) = self.last_key.take() {
-			old::ContractInfoOf::<T>::iter_from(last_key.to_vec())
+		let mut iter = if let Some(last_account) = self.last_account.take() {
+			old::ContractInfoOf::<T>::iter_from(old::ContractInfoOf::<T>::hashed_key_for(
+				last_account,
+			))
 		} else {
 			old::ContractInfoOf::<T>::iter()
 		};
@@ -135,9 +136,6 @@ impl<T: Config> MigrationStep for Migration<T> {
 			let min_balance = Pallet::<T>::min_balance();
 			log::debug!(target: LOG_TARGET, "Account: 0x{} ", HexDisplay::from(&account.encode()));
 
-			// Store last key for next migration step
-			self.last_key = Some(iter.last_raw_key().to_vec().try_into().unwrap());
-
 			// Get the new deposit account address
 			let deposit_account: DepositAccount<T> =
 				DepositAccount(T::AddressGenerator::deposit_address(&account));
@@ -223,6 +221,10 @@ impl<T: Config> MigrationStep for Migration<T> {
 			};
 
 			ContractInfoOf::<T>::insert(&account, new_contract_info);
+
+			// Store last key for next migration step
+			self.last_account = Some(account);
+
 			(IsFinished::No, T::WeightInfo::v10_migration_step())
 		} else {
 			log::debug!(target: LOG_TARGET, "Done Migrating contract info");
diff --git a/frame/contracts/src/migration/v9.rs b/frame/contracts/src/migration/v9.rs
index d50f79c72c3be..e6c6642955642 100644
--- a/frame/contracts/src/migration/v9.rs
+++ b/frame/contracts/src/migration/v9.rs
@@ -23,12 +23,10 @@ use crate::{
 	CodeHash, Config, Determinism, Pallet, Weight, LOG_TARGET,
 };
 use codec::{Decode, Encode};
-use frame_support::{
-	codec, pallet_prelude::*, storage_alias, BoundedVec, DefaultNoBound, Identity,
-};
+use frame_support::{codec, pallet_prelude::*, storage_alias, DefaultNoBound, Identity};
 #[cfg(feature = "try-runtime")]
 use sp_runtime::TryRuntimeError;
-use sp_std::{marker::PhantomData, prelude::*};
+use sp_std::prelude::*;
 
 mod old {
 	use super::*;
@@ -79,8 +77,7 @@ type CodeStorage<T: Config> = StorageMap<Pallet<T>, Identity, CodeHash<T>, Prefa
 
 #[derive(Encode, Decode, MaxEncodedLen, DefaultNoBound)]
 pub struct Migration<T: Config> {
-	last_key: Option<BoundedVec<u8, ConstU32<256>>>,
-	_phantom: PhantomData<T>,
+	last_code_hash: Option<CodeHash<T>>,
 }
 
 impl<T: Config> MigrationStep for Migration<T> {
@@ -91,8 +88,8 @@ impl<T: Config> MigrationStep for Migration<T> {
 	}
 
 	fn step(&mut self) -> (IsFinished, Weight) {
-		let mut iter = if let Some(last_key) = self.last_key.take() {
-			old::CodeStorage::<T>::iter_from(last_key.to_vec())
+		let mut iter = if let Some(last_key) = self.last_code_hash.take() {
+			old::CodeStorage::<T>::iter_from(old::CodeStorage::<T>::hashed_key_for(last_key))
 		} else {
 			old::CodeStorage::<T>::iter()
 		};
@@ -108,7 +105,7 @@ impl<T: Config> MigrationStep for Migration<T> {
 				determinism: Determinism::Enforced,
 			};
 			CodeStorage::<T>::insert(key, module);
-			self.last_key = Some(iter.last_raw_key().to_vec().try_into().unwrap());
+			self.last_code_hash = Some(key);
 			(IsFinished::No, T::WeightInfo::v9_migration_step(len))
 		} else {
 			log::debug!(target: LOG_TARGET, "No more contracts code to migrate");