Add Score to Bags List (#11357)

* Add Score to Bags List

* fix ordering

* make compile

* in progress migration

* make migration compile

* remove old check

* remove runtime specific migration

* fix warning

* Apply suggestions from code review

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* improve migration

* fix

* fix merge

* fmt

* Update migrations.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

bin/node/runtime/src/lib.rs

@@ -1572,7 +1572,7 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
-	pallet_bags_list::migrations::CheckCounterPrefix<Runtime>,
+	(),
 >;
 
 /// MMR helper types.

frame/bags-list/fuzzer/src/main.rs

@@ -70,13 +70,20 @@ fn main() {
 				},
 				Action::Update => {
 					let already_contains = BagsList::contains(&id);
-					BagsList::on_update(&id, vote_weight).unwrap();
 					if already_contains {
+						BagsList::on_update(&id, vote_weight).unwrap();
 						assert!(BagsList::contains(&id));
+					} else {
+						BagsList::on_update(&id, vote_weight).unwrap_err();
 					}
 				},
 				Action::Remove => {
-					BagsList::on_remove(&id).unwrap();
+					let already_contains = BagsList::contains(&id);
+					if already_contains {
+						BagsList::on_remove(&id).unwrap();
+					} else {
+						BagsList::on_remove(&id).unwrap_err();
+					}
 					assert!(!BagsList::contains(&id));
 				},
 			}

frame/bags-list/src/lib.rs

@@ -82,7 +82,10 @@ macro_rules! log {
 	($level:tt, $patter:expr $(, $values:expr)* $(,)?) => {
 		log::$level!(
 			target: crate::LOG_TARGET,
-			concat!("[{:?}] 👜", $patter), <frame_system::Pallet<T>>::block_number() $(, $values)*
+			concat!("[{:?}] 👜 [{}]", $patter),
+			<frame_system::Pallet<T>>::block_number(),
+			<crate::Pallet::<T, I> as frame_support::traits::PalletInfoAccess>::name()
+			$(, $values)*
 		)
 	};
 }
@@ -189,6 +192,8 @@ pub mod pallet {
 	pub enum Event<T: Config<I>, I: 'static = ()> {
 		/// Moved an account from one bag to another.
 		Rebagged { who: T::AccountId, from: T::Score, to: T::Score },
+		/// Updated the score of some account to the given amount.
+		ScoreUpdated { who: T::AccountId, new_score: T::Score },
 	}
 
 	#[pallet::error]
@@ -212,13 +217,16 @@ pub mod pallet {
 		///
 		/// Anyone can call this function about any potentially dislocated account.
 		///
-		/// Will never return an error; if `dislocated` does not exist or doesn't need a rebag, then
-		/// it is a noop and fees are still collected from `origin`.
+		/// Will always update the stored score of `dislocated` to the correct score, based on
+		/// `ScoreProvider`.
+		///
+		/// If `dislocated` does not exists, it returns an error.
 		#[pallet::weight(T::WeightInfo::rebag_non_terminal().max(T::WeightInfo::rebag_terminal()))]
 		pub fn rebag(origin: OriginFor<T>, dislocated: T::AccountId) -> DispatchResult {
 			ensure_signed(origin)?;
 			let current_score = T::ScoreProvider::score(&dislocated);
-			let _ = Pallet::<T, I>::do_rebag(&dislocated, current_score);
+			let _ = Pallet::<T, I>::do_rebag(&dislocated, current_score)
+				.map_err::<Error<T, I>, _>(Into::into)?;
 			Ok(())
 		}
 
@@ -265,6 +273,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		if let Some((from, to)) = maybe_movement {
 			Self::deposit_event(Event::<T, I>::Rebagged { who: account.clone(), from, to });
 		};
+		Self::deposit_event(Event::<T, I>::ScoreUpdated { who: account.clone(), new_score });
 		Ok(maybe_movement)
 	}
 
@@ -302,6 +311,10 @@ impl<T: Config<I>, I: 'static> SortedListProvider<T::AccountId> for Pallet<T, I>
 		List::<T, I>::insert(id, score)
 	}
 
+	fn get_score(id: &T::AccountId) -> Result<T::Score, ListError> {
+		List::<T, I>::get_score(id)
+	}
+
 	fn on_update(id: &T::AccountId, new_score: T::Score) -> Result<(), ListError> {
 		Pallet::<T, I>::do_rebag(id, new_score).map(|_| ())
 	}
@@ -359,3 +372,22 @@ impl<T: Config<I>, I: 'static> SortedListProvider<T::AccountId> for Pallet<T, I>
 		}
 	}
 }
