From e80f794dffeea3f63c3a792f505a805204069c53 Mon Sep 17 00:00:00 2001
From: akildemir <aeg_asd@hotmail.com>
Date: Sat, 18 Nov 2023 10:53:28 +0300
Subject: [PATCH 1/2] add gene_session & new_session public functions

---
 frame/grandpa/src/lib.rs | 49 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/frame/grandpa/src/lib.rs b/frame/grandpa/src/lib.rs
index 64f0490f6f6c8..f6fddccda422b 100644
--- a/frame/grandpa/src/lib.rs
+++ b/frame/grandpa/src/lib.rs
@@ -507,6 +507,55 @@ impl<T: Config> Pallet<T> {
 		}
 	}
 
+	/// This function is added later to be used in place of overly specific trait
+ 	/// implementation of OneSessionHandler::on_genesis_session.
+ 	pub fn genesis_session(authorities: AuthorityList) {
+		Self::initialize(&authorities);
+	}
+
+	/// This function is added later to be used in place of overly specific trait
+	/// implementation of OneSessionHandler::on_new_session.
+	pub fn new_session(changed: bool, session_index: u32, next_authorities: AuthorityList) {
+		// Always issue a change if `session` says that the validators have changed.
+		// Even if their session keys are the same as before, the underlying economic
+		// identities have changed.
+		let current_set_id = if changed || <Stalled<T>>::exists() {
+
+			let res = if let Some((further_wait, median)) = <Stalled<T>>::take() {
+				Self::schedule_change(next_authorities, further_wait, Some(median))
+			} else {
+				Self::schedule_change(next_authorities, Zero::zero(), None)
+			};
+
+			if res.is_ok() {
+				let current_set_id = CurrentSetId::<T>::mutate(|s| {
+					*s += 1;
+					*s
+				});
+
+				let max_set_id_session_entries = T::MaxSetIdSessionEntries::get().max(1);
+				if current_set_id >= max_set_id_session_entries {
+					SetIdSession::<T>::remove(current_set_id - max_set_id_session_entries);
+				}
+
+				current_set_id
+			} else {
+				// either the session module signalled that the validators have changed
+				// or the set was stalled. but since we didn't successfully schedule
+				// an authority set change we do not increment the set id.
+				Self::current_set_id()
+			}
+		} else {
+			// nothing's changed, neither economic conditions nor session keys. update the pointer
+			// of the current set.
+			Self::current_set_id()
+		};
+
+		// update the mapping to note that the current set corresponds to the
+		// latest equivalent session (i.e. now).
+		SetIdSession::<T>::insert(current_set_id, &session_index);
+	}
+
 	/// Deposit one of this module's logs.
 	fn deposit_log(log: ConsensusLog<BlockNumberFor<T>>) {
 		let log = DigestItem::Consensus(GRANDPA_ENGINE_ID, log.encode());

From a7910c652a5056d4fe3ef1380a58abcee9d7d561 Mon Sep 17 00:00:00 2001
From: akildemir <aeg_asd@hotmail.com>
Date: Sat, 18 Nov 2023 21:51:57 +0300
Subject: [PATCH 2/2] use new functions from within SessionHandler

---
 frame/grandpa/src/lib.rs | 43 +++-------------------------------------
 1 file changed, 3 insertions(+), 40 deletions(-)

diff --git a/frame/grandpa/src/lib.rs b/frame/grandpa/src/lib.rs
index f6fddccda422b..a9f0ccfeb7482 100644
--- a/frame/grandpa/src/lib.rs
+++ b/frame/grandpa/src/lib.rs
@@ -610,53 +610,16 @@ where
 		I: Iterator<Item = (&'a T::AccountId, AuthorityId)>,
 	{
 		let authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
-		Self::initialize(&authorities);
+		Self::genesis_session(authorities);
 	}
 
 	fn on_new_session<'a, I: 'a>(changed: bool, validators: I, _queued_validators: I)
 	where
 		I: Iterator<Item = (&'a T::AccountId, AuthorityId)>,
 	{
-		// Always issue a change if `session` says that the validators have changed.
-		// Even if their session keys are the same as before, the underlying economic
-		// identities have changed.
-		let current_set_id = if changed || <Stalled<T>>::exists() {
-			let next_authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
-
-			let res = if let Some((further_wait, median)) = <Stalled<T>>::take() {
-				Self::schedule_change(next_authorities, further_wait, Some(median))
-			} else {
-				Self::schedule_change(next_authorities, Zero::zero(), None)
-			};
-
-			if res.is_ok() {
-				let current_set_id = CurrentSetId::<T>::mutate(|s| {
-					*s += 1;
-					*s
-				});
-
-				let max_set_id_session_entries = T::MaxSetIdSessionEntries::get().max(1);
-				if current_set_id >= max_set_id_session_entries {
-					SetIdSession::<T>::remove(current_set_id - max_set_id_session_entries);
-				}
-
-				current_set_id
-			} else {
-				// either the session module signalled that the validators have changed
-				// or the set was stalled. but since we didn't successfully schedule
-				// an authority set change we do not increment the set id.
-				Self::current_set_id()
-			}
-		} else {
-			// nothing's changed, neither economic conditions nor session keys. update the pointer
-			// of the current set.
-			Self::current_set_id()
-		};
-
-		// update the mapping to note that the current set corresponds to the
-		// latest equivalent session (i.e. now).
+		let authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
 		let session_index = <pallet_session::Pallet<T>>::current_index();
-		SetIdSession::<T>::insert(current_set_id, &session_index);
+		Self::new_session(changed, session_index, authorities);
 	}
 
 	fn on_disabled(i: u32) {