use CountedMap in pallet-bags-list (#10179)

* use CountedMap in pallet-bags-list

* Fix build

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

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* add a check as well

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

bin/node/runtime/src/lib.rs

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

frame/bags-list/Cargo.toml

@@ -64,9 +64,10 @@ runtime-benchmarks = [
 	"frame-election-provider-support/runtime-benchmarks",
 ]
 fuzz = [
-  "sp-core",
+	"sp-core",
 	"sp-io",
-  "pallet-balances",
+	"pallet-balances",
 	"sp-tracing",
 ]
+try-runtime = [ "frame-support/try-runtime" ]
 

frame/bags-list/src/lib.rs

@@ -59,6 +59,7 @@ use sp_std::prelude::*;
 mod benchmarks;
 
 mod list;
+pub mod migrations;
 #[cfg(any(test, feature = "fuzz"))]
 pub mod mock;
 #[cfg(test)]
@@ -151,17 +152,12 @@ pub mod pallet {
 		type BagThresholds: Get<&'static [VoteWeight]>;
 	}
 
-	/// How many ids are registered.
-	// NOTE: This is merely a counter for `ListNodes`. It should someday be replaced by the
-	// `CountedMap` storage.
-	#[pallet::storage]
-	pub(crate) type CounterForListNodes<T> = StorageValue<_, u32, ValueQuery>;
-
 	/// A single node, within some bag.
 	///
 	/// Nodes store links forward and back within their respective bags.
 	#[pallet::storage]
-	pub(crate) type ListNodes<T: Config> = StorageMap<_, Twox64Concat, T::AccountId, list::Node<T>>;
+	pub(crate) type ListNodes<T: Config> =
+		CountedStorageMap<_, Twox64Concat, T::AccountId, list::Node<T>>;
 
 	/// A bag stored in storage.
 	///
@@ -240,7 +236,7 @@ impl<T: Config> SortedListProvider<T::AccountId> for Pallet<T> {
 	}
 
 	fn count() -> u32 {
-		CounterForListNodes::<T>::get()
+		ListNodes::<T>::count()
 	}
 
 	fn contains(id: &T::AccountId) -> bool {

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

@@ -18,7 +18,7 @@
 //! Implementation of a "bags list": a semi-sorted list where ordering granularity is dictated by
 //! configurable thresholds that delineate the boundaries of bags. It uses a pattern of composite
 //! data structures, where multiple storage items are masked by one outer API. See [`ListNodes`],
-//! [`CounterForListNodes`] and [`ListBags`] for more information.
+//! [`ListBags`] for more information.
 //!
 //! The outer API of this module is the [`List`] struct. It wraps all acceptable operations on top
 //! of the aggregate linked list. All operations with the bags list should happen through this
@@ -77,17 +77,18 @@ pub struct List<T: Config>(PhantomData<T>);
 
 impl<T: Config> List<T> {
 	/// Remove all data associated with the list from storage. Parameter `items` is the number of
-	/// items to clear from the list. WARNING: `None` will clear all items and should generally not
-	/// be used in production as it could lead to an infinite number of storage accesses.
+	/// items to clear from the list.
+	///
+	/// ## WARNING
+	///
+	/// `None` will clear all items and should generally not be used in production as it could lead
+	/// to a very large number of storage accesses.
 	pub(crate) fn clear(maybe_count: Option<u32>) -> u32 {
 		crate::ListBags::<T>::remove_all(maybe_count);
+		let pre = crate::ListNodes::<T>::count();
 		crate::ListNodes::<T>::remove_all(maybe_count);
-		if let Some(count) = maybe_count {
-			crate::CounterForListNodes::<T>::mutate(|items| *items - count);
-			count
-		} else {
-			crate::CounterForListNodes::<T>::take()
-		}
+		let post = crate::ListNodes::<T>::count();
+		pre.saturating_sub(post)
 	}
 
 	/// Regenerate all of the data from the given ids.
@@ -274,17 +275,13 @@ impl<T: Config> List<T> {
 		// new inserts are always the tail, so we must write the bag.
 		bag.put();
 
-		crate::CounterForListNodes::<T>::mutate(|prev_count| {
-			*prev_count = prev_count.saturating_add(1)
-		});
-
 		crate::log!(
 			debug,
 			"inserted {:?} with weight {} into bag {:?}, new count is {}",
 			id,
 			weight,
 			bag_weight,
-			crate::CounterForListNodes::<T>::get(),
+			crate::ListNodes::<T>::count(),
 		);
 
 		Ok(())
@@ -331,10 +328,6 @@ impl<T: Config> List<T> {
 			bag.put();
 		}
 
-		crate::CounterForListNodes::<T>::mutate(|prev_count| {
-			*prev_count = prev_count.saturating_sub(count)
-		});
-
 		count
 	}
 
@@ -390,7 +383,7 @@ impl<T: Config> List<T> {
 	/// is being used, after all other staking data (such as counter) has been updated. It checks:
 	///
 	/// * there are no duplicate ids,
-	/// * length of this list is in sync with `CounterForListNodes`,
+	/// * length of this list is in sync with `ListNodes::count()`,
 	/// * and sanity-checks all bags and nodes. This will cascade down all the checks and makes sure
 	/// all bags and nodes are checked per *any* update to `List`.
 	#[cfg(feature = "std")]
@@ -403,7 +396,7 @@ impl<T: Config> List<T> {
 		);
 
 		let iter_count = Self::iter().count() as u32;
-		let stored_count = crate::CounterForListNodes::<T>::get();
+		let stored_count = crate::ListNodes::<T>::count();
 		let nodes_count = crate::ListNodes::<T>::iter().count() as u32;
 		ensure!(iter_count == stored_count, "iter_count != stored_count");
 		ensure!(stored_count == nodes_count, "stored_count != nodes_count");

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

@@ -18,7 +18,7 @@
 use super::*;
 use crate::{
 	mock::{test_utils::*, *},
-	CounterForListNodes, ListBags, ListNodes,
+	ListBags, ListNodes,
 };
 use frame_election_provider_support::SortedListProvider;
 use frame_support::{assert_ok, assert_storage_noop};
@@ -29,7 +29,7 @@ fn basic_setup_works() {
 		// syntactic sugar to create a raw node
 		let node = |id, prev, next, bag_upper| Node::<Runtime> { id, prev, next, bag_upper };
 
-		assert_eq!(CounterForListNodes::<Runtime>::get(), 4);
+		assert_eq!(ListNodes::<Runtime>::count(), 4);
 		assert_eq!(ListNodes::<Runtime>::iter().count(), 4);
 		assert_eq!(ListBags::<Runtime>::iter().count(), 2);
 
@@ -249,10 +249,10 @@ mod list {
 
 	#[test]
 	fn remove_works() {
-		use crate::{CounterForListNodes, ListBags, ListNodes};
+		use crate::{ListBags, ListNodes};
 		let ensure_left = |id, counter| {
 			assert!(!ListNodes::<Runtime>::contains_key(id));
-			assert_eq!(CounterForListNodes::<Runtime>::get(), counter);
+			assert_eq!(ListNodes::<Runtime>::count(), counter);
 			assert_eq!(ListNodes::<Runtime>::iter().count() as u32, counter);
 		};
 
@@ -357,10 +357,19 @@ mod list {
 			assert_eq!(List::<Runtime>::sanity_check(), Err("duplicate identified"));
 		});
 
-		// ensure count is in sync with `CounterForListNodes`.
+		// ensure count is in sync with `ListNodes::count()`.
 		ExtBuilder::default().build_and_execute_no_post_check(|| {
-			crate::CounterForListNodes::<Runtime>::mutate(|counter| *counter += 1);
-			assert_eq!(crate::CounterForListNodes::<Runtime>::get(), 5);
+			assert_eq!(crate::ListNodes::<Runtime>::count(), 4);
+			// we do some wacky stuff here to get access to the counter, since it is (reasonably)
+			// not exposed as mutable in any sense.
+			frame_support::generate_storage_alias!(
+				BagsList,
+				CounterForListNodes
+				=> Value<u32, frame_support::pallet_prelude::ValueQuery>
+			);
+			CounterForListNodes::mutate(|counter| *counter += 1);
+			assert_eq!(crate::ListNodes::<Runtime>::count(), 5);
+
 			assert_eq!(List::<Runtime>::sanity_check(), Err("iter_count != stored_count"));
 		});
 	}

frame/bags-list/src/migrations.rs

@@ -0,0 +1,49 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! The migrations of this pallet.
+
+use frame_support::traits::OnRuntimeUpgrade;
+
+/// A struct that does not migration, but only checks that the counter prefix exists and is correct.
+pub struct CheckCounterPrefix<T: crate::Config>(sp_std::marker::PhantomData<T>);
+impl<T: crate::Config> OnRuntimeUpgrade for CheckCounterPrefix<T> {
+	fn on_runtime_upgrade() -> frame_support::weights::Weight {
+		0
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<(), &'static str> {
+		use frame_support::ensure;
+		// The old explicit storage item.
+		frame_support::generate_storage_alias!(BagsList, CounterForListNodes => Value<u32>);
+
+		// ensure that a value exists in the counter struct.
+		ensure!(
+			crate::ListNodes::<T>::count() == CounterForListNodes::get().unwrap(),
+			"wrong list node counter"
+		);
+
+		crate::log!(
+			info,
+			"checked bags-list prefix to be correct and have {} nodes",
+			crate::ListNodes::<T>::count()
+		);
+
+		Ok(())
+	}
+}

frame/bags-list/src/tests.rs

@@ -340,7 +340,7 @@ mod sorted_list_provider {
 		let ensure_left = |id, counter| {
 			assert!(!ListNodes::<Runtime>::contains_key(id));
 			assert_eq!(BagsList::count(), counter);
-			assert_eq!(CounterForListNodes::<Runtime>::get(), counter);
+			assert_eq!(ListNodes::<Runtime>::count(), counter);
 			assert_eq!(ListNodes::<Runtime>::iter().count() as u32, counter);
 		};
 

frame/support/src/storage/types/counted_map.rs

@@ -31,6 +31,7 @@ use crate::{
 	Never,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen, Ref};
+use sp_arithmetic::traits::Bounded;
 use sp_runtime::traits::Saturating;
 use sp_std::prelude::*;
 
@@ -262,9 +263,10 @@ where
 	}
 
 	/// Remove all value of the storage.
-	pub fn remove_all() {
-		CounterFor::<Prefix>::set(0u32);
-		<Self as MapWrapper>::Map::remove_all(None);
+	pub fn remove_all(maybe_limit: Option<u32>) {
+		let leftover = Self::count().saturating_sub(maybe_limit.unwrap_or_else(Bounded::max_value));
+		CounterFor::<Prefix>::set(leftover);
+		<Self as MapWrapper>::Map::remove_all(maybe_limit);
 	}
 
 	/// Iter over all value of the storage.
@@ -676,7 +678,7 @@ mod test {
 			assert_eq!(A::count(), 2);
 
 			// Remove all.
-			A::remove_all();
+			A::remove_all(None);
 
 			assert_eq!(A::count(), 0);
 			assert_eq!(A::initialize_counter(), 0);
@@ -907,7 +909,7 @@ mod test {
 			assert_eq!(B::count(), 2);
 
 			// Remove all.
-			B::remove_all();
+			B::remove_all(None);
 
 			assert_eq!(B::count(), 0);
 			assert_eq!(B::initialize_counter(), 0);