+
+impl<T: Config<I>, I: 'static> ScoreProvider<T::AccountId> for Pallet<T, I> {
+	type Score = <Pallet<T, I> as SortedListProvider<T::AccountId>>::Score;
+
+	fn score(id: &T::AccountId) -> T::Score {
+		Node::<T, I>::get(id).map(|node| node.score()).unwrap_or_default()
+	}
+
+	#[cfg(any(feature = "runtime-benchmarks", test))]
+	fn set_score_of(id: &T::AccountId, new_score: T::Score) {
+		ListNodes::<T, I>::mutate(id, |maybe_node| {
+			if let Some(node) = maybe_node.as_mut() {
+				node.set_score(new_score)
+			} else {
+				panic!("trying to mutate {:?} which does not exists", id);
+			}
+		})
+	}
+}

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

@@ -28,6 +28,7 @@ use crate::Config;
 use codec::{Decode, Encode, MaxEncodedLen};
 use frame_election_provider_support::ScoreProvider;
 use frame_support::{
+	ensure,
 	traits::{Defensive, Get},
 	DefaultNoBound, PalletError,
 };
@@ -38,7 +39,7 @@ use sp_std::{
 	collections::{btree_map::BTreeMap, btree_set::BTreeSet},
 	iter,
 	marker::PhantomData,
-	vec::Vec,
+	prelude::*,
 };
 
 #[derive(Debug, PartialEq, Eq, Encode, Decode, MaxEncodedLen, TypeInfo, PalletError)]
@@ -227,6 +228,11 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 		crate::ListNodes::<T, I>::contains_key(id)
 	}
 
+	/// Get the score of the given node,
+	pub fn get_score(id: &T::AccountId) -> Result<T::Score, ListError> {
+		Node::<T, I>::get(id).map(|node| node.score()).ok_or(ListError::NodeNotFound)
+	}
+
 	/// Iterate over all nodes in all bags in the list.
 	///
 	/// Full iteration can be expensive; it's recommended to limit the number of items with
@@ -310,7 +316,7 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 		let bag_score = notional_bag_for::<T, I>(score);
 		let mut bag = Bag::<T, I>::get_or_make(bag_score);
 		// unchecked insertion is okay; we just got the correct `notional_bag_for`.
-		bag.insert_unchecked(id.clone());
+		bag.insert_unchecked(id.clone(), score);
 
 		// new inserts are always the tail, so we must write the bag.
 		bag.put();
@@ -381,16 +387,18 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 	/// If the node was in the correct bag, no effect. If the node was in the incorrect bag, they
 	/// are moved into the correct bag.
 	///
