Add ClassAccount storage to unique pallet (#9940)

* add ClassAccount to uniques storage

* add tests for Class and ClassAccount storage

* fix format

* fix description

* add migration

* remove extra iteration

* Update frame/uniques/src/migration.rs

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

* cargo run --quiet --release --features=runtime-benchmarks --manifest-path=bin/node/cli/Cargo.toml -- benchmark --chain=dev --steps=50 --repeat=20 --pallet=pallet_uniques --extrinsic=* --execution=wasm --wasm-execution=compiled --heap-pages=4096 --output=./frame/uniques/src/weights.rs --template=./.maintain/frame-weight-template.hbs

* fix format

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Parity Bot <admin@parity.io>
Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

frame/uniques/Cargo.toml

@@ -20,6 +20,7 @@ sp-runtime = { version = "4.0.0", default-features = false, path = "../../primit
 frame-support = { version = "4.0.0-dev", default-features = false, path = "../support" }
 frame-system = { version = "4.0.0-dev", default-features = false, path = "../system" }
 frame-benchmarking = { version = "4.0.0-dev", default-features = false, path = "../benchmarking", optional = true }
+log = { version = "0.4.14", default-features = false }
 
 [dev-dependencies]
 sp-std = { version = "4.0.0", path = "../../primitives/std" }

frame/uniques/src/functions.rs

@@ -81,6 +81,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			},
 		);
 
+		ClassAccount::<T, I>::insert(&owner, &class, ());
 		Self::deposit_event(event);
 		Ok(())
 	}
@@ -108,6 +109,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			InstanceMetadataOf::<T, I>::remove_prefix(&class, None);
 			ClassMetadataOf::<T, I>::remove(&class);
 			Attribute::<T, I>::remove_prefix((&class,), None);
+			ClassAccount::<T, I>::remove(&class_details.owner, &class);
 			T::Currency::unreserve(&class_details.owner, class_details.total_deposit);
 
 			Self::deposit_event(Event::Destroyed { class });

frame/uniques/src/lib.rs

@@ -40,6 +40,8 @@ mod impl_nonfungibles;
 mod types;
 pub use types::*;
 
+mod migration;
+
 use codec::{Decode, Encode, HasCompact};
 use frame_support::traits::{BalanceStatus::Reserved, Currency, ReservableCurrency};
 use frame_system::Config as SystemConfig;
@@ -141,6 +143,19 @@ pub mod pallet {
 		OptionQuery,
 	>;
 
+	#[pallet::storage]
+	/// The classes owned by any given account; set out this way so that classes owned by a single
+	/// account can be enumerated.
+	pub(super) type ClassAccount<T: Config<I>, I: 'static = ()> = StorageDoubleMap<
+		_,
+		Blake2_128Concat,
+		T::AccountId,
+		Blake2_128Concat,
+		T::ClassId,
+		(),
+		OptionQuery,
+	>;
+
 	#[pallet::storage]
 	/// The assets in existence and their ownership details.
 	pub(super) type Asset<T: Config<I>, I: 'static = ()> = StorageDoubleMap<
@@ -302,7 +317,11 @@ pub mod pallet {
 	}
 
 	#[pallet::hooks]
-	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {}
+	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
+		fn on_runtime_upgrade() -> frame_support::weights::Weight {
+			migration::migrate_to_v1::<T, I, Self>()
+		}
+	}
 
 	impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		/// Get the owner of the asset instance, if the asset exists.
@@ -731,6 +750,8 @@ pub mod pallet {
 					details.total_deposit,
 					Reserved,
 				)?;
+				ClassAccount::<T, I>::remove(&details.owner, &class);
+				ClassAccount::<T, I>::insert(&owner, &class, ());
 				details.owner = owner.clone();
 
 				Self::deposit_event(Event::OwnerChanged { class, new_owner: owner });
