Suggestions for #9081 (Store voters in unsorted bags)

bin/node/runtime/voter-bags/src/main.rs

@@ -62,7 +62,7 @@ struct Opt {
 }
 
 fn main() -> Result<(), std::io::Error> {
-	let Opt {n_bags, output, runtime } = Opt::from_args();
+	let Opt { n_bags, output, runtime } = Opt::from_args();
 	let mut ext = sp_io::TestExternalities::new_empty();
 	ext.execute_with(|| runtime.generate_thresholds()(n_bags, &output))
 }

frame/staking/src/mock/voter_bags.rs

@@ -31,7 +31,7 @@ pub const EXISTENTIAL_WEIGHT: u64 = 1;
 pub const CONSTANT_RATIO: f64 = 2.0000000000000000;
 
 /// Upper thresholds delimiting the bag list.
-pub const THRESHOLDS: [u64; 65] = [5, 15, 25, 50, 100, 1000, 2000, 3000]
+// pub const THRESHOLDS: [u64; 65] = [5, 15, 25, 50, 100, 1000, 2000, 3000]; // TODO: switch to this
 pub const THRESHOLDS: [u64; 65] = [
 	                         1,
 	                         2,

frame/staking/src/tests.rs

@@ -153,7 +153,7 @@ fn basic_setup_works() {
 		// for these stash ids, see
 		// https://github.com/paritytech/substrate/
 		//   blob/631d4cdbcad438248c2597213918d8207d85bf6e/frame/staking/src/mock.rs#L435-L441
-		for genesis_stash_account_id in [11, 21, 31, 101] {
+		for genesis_stash_account_id in &[11, 21, 31, 101] {
 			let node = crate::voter_bags::Node::<Test>::from_id(&genesis_stash_account_id)
 				.expect(&format!("node was created for account {}", genesis_stash_account_id));
 			assert!(!node.is_misplaced(&weight_of));

frame/staking/src/voter_bags.rs

@@ -21,12 +21,8 @@
 //! - It's efficient to iterate over the top* N voters by stake, where the precise ordering of
 //!   voters doesn't particularly matter.
 
-use crate::{
-	AccountIdOf, Config, Nominations, Nominators, Pallet, Validators, VoteWeight, VoterBagFor,
-	VotingDataOf, slashing::SlashingSpans,
-};
-use codec::{Encode, Decode};
-use frame_support::{DefaultNoBound, traits::Get};
+use codec::{Decode, Encode};
+use frame_support::{traits::Get, DefaultNoBound};
 use sp_runtime::SaturatedConversion;
 use sp_std::{
 	boxed::Box,
@@ -35,6 +31,11 @@ use sp_std::{
 	marker::PhantomData,
 };
 
+use crate::{
+	slashing::SlashingSpans, AccountIdOf, Config, Nominations, Nominators, Pallet, Validators,
+	VoteWeight, VoterBagFor, VotingDataOf,
+};
+
 /// [`Voter`] parametrized by [`Config`] instead of by `AccountId`.
 pub type VoterOf<T> = Voter<AccountIdOf<T>>;
 
@@ -438,9 +439,9 @@ impl<T: Config> Bag<T> {
 		node.put();
 
 		// update the previous tail
-		if let Some(mut tail) = self.tail() {
-			tail.next = Some(id.clone());
-			tail.put();
+		if let Some(mut old_tail) = self.tail() {
+			old_tail.next = Some(id.clone());
+			old_tail.put();
 		}
 
 		// update the internal bag links
@@ -458,19 +459,56 @@ impl<T: Config> Bag<T> {
 	/// the first place. Generally, use [`VoterList::remove`] instead.
 	///
 	/// Storage note: this modifies storage, but only for adjacent nodes. You still need to call
-	/// `self.put()` and `node.put()` after use.
+	/// `self.put()`, `VoterNodes::remove(voter_id)` and `VoterBagFor::remove(voter_id)`
+	/// to update storage for the bag and `node`.
 	fn remove_node(&mut self, node: &Node<T>) {
-		// TODO: we could merge this function here.
-		// node.excise();
+		// Excise `node`.
 		if let Some(mut prev) = node.prev() {
-			prev.next = self.next.clone();
+			debug_assert!(
+				self.head.as_ref() != Some(&node.voter.id),
+				"node is the head, but has Some prev"
+			);
+			debug_assert!(
+				prev.prev().is_some() || self.head.as_ref() == Some(&prev.voter.id),
+				"node.prev.prev should be Some OR node.prev should be the head"
+			);
+
+			prev.next = node.next.clone();
+			debug_assert!(
+				prev.next().and_then(|prev_next|
+					// prev.next.prev should point at node prior to being reassigned
+					Some(prev_next.prev().unwrap().voter.id == node.voter.id)
+				)
+				// unless prev.next is None, in which case node has to be the tail
+				.unwrap_or(self.tail.as_ref() == Some(&node.voter.id)),
+				"prev.next.prev should point at node prior to being reassigned OR node should be the tail"
+			);
+
 			prev.put();
 		}
 		if let Some(mut next) = node.next() {
-			next.prev = self.prev.clone();
+			debug_assert!(
+				self.tail.as_ref() != Some(&node.voter.id),
+				"node is the tail, but has Some next"
+			);
+			debug_assert!(
+				next.next().is_some() || self.tail.as_ref() == Some(&next.voter.id),
+				"node.next.next should be Some OR node.next should be the head"
+			);
+
+			next.prev = node.prev.clone();
+			debug_assert!(
+				next.prev().and_then(|next_prev|
+					// next.prev.next should point at next after being reassigned
+					Some(next_prev.next().unwrap().voter.id == next.voter.id)
+				)
+				// unless next.prev is None, in which case node has to be the head
+				.unwrap_or_else(|| self.head.as_ref() == Some(&node.voter.id)),
+				"next.prev.next should point at next after being reassigned OR node should be the head"
+			);
+
 			next.put();
 		}
-		// IDEA: debug_assert! prev.next.prev == self
 
 		// clear the bag head/tail pointers as necessary
 		if self.head.as_ref() == Some(&node.voter.id) {
@@ -498,7 +536,10 @@ pub struct Node<T: Config> {
 impl<T: Config> Node<T> {
 	/// Get a node by bag idx and account id.
 	pub fn get(bag_upper: VoteWeight, account_id: &AccountIdOf<T>) -> Option<Node<T>> {
-		// debug_assert!(bag_upper is in Threshold)
+		// debug_assert!( // TODO: figure out why this breaks test take_works
+		// 	T::VoterBagThresholds::get().contains(&bag_upper) || bag_upper == VoteWeight::MAX,
+		// 	"it is a logic error to attempt to get a bag which is not in the thresholds list"
+		// );
 		crate::VoterNodes::<T>::try_get(account_id).ok().map(|mut node| {
 			node.bag_upper = bag_upper;
 			node
@@ -566,6 +607,7 @@ impl<T: Config> Node<T> {
 	/// Remove this node from the linked list.
 	///
 	/// Modifies storage, but only modifies the adjacent nodes. Does not modify `self` or any bag.
+	#[allow(dead_code)]// TODO: do we keep? (equivalent code in `fn remove_node`)
 	fn excise(&self) {
 		if let Some(mut prev) = self.prev() {
 			prev.next = self.next.clone();
@@ -855,10 +897,11 @@ pub mod make_bags {
 
 #[cfg(test)]
 mod tests {
-	use crate::mock::{ExtBuilder, Staking, Test};
 	use frame_support::traits::Currency;
 	use substrate_test_utils::assert_eq_uvec;
+
 	use super::*;
+	use crate::mock::*;
 
 	const GENESIS_VOTER_IDS: [u64; 5] = [11, 21, 31, 41, 101];
 
@@ -890,6 +933,9 @@ mod tests {
 			// initialize the voters' deposits
 			let existential_deposit = <Test as Config>::Currency::minimum_balance();
 			let mut balance = existential_deposit + 1;
+			assert_eq!(VoterBagThresholds::get()[1] as u128, balance);
+			assert_eq!(balance, 2);
+
 			for voter_id in voters.iter().rev() {
 				<Test as Config>::Currency::make_free_balance_be(voter_id, balance);
 				let controller = Staking::bonded(voter_id).unwrap();
@@ -906,4 +952,58 @@ mod tests {
 			assert_eq!(voters, have_voters);
 		});
 	}
+
+	/// This tests that we can `take` x voters, even if that quantity ends midway through a list.
+	#[test]
+	fn take_works() {
+		ExtBuilder::default().validator_pool(true).build_and_execute(|| {
+			// initialize the voters' deposits
+			let existential_deposit = <Test as Config>::Currency::minimum_balance();
+			let mut balance = existential_deposit + 1;
+			assert_eq!(VoterBagThresholds::get()[1] as u128, balance);
+			assert_eq!(balance, 2);
+
+			for (idx, voter_id) in GENESIS_VOTER_IDS.iter().enumerate() {
+				if idx % 3 == 0 {
+					// This increases the balance by a constant factor of 2, which is
+					// is the factor used to generate the mock bags. Thus this will
+					// increase the balance by 1 bag.
+					//
+					// This will create 2 bags, the lower threshold bag having
+					// 3 voters with balance 4, and the higher threshold bag having
+					// 2 voters with balance 8.
+					balance *= 2;
+				}
+
+				<Test as Config>::Currency::make_free_balance_be(voter_id, balance);
+				let controller = Staking::bonded(voter_id).unwrap();
+				let mut ledger = Staking::ledger(&controller).unwrap();
+				ledger.total = balance;
+				ledger.active = balance;
+				Staking::update_ledger(&controller, &ledger);
+				Staking::do_rebag(voter_id);
+			}
+
+			let bag_thresh4 = <Staking as crate::Store>::VoterBags::get(&4).unwrap().iter()
+				.map(|node| node.voter.id).collect::<Vec<_>>();
+			assert_eq!(bag_thresh4, vec![11, 21, 31]);
+
+			let bag_thresh8 = <Staking as crate::Store>::VoterBags::get(&8).unwrap().iter()
+				.map(|node| node.voter.id).collect::<Vec<_>>();
+			assert_eq!(bag_thresh8, vec![41, 101]);
+
+
+			let voters: Vec<_> = VoterList::<Test>::iter()
+				// take 4/5 from [41, 101],[11, 21, 31], demonstrating that we can do a
+				// take that stops mid bag.
+				.take(4)
+				.map(|node| node.voter.id)
+				.collect();
+
+			assert_eq!(voters, vec![41, 101, 11, 21]);
+		});
+	}
+
+	// TODO:
+	// - storage is cleaned up when a voter is removed
 }