-	/// Returns `Some((old_idx, new_idx))` if the node moved, otherwise `None`.
+	/// Returns `Some((old_idx, new_idx))` if the node moved, otherwise `None`. In both cases, the
+	/// node's score is written to the `score` field. Thus, this is not a noop, even if `None`.
 	///
 	/// This operation is somewhat more efficient than simply calling [`self.remove`] followed by
 	/// [`self.insert`]. However, given large quantities of nodes to move, it may be more efficient
 	/// to call [`self.remove_many`] followed by [`self.insert_many`].
 	pub(crate) fn update_position_for(
-		node: Node<T, I>,
+		mut node: Node<T, I>,
 		new_score: T::Score,
 	) -> Option<(T::Score, T::Score)> {
-		node.is_misplaced(new_score).then(move || {
+		node.score = new_score;
+		if node.is_misplaced(new_score) {
 			let old_bag_upper = node.bag_upper;
 
 			if !node.is_terminal() {
@@ -402,12 +410,9 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 				bag.remove_node_unchecked(&node);
 				bag.put();
 			} else {
-				crate::log!(
-					error,
-					"Node {:?} did not have a bag; ListBags is in an inconsistent state",
-					node.id,
+				frame_support::defensive!(
+					"Node did not have a bag; BagsList is in an inconsistent state"
 				);
-				debug_assert!(false, "every node must have an extant bag associated with it");
 			}
 
 			// put the node into the appropriate new bag.
@@ -418,8 +423,12 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 			bag.insert_node_unchecked(node);
 			bag.put();
 
-			(old_bag_upper, new_bag_upper)
-		})
+			Some((old_bag_upper, new_bag_upper))
+		} else {
+			// just write the new score.
+			node.put();
+			None
+		}
 	}
 
 	/// Put `heavier_id` to the position directly in front of `lighter_id`. Both ids must be in the
@@ -428,8 +437,6 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 		lighter_id: &T::AccountId,
 		heavier_id: &T::AccountId,
 	) -> Result<(), ListError> {
-		use frame_support::ensure;
-
 		let lighter_node = Node::<T, I>::get(&lighter_id).ok_or(ListError::NodeNotFound)?;
 		let heavier_node = Node::<T, I>::get(&heavier_id).ok_or(ListError::NodeNotFound)?;
 
@@ -510,7 +517,6 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 	/// all bags and nodes are checked per *any* update to `List`.
 	#[cfg(feature = "std")]
 	pub(crate) fn sanity_check() -> Result<(), &'static str> {
-		use frame_support::ensure;
 		let mut seen_in_list = BTreeSet::new();
 		ensure!(
 			Self::iter().map(|node| node.id).all(|id| seen_in_list.insert(id)),
@@ -523,7 +529,7 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 		ensure!(iter_count == stored_count, "iter_count != stored_count");
 		ensure!(stored_count == nodes_count, "stored_count != nodes_count");
 
-		crate::log!(debug, "count of nodes: {}", stored_count);
+		crate::log!(trace, "count of nodes: {}", stored_count);
 
 		let active_bags = {
 			let thresholds = T::BagThresholds::get().iter().copied();
@@ -544,7 +550,7 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 			active_bags.clone().fold(0u32, |acc, cur| acc + cur.iter().count() as u32);
 		ensure!(nodes_count == nodes_in_bags_count, "stored_count != nodes_in_bags_count");
 
-		crate::log!(debug, "count of active bags {}", active_bags.count());
+		crate::log!(trace, "count of active bags {}", active_bags.count());
 
 		// check that all nodes are sane. We check the `ListNodes` storage item directly in case we
 		// have some "stale" nodes that are not in a bag.
@@ -667,7 +673,7 @@ impl<T: Config<I>, I: 'static> Bag<T, I> {
 	///
 	/// Storage note: this modifies storage, but only for the nodes. You still need to call
 	/// `self.put()` after use.
-	fn insert_unchecked(&mut self, id: T::AccountId) {
+	fn insert_unchecked(&mut self, id: T::AccountId, score: T::Score) {
 		// insert_node will overwrite `prev`, `next` and `bag_upper` to the proper values. As long
 		// as this bag is the correct one, we're good. All calls to this must come after getting the
 		// correct [`notional_bag_for`].
@@ -676,6 +682,7 @@ impl<T: Config<I>, I: 'static> Bag<T, I> {
 			prev: None,
 			next: None,
 			bag_upper: Zero::zero(),
+			score,
 			_phantom: PhantomData,
 		});
 	}
@@ -784,11 +791,6 @@ impl<T: Config<I>, I: 'static> Bag<T, I> {
 		Ok(())
 	}
 
-	#[cfg(not(feature = "std"))]
-	fn sanity_check(&self) -> Result<(), &'static str> {
-		Ok(())
-	}
-
 	/// Iterate over the nodes in this bag (public for tests).
 	#[cfg(feature = "std")]
 	#[allow(dead_code)]
@@ -809,12 +811,13 @@ impl<T: Config<I>, I: 'static> Bag<T, I> {
 #[scale_info(skip_type_params(T, I))]
 #[cfg_attr(feature = "std", derive(frame_support::DebugNoBound, Clone, PartialEq))]
 pub struct Node<T: Config<I>, I: 'static = ()> {
-	id: T::AccountId,
-	prev: Option<T::AccountId>,
-	next: Option<T::AccountId>,
-	bag_upper: T::Score,
+	pub(crate) id: T::AccountId,
+	pub(crate) prev: Option<T::AccountId>,
+	pub(crate) next: Option<T::AccountId>,
+	pub(crate) bag_upper: T::Score,
+	pub(crate) score: T::Score,
 	#[codec(skip)]
-	_phantom: PhantomData<I>,
+	pub(crate) _phantom: PhantomData<I>,
 }
 
 impl<T: Config<I>, I: 'static> Node<T, I> {
@@ -877,13 +880,23 @@ impl<T: Config<I>, I: 'static> Node<T, I> {
 		&self.id
 	}
 
+	/// Get the current vote weight of the node.
+	pub(crate) fn score(&self) -> T::Score {
+		self.score
+	}
+
 	/// Get the underlying voter (public fo tests).
 	#[cfg(feature = "std")]
 	#[allow(dead_code)]
 	pub fn std_id(&self) -> &T::AccountId {
 		&self.id
 	}
 
+	#[cfg(any(feature = "runtime-benchmarks", test))]
+	pub fn set_score(&mut self, s: T::Score) {
+		self.score = s
+	}
+
 	/// The bag this nodes belongs to (public for benchmarks).
 	#[cfg(feature = "runtime-benchmarks")]
 	#[allow(dead_code)]