@@ -906,13 +927,17 @@ pub mod pallet {
 
 			Class::<T, I>::try_mutate(class, |maybe_asset| {
 				let mut asset = maybe_asset.take().ok_or(Error::<T, I>::Unknown)?;
-				asset.owner = T::Lookup::lookup(owner)?;
+				let old_owner = asset.owner;
+				let new_owner = T::Lookup::lookup(owner)?;
+				asset.owner = new_owner.clone();
 				asset.issuer = T::Lookup::lookup(issuer)?;
 				asset.admin = T::Lookup::lookup(admin)?;
 				asset.freezer = T::Lookup::lookup(freezer)?;
 				asset.free_holding = free_holding;
 				asset.is_frozen = is_frozen;
 				*maybe_asset = Some(asset);
+				ClassAccount::<T, I>::remove(&old_owner, &class);
+				ClassAccount::<T, I>::insert(&new_owner, &class, ());
 
 				Self::deposit_event(Event::AssetStatusChanged { class });
 				Ok(())

frame/uniques/src/migration.rs

@@ -0,0 +1,56 @@
+// 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.
+
+//! Various pieces of common functionality.
+use super::*;
+use frame_support::{
+	traits::{Get, GetStorageVersion, PalletInfoAccess, StorageVersion},
+	weights::Weight,
+};
+
+pub fn migrate_to_v1<T: Config<I>, I: 'static, P: GetStorageVersion + PalletInfoAccess>(
+) -> frame_support::weights::Weight {
+	let on_chain_storage_version = <P as GetStorageVersion>::on_chain_storage_version();
+	log::info!(
+		target: "runtime::uniques",
+		"Running migration storage v1 for uniques with storage version {:?}",
+		on_chain_storage_version,
+	);
+
+	if on_chain_storage_version < 1 {
+		let mut count = 0;
+		for (class, detail) in Class::<T, I>::iter() {
+			ClassAccount::<T, I>::insert(&detail.owner, &class, ());
+			count += 1;
+		}
+		StorageVersion::new(1).put::<P>();
+		log::info!(
+			target: "runtime::uniques",
+			"Running migration storage v1 for uniques with storage version {:?} was complete",
+			on_chain_storage_version,
+		);
+		// calculate and return migration weights
+		T::DbWeight::get().reads_writes(count as Weight + 1, count as Weight + 1)
+	} else {
+		log::warn!(
+			target: "runtime::uniques",
+			"Attempted to apply migration to v1 but failed because storage version is {:?}",
+			on_chain_storage_version,
+		);
+		T::DbWeight::get().reads(1)
+	}
+}

frame/uniques/src/tests.rs

@@ -48,6 +48,15 @@ fn assets() -> Vec<(u64, u32, u32)> {
 	r
 }
 
+fn classes() -> Vec<(u64, u32)> {
+	let mut r: Vec<_> = ClassAccount::<Test>::iter().map(|x| (x.0, x.1)).collect();
+	r.sort();
+	let mut s: Vec<_> = Class::<Test>::iter().map(|x| (x.1.owner, x.0)).collect();
+	s.sort();
+	assert_eq!(r, s);
+	r
+}
+
 macro_rules! bvec {
 	($( $x:tt )*) => {
 		vec![$( $x )*].try_into().unwrap()
@@ -73,10 +82,12 @@ fn basic_setup_works() {
 fn basic_minting_should_work() {
 	new_test_ext().execute_with(|| {
 		assert_ok!(Uniques::force_create(Origin::root(), 0, 1, true));
+		assert_eq!(classes(), vec![(1, 0)]);
 		assert_ok!(Uniques::mint(Origin::signed(1), 0, 42, 1));
 		assert_eq!(assets(), vec![(1, 0, 42)]);
 
 		assert_ok!(Uniques::force_create(Origin::root(), 1, 2, true));
+		assert_eq!(classes(), vec![(1, 0), (2, 1)]);
 		assert_ok!(Uniques::mint(Origin::signed(2), 1, 69, 1));
 		assert_eq!(assets(), vec![(1, 0, 42), (1, 1, 69)]);
 	});
@@ -88,7 +99,7 @@ fn lifecycle_should_work() {
 		Balances::make_free_balance_be(&1, 100);
 		assert_ok!(Uniques::create(Origin::signed(1), 0, 1));
 		assert_eq!(Balances::reserved_balance(&1), 2);
-
+		assert_eq!(classes(), vec![(1, 0)]);
 		assert_ok!(Uniques::set_class_metadata(Origin::signed(1), 0, bvec![0, 0], false));
 		assert_eq!(Balances::reserved_balance(&1), 5);
 		assert!(ClassMetadataOf::<Test>::contains_key(0));
@@ -120,6 +131,7 @@ fn lifecycle_should_work() {
 		assert!(!ClassMetadataOf::<Test>::contains_key(0));
 		assert!(!InstanceMetadataOf::<Test>::contains_key(0, 42));
 		assert!(!InstanceMetadataOf::<Test>::contains_key(0, 69));
+		assert_eq!(classes(), vec![]);
 		assert_eq!(assets(), vec![]);
 	});
 }
@@ -142,6 +154,7 @@ fn mint_should_work() {
 		assert_ok!(Uniques::force_create(Origin::root(), 0, 1, true));
 		assert_ok!(Uniques::mint(Origin::signed(1), 0, 42, 1));
 		assert_eq!(Uniques::owner(0, 42).unwrap(), 1);
+		assert_eq!(classes(), vec![(1, 0)]);
 		assert_eq!(assets(), vec![(1, 0, 42)]);
 	});
 }
@@ -204,7 +217,9 @@ fn transfer_owner_should_work() {
 		Balances::make_free_balance_be(&2, 100);
 		Balances::make_free_balance_be(&3, 100);
 		assert_ok!(Uniques::create(Origin::signed(1), 0, 1));
+		assert_eq!(classes(), vec![(1, 0)]);
 		assert_ok!(Uniques::transfer_ownership(Origin::signed(1), 0, 2));
+		assert_eq!(classes(), vec![(2, 0)]);
 		assert_eq!(Balances::total_balance(&1), 98);
 		assert_eq!(Balances::total_balance(&2), 102);
 		assert_eq!(Balances::reserved_balance(&1), 0);
@@ -220,6 +235,7 @@ fn transfer_owner_should_work() {
 		assert_ok!(Uniques::mint(Origin::signed(1), 0, 42, 1));
 		assert_ok!(Uniques::set_metadata(Origin::signed(2), 0, 42, bvec![0u8; 20], false));
 		assert_ok!(Uniques::transfer_ownership(Origin::signed(2), 0, 3));
+		assert_eq!(classes(), vec![(3, 0)]);
 		assert_eq!(Balances::total_balance(&2), 57);
 		assert_eq!(Balances::total_balance(&3), 145);
 		assert_eq!(Balances::reserved_balance(&2), 0);