Add extrinsic to improve position in a bag of bags-list

frame/bags-list/src/benchmarks.rs

@@ -19,7 +19,7 @@
 
 use super::*;
 use crate::list::List;
-use frame_benchmarking::{account, whitelisted_caller};
+use frame_benchmarking::{account, whitelist_account, whitelisted_caller};
 use frame_election_provider_support::VoteWeightProvider;
 use frame_support::{assert_ok, traits::Get};
 use frame_system::RawOrigin as SystemOrigin;
@@ -135,9 +135,41 @@ frame_benchmarking::benchmarks! {
 		);
 	}
 
-	impl_benchmark_test_suite!(
-		Pallet,
-		crate::mock::ExtBuilder::default().build(),
-		crate::mock::Runtime,
-	)
+	put_in_front_of {
+		// The most expensive case for `put_in_front_of`:
+		//
+		// - both heavier's `prev` and `next` are nodes that will need to be read and written.
+		// - `lighter` is the bag's `head`, so the bag will need to be read and written.
+
+		let bag_thresh = T::BagThresholds::get()[0];
+
+		// insert the nodes in order
+		let lighter: T::AccountId = account("lighter", 0, 0);
+		assert_ok!(List::<T>::insert(lighter.clone(), bag_thresh));
+
+		let heavier_prev: T::AccountId = account("heavier_prev", 0, 0);
+		assert_ok!(List::<T>::insert(heavier_prev.clone(), bag_thresh));
+
+		let heavier: T::AccountId = account("heavier", 0, 0);
+		assert_ok!(List::<T>::insert(heavier.clone(), bag_thresh));
+
+		let heavier_next: T::AccountId = account("heavier_next", 0, 0);
+		assert_ok!(List::<T>::insert(heavier_next.clone(), bag_thresh));
+
+		T::VoteWeightProvider::set_vote_weight_of(&lighter, bag_thresh - 1);
+		T::VoteWeightProvider::set_vote_weight_of(&heavier, bag_thresh);
+
+		assert_eq!(
+			List::<T>::iter().map(|n| n.id().clone()).collect::<Vec<_>>(),
+			vec![lighter.clone(), heavier_prev.clone(), heavier.clone(), heavier_next.clone()]
+		);
+
+		whitelist_account!(heavier);
+	}: _(SystemOrigin::Signed(heavier.clone()), lighter.clone())
+	verify {
+		assert_eq!(
+			List::<T>::iter().map(|n| n.id().clone()).collect::<Vec<_>>(),
+			vec![heavier, lighter, heavier_prev, heavier_next]
+		)
+	}
 }

frame/bags-list/src/lib.rs

@@ -173,6 +173,17 @@ pub mod pallet {
 		Rebagged { who: T::AccountId, from: VoteWeight, to: VoteWeight },
 	}
 
+	#[pallet::error]
+	#[cfg_attr(test, derive(PartialEq))]
+	pub enum Error<T> {
+		/// Attempted to place node in front of a node in another bag.
+		NotInSameBag,
+		/// Id not found in list.
+		IdNotFound,
+		/// An Id does not have a greater vote weight than another Id.
+		NotHeavier,
+	}
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {
 		/// Declare that some `dislocated` account has, through rewards or penalties, sufficiently
@@ -190,6 +201,20 @@ pub mod pallet {
 			let _ = Pallet::<T>::do_rebag(&dislocated, current_weight);
 			Ok(())
 		}
+
+		/// Move the caller's Id directly in front of `lighter`.
+		///
+		/// The dispatch origin for this call must be _Signed_ and can only be called by the Id of
+		/// the account going in front of `lighter`.
+		///
+		/// Only works if
+		/// - both nodes are within the same bag,
+		/// - and `origin` has a greater `VoteWeight` than `lighter`.
+		#[pallet::weight(T::WeightInfo::put_in_front_of())]
+		pub fn put_in_front_of(origin: OriginFor<T>, lighter: T::AccountId) -> DispatchResult {
+			let heavier = ensure_signed(origin)?;
+			List::<T>::put_in_front_of(&lighter, &heavier).map_err(Into::into)
+		}
 	}
 
 	#[pallet::hooks]

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

