feat(assets): add force transfer all functionality (#4390)

ComposableFi · Dec 26, 2023 · e0c1bcf · e0c1bcf
2 parents b8417ac + 07b75ab
commit e0c1bcf
Show file tree

Hide file tree

Showing 4 changed files with 172 additions and 0 deletions.
diff --git a/code/parachain/frame/assets/src/benchmarking.rs b/code/parachain/frame/assets/src/benchmarking.rs
@@ -60,6 +60,23 @@ benchmarks! {
 		T::NativeCurrency::mint_into(&caller, amount).expect("always can mint in test");
 	}: _(RawOrigin::Root, from, dest, amount, false)
 
+	force_transfer_all {
+		let caller: T::AccountId = FROM_ACCOUNT.into();
+		let asset_id: T::AssetId = ASSET_ID.into();
+		let from = T::Lookup::unlookup(FROM_ACCOUNT.into());
+		let dest = T::Lookup::unlookup(TO_ACCOUNT.into());
+		let amount: T::Balance = TRANSFER_AMOUNT.into();
+		T::MultiCurrency::mint_into(asset_id, &caller, amount).expect("always can mint in test");
+	}: _(RawOrigin::Root, asset_id, from, dest, false)
+
+	force_transfer_all_native {
+		let caller: T::AccountId = FROM_ACCOUNT.into();
+		let from = T::Lookup::unlookup(FROM_ACCOUNT.into());
+		let dest = T::Lookup::unlookup(TO_ACCOUNT.into());
+		let amount: T::Balance = TRANSFER_AMOUNT.into();
+		T::NativeCurrency::mint_into(&caller, amount).expect("always can mint in test");
+	}: _(RawOrigin::Root, from, dest, false)
+
 	transfer_all {
 		let caller: T::AccountId = whitelisted_caller();
 		let asset_id: T::AssetId = ASSET_ID.into();

diff --git a/code/parachain/frame/assets/src/lib.rs b/code/parachain/frame/assets/src/lib.rs
@@ -309,6 +309,73 @@ pub mod pallet {
 			)?;
 			Ok(().into())
 		}
+
+		/// Transfer all free balance of the `asset` from `source` to `dest`.
+		///
+		/// # Errors
+		///  - When `origin` is not signed.
+		///  - If the `dest` cannot be looked up.
+		#[pallet::weight(T::WeightInfo::force_transfer_all())]
+		#[pallet::call_index(10)]
+		pub fn force_transfer_all(
+			origin: OriginFor<T>,
+			asset: T::AssetId,
+			source: <T::Lookup as StaticLookup>::Source,
+			dest: <T::Lookup as StaticLookup>::Source,
+			keep_alive: bool,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+			let source = T::Lookup::lookup(source)?;
+			let dest = T::Lookup::lookup(dest)?;
+			let keep_alive =
+				if keep_alive { Preservation::Preserve } else { Preservation::Expendable };
+			let reducible_balance = <Self as Inspect<T::AccountId>>::reducible_balance(
+				asset,
+				&source,
+				keep_alive,
+				Fortitude::Polite,
+			);
+			<Self as Mutate<T::AccountId>>::transfer(
+				asset,
+				&source,
+				&dest,
+				reducible_balance,
+				keep_alive,
+			)?;
+			Ok(())
+		}
+
+		/// Transfer all free balance of the native asset from `source` to `dest`.
+		///
+		/// # Errors
+		///  - When `origin` is not signed.
+		///  - If the `dest` cannot be looked up.
+		#[pallet::weight(T::WeightInfo::force_transfer_all_native())]
+		#[pallet::call_index(11)]
+		pub fn force_transfer_all_native(
+			origin: OriginFor<T>,
+			source: <T::Lookup as StaticLookup>::Source,
+			dest: <T::Lookup as StaticLookup>::Source,
+			keep_alive: bool,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+			let source = T::Lookup::lookup(source)?;
+			let dest = T::Lookup::lookup(dest)?;
+			let keep_alive =
+				if keep_alive { Preservation::Preserve } else { Preservation::Expendable };
+			let reducible_balance = <Self as NativeInspect<T::AccountId>>::reducible_balance(
+				&source,
+				keep_alive,
+				Fortitude::Polite,
+			);
+			<Self as NativeMutate<T::AccountId>>::transfer(
+				&source,
+				&dest,
+				reducible_balance,
+				keep_alive,
+			)?;
+			Ok(())
+		}
 	}
 
 	pub(crate) fn valid_asset_id<T: Config>(asset_id: T::AssetId) -> Option<T::AssetId> {

diff --git a/code/parachain/frame/assets/src/tests/extrinsics.rs b/code/parachain/frame/assets/src/tests/extrinsics.rs
@@ -96,6 +96,68 @@ fn test_force_transfer_native() {
 	});
 }
 
