From a2f073e643f014a9c4c0eb50042070ba2100091d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Sun, 19 Nov 2023 21:34:43 +0100
Subject: [PATCH 01/16] Deprecate RewardDestination default

---
 substrate/frame/staking/src/ledger.rs       |  4 ++--
 substrate/frame/staking/src/lib.rs          |  6 ------
 substrate/frame/staking/src/pallet/impls.rs | 15 +++++++++-----
 substrate/frame/staking/src/pallet/mod.rs   |  2 +-
 substrate/frame/staking/src/tests.rs        | 22 ++++++++++-----------
 5 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/substrate/frame/staking/src/ledger.rs b/substrate/frame/staking/src/ledger.rs
index 84bb4d167dcb0..5947adb9028b3 100644
--- a/substrate/frame/staking/src/ledger.rs
+++ b/substrate/frame/staking/src/ledger.rs
@@ -126,7 +126,7 @@ impl<T: Config> StakingLedger<T> {
 	/// default reward destination.
 	pub(crate) fn reward_destination(
 		account: StakingAccount<T::AccountId>,
-	) -> RewardDestination<T::AccountId> {
+	) -> Option<RewardDestination<T::AccountId>> {
 		let stash = match account {
 			StakingAccount::Stash(stash) => Some(stash),
 			StakingAccount::Controller(controller) =>
@@ -137,7 +137,7 @@ impl<T: Config> StakingLedger<T> {
 			<Payee<T>>::get(stash)
 		} else {
 			defensive!("fetched reward destination from unbonded stash {}", stash);
-			RewardDestination::default()
+			None
 		}
 	}
 
diff --git a/substrate/frame/staking/src/lib.rs b/substrate/frame/staking/src/lib.rs
index 9e4697e845b61..ed3f335b8c3c5 100644
--- a/substrate/frame/staking/src/lib.rs
+++ b/substrate/frame/staking/src/lib.rs
@@ -412,12 +412,6 @@ pub enum RewardDestination<AccountId> {
 	None,
 }
 
-impl<AccountId> Default for RewardDestination<AccountId> {
-	fn default() -> Self {
-		RewardDestination::Staked
-	}
-}
-
 /// Preference of what happens regarding validation.
 #[derive(PartialEq, Eq, Clone, Encode, Decode, RuntimeDebug, TypeInfo, Default, MaxEncodedLen)]
 pub struct ValidatorPrefs {
diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index 40f30735258f6..b732a360f64d9 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -75,7 +75,7 @@ impl<T: Config> Pallet<T> {
 		StakingLedger::<T>::get(account)
 	}
 
-	pub fn payee(account: StakingAccount<T::AccountId>) -> RewardDestination<T::AccountId> {
+	pub fn payee(account: StakingAccount<T::AccountId>) -> Option<RewardDestination<T::AccountId>> {
 		StakingLedger::<T>::reward_destination(account)
 	}
 
@@ -336,8 +336,14 @@ impl<T: Config> Pallet<T> {
 			return None
 		}
 
-		let dest = Self::payee(StakingAccount::Stash(stash.clone()));
-		let maybe_imbalance = match dest {
+		let dest = match Self::payee(StakingAccount::Stash(stash.clone()))
+			.defensive_proof("payee must exist for a ledger requesting a payout")
+		{
+			Some(dest) => dest,
+			None => return None,
+		};
+
+		let maybe_imbalance = match dest.clone() {
 			RewardDestination::Controller => Self::bonded(stash)
 				.map(|controller| T::Currency::deposit_creating(&controller, amount)),
 			RewardDestination::Stash => T::Currency::deposit_into_existing(stash, amount).ok(),
@@ -358,8 +364,7 @@ impl<T: Config> Pallet<T> {
 				Some(T::Currency::deposit_creating(&dest_account, amount)),
 			RewardDestination::None => None,
 		};
-		maybe_imbalance
-			.map(|imbalance| (imbalance, Self::payee(StakingAccount::Stash(stash.clone()))))
+		maybe_imbalance.map(|imbalance| (imbalance, dest))
 	}
 
 	/// Plan a new session potentially trigger a new era.
diff --git a/substrate/frame/staking/src/pallet/mod.rs b/substrate/frame/staking/src/pallet/mod.rs
index 18ad3e4a6cf1d..6cbd3b268e5f3 100644
--- a/substrate/frame/staking/src/pallet/mod.rs
+++ b/substrate/frame/staking/src/pallet/mod.rs
@@ -336,7 +336,7 @@ pub mod pallet {
 	/// TWOX-NOTE: SAFE since `AccountId` is a secure hash.
 	#[pallet::storage]
 	pub type Payee<T: Config> =
-		StorageMap<_, Twox64Concat, T::AccountId, RewardDestination<T::AccountId>, ValueQuery>;
+		StorageMap<_, Twox64Concat, T::AccountId, RewardDestination<T::AccountId>, OptionQuery>;
 
 	/// The map from (wannabe) validator stash key to the preferences of that validator.
 	///
diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index bac2530b19bb4..ea6ad87f74cfc 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -774,7 +774,7 @@ fn double_staking_should_fail() {
 		let (stash, controller) = testing_utils::create_unique_stash_controller::<Test>(
 			0,
 			arbitrary_value,
-			RewardDestination::default(),
+			RewardDestination::Staked,
 			false,
 		)
 		.unwrap();
@@ -784,7 +784,7 @@ fn double_staking_should_fail() {
 			Staking::bond(
 				RuntimeOrigin::signed(stash),
 				arbitrary_value.into(),
-				RewardDestination::default()
+				RewardDestination::Staked,
 			),
 			Error::<Test>::AlreadyBonded,
 		);
@@ -808,7 +808,7 @@ fn double_controlling_attempt_should_fail() {
 		let (stash, _) = testing_utils::create_unique_stash_controller::<Test>(
 			0,
 			arbitrary_value,
-			RewardDestination::default(),
+			RewardDestination::Staked,
 			false,
 		)
 		.unwrap();
@@ -818,7 +818,7 @@ fn double_controlling_attempt_should_fail() {
 			Staking::bond(
 				RuntimeOrigin::signed(stash),
 				arbitrary_value.into(),
-				RewardDestination::default()
+				RewardDestination::Staked,
 			),
 			Error::<Test>::AlreadyBonded,
 		);
@@ -1066,8 +1066,8 @@ fn reward_destination_works() {
 		mock::start_active_era(1);
 		mock::make_all_reward_payment(0);
 
-		// Check that RewardDestination is Staked (default)
-		assert_eq!(Staking::payee(11.into()), RewardDestination::Staked);
+		// Check that RewardDestination is Staked
+		assert_eq!(Staking::payee(11.into()), Some(RewardDestination::Staked));
 		// Check that reward went to the stash account of validator
 		assert_eq!(Balances::free_balance(11), 1000 + total_payout_0);
 		// Check that amount at stake increased accordingly
@@ -1096,7 +1096,7 @@ fn reward_destination_works() {
 		mock::make_all_reward_payment(1);
 
 		// Check that RewardDestination is Stash
-		assert_eq!(Staking::payee(11.into()), RewardDestination::Stash);
+		assert_eq!(Staking::payee(11.into()), Some(RewardDestination::Stash));
 		// Check that reward went to the stash account
 		assert_eq!(Balances::free_balance(11), 1000 + total_payout_0 + total_payout_1);
 		// Record this value
@@ -1130,7 +1130,7 @@ fn reward_destination_works() {
 		mock::make_all_reward_payment(2);
 
 		// Check that RewardDestination is Controller
-		assert_eq!(Staking::payee(11.into()), RewardDestination::Controller);
+		assert_eq!(Staking::payee(11.into()), Some(RewardDestination::Controller));
 		// Check that reward went to the controller account
 		assert_eq!(Balances::free_balance(11), recorded_stash_balance + total_payout_2);
 		// Check that amount at stake is NOT increased
@@ -2204,7 +2204,7 @@ fn reward_validator_slashing_validator_does_not_overflow() {
 		let _ = Balances::make_free_balance_be(&2, stake);
 
 		// only slashes out of bonded stake are applied. without this line, it is 0.
-		Staking::bond(RuntimeOrigin::signed(2), stake - 1, RewardDestination::default()).unwrap();
+		Staking::bond(RuntimeOrigin::signed(2), stake - 1, RewardDestination::Staked).unwrap();
 		// Override exposure of 11
 		EraInfo::<Test>::set_exposure(
 			0,
@@ -5054,7 +5054,7 @@ mod election_data_provider {
 			assert_eq!(MinNominatorBond::<Test>::get(), 1);
 			assert_eq!(<Test as Config>::VoterList::count(), 4);
 
-			assert_ok!(Staking::bond(RuntimeOrigin::signed(4), 5, Default::default(),));
+			assert_ok!(Staking::bond(RuntimeOrigin::signed(4), 5, RewardDestination::Staked,));
 			assert_ok!(Staking::nominate(RuntimeOrigin::signed(4), vec![1]));
 			assert_eq!(<Test as Config>::VoterList::count(), 5);
 
@@ -6868,7 +6868,7 @@ mod ledger {
 			assert_ok!(ledger.clone().bond(reward_dest));
 			assert!(StakingLedger::<Test>::is_bonded(StakingAccount::Stash(42)));
 			assert!(<Bonded<Test>>::get(&42).is_some());
-			assert_eq!(<Payee<Test>>::get(&42), reward_dest);
+			assert_eq!(<Payee<Test>>::get(&42), Some(reward_dest));
 
 			// cannot bond again.
 			assert!(ledger.clone().bond(reward_dest).is_err());

From eb005955f29a0d1eededfbdff60796c914850cca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Sun, 19 Nov 2023 21:54:09 +0100
Subject: [PATCH 02/16] Fixes nomination-pools tests

---
 substrate/frame/nomination-pools/test-staking/src/lib.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/substrate/frame/nomination-pools/test-staking/src/lib.rs b/substrate/frame/nomination-pools/test-staking/src/lib.rs
index 9108da510a3a5..865b7a71e6884 100644
--- a/substrate/frame/nomination-pools/test-staking/src/lib.rs
+++ b/substrate/frame/nomination-pools/test-staking/src/lib.rs
@@ -214,7 +214,10 @@ fn pool_slash_e2e() {
 			]
 		);
 
-		assert_eq!(Payee::<Runtime>::get(POOL1_BONDED), RewardDestination::Account(POOL1_REWARD));
+		assert_eq!(
+			Payee::<Runtime>::get(POOL1_BONDED),
+			Some(RewardDestination::Account(POOL1_REWARD))
+		);
 
 		// have two members join
 		assert_ok!(Pools::join(RuntimeOrigin::signed(20), 20, 1));

From 34b6cbf7e9518bb47cd21485549563eba8ebcdc9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Sun, 19 Nov 2023 23:27:47 +0100
Subject: [PATCH 03/16] Fixes staking benchmarks

---
 substrate/frame/staking/src/benchmarking.rs | 26 ++++++++++-----------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/substrate/frame/staking/src/benchmarking.rs b/substrate/frame/staking/src/benchmarking.rs
index 05c6bc3970977..1ca023a8a4b7d 100644
--- a/substrate/frame/staking/src/benchmarking.rs
+++ b/substrate/frame/staking/src/benchmarking.rs
@@ -175,7 +175,7 @@ impl<T: Config> ListScenario<T> {
 		let (origin_stash1, origin_controller1) = create_stash_controller_with_balance::<T>(
 			USER_SEED + 2,
 			origin_weight,
-			Default::default(),
+			RewardDestination::Staked,
 		)?;
 		Staking::<T>::nominate(
 			RawOrigin::Signed(origin_controller1.clone()).into(),
@@ -186,7 +186,7 @@ impl<T: Config> ListScenario<T> {
 		let (_origin_stash2, origin_controller2) = create_stash_controller_with_balance::<T>(
 			USER_SEED + 3,
 			origin_weight,
-			Default::default(),
+			RewardDestination::Staked,
 		)?;
 		Staking::<T>::nominate(
 			RawOrigin::Signed(origin_controller2).into(),
@@ -206,7 +206,7 @@ impl<T: Config> ListScenario<T> {
 		let (_dest_stash1, dest_controller1) = create_stash_controller_with_balance::<T>(
 			USER_SEED + 1,
 			dest_weight,
-			Default::default(),
+			RewardDestination::Staked,
 		)?;
 		Staking::<T>::nominate(
 			RawOrigin::Signed(dest_controller1).into(),
@@ -290,7 +290,7 @@ benchmarks! {
 	withdraw_unbonded_update {
 		// Slashing Spans
 		let s in 0 .. MAX_SPANS;
-		let (stash, controller) = create_stash_controller::<T>(0, 100, Default::default())?;
+		let (stash, controller) = create_stash_controller::<T>(0, 100, RewardDestination::Staked)?;
 		add_slashing_spans::<T>(&stash, s);
 		let amount = T::Currency::minimum_balance() * 5u32.into(); // Half of total
 		Staking::<T>::unbond(RawOrigin::Signed(controller.clone()).into(), amount)?;
@@ -339,7 +339,7 @@ benchmarks! {
 		let (stash, controller) = create_stash_controller::<T>(
 			MaxNominationsOf::<T>::get() - 1,
 			100,
-			Default::default(),
+			RewardDestination::Staked,
 		)?;
 		// because it is chilled.
 		assert!(!T::VoterList::contains(&stash));
@@ -367,7 +367,7 @@ benchmarks! {
 		let (stash, controller) = create_stash_controller::<T>(
 			MaxNominationsOf::<T>::get() - 1,
 			100,
-			Default::default(),
+			RewardDestination::Staked,
 		)?;
 		let stash_lookup = T::Lookup::unlookup(stash.clone());
 
@@ -382,7 +382,7 @@ benchmarks! {
 			let (n_stash, n_controller) = create_stash_controller::<T>(
 				MaxNominationsOf::<T>::get() + i,
 				100,
-				Default::default(),
+				RewardDestination::Staked,
 			)?;
 
 			// bake the nominations; we first clone them from the rest of the validators.
@@ -430,7 +430,7 @@ benchmarks! {
 		let (stash, controller) = create_stash_controller_with_balance::<T>(
 			SEED + MaxNominationsOf::<T>::get() + 1, // make sure the account does not conflict with others
 			origin_weight,
-			Default::default(),
+			RewardDestination::Staked,
 		).unwrap();
 
 		assert!(!Nominators::<T>::contains_key(&stash));
@@ -464,16 +464,16 @@ benchmarks! {
 	}
 
 	set_payee {
-		let (stash, controller) = create_stash_controller::<T>(USER_SEED, 100, Default::default())?;
-		assert_eq!(Payee::<T>::get(&stash), RewardDestination::Staked);
+		let (stash, controller) = create_stash_controller::<T>(USER_SEED, 100, RewardDestination::Staked)?;
+		assert_eq!(Payee::<T>::get(&stash), Some(RewardDestination::Staked));
 		whitelist_account!(controller);
 	}: _(RawOrigin::Signed(controller), RewardDestination::Controller)
 	verify {
-		assert_eq!(Payee::<T>::get(&stash), RewardDestination::Controller);
+		assert_eq!(Payee::<T>::get(&stash), Some(RewardDestination::Controller));
 	}
 
 	set_controller {
-		let (stash, ctlr) = create_unique_stash_controller::<T>(9000, 100, Default::default(), false)?;
+		let (stash, ctlr) = create_unique_stash_controller::<T>(9000, 100, RewardDestination::Staked, false)?;
 		// ensure `ctlr` is the currently stored controller.
 		assert!(!Ledger::<T>::contains_key(&stash));
 		assert!(Ledger::<T>::contains_key(&ctlr));
@@ -772,7 +772,7 @@ benchmarks! {
 	#[extra]
 	do_slash {
 		let l in 1 .. T::MaxUnlockingChunks::get() as u32;
-		let (stash, controller) = create_stash_controller::<T>(0, 100, Default::default())?;
+		let (stash, controller) = create_stash_controller::<T>(0, 100, RewardDestination::Staked)?;
 		let mut staking_ledger = Ledger::<T>::get(controller.clone()).unwrap();
 		let unlock_chunk = UnlockChunk::<BalanceOf<T>> {
 			value: 1u32.into(),

From ace469291ebf64f1e8406623c4d655a359f68351 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Mon, 27 Nov 2023 15:01:03 +0100
Subject: [PATCH 04/16] Adds try-state checks

---
 substrate/frame/staking/src/pallet/impls.rs | 17 +++++++++
 substrate/frame/staking/src/tests.rs        | 42 +++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index f463239c40cb0..909a77e5dc136 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1830,6 +1830,7 @@ impl<T: Config> Pallet<T> {
 			"VoterList contains non-staker"
 		);
 
+		Self::check_payees()?;
 		Self::check_nominators()?;
 		Self::check_exposures()?;
 		Self::check_paged_exposures()?;
@@ -1837,6 +1838,22 @@ impl<T: Config> Pallet<T> {
 		Self::check_count()
 	}
 
+	/// Invariants:
+	/// * A bonded ledger should always have an assigned `Payee`.
+	/// * The number of entries in `Payee` and of bonded staking ledgers *must* match.
+	fn check_payees() -> Result<(), TryRuntimeError> {
+		for (acc, _) in Ledger::<T>::iter() {
+			ensure!(Payee::<T>::get(acc).is_some(), "bonded ledger does not have payee set");
+		}
+
+		ensure!(
+			Ledger::<T>::iter().count() == Payee::<T>::iter().count(),
+			"number of entries in payee storage items not match the number of bonded ledgers",
+		);
+
+		Ok(())
+	}
+
 	fn check_count() -> Result<(), TryRuntimeError> {
 		ensure!(
 			<T as Config>::VoterList::count() ==
diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index ac897245dcc48..017ba33ad2db9 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -751,7 +751,11 @@ fn nominators_also_get_slashed_pro_rata() {
 	});
 }
 
+// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
+// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+// asserts in the tests are checked before.
 #[test]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn double_staking_should_fail() {
 	// should test (in the same order):
 	// * an account already bonded as stash cannot be be stashed again.
@@ -786,7 +790,11 @@ fn double_staking_should_fail() {
 	});
 }
 
+// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
+// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+// asserts in the tests are checked before.
 #[test]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn double_controlling_attempt_should_fail() {
 	// should test (in the same order):
 	// * an account already bonded as controller CANNOT be reused as the controller of another
@@ -1725,7 +1733,11 @@ fn rebond_emits_right_value_in_event() {
 	});
 }
 
+// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
+// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+// asserts in the tests are checked before.
 #[test]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn reward_to_stake_works() {
 	ExtBuilder::default()
 		.nominate(false)
@@ -5401,6 +5413,28 @@ fn count_check_works() {
 	})
 }
 
+#[test]
+#[should_panic]
+fn check_payee_invariant1_works() {
+	// A bonded ledger should always have an assigned `Payee` This test should panic as we verify
+	// that a bad state will panic due to the `try_state` checks in the `post_checks` in `mock`.
+	ExtBuilder::default().build_and_execute(|| {
+		let rogue_ledger = StakingLedger::<Test>::new(123456, 20);
+		Ledger::<Test>::insert(123456, rogue_ledger);
+	})
+}
+
+#[test]
+#[should_panic]
+fn check_payee_invariant2_works() {
+	// The number of entries in both `Payee` and of bonded staking ledgers should match. This test
+	// should panic as we verify that a bad state will panic due to the `try_state` checks in the
+	// `post_checks` in `mock`.
+	ExtBuilder::default().build_and_execute(|| {
+		Payee::<Test>::insert(1111, RewardDestination::Staked);
+	})
+}
+
 #[test]
 fn min_bond_checks_work() {
 	ExtBuilder::default()
@@ -6718,7 +6752,11 @@ mod staking_interface {
 mod ledger {
 	use super::*;
 
+	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
+	// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+	// asserts in the tests are checked before.
 	#[test]
+	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn paired_account_works() {
 		ExtBuilder::default().build_and_execute(|| {
 			assert_ok!(Staking::bond(
@@ -6753,7 +6791,11 @@ mod ledger {
 		})
 	}
 
+	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
+	// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+	// asserts in the tests are checked before.
 	#[test]
+	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn get_ledger_works() {
 		ExtBuilder::default().build_and_execute(|| {
 			// stash does not exist

From f3c8ca7e8c3e79262ecf4855257aa2b89a5bbc49 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Mon, 27 Nov 2023 20:20:00 +0100
Subject: [PATCH 05/16] fixes benchmarks

---
 substrate/frame/staking/src/benchmarking.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/frame/staking/src/benchmarking.rs b/substrate/frame/staking/src/benchmarking.rs
index a864a59774ebd..7513f30b251c7 100644
--- a/substrate/frame/staking/src/benchmarking.rs
+++ b/substrate/frame/staking/src/benchmarking.rs
@@ -481,7 +481,7 @@ benchmarks! {
 		whitelist_account!(controller);
 	}: _(RawOrigin::Signed(controller.clone()), controller.clone())
 	verify {
-		assert_eq!(Payee::<T>::get(&stash), RewardDestination::Account(controller));
+		assert_eq!(Payee::<T>::get(&stash), Some(RewardDestination::Account(controller)));
 	}
 
 	set_controller {

From d3e40581f6778db0caf180a85134cf41e5307e8a Mon Sep 17 00:00:00 2001
From: command-bot <>
Date: Mon, 27 Nov 2023 20:16:48 +0000
Subject: [PATCH 06/16] ".git/.scripts/commands/fmt/fmt.sh"

---
 substrate/frame/staking/src/tests.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index 017ba33ad2db9..6a2ad72218f44 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -6752,9 +6752,9 @@ mod staking_interface {
 mod ledger {
 	use super::*;
 
-	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-	// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
-	// asserts in the tests are checked before.
+	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger
+	// the `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus
+	// the asserts in the tests are checked before.
 	#[test]
 	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn paired_account_works() {
@@ -6791,9 +6791,9 @@ mod ledger {
 		})
 	}
 
-	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-	// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
-	// asserts in the tests are checked before.
+	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger
+	// the `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus
+	// the asserts in the tests are checked before.
 	#[test]
 	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn get_ledger_works() {

From f7469d0b89de7018defb6ae6808dce3997caf9be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Tue, 28 Nov 2023 11:35:07 +0100
Subject: [PATCH 07/16] Update substrate/frame/staking/src/tests.rs

Co-authored-by: Ross Bulat <ross@parity.io>
---
 substrate/frame/staking/src/tests.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index 6a2ad72218f44..772a32e6b23d6 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -752,7 +752,7 @@ fn nominators_also_get_slashed_pro_rata() {
 }
 
 // note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
+// `try_state` checks. The check that panics runs at the end of the test (`post_check`), thus the
 // asserts in the tests are checked before.
 #[test]
 #[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]

From a81406b725fabd6d42b6df9fc64440b3fc32ce5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Tue, 28 Nov 2023 11:35:15 +0100
Subject: [PATCH 08/16] Update substrate/frame/staking/src/pallet/impls.rs

Co-authored-by: Ross Bulat <ross@parity.io>
---
 substrate/frame/staking/src/pallet/impls.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index 909a77e5dc136..ef7f8d35f70b7 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1848,7 +1848,7 @@ impl<T: Config> Pallet<T> {
 
 		ensure!(
 			Ledger::<T>::iter().count() == Payee::<T>::iter().count(),
-			"number of entries in payee storage items not match the number of bonded ledgers",
+			"number of entries in payee storage items does not match the number of bonded ledgers",
 		);
 
 		Ok(())

From dfee4d4d3548cc48ee9f50f25f558658276e87fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Tue, 28 Nov 2023 11:41:47 +0100
Subject: [PATCH 09/16] fixes review comments

---
 substrate/frame/staking/src/pallet/impls.rs | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index ef7f8d35f70b7..25ae89dace6ea 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -336,15 +336,9 @@ impl<T: Config> Pallet<T> {
 		if amount.is_zero() {
 			return None
 		}
+		let dest = Self::payee(StakingAccount::Stash(stash.clone()))?;
 
-		let dest = match Self::payee(StakingAccount::Stash(stash.clone()))
-			.defensive_proof("payee must exist for a ledger requesting a payout")
-		{
-			Some(dest) => dest,
-			None => return None,
-		};
-
-		let maybe_imbalance = match dest.clone() {
+		let maybe_imbalance = match dest {
 			RewardDestination::Stash => T::Currency::deposit_into_existing(stash, amount).ok(),
 			RewardDestination::Staked => Self::ledger(Stash(stash.clone()))
 				.and_then(|mut ledger| {
@@ -359,7 +353,7 @@ impl<T: Config> Pallet<T> {
 					Ok(r)
 				})
 				.unwrap_or_default(),
-			RewardDestination::Account(dest_account) =>
+			RewardDestination::Account(ref dest_account) =>
 				Some(T::Currency::deposit_creating(&dest_account, amount)),
 			RewardDestination::None => None,
 			#[allow(deprecated)]

From 0d6bf8b9069145f8129f512ac34d891dbf7c7677 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Sat, 2 Dec 2023 21:46:46 +0100
Subject: [PATCH 10/16] Adds logging when check state invariant fails

---
 substrate/frame/staking/src/pallet/impls.rs | 5 ++++-
 substrate/frame/staking/src/tests.rs        | 4 ++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index 25ae89dace6ea..efc9453df79fb 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1837,7 +1837,10 @@ impl<T: Config> Pallet<T> {
 	/// * The number of entries in `Payee` and of bonded staking ledgers *must* match.
 	fn check_payees() -> Result<(), TryRuntimeError> {
 		for (acc, _) in Ledger::<T>::iter() {
-			ensure!(Payee::<T>::get(acc).is_some(), "bonded ledger does not have payee set");
+			if Payee::<T>::get(&acc).is_none() {
+				log!(error, "ledger {} does not have payee set", acc);
+				return Err("bonded ledger does not have payee set".into());
+			}
 		}
 
 		ensure!(
diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index 772a32e6b23d6..5bc9bb40f3a81 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -5414,7 +5414,7 @@ fn count_check_works() {
 }
 
 #[test]
-#[should_panic]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn check_payee_invariant1_works() {
 	// A bonded ledger should always have an assigned `Payee` This test should panic as we verify
 	// that a bad state will panic due to the `try_state` checks in the `post_checks` in `mock`.
@@ -5425,7 +5425,7 @@ fn check_payee_invariant1_works() {
 }
 
 #[test]
-#[should_panic]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"number of entries in payee storage items does not match the number of bonded ledgers\")"]
 fn check_payee_invariant2_works() {
 	// The number of entries in both `Payee` and of bonded staking ledgers should match. This test
 	// should panic as we verify that a bad state will panic due to the `try_state` checks in the

From a78e96d44d946a21791393ef923dab81fc13590e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Mon, 4 Dec 2023 10:12:12 +0100
Subject: [PATCH 11/16] account cannot be formatted

---
 substrate/frame/staking/src/pallet/impls.rs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index efc9453df79fb..f7aff0b37c9c5 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1837,10 +1837,7 @@ impl<T: Config> Pallet<T> {
 	/// * The number of entries in `Payee` and of bonded staking ledgers *must* match.
 	fn check_payees() -> Result<(), TryRuntimeError> {
 		for (acc, _) in Ledger::<T>::iter() {
-			if Payee::<T>::get(&acc).is_none() {
-				log!(error, "ledger {} does not have payee set", acc);
-				return Err("bonded ledger does not have payee set".into());
-			}
+			ensure!(Payee::<T>::get(&acc).is_some(), "bonded ledger does not have payee set");
 		}
 
 		ensure!(

From 904ca63cdd5ade6bb683d6fd8f9228d7a438fb52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Tue, 12 Dec 2023 15:00:37 +0100
Subject: [PATCH 12/16] Small fixes; skips try-state checks in tests if
 explicitly set

---
 substrate/frame/staking/src/mock.rs       | 13 ++++++++++++-
 substrate/frame/staking/src/pallet/mod.rs |  2 +-
 substrate/frame/staking/src/tests.rs      |  4 ++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/substrate/frame/staking/src/mock.rs b/substrate/frame/staking/src/mock.rs
index 5332dbfdd5b2d..961dd9fa3f300 100644
--- a/substrate/frame/staking/src/mock.rs
+++ b/substrate/frame/staking/src/mock.rs
@@ -345,6 +345,11 @@ where
 pub(crate) type StakingCall = crate::Call<Test>;
 pub(crate) type TestCall = <Test as frame_system::Config>::RuntimeCall;
 
+parameter_types! {
+	// if true, skips the try-state for the test running.
+	pub static SkipTryStateCheck: bool = false;
+}
+
 pub struct ExtBuilder {
 	nominate: bool,
 	validator_count: u32,
@@ -582,11 +587,17 @@ impl ExtBuilder {
 		let mut ext = self.build();
 		ext.execute_with(test);
 		ext.execute_with(|| {
-			Staking::do_try_state(System::block_number()).unwrap();
+			if !SkipTryStateCheck::get() {
+				Staking::do_try_state(System::block_number()).unwrap();
+			}
 		});
 	}
 }
 
+pub(crate) fn skip_try_state_checks() {
+	SkipTryStateCheck::set(true);
+}
+
 pub(crate) fn active_era() -> EraIndex {
 	Staking::active_era().unwrap().index
 }
diff --git a/substrate/frame/staking/src/pallet/mod.rs b/substrate/frame/staking/src/pallet/mod.rs
index 0ded1a61c005c..0689418d02bd9 100644
--- a/substrate/frame/staking/src/pallet/mod.rs
+++ b/substrate/frame/staking/src/pallet/mod.rs
@@ -1948,7 +1948,7 @@ pub mod pallet {
 						// `Controller` variant, skip deprecating this account.
 						let payee_deprecated = Payee::<T>::get(&ledger.stash) == {
 							#[allow(deprecated)]
-							RewardDestination::Controller
+							Some(RewardDestination::Controller)
 						};
 
 						if ledger.stash != *controller && !payee_deprecated {
diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index ab124a5169bd8..60e962c858f50 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -7084,6 +7084,10 @@ mod ledger {
 			// Ledger is still keyed by controller.
 			let ledger_updated = Ledger::<Test>::get(ctlr).unwrap();
 			assert_eq!(ledger_updated.stash, stash);
+
+			// this test manually inserts a ledger in storage. Let's skip the try-state assertions
+			// at the end of the test, as we expect them to fail.
+			skip_try_state_checks();
 		})
 	}
 }

From 8778b199820f7b31b4247612c6ff6acc742db957 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Mon, 15 Jan 2024 22:53:06 +0100
Subject: [PATCH 13/16] Adds try_state ext builder method to skip try state
 checks on tests and remove panicking

---
 substrate/frame/staking/src/mock.rs  |  8 +++----
 substrate/frame/staking/src/tests.rs | 35 +++++-----------------------
 2 files changed, 10 insertions(+), 33 deletions(-)

diff --git a/substrate/frame/staking/src/mock.rs b/substrate/frame/staking/src/mock.rs
index 961dd9fa3f300..73b6f8ccf27f6 100644
--- a/substrate/frame/staking/src/mock.rs
+++ b/substrate/frame/staking/src/mock.rs
@@ -459,6 +459,10 @@ impl ExtBuilder {
 		self.balance_factor = factor;
 		self
 	}
+	pub fn try_state(self, enable: bool) -> Self {
+		SkipTryStateCheck::set(!enable);
+		self
+	}
 	fn build(self) -> sp_io::TestExternalities {
 		sp_tracing::try_init_simple();
 		let mut storage = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
@@ -594,10 +598,6 @@ impl ExtBuilder {
 	}
 }
 
-pub(crate) fn skip_try_state_checks() {
-	SkipTryStateCheck::set(true);
-}
-
 pub(crate) fn active_era() -> EraIndex {
 	Staking::active_era().unwrap().index
 }
diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index 60e962c858f50..c0ecd4213cdb6 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -751,17 +751,13 @@ fn nominators_also_get_slashed_pro_rata() {
 	});
 }
 
-// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-// `try_state` checks. The check that panics runs at the end of the test (`post_check`), thus the
-// asserts in the tests are checked before.
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn double_staking_should_fail() {
 	// should test (in the same order):
 	// * an account already bonded as stash cannot be be stashed again.
 	// * an account already bonded as stash cannot nominate.
 	// * an account already bonded as controller can nominate.
-	ExtBuilder::default().build_and_execute(|| {
+	ExtBuilder::default().try_state(false).build_and_execute(|| {
 		let arbitrary_value = 5;
 		let (stash, controller) = testing_utils::create_unique_stash_controller::<Test>(
 			0,
@@ -790,16 +786,12 @@ fn double_staking_should_fail() {
 	});
 }
 
-// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
-// asserts in the tests are checked before.
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn double_controlling_attempt_should_fail() {
 	// should test (in the same order):
 	// * an account already bonded as controller CANNOT be reused as the controller of another
 	//   account.
-	ExtBuilder::default().build_and_execute(|| {
+	ExtBuilder::default().try_state(false).build_and_execute(|| {
 		let arbitrary_value = 5;
 		let (stash, _) = testing_utils::create_unique_stash_controller::<Test>(
 			0,
@@ -1733,17 +1725,14 @@ fn rebond_emits_right_value_in_event() {
 	});
 }
 
-// note: this test will panic due to manual changes in the `Ledger` storage that will trigger the
-// `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus the
-// asserts in the tests are checked before.
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 fn reward_to_stake_works() {
 	ExtBuilder::default()
 		.nominate(false)
 		.set_status(31, StakerStatus::Idle)
 		.set_status(41, StakerStatus::Idle)
 		.set_stake(21, 2000)
+		.try_state(false)
 		.build_and_execute(|| {
 			assert_eq!(Staking::validator_count(), 2);
 			// Confirm account 10 and 20 are validators
@@ -6752,13 +6741,9 @@ mod staking_interface {
 mod ledger {
 	use super::*;
 
-	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger
-	// the `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus
-	// the asserts in the tests are checked before.
 	#[test]
-	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn paired_account_works() {
-		ExtBuilder::default().build_and_execute(|| {
+		ExtBuilder::default().try_state(false).build_and_execute(|| {
 			assert_ok!(Staking::bond(
 				RuntimeOrigin::signed(10),
 				100,
@@ -6791,13 +6776,9 @@ mod ledger {
 		})
 	}
 
-	// note: this test will panic due to manual changes in the `Ledger` storage that will trigger
-	// the `try_state` checks. The check tat panics runs at the end of the test (`post_check`), thus
-	// the asserts in the tests are checked before.
 	#[test]
-	#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
 	fn get_ledger_works() {
-		ExtBuilder::default().build_and_execute(|| {
+		ExtBuilder::default().try_state(false).build_and_execute(|| {
 			// stash does not exist
 			assert!(StakingLedger::<Test>::get(StakingAccount::Stash(42)).is_err());
 
@@ -7044,7 +7025,7 @@ mod ledger {
 
 	#[test]
 	fn deprecate_controller_batch_skips_unmigrated_controller_payees() {
-		ExtBuilder::default().build_and_execute(|| {
+		ExtBuilder::default().try_state(false).build_and_execute(|| {
 			// Given:
 
 			let stash: u64 = 1000;
@@ -7084,10 +7065,6 @@ mod ledger {
 			// Ledger is still keyed by controller.
 			let ledger_updated = Ledger::<Test>::get(ctlr).unwrap();
 			assert_eq!(ledger_updated.stash, stash);
-
-			// this test manually inserts a ledger in storage. Let's skip the try-state assertions
-			// at the end of the test, as we expect them to fail.
-			skip_try_state_checks();
 		})
 	}
 }

From 8028f832d54f535a2a9fcc8ea2c0871951cf897b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Tue, 16 Jan 2024 09:19:58 +0100
Subject: [PATCH 14/16] restart CI


From 16d02404c8f773ddc6326c52bcf8868a92de9e6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Fri, 26 Jan 2024 12:32:19 +0100
Subject: [PATCH 15/16] Fixes try-state check introduced in this PR

---
 substrate/frame/staking/src/pallet/impls.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/substrate/frame/staking/src/pallet/impls.rs b/substrate/frame/staking/src/pallet/impls.rs
index 70f53f82e9b90..8b45430c688ec 100644
--- a/substrate/frame/staking/src/pallet/impls.rs
+++ b/substrate/frame/staking/src/pallet/impls.rs
@@ -1836,12 +1836,13 @@ impl<T: Config> Pallet<T> {
 	/// * A bonded ledger should always have an assigned `Payee`.
 	/// * The number of entries in `Payee` and of bonded staking ledgers *must* match.
 	fn check_payees() -> Result<(), TryRuntimeError> {
-		for (acc, _) in Ledger::<T>::iter() {
-			ensure!(Payee::<T>::get(&acc).is_some(), "bonded ledger does not have payee set");
+		for (stash, _) in Bonded::<T>::iter() {
+			ensure!(Payee::<T>::get(&stash).is_some(), "bonded ledger does not have payee set");
 		}
 
 		ensure!(
-			Ledger::<T>::iter().count() == Payee::<T>::iter().count(),
+			(Ledger::<T>::iter().count() == Payee::<T>::iter().count()) &&
+				(Ledger::<T>::iter().count() == Bonded::<T>::iter().count()),
 			"number of entries in payee storage items does not match the number of bonded ledgers",
 		);
 

From be10ca0dd2bcd82e6c67b5f43ce8d034f667c312 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gon=C3=A7alo=20Pestana?= <g6pestana@gmail.com>
Date: Fri, 26 Jan 2024 14:18:59 +0100
Subject: [PATCH 16/16] fixes expected try-state fail msg in tests

---
 substrate/frame/staking/src/tests.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs
index 12e35f7fee795..85ee7dd3bf59d 100644
--- a/substrate/frame/staking/src/tests.rs
+++ b/substrate/frame/staking/src/tests.rs
@@ -5417,7 +5417,7 @@ fn count_check_works() {
 }
 
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"bonded ledger does not have payee set\")"]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: Other(\"number of entries in payee storage items does not match the number of bonded ledgers\")"]
 fn check_payee_invariant1_works() {
 	// A bonded ledger should always have an assigned `Payee` This test should panic as we verify
 	// that a bad state will panic due to the `try_state` checks in the `post_checks` in `mock`.