From ace0f6e9e12c4b614f7dce95dfc7b1aeb34b9bf9 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Tue, 6 Jun 2023 15:56:02 +0300
Subject: [PATCH] remove unnecessary types from v0 mod

---
 frame/im-online/src/lib.rs       |  9 +++---
 frame/im-online/src/migration.rs | 47 +++++++++++++++-----------------
 2 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/frame/im-online/src/lib.rs b/frame/im-online/src/lib.rs
index 0c357adc57aec..198c428c53b98 100644
--- a/frame/im-online/src/lib.rs
+++ b/frame/im-online/src/lib.rs
@@ -84,13 +84,17 @@ pub mod weights;
 
 use codec::{Decode, Encode, MaxEncodedLen};
 use frame_support::{
+	pallet_prelude::*,
 	traits::{
 		EstimateNextSessionRotation, Get, OneSessionHandler, ValidatorSet,
 		ValidatorSetWithIdentification,
 	},
 	BoundedSlice, WeakBoundedVec,
 };
-use frame_system::offchain::{SendTransactionTypes, SubmitTransaction};
+use frame_system::{
+	offchain::{SendTransactionTypes, SubmitTransaction},
+	pallet_prelude::*,
+};
 pub use pallet::*;
 use scale_info::TypeInfo;
 use sp_application_crypto::RuntimeAppPublic;
@@ -106,9 +110,6 @@ use sp_staking::{
 use sp_std::prelude::*;
 pub use weights::WeightInfo;
 
-use frame_support::pallet_prelude::*;
-use frame_system::pallet_prelude::*;
-
 pub mod sr25519 {
 	mod app_sr25519 {
 		use sp_application_crypto::{app_crypto, key_types::IM_ONLINE, sr25519};
diff --git a/frame/im-online/src/migration.rs b/frame/im-online/src/migration.rs
index 94ab7841b1c3d..db6461c618764 100644
--- a/frame/im-online/src/migration.rs
+++ b/frame/im-online/src/migration.rs
@@ -28,7 +28,7 @@ use sp_runtime::TryRuntimeError;
 /// The log target.
 const TARGET: &'static str = "runtime::im-online::migration::v1";
 
-/// The original data layout of the im-online pallet without a specific version number.
+/// The original data layout of the im-online pallet (`ReceivedHeartbeats` storage item).
 mod v0 {
 	use super::*;
 	use frame_support::traits::WrapperOpaque;
@@ -55,13 +55,6 @@ mod v0 {
 			WeakBoundedVec<WeakBoundedVec<u8, MultiAddrEncodingLimit>, AddressesLimit>,
 	}
 
-	#[storage_alias]
-	pub(crate) type HeartbeatAfter<T: Config> = StorageValue<Pallet<T>, T::BlockNumber, ValueQuery>;
-
-	#[storage_alias]
-	pub(crate) type Keys<T: Config> =
-		StorageValue<Pallet<T>, WeakBoundedVec<T::AuthorityId, T::MaxKeys>, ValueQuery>;
-
 	#[storage_alias]
 	pub(crate) type ReceivedHeartbeats<T: Config> = StorageDoubleMap<
 		Pallet<T>,
@@ -71,23 +64,12 @@ mod v0 {
 		AuthIndex,
 		WrapperOpaque<BoundedOpaqueNetworkState<u32, u32, T::MaxPeerInHeartbeats>>,
 	>;
-
-	#[storage_alias]
-	pub(crate) type AuthoredBlocks<T: pallet::Config> = StorageDoubleMap<
-		Pallet<T>,
-		Twox64Concat,
-		SessionIndex,
-		Twox64Concat,
-		ValidatorId<T>,
-		u32,
-		ValueQuery,
-	>;
 }
 
 pub mod v1 {
 	use super::*;
 
-	/// Migration for moving im-online from V0 to V1 storage.
+	/// Simple migration that replaces `ReceivedHeartbeats` values with `()`.
 	pub struct Migration<T>(sp_std::marker::PhantomData<T>);
 
 	impl<T: Config> OnRuntimeUpgrade for Migration<T> {
@@ -105,17 +87,32 @@ pub mod v1 {
 		}
 
 		fn on_runtime_upgrade() -> Weight {
-			ReceivedHeartbeats::<T>::translate::<_, _>(
+			let mut weight = T::DbWeight::get().reads(1);
+			if StorageVersion::get::<Pallet<T>>() != 0 {
+				log::warn!(
+					target: TARGET,
+					"Skipping migration because current storage version is not 0"
+				);
+				return weight
+			}
+
+			let count = v0::ReceivedHeartbeats::<T>::iter().count();
+			weight.saturating_accrue(
+				T::DbWeight::get().reads(v0::ReceivedHeartbeats::<T>::iter().count() as u64),
+			);
+			weight.saturating_accrue(
+				T::DbWeight::get().writes(v0::ReceivedHeartbeats::<T>::iter().count() as u64),
+			);
+
+			v0::ReceivedHeartbeats::<T>::translate::<_, _>(
 				|k: T::SessionIndex, T::AccountId, state: _| {
-					log::info!(target: TARGET, "Migrated received heartbeat for {:?}...", k);
+					log::trace!(target: TARGET, "Migrated received heartbeat for {:?}...", k);
 					Some(())
 				},
 			);
 
 			StorageVersion::new(1).put::<Pallet<T>>();
-
-			let count = ReceivedHeartbeats::<T>::iter().count();
-			T::DbWeight::get().reads_writes(count as Weight + 1, count as Weight + 1)
+			weight.saturating_add(T::DbWeight::get().writes(1))
 		}
 
 		#[cfg(feature = "try-runtime")]