+#[test]
+fn test_force_transfer_all_alive() {
+	new_test_ext().execute_with(|| {
+		Pallet::<Test>::force_transfer_all(
+			RuntimeOrigin::root(),
+			ASSET_ID,
+			FROM_ACCOUNT,
+			TO_ACCOUNT,
+			true,
+		)
+		.expect("force_transfer should work");
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &FROM_ACCOUNT), 1);
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &TO_ACCOUNT), INIT_AMOUNT * 2 - 1);
+	});
+}
+
+#[test]
+fn test_force_transfer_all_native_alive() {
+	new_test_ext().execute_with(|| {
+		Pallet::<Test>::force_transfer_all_native(
+			RuntimeOrigin::root(),
+			FROM_ACCOUNT,
+			TO_ACCOUNT,
+			true,
+		)
+		.expect("force_transfer_native should work");
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &FROM_ACCOUNT), 1);
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &TO_ACCOUNT), INIT_AMOUNT * 2 - 1);
+	});
+}
+
+#[test]
+fn test_force_transfer_all() {
+	new_test_ext().execute_with(|| {
+		Pallet::<Test>::force_transfer_all(
+			RuntimeOrigin::root(),
+			ASSET_ID,
+			FROM_ACCOUNT,
+			TO_ACCOUNT,
+			false,
+		)
+		.expect("force_transfer should work");
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &FROM_ACCOUNT), 0);
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &TO_ACCOUNT), INIT_AMOUNT * 2);
+	});
+}
+
+#[test]
+fn test_force_transfer_all_native() {
+	new_test_ext().execute_with(|| {
+		Pallet::<Test>::force_transfer_all_native(
+			RuntimeOrigin::root(),
+			FROM_ACCOUNT,
+			TO_ACCOUNT,
+			false,
+		)
+		.expect("force_transfer_native should work");
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &FROM_ACCOUNT), 0);
+		assert_eq!(Pallet::<Test>::total_balance(ASSET_ID, &TO_ACCOUNT), INIT_AMOUNT * 2);
+	});
+}
+
 #[test]
 fn test_transfer_all() {
 	new_test_ext().execute_with(|| {

diff --git a/code/parachain/frame/assets/src/weights.rs b/code/parachain/frame/assets/src/weights.rs
@@ -14,6 +14,8 @@ pub trait WeightInfo {
 	fn force_transfer_native() -> Weight;
 	fn transfer_all() -> Weight;
 	fn transfer_all_native() -> Weight;
+	fn force_transfer_all() -> Weight;
+	fn force_transfer_all_native() -> Weight;
 	fn mint_initialize() -> Weight;
 	fn set_administrator() -> Weight;
 	fn mint_into() -> Weight;
@@ -34,6 +36,18 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().writes(4_u64))
 	}
 
+	fn force_transfer_all() -> Weight {
+		Weight::from_parts(83_205_000_u64, 0)
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(4_u64))
+	}
+
+	fn force_transfer_all_native() -> Weight {
+		Weight::from_parts(83_205_000_u64, 0)
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(4_u64))
+	}
+
 	fn transfer() -> Weight {
 		Weight::from_parts(83_205_000_u64, 0)
 			.saturating_add(T::DbWeight::get().reads(4_u64))
@@ -106,6 +120,18 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().writes(4_u64))
 	}
 
+	fn force_transfer_all() -> Weight {
+		Weight::from_parts(83_205_000_u64, 0)
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+
+	fn force_transfer_all_native() -> Weight {
+		Weight::from_parts(83_205_000_u64, 0)
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+
 	fn transfer() -> Weight {
 		Weight::from_parts(83_205_000_u64, 0)
 			.saturating_add(RocksDbWeight::get().reads(4_u64))