Enable Merklized Distributions in Assets Pallet

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

@@ -541,4 +541,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-rococo/src/weights/pallet_assets_local.rs

@@ -538,4 +538,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-rococo/src/weights/pallet_assets_pool.rs

@@ -538,4 +538,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_foreign.rs

@@ -547,4 +547,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_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,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()
+	}
 }

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

@@ -139,6 +139,29 @@ fn assert_event<T: Config<I>, I: 'static>(generic_event: <T as Config<I>>::Runti
 	frame_system::Pallet::<T>::assert_has_event(generic_event.into());
 }
 
+// fn generate_merkle_trie<T: Config<I>, I: 'static>(items: u32) -> (DistributionHashOf<T, I>,
+// DistributionProofOf<T, I>) { 	use codec::Encode;
+// 	let flat_distribution = Vec::<Vec<u8>>::with_capacity(items as usize);
+// 	for i in 0..items {
+// 		let account: T::AccountId = account("target", i, SEED);
+// 		let balance: T::Balance = i.into();
+
+// 		flat_distribution.push((account, balance).encode());
+// 	}
+
+// 	let root = binary_merkle_tree::merkle_root::<<T as frame_system::Config>::Hashing, _>(
+// 		flat_distribution.clone(),
+// 	);
+
+// 	let proof = binary_merkle_tree::merkle_proof::<
+// 		<T as frame_system::Config>::Hashing,
+// 		_,
+// 		_,
+// 	>(flat_distribution.clone(), items);
+
+// 	return (root, proof)
+// }
+
 benchmarks_instance_pallet! {
 	create {
 		let asset_id = default_asset_id::<T, I>();
@@ -564,5 +587,61 @@ 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 = MerklizedDistribution::<T, I>::count();
+	}: _(SystemOrigin::Signed(caller.clone()), asset_id.clone(), DistributionHashOf::<T, I>::default())
+	verify {
+		let count = MerklizedDistribution::<T, I>::count();
+		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());
+	}
+
+	// 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 = MerklizedDistribution::<T, I>::count();
+		Assets::<T, I>::mint_distribution(
+			SystemOrigin::Signed(caller.clone()).into(),
+			asset_id.clone(),
+			DistributionHashOf::<T, I>::default(),
+		)?;
+		let count = MerklizedDistribution::<T, I>::count();
+		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 = MerklizedDistribution::<T, I>::count();
+		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,128 @@ 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 = MerklizedDistribution::<T, I>::count();
+		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
+
+		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>,
+	) -> DispatchResult {
+		let proof =
+			codec::Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadProof)?;
+
+		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)
+			.map_err(|()| Error::<T, I>::BadProof)?;
+		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 {
+			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