@@ -377,6 +377,83 @@ impl<T: Config> List<T> {
 		})
 	}
 
+	/// Put `heavier_id` to the position directly in front of `lighter_id`. Both ids must be in the
+	/// same bag and the `weight_of` `lighter_id` must be less than that of `heavier_id`.
+	pub(crate) fn put_in_front_of(
+		lighter_id: &T::AccountId,
+		heavier_id: &T::AccountId,
+	) -> Result<(), crate::pallet::Error<T>> {
+		use crate::pallet;
+		use frame_support::ensure;
+
+		let lighter_node = Node::<T>::get(&lighter_id).ok_or(pallet::Error::IdNotFound)?;
+		let heavier_node = Node::<T>::get(&heavier_id).ok_or(pallet::Error::IdNotFound)?;
+
+		ensure!(lighter_node.bag_upper == heavier_node.bag_upper, pallet::Error::NotInSameBag);
+
+		// this is the most expensive check, so we do it last.
+		ensure!(
+			T::VoteWeightProvider::vote_weight(&heavier_id) >
+				T::VoteWeightProvider::vote_weight(&lighter_id),
+			pallet::Error::NotHeavier
+		);
+
+		// remove the heavier node from this list. Note that this removes the node from storage and
+		// decrements the node counter.
+		Self::remove(&heavier_id);
+
+		// re-fetch `lighter_node` from storage since it may have been updated when `heavier_node`
+		// was removed.
+		let lighter_node = Node::<T>::get(&lighter_id).ok_or(pallet::Error::IdNotFound)?;
+		// insert `heavier_node` directly in front of `lighter_node`. This will update both nodes
+		// in storage and update the node counter.
+		Self::insert_at_unchecked(lighter_node, heavier_node);
+
+		Ok(())
+	}
+
+	/// Insert `node` directly in front of `at`.
+	///
+	/// WARNINGS:
+	/// - this is a naive function in that it does not check if `node` belongs to the same bag as
+	/// `at`. It is expected that the call site will check preconditions.
+	/// - this will panic if `at.bag_upper` is not a bag that already exists in storage.
+	fn insert_at_unchecked(mut at: Node<T>, mut node: Node<T>) {
+		// connect `node` to its new `prev`.
+		node.prev = at.prev.clone();
+		if let Some(mut prev) = at.prev() {
+			prev.next = Some(node.id().clone());
+			prev.put()
+		}
+
+		// connect `node` and `at`.
+		node.next = Some(at.id().clone());
+		at.prev = Some(node.id().clone());
+
+		if at.is_terminal() || node.is_terminal() {
+			// `at` is the tail and/or `node` is the head, so we make sure the bag is updated. Note,
+			// since `node` is always in front of `at` we know that 1) there is always at least 2
+			// nodes in the bag, and 2) only `node` could be the head and only `at` could be the
+			// tail.
+			let mut bag = Bag::<T>::get(at.bag_upper)
+				.expect("given nodes must always have a valid bag. qed.");
+
+			if node.prev == None {
+				bag.head = Some(node.id().clone())
+			}
+
+			if at.next == None {
+				bag.tail = Some(at.id().clone())
+			}
+
+			bag.put()
+		};
+
+		// write the updated nodes to storage.
+		at.put();
+		node.put();
+	}
+
 	/// Sanity check the list.
 	///
 	/// This should be called from the call-site, whenever one of the mutating apis (e.g. `insert`)

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

@@ -383,6 +383,123 @@ mod list {
 			assert!(non_existent_ids.iter().all(|id| !List::<Runtime>::contains(id)));
 		})
 	}
