Speed up try runtime checks for pallet-bags-list

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

@@ -543,7 +543,16 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 			thresholds.into_iter().filter_map(|t| Bag::<T, I>::get(t))
 		};
 
-		let _ = active_bags.clone().try_for_each(|b| b.do_try_state())?;
+		// build map of bags and the corresponding nodes to avoid multiple lookups
+		let mut bags_map = BTreeMap::<T::Score, Vec<T::AccountId>>::new();
+
+		let _ = active_bags.clone().try_for_each(|b| {
+			bags_map.insert(
+				b.bag_upper,
+				b.iter().map(|n: Node<T, I>| n.id().clone()).collect::<Vec<_>>(),
+			);
+			b.do_try_state()
+		})?;
 
 		let nodes_in_bags_count =
 			active_bags.clone().fold(0u32, |acc, cur| acc + cur.iter().count() as u32);
@@ -554,6 +563,13 @@ impl<T: Config<I>, I: 'static> List<T, I> {
 		// 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.
 		for (_id, node) in crate::ListNodes::<T, I>::iter() {
+			// check that the node is in the correct bag
+			let expected_bag = bags_map
+				.get(&node.bag_upper)
+				.ok_or("bag not found for the node in active bags")?;
+			frame_support::ensure!(expected_bag.contains(node.id()), "node not found in the bag");
+
+			// verify node state
 			node.do_try_state()?
 		}
 
@@ -790,12 +806,6 @@ impl<T: Config<I>, I: 'static> Bag<T, I> {
 	pub fn std_iter(&self) -> impl Iterator<Item = Node<T, I>> {
 		sp_std::iter::successors(self.head(), |prev| prev.next())
 	}
-
-	/// Check if the bag contains a node with `id`.
-	#[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
-	fn contains(&self, id: &T::AccountId) -> bool {
-		self.iter().any(|n| n.id() == id)
-	}
 }
 
 /// A Node is the fundamental element comprising the doubly-linked list described by `Bag`.
@@ -900,11 +910,8 @@ impl<T: Config<I>, I: 'static> Node<T, I> {
 	#[cfg(any(test, feature = "try-runtime", feature = "fuzz"))]
 	fn do_try_state(&self) -> Result<(), TryRuntimeError> {
 		let expected_bag = Bag::<T, I>::get(self.bag_upper).ok_or("bag not found for node")?;
-
 		let id = self.id();
 
-		frame_support::ensure!(expected_bag.contains(id), "node does not exist in the bag");
-
 		let non_terminal_check = !self.is_terminal() &&
 			expected_bag.head.as_ref() != Some(id) &&
 			expected_bag.tail.as_ref() != Some(id);