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

@@ -178,6 +178,21 @@ pub mod pallet {
 		Rebagged(T::AccountId, VoteWeight, 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,
+		/// Heavier weight Id is already in a higher position than a lesser weight Id.
+		AlreadyHigherPosition,
+		/// Bag could not be found. This is a system logic error that should never happen.
+		BagNotFound,
+	}
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {
 		/// Declare that some `dislocated` account has, through rewards or penalties, sufficiently
@@ -195,6 +210,25 @@ 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,
+				Box::new(T::VoteWeightProvider::vote_weight),
+			)
+			.map_err(Into::into)
+		}
 	}
 
 	#[pallet::hooks]

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

@@ -384,6 +384,85 @@ impl<T: Config> List<T> {
 		})
 	}
 
+	/// Put `heavier_id` to the position directly in front of `lighter_id`. Both ids must be in
+	/// 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,
+		weight_of: Box<dyn Fn(&T::AccountId) -> VoteWeight>,
+	) -> 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!(weight_of(&heavier_id) > weight_of(&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. 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();
+
+		// account for `node` being added to the list.
+		crate::CounterForListNodes::<T>::mutate(|prev_count| {
+			*prev_count = prev_count.saturating_add(1)
+		});
+	}
+
 	/// 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

@@ -109,42 +109,34 @@ fn remove_last_node_in_bags_cleans_bag() {
 
 #[test]
 fn migrate_works() {
-	ExtBuilder::default()
-		.add_ids(vec![(710, 15), (711, 16), (712, 2_000)])
-		.build_and_execute(|| {
-			// given
-			assert_eq!(
-				List::<Runtime>::get_bags(),
-				vec![
-					(10, vec![1]),
-					(20, vec![710, 711]),
-					(1_000, vec![2, 3, 4]),
-					(2_000, vec![712])
-				]
-			);
-			let old_thresholds = <Runtime as Config>::BagThresholds::get();
-			assert_eq!(old_thresholds, vec![10, 20, 30, 40, 50, 60, 1_000, 2_000, 10_000]);
+	ExtBuilder::default().add_aux_data().build_and_execute(|| {
+		// given
+		assert_eq!(
+			List::<Runtime>::get_bags(),
+			vec![(10, vec![1]), (20, vec![710, 711]), (1_000, vec![2, 3, 4]), (2_000, vec![712])]
+		);
+		let old_thresholds = <Runtime as Config>::BagThresholds::get();
+		assert_eq!(old_thresholds, vec![10, 20, 30, 40, 50, 60, 1_000, 2_000, 10_000]);
 
-			// when the new thresholds adds `15` and removes `2_000`
-			const NEW_THRESHOLDS: &'static [VoteWeight] =
-				&[10, 15, 20, 30, 40, 50, 60, 1_000, 10_000];
-			BagThresholds::set(NEW_THRESHOLDS);
-			// and we call
-			List::<Runtime>::migrate(old_thresholds);
+		// when the new thresholds adds `15` and removes `2_000`
+		const NEW_THRESHOLDS: &'static [VoteWeight] = &[10, 15, 20, 30, 40, 50, 60, 1_000, 10_000];
+		BagThresholds::set(NEW_THRESHOLDS);
+		// and we call
+		List::<Runtime>::migrate(old_thresholds);
 
-			// then
-			assert_eq!(
-				List::<Runtime>::get_bags(),
-				vec![
-					(10, vec![1]),
-					(15, vec![710]), // nodes in range 11 ..= 15 move from bag 20 to bag 15
-					(20, vec![711]),
-					(1_000, vec![2, 3, 4]),
-					// nodes in range 1_001 ..= 2_000 move from bag 2_000 to bag 10_000
-					(10_000, vec![712]),
-				]
-			);
-		});
+		// then
+		assert_eq!(
+			List::<Runtime>::get_bags(),
+			vec![
+				(10, vec![1]),
+				(15, vec![710]), // nodes in range 11 ..= 15 move from bag 20 to bag 15
+				(20, vec![711]),
+				(1_000, vec![2, 3, 4]),
+				// nodes in range 1_001 ..= 2_000 move from bag 2_000 to bag 10_000
+				(10_000, vec![712]),
+			]
+		);
+	});
 }
 
 mod list {
@@ -374,6 +366,127 @@ mod list {
 			assert!(non_existent_ids.iter().all(|id| !List::<Runtime>::contains(id)));
 		})
 	}
+
+	#[test]
+	#[should_panic = "given nodes must always have a valid. qed."]
+	fn put_in_front_of_panics_if_bag_not_found() {
+		ExtBuilder::default().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![(10, vec![1]), (1_000, vec![2, 3, 4])]);
+
+			let weight_fn = Box::new(<Runtime as Config>::VoteWeightProvider::vote_weight);
+
+			// then
+			assert_storage_noop!(assert_eq!(
+				List::<Runtime>::put_in_front_of(&10, &11, weight_fn).unwrap_err(),
+				crate::pallet::Error::<Runtime>::BagNotFound
+			));
+		});
+	}
+
+	#[test]
+	fn insert_at_unchecked_at_is_only_node() {
+		// Note that `insert_at_unchecked` tests should fail post checks because it does not
+		// increment the node counter.
+		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_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(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_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: 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_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(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 {