+
+	#[test]
+	#[should_panic = "given nodes must always have a valid bag. qed."]
+	fn put_in_front_of_panics_if_bag_not_found() {
+		ExtBuilder::default().skip_genesis_ids().build_and_execute_no_post_check(|| {
+			let node_10_no_bag = Node::<Runtime> { id: 10, prev: None, next: None, bag_upper: 15 };
+			let node_11_no_bag = Node::<Runtime> { id: 11, prev: None, next: None, bag_upper: 15 };
+
+			// given
+			ListNodes::<Runtime>::insert(10, node_10_no_bag);
+			ListNodes::<Runtime>::insert(11, node_11_no_bag);
+			StakingMock::set_vote_weight_of(&10, 14);
+			StakingMock::set_vote_weight_of(&11, 15);
+			assert!(!ListBags::<Runtime>::contains_key(15));
+			assert_eq!(List::<Runtime>::get_bags(), vec![]);
+
+			// then .. this panics
+			let _ = List::<Runtime>::put_in_front_of(&10, &11);
+		});
+	}
+
+	#[test]
+	fn insert_at_unchecked_at_is_only_node() {
+		// Note that this `insert_at_unchecked` test should fail post checks because node 42 does
+		// not get re-assigned the correct bagu pper. This is because `insert_at_unchecked` assumes
+		// both nodes are already in the same bag with the correct bag upper.
+		ExtBuilder::default().build_and_execute_no_post_check(|| {
+			// given
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			// implicitly also test that `node`'s `prev`/`next` are correctly re-assigned.
+			let node_42 =
+				Node::<Runtime> { id: 42, prev: Some(1), next: Some(2), bag_upper: 1_000 };
+			assert!(!crate::ListNodes::<Runtime>::contains_key(42));
+
+			let node_1 = crate::ListNodes::<Runtime>::get(&1).unwrap();
+
+			// when
+			List::<Runtime>::insert_at_unchecked(node_1, node_42);
+
+			// then
+			assert_eq!(
+				List::<Runtime>::get_bags(),
+				vec![(10, vec![42, 1]), (1_000, vec![2, 3, 4])]
+			);
+		})
+	}
+
+	#[test]
+	fn insert_at_unchecked_at_is_head() {
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			// implicitly also test that `node`'s `prev`/`next` are correctly re-assigned.
+			let node_42 = Node::<Runtime> { id: 42, prev: Some(4), next: None, bag_upper: 1_000 };
+			assert!(!crate::ListNodes::<Runtime>::contains_key(42));
+
+			let node_2 = crate::ListNodes::<Runtime>::get(&2).unwrap();
+
+			// when
+			List::<Runtime>::insert_at_unchecked(node_2, node_42);
+
+			// then
+			assert_eq!(
+				List::<Runtime>::get_bags(),
+				vec![(10, vec![1]), (1_000, vec![42, 2, 3, 4])]
+			);
+		})
+	}
+
+	#[test]
+	fn insert_at_unchecked_at_is_non_terminal() {
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			// implicitly also test that `node`'s `prev`/`next` are correctly re-assigned.
+			let node_42 = Node::<Runtime> { id: 42, prev: None, next: Some(2), bag_upper: 1_000 };
+			assert!(!crate::ListNodes::<Runtime>::contains_key(42));
+
+			let node_3 = crate::ListNodes::<Runtime>::get(&3).unwrap();
+
+			// when
+			List::<Runtime>::insert_at_unchecked(node_3, node_42);
+
+			// then
+			assert_eq!(
+				List::<Runtime>::get_bags(),
+				vec![(10, vec![1]), (1_000, vec![2, 42, 3, 4])]
+			);
+		})
+	}
+
+	#[test]
+	fn insert_at_unchecked_at_is_tail() {
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(List::<Runtime>::get_bags(), vec![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			// implicitly also test that `node`'s `prev`/`next` are correctly re-assigned.
+			let node_42 =
+				Node::<Runtime> { id: 42, prev: Some(42), next: Some(42), bag_upper: 1_000 };
+			assert!(!crate::ListNodes::<Runtime>::contains_key(42));
+
+			let node_4 = crate::ListNodes::<Runtime>::get(&4).unwrap();
+
+			// when
+			List::<Runtime>::insert_at_unchecked(node_4, node_42);
+
+			// then
+			assert_eq!(
+				List::<Runtime>::get_bags(),
+				vec![(10, vec![1]), (1_000, vec![2, 3, 42, 4])]
+			);
+		})
+	}
 }
 
 mod bags {

frame/bags-list/src/mock.rs

@@ -21,28 +21,26 @@ use super::*;
 use crate::{self as bags_list};
 use frame_election_provider_support::VoteWeight;
 use frame_support::parameter_types;
+use std::collections::HashMap;
 
 pub type AccountId = u32;
 pub type Balance = u32;
 
 parameter_types! {
+	// Set the vote weight for any id who's weight has _not_ been set with `set_vote_weight_of`.
 	pub static NextVoteWeight: VoteWeight = 0;
+	pub static NextVoteWeightMap: HashMap<AccountId, VoteWeight> = Default::default();
 }
 
 pub struct StakingMock;
 impl frame_election_provider_support::VoteWeightProvider<AccountId> for StakingMock {
 	fn vote_weight(id: &AccountId) -> VoteWeight {
-		match id {
-			710 => 15,
-			711 => 16,
-			712 => 2_000, // special cases used for migrate test
-			_ => NextVoteWeight::get(),
-		}
+		*NextVoteWeightMap::get().get(id).unwrap_or(&NextVoteWeight::get())
 	}
+
 	#[cfg(any(feature = "runtime-benchmarks", test))]
-	fn set_vote_weight_of(_: &AccountId, weight: VoteWeight) {
-		// we don't really keep a mapping, just set weight for everyone.
-		NextVoteWeight::set(weight)
+	fn set_vote_weight_of(id: &AccountId, weight: VoteWeight) {
+		NEXT_VOTE_WEIGHT_MAP.with(|m| m.borrow_mut().insert(id.clone(), weight));
 	}
 }
 
@@ -103,9 +101,17 @@ pub(crate) const GENESIS_IDS: [(AccountId, VoteWeight); 4] =
 #[derive(Default)]
 pub struct ExtBuilder {
 	ids: Vec<(AccountId, VoteWeight)>,
+	skip_genesis_ids: bool,
 }
 
 impl ExtBuilder {
+	/// Skip adding the default genesis ids to the list.
+	#[cfg(test)]
+	pub(crate) fn skip_genesis_ids(mut self) -> Self {
+		self.skip_genesis_ids = true;
+		self
+	}
+
 	/// Add some AccountIds to insert into `List`.
 	#[cfg(test)]
 	pub(crate) fn add_ids(mut self, ids: Vec<(AccountId, VoteWeight)>) -> Self {
@@ -117,10 +123,17 @@ impl ExtBuilder {
 		sp_tracing::try_init_simple();
 		let storage = frame_system::GenesisConfig::default().build_storage::<Runtime>().unwrap();
 
+		let ids_with_weight: Vec<_> = if self.skip_genesis_ids {
+			self.ids.iter().collect()
+		} else {
+			GENESIS_IDS.iter().chain(self.ids.iter()).collect()
+		};
+
 		let mut ext = sp_io::TestExternalities::from(storage);
 		ext.execute_with(|| {
-			for (id, weight) in GENESIS_IDS.iter().chain(self.ids.iter()) {
+			for (id, weight) in ids_with_weight {
 				frame_support::assert_ok!(List::<Runtime>::insert(*id, *weight));
+				StakingMock::set_vote_weight_of(id, *weight);
 			}
 		});