Enable Merklized Distributions in Assets Pallet

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs

@@ -541,4 +541,24 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn trie_hash(_h: u32) -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs

@@ -538,4 +538,24 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn trie_hash(_h: u32) -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs

@@ -538,4 +538,24 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn trie_hash(_h: u32) -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs

@@ -547,4 +547,24 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn trie_hash(_h: u32) -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs

@@ -545,4 +545,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs

@@ -539,4 +539,24 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn trie_hash(_h: u32) -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn destroy_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

substrate/frame/assets/Cargo.toml

@@ -32,6 +32,7 @@ sp-core = { workspace = true }
 [dev-dependencies]
 sp-io = { workspace = true, default-features = true }
 pallet-balances = { workspace = true, default-features = true }
+binary-merkle-tree = { workspace = true }
 
 [features]
 default = ["std"]

substrate/frame/assets/src/benchmarking.rs

@@ -26,7 +26,7 @@ use frame_benchmarking::v1::{
 };
 use frame_support::traits::{EnsureOrigin, Get, UnfilteredDispatchable};
 use frame_system::RawOrigin as SystemOrigin;
-use sp_runtime::traits::Bounded;
+use sp_runtime::traits::{Bounded, Hash};
 
 use crate::Pallet as Assets;
 
@@ -564,5 +564,76 @@ benchmarks_instance_pallet! {
 		assert_last_event::<T, I>(Event::Transferred { asset_id: asset_id.into(), from: caller, to: target, amount }.into());
 	}
 
+	// This function is O(1), so placing any hash as a merkle root should work.
+	mint_distribution {
+		let (asset_id, caller, _) = create_default_asset::<T, I>(true);
+		let before_count = CountForMerklizedDistribution::<T, I>::get();
+	}: _(SystemOrigin::Signed(caller.clone()), asset_id.clone(), DistributionHashOf::<T, I>::default())
+	verify {
+		let count = CountForMerklizedDistribution::<T, I>::get();
+		assert_eq!(count, before_count + 1);
+		assert_last_event::<T, I>(Event::DistributionIssued { distribution_id: before_count, asset_id: asset_id.into(), merkle_root: DistributionHashOf::<T, I>::default() }.into());
+	}
+
+	// Calculate the cost of `h` hashes.
+	trie_hash {
+		let h in 0 .. 1_000;
+		let mut hash = T::Hash::default();
+	}: {
+		for _ in 0..h {
+			// Our hashes will be composed of two sub-hashes.
+			hash = T::Hashing::hash_of(&(hash, hash))
+		}
+	} verify {
+		if h > 0 {
+			assert!(hash != T::Hash::default());
+		}
+	}
+
+	// This function is O(1), so ending any distribution should work.
+	end_distribution {
+		let (asset_id, caller, _) = create_default_asset::<T, I>(true);
+		let before_count = CountForMerklizedDistribution::<T, I>::get();
+		Assets::<T, I>::mint_distribution(
+			SystemOrigin::Signed(caller.clone()).into(),
+			asset_id.clone(),
+			DistributionHashOf::<T, I>::default(),
+		)?;
+		let count = CountForMerklizedDistribution::<T, I>::get();
+		assert_eq!(count, before_count + 1);
+		assert_last_event::<T, I>(Event::DistributionIssued { distribution_id: before_count, asset_id: asset_id.into(), merkle_root: DistributionHashOf::<T, I>::default() }.into());
+	}: _(SystemOrigin::Signed(caller.clone()), before_count)
+	verify {
+		assert_last_event::<T, I>(Event::DistributionEnded { distribution_id: before_count }.into());
+	}
+
+	// This function is O(N), where N is the number of items destroyed in one extrinsic call.
+	// This benchmark cheats a little to avoid having to do hundreds or thousands of merkle proofs.
+	// Instead we call low level storage to populate the `MerklizedDistributionTracker` for our needs.
+	// If the logic of the `do_destroy_distribution` function changes, then this also needs to be updated.
+	destroy_distribution {
+		let c in 0 .. T::RemoveItemsLimit::get();
+		let (asset_id, caller, _) = create_default_asset::<T, I>(true);
+		let before_count = CountForMerklizedDistribution::<T, I>::get();
+		Assets::<T, I>::mint_distribution(
+			SystemOrigin::Signed(caller.clone()).into(),
+			asset_id.clone(),
+			DistributionHashOf::<T, I>::default(),
+		)?;
+		Assets::<T, I>::end_distribution(
+			SystemOrigin::Signed(caller.clone()).into(),
+			before_count,
+		)?;
+		for i in 0..c {
+			let account_id: T::AccountId = account("target", i, SEED);
+			MerklizedDistributionTracker::<T, I>::insert(before_count, account_id, ());
+		}
+		assert_eq!(MerklizedDistributionTracker::<T, I>::iter().count() as u32, c);
+	}: _(SystemOrigin::Signed(caller.clone()), before_count)
+	verify {
+		assert_last_event::<T, I>(Event::DistributionCleaned { distribution_id: before_count }.into());
+		assert_eq!(MerklizedDistributionTracker::<T, I>::iter().count() as u32, 0);
+	}
+
 	impl_benchmark_test_suite!(Assets, crate::mock::new_test_ext(), crate::mock::Test)
 }

