update to check head pre & post conditions

frame/bags-list/src/lib.rs

@@ -140,8 +140,8 @@ pub mod pallet {
 		///   the procedure given above, then the constant ratio is equal to 2.
 		/// - If `BagThresholds::get().len() == 200`, and the thresholds are determined according to
 		///   the procedure given above, then the constant ratio is approximately equal to 1.248.
-		/// - If the threshold list begins `[1, 2, 3, ...]`, then an id with weight 0 or 1 will
-		///   fall into bag 0, an id with weight 2 will fall into bag 1, etc.
+		/// - If the threshold list begins `[1, 2, 3, ...]`, then an id with weight 0 or 1 will fall
+		///   into bag 0, an id with weight 2 will fall into bag 1, etc.
 		///
 		/// # Migration
 		///
@@ -265,7 +265,23 @@ impl<T: Config> SortedListProvider<T::AccountId> for Pallet<T> {
 	}
 
 	#[cfg(feature = "runtime-benchmarks")]
-	fn is_in_bag(id: &T::AccountId, weight: VoteWeight) -> bool {
+	fn is_in_bag(id: &T::AccountId, weight: VoteWeight, _: bool) -> bool {
+		use frame_support::traits::Get;
+		let info = T::BagThresholds::get()
+			.into_iter()
+			.chain(std::iter::once(&VoteWeight::MAX)) // assumes this is not an explicit threshold
+			.filter_map(|t| {
+				list::Bag::<T>::get(*t)
+					.map(|bag| (*t, bag.iter().map(|n| n.id().clone()).collect::<Vec<_>>()))
+			})
+			.collect::<Vec<_>>();
+		println!("bags info {:#?}", info);
+
 		list::Bag::<T>::get(list::notional_bag_for::<T>(weight)).unwrap().contains(id)
 	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn is_bag_head(id: &T::AccountId, weight: VoteWeight, _: bool) -> bool {
+		list::Bag::<T>::get(list::notional_bag_for::<T>(weight)).unwrap().head_id() == Some(id)
+	}
 }

frame/bags-list/src/list/mod.rs

@@ -111,10 +111,10 @@ impl<T: Config> List<T> {
 	/// Postconditions:
 	///
 	/// - All `bag_upper` currently in storage are members of `T::BagThresholds`.
-	/// - No id is changed unless required to by the difference between the old threshold list
-	///   and the new.
-	/// - ids whose bags change at all are implicitly rebagged into the appropriate bag in the
-	///   new threshold set.
+	/// - No id is changed unless required to by the difference between the old threshold list and
+	///   the new.
+	/// - ids whose bags change at all are implicitly rebagged into the appropriate bag in the new
+	///   threshold set.
 	#[allow(dead_code)]
 	pub fn migrate(old_thresholds: &[VoteWeight]) -> u32 {
 		// we can't check all preconditions, but we can check one
@@ -590,6 +590,11 @@ impl<T: Config> Bag<T> {
 	pub(crate) fn contains(&self, id: &T::AccountId) -> bool {
 		self.iter().any(|n| n.id() == id)
 	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	pub(crate) fn head_id(&self) -> Option<&T::AccountId> {
+		self.head.as_ref()
+	}
 }
 
 /// A Node is the fundamental element comprising the doubly-linked list described by `Bag`.

frame/bags-list/src/list/tests.rs

@@ -136,7 +136,8 @@ fn migrate_works() {
 					(15, vec![710]), // nodes in range 11 ..= 15 move from bag 20 to bag 15
 					(20, vec![711]),
 					(1000, vec![2, 3, 4]),
-					(10_000, vec![712]), // nodes in range 1_001 ..= 2_000 move from bag 2_000 to bag 10_000
+					(10_000, vec![712]), /* nodes in range 1_001 ..= 2_000 move from bag 2_000
+					                      * to bag 10_000 */
 				]
 			);
 		});

frame/election-provider-support/src/lib.rs

@@ -332,11 +332,18 @@ pub trait SortedListProvider<AccountId> {
 	/// Sanity check internal state of list. Only meant for debug compilation.
 	fn sanity_check() -> Result<(), &'static str>;
 
-	/// If the voter is in the notional bag for the given weight.
+	/// Wether the voter is in the notional bag for the given weight. Only relevant for
+	/// implementations that use bags.
 	#[cfg(any(feature = "runtime-benchmarks", test))]
-	fn is_in_bag(_: &AccountId, _: VoteWeight) -> bool {
-		// TODO problematic if this always returns true
-		true // default to true for impls that don't have a bag.
+	fn is_in_bag(_: &AccountId, _: VoteWeight, mock: bool) -> bool {
+		mock // default to returning mock value for values that are not relevant
+	}
+
+	/// Wether the voter is the head of the notional bag for the given weight. Only relevant for
+	/// implementations that use bags.
+	#[cfg(any(feature = "runtime-benchmarks", test))]
+	fn is_bag_head(_: &AccountId, _: VoteWeight, mock: bool) -> bool {
+		mock // default to returning mock value for values that are not relevant
 	}
 }
 

frame/generate-bags/src/lib.rs

@@ -37,8 +37,8 @@
 //!    ```
 //!
 //! 2. Write a little program to generate the definitions. This can be a near-identical copy of
-//!    `substrate/frame/bags-list/generate-bags`. This program exists only to hook together the runtime
-//!    definitions with the various calculations here.
+//!    `substrate/frame/bags-list/generate-bags`. This program exists only to hook together the
+//! runtime    definitions with the various calculations here.
 //!
 //! 3. Run that program:
 //!

frame/staking/src/benchmarking.rs

@@ -139,16 +139,16 @@ struct RebagScenario<T: Config> {
 	/// Controller of the Stash that is expected to be rebagged.
 	origin_controller1: T::AccountId,
 	origin_stash2: T::AccountId,
-	origin_bag_thresh: BalanceOf<T>,
-	dest_bag_thresh: BalanceOf<T>,
+	origin_thresh_as_vote: VoteWeight,
+	dest_thresh_as_vote: VoteWeight,
 }
 
 impl<T: Config> RebagScenario<T> {
 	/// An expensive rebag scenario:
 	///
 	/// - the node to be rebagged (r) is the head of a bag that has at least one other node. The bag
-	///   itself will need to be read and written to update its head. The node pointed to by
-	///   r.next will need to be read and written as it will need to have its prev pointer updated.
+	///   itself will need to be read and written to update its head. The node pointed to by r.next
+	///   will need to be read and written as it will need to have its prev pointer updated.
 	///
 	/// - the destination bag has at least one node, which will need its next pointer updated.
 	fn new(
@@ -198,27 +198,75 @@ impl<T: Config> RebagScenario<T> {
 		)?;
 		Staking::<T>::nominate(RawOrigin::Signed(dest_controller1).into(), validators)?;
 
+		let total_issuance = T::Currency::total_issuance();
+		let origin_thresh_as_vote = T::CurrencyToVote::to_vote(origin_bag_thresh, total_issuance);
+		let dest_thresh_as_vote = T::CurrencyToVote::to_vote(dest_bag_thresh, total_issuance);
+
 		Ok(RebagScenario {
 			dest_stash1,
 			origin_stash1,
 			origin_controller1,
 			origin_stash2,
-			origin_bag_thresh,
-			dest_bag_thresh,
+			origin_thresh_as_vote,
+			dest_thresh_as_vote,
 		})
 	}
 
-	fn assert_preconditions(&self) {
-		let total_issuance = T::Currency::total_issuance();
-		let dest_thresh_as_vote = T::CurrencyToVote::to_vote(self.dest_bag_thresh, total_issuance);
-		let origin_thresh_as_vote = T::CurrencyToVote::to_vote(self.origin_bag_thresh, total_issuance);
+	fn check_preconditions(&self) {
+		let RebagScenario {
+			dest_stash1,
+			origin_stash2,
+			dest_thresh_as_vote,
+			origin_thresh_as_vote,
+			..
+		} = self;
+		// destination stash is not in the origin bag.
+		assert!(!T::SortedListProvider::is_in_bag(&dest_stash1, *origin_thresh_as_vote, false));
+		// origin stash2 is in the origin bag
+		assert!(T::SortedListProvider::is_in_bag(&origin_stash2, *origin_thresh_as_vote, true));
+		// head checks implicitly check that dest stash1 and origin stash1 are in the correct bags.
+		self.check_head_preconditions();
+	}
 
+	fn check_postconditions(&self) {
+		let RebagScenario {
+			dest_stash1,
+			origin_stash1,
+			dest_thresh_as_vote,
+			origin_thresh_as_vote,
+			..
+		} = self;
 		// destination stash is not in the origin bag
-		assert!(!T::SortedListProvider::is_in_bag(&self.dest_stash1, origin_thresh_as_vote));
+		assert!(!T::SortedListProvider::is_in_bag(&dest_stash1, *origin_thresh_as_vote, false));
 		// and is in the destination bag.
-		assert!(T::SortedListProvider::is_in_bag(&self.dest_stash1, dest_thresh_as_vote));
-		// the origin stash is in the origin bag.
-		assert!(T::SortedListProvider::is_in_bag(&self.origin_stash1, origin_thresh_as_vote));
+		assert!(T::SortedListProvider::is_in_bag(&dest_stash1, *dest_thresh_as_vote, true));
+		// the origin stash is in the destination bag.
+		assert!(T::SortedListProvider::is_in_bag(&origin_stash1, *dest_thresh_as_vote, true));
+		self.check_head_postconditions();
+	}
+
+	fn check_head_preconditions(&self) {
+		let RebagScenario {
+			dest_stash1,
+			origin_stash1,
+			dest_thresh_as_vote,
+			origin_thresh_as_vote,
+			..
+		} = self;
+		assert!(T::SortedListProvider::is_bag_head(&dest_stash1, *dest_thresh_as_vote, true));
+		assert!(T::SortedListProvider::is_bag_head(&origin_stash1, *origin_thresh_as_vote, true));
+	}
+
+	fn check_head_postconditions(&self) {
+		let RebagScenario {
+			dest_stash1,
+			origin_stash2,
+			dest_thresh_as_vote,
+			origin_thresh_as_vote,
+			..
+		} = self;
+		assert!(T::SortedListProvider::is_bag_head(&dest_stash1, *dest_thresh_as_vote, true));
+		assert!(T::SortedListProvider::is_bag_head(&origin_stash2, *origin_thresh_as_vote, true));
 	}
 }
 
@@ -262,13 +310,15 @@ benchmarks! {
 		let original_bonded: BalanceOf<T> = ledger.active;
 		whitelist_account!(stash);
 
-		scenario.assert_preconditions();
+		scenario.check_preconditions();
 
 	}: _(RawOrigin::Signed(stash.clone()), max_additional)
 	verify {
 		let ledger = Ledger::<T>::get(&controller).ok_or("ledger not created after")?;
 		let new_bonded: BalanceOf<T> = ledger.active;
 		assert!(original_bonded < new_bonded);
+
+		scenario.check_postconditions();
 	}
 
 	unbond {
@@ -293,14 +343,16 @@ benchmarks! {
 		let ledger = Ledger::<T>::get(&controller).ok_or("ledger not created before")?;
 		let original_bonded: BalanceOf<T> = ledger.active;
 
-		scenario.assert_preconditions();
+		scenario.check_preconditions();
 
 		whitelist_account!(controller);
 	}: _(RawOrigin::Signed(controller.clone()), amount)
 	verify {
 		let ledger = Ledger::<T>::get(&controller).ok_or("ledger not created after")?;
 		let new_bonded: BalanceOf<T> = ledger.active;
 		assert!(original_bonded > new_bonded);
+
+		scenario.check_postconditions();
 	}
 
 	// Withdraw only updates the ledger
@@ -681,11 +733,7 @@ benchmarks! {
 		Ledger::<T>::insert(controller.clone(), staking_ledger.clone());
 		let original_bonded: BalanceOf<T> = staking_ledger.active;
 
-		scenario.assert_preconditions();
-		assert_eq!(
-			T::SortedListProvider::iter().collect::<Vec<_>>(),
-			vec![scenario.dest_stash1.clone(), stash.clone(), scenario.origin_stash2.clone()]
-		);
+		scenario.check_preconditions();
 
 		whitelist_account!(controller);
 	}: _(RawOrigin::Signed(controller.clone()), rebond_amount)
@@ -694,11 +742,7 @@ benchmarks! {
 		let new_bonded: BalanceOf<T> = ledger.active;
 		assert!(original_bonded < new_bonded);
 
-		// the ordering in the list doesn't change
-		assert_eq!(
-			T::SortedListProvider::iter().collect::<Vec<_>>(),
-			vec![scenario.dest_stash1.clone(), stash.clone(), scenario.origin_stash2.clone()]
-		);
+		scenario.check_postconditions();
 	}
 
 	set_history_depth {

frame/staking/src/pallet/mod.rs

@@ -768,13 +768,14 @@ pub mod pallet {
 					Error::<T>::InsufficientBond
 				);
 
+				// ledger must be updated prior to call to `Self::weight_of`.
+				Self::update_ledger(&controller, &ledger);
 				// update this staker in the sorted list, if they exist in it.
 				if T::SortedListProvider::contains(&stash) {
 					T::SortedListProvider::on_update(&stash, Self::weight_of(&ledger.stash));
 				}
 
 				Self::deposit_event(Event::<T>::Bonded(stash.clone(), extra));
-				Self::update_ledger(&controller, &ledger);
 			}
 			Ok(())
 		}