substrate/frame/assets/src/functions.rs

@@ -438,6 +438,137 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// Creates a distribution in storage for asset `id`, which can be claimed via
+	/// `do_claim_distribution`.
+	pub(super) fn do_mint_distribution(
+		id: T::AssetId,
+		merkle_root: DistributionHashOf<T, I>,
+		maybe_check_issuer: Option<T::AccountId>,
+	) -> DispatchResult {
+		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
+		ensure!(details.status == AssetStatus::Live, Error::<T, I>::AssetNotLive);
+
+		if let Some(check_issuer) = maybe_check_issuer {
+			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
+		}
+
+		let info = DistributionInfo {
+			asset_id: id.clone(),
+			merkle_root: merkle_root.clone(),
+			active: true,
+		};
+
+		let distribution_id: u32 = CountForMerklizedDistribution::<T, I>::get();
+		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
+		CountForMerklizedDistribution::<T, I>::put(
+			distribution_id.checked_add(1).ok_or(ArithmeticError::Overflow)?,
+		);
+
+		Self::deposit_event(Event::DistributionIssued {
+			distribution_id,
+			asset_id: id,
+			merkle_root,
+		});
+
+		Ok(())
+	}
+
+	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to
+	/// whom.
+	pub(super) fn do_claim_distribution(
+		distribution_id: DistributionCounter,
+		merkle_proof: Vec<u8>,
+		hashes: u32,
+	) -> DispatchResult {
+		let proof =
+			codec::Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadProof)?;
+
+		let expected_hashes = T::VerifyExistenceProof::proof_to_hashes(&proof)
+			.map_err(|_| Error::<T, I>::BadProof)?;
+
+		ensure!(hashes >= expected_hashes, Error::<T, I>::TooManyHashes);
+
+		let DistributionInfo { asset_id, merkle_root, active } =
+			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		ensure!(active, Error::<T, I>::DistributionEnded);
+
+		let leaf = T::VerifyExistenceProof::verify_proof(proof, &merkle_root)?;
+		let (beneficiary, amount) =
+			codec::Decode::decode(&mut &leaf[..]).map_err(|_| Error::<T, I>::CannotDecodeLeaf)?;
+
+		ensure!(
+			!MerklizedDistributionTracker::<T, I>::contains_key(distribution_id, &beneficiary),
+			Error::<T, I>::AlreadyClaimed
+		);
+
+		Self::do_mint(asset_id, &beneficiary, amount, None)?;
+		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
+
+		Ok(())
+	}
+
+	/// Ends the asset distribution of `distribution_id`.
+	pub(super) fn do_end_distribution(
+		distribution_id: DistributionCounter,
+		maybe_check_issuer: Option<T::AccountId>,
+	) -> DispatchResult {
+		let mut info =
+			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
+		let details = Asset::<T, I>::get(&info.asset_id).ok_or(Error::<T, I>::Unknown)?;
+
+		if let Some(check_issuer) = maybe_check_issuer {
+			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
+		}
+
+		info.active = false;
+
+		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
+
+		Self::deposit_event(Event::DistributionEnded { distribution_id });
+
+		Ok(())
+	}
+
+	/// Cleans up the distribution tracker of `distribution_id`.
+	/// Iterates and cleans up data in the `MerklizedDistributionTracker` map `RemoveItemsLimit` at
+	/// a time. This function may need to be called multiple times to complete successfully.
+	pub(super) fn do_destroy_distribution(
+		distribution_id: DistributionCounter,
+	) -> DispatchResultWithPostInfo {
+		let info =
+			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		ensure!(!info.active, Error::<T, I>::DistributionActive);
+
+		let mut refund_count = 0u32;
+		let distribution_iterator =
+			MerklizedDistributionTracker::<T, I>::iter_key_prefix(&distribution_id);
+
+		let mut all_refunded = true;
+		for who in distribution_iterator {
+			if refund_count >= T::RemoveItemsLimit::get() {
+				// Not everyone was able to be refunded this time around.
+				all_refunded = false;
+				break
+			}
+
+			MerklizedDistributionTracker::<T, I>::remove(&distribution_id, &who);
+			refund_count += 1;
+		}
+
+		if all_refunded {
+			MerklizedDistribution::<T, I>::remove(&distribution_id);
+			Self::deposit_event(Event::<T, I>::DistributionCleaned { distribution_id });
+			// Refund weight only the amount we actually used.
+			Ok(Some(T::WeightInfo::destroy_distribution(refund_count)).into())
+		} else {
+			Self::deposit_event(Event::<T, I>::DistributionPartiallyCleaned { distribution_id });
+			// No weight to refund since we did not finish the loop.
+			Ok(().into())
+		}
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need