restore set_balance as it was

pallets/restricted-tokens/src/benchmarking.rs

@@ -332,13 +332,16 @@ benchmarks! {
 	// We let the other die to have clean-up logic in weight
 	set_balance_native {
 		let free = as_balance::<T>(300);
+		let reserved = as_balance::<T>(200);
 		let currency: <T as Config>::CurrencyId = CurrencyId::Native.into();
 		let recv = get_account::<T>("receiver", false);
 		let recv_lookup: <T::Lookup as StaticLookup>::Source = T::Lookup::unlookup(recv.clone());
 
 		make_free_balance::<T>(currency, &recv, free + free);
-	}:set_balance(RawOrigin::Root, recv_lookup, currency, free)
+		reserve_balance::<T>(currency, &recv, reserved + reserved);
+	}:set_balance(RawOrigin::Root, recv_lookup, currency, free, reserved)
 	verify {
+		assert!(<pallet_balances::Pallet<T> as fungible::InspectHold<T::AccountId>>::total_balance_on_hold(&recv) == reserved);
 		assert!(<pallet_balances::Pallet<T> as fungible::Inspect<T::AccountId>>::reducible_balance(&recv, Preservation::Protect, Fortitude::Polite) == free - <pallet_balances::Pallet<T> as fungible::Inspect<T::AccountId>>::minimum_balance());
 		assert!(<pallet_balances::Pallet<T> as fungible::Inspect<T::AccountId>>::balance(&recv) == (free));
 	}
@@ -349,13 +352,16 @@ benchmarks! {
 	// We let the other die to have clean-up logic in weight
 	set_balance_other {
 		let free = as_balance::<T>(300);
+		let reserved = as_balance::<T>(200);
 		let currency: <T as Config>::CurrencyId = get_non_native_currency::<T>();
 		let recv = get_account_maybe_permission::<T>("receiver", currency.clone());
 		let recv_loopup: <T::Lookup as StaticLookup>::Source = T::Lookup::unlookup(recv.clone());
 
 		make_free_balance::<T>(currency, &recv, free + free);
-	}:set_balance(RawOrigin::Root, recv_loopup, currency.clone(), free)
+		reserve_balance::<T>(currency, &recv, reserved + reserved);
+	}:set_balance(RawOrigin::Root, recv_loopup, currency.clone(), free, reserved)
 	verify {
+		assert!(<pallet_balances::Pallet<T> as fungible::InspectHold<T::AccountId>>::total_balance_on_hold(&recv) == reserved);
 		assert!(<orml_tokens::Pallet<T> as fungibles::Inspect<T::AccountId>>::reducible_balance(currency, &recv, Preservation::Protect, Fortitude::Polite) == free - <orml_tokens::Pallet<T> as fungibles::Inspect<T::AccountId>>::minimum_balance(currency));
 		assert!(<orml_tokens::Pallet<T> as fungibles::Inspect<T::AccountId>>::balance(currency, &recv) == (free));
 	}

pallets/restricted-tokens/src/lib.rs

@@ -67,7 +67,7 @@ pub mod pallet {
 	use frame_support::{
 		pallet_prelude::TypeInfo,
 		sp_runtime::{
-			traits::{AtLeast32BitUnsigned, StaticLookup},
+			traits::{AtLeast32BitUnsigned, EnsureAdd, StaticLookup},
 			FixedPointOperand,
 		},
 		traits::tokens::{Fortitude, Precision, Preservation},
@@ -238,6 +238,7 @@ pub mod pallet {
 			currency_id: T::CurrencyId,
 			who: T::AccountId,
 			free: T::Balance,
+			reserved: T::Balance,
 		},
 	}
 
@@ -501,22 +502,53 @@ pub mod pallet {
 			who: <T::Lookup as StaticLookup>::Source,
 			currency_id: T::CurrencyId,
 			#[pallet::compact] new_free: T::Balance,
+			#[pallet::compact] new_reserved: T::Balance,
 		) -> DispatchResultWithPostInfo {
 			ensure_root(origin)?;
 			let who = T::Lookup::lookup(who)?;
+			let new_total = new_free.ensure_add(new_reserved)?;
 
 			let token = if T::NativeToken::get() == currency_id {
+				let old_reserved = <Self as fungible::InspectHold<T::AccountId>>::balance_on_hold(
+					&HoldReason::NativeIndex.into(),
+					&who,
+				);
+
+				<Self as fungible::MutateHold<T::AccountId>>::release(
+					&HoldReason::NativeIndex.into(),
+					&who,
+					old_reserved,
+					Precision::Exact,
+				)?;
 				let to_burn = <Self as fungible::Inspect<T::AccountId>>::balance(&who);
 				<Self as fungible::Mutate<T::AccountId>>::burn_from(
 					&who,
 					to_burn,
 					Precision::Exact,
 					Fortitude::Force,
 				)?;
-				<Self as fungible::Mutate<T::AccountId>>::mint_into(&who, new_free)?;
+				<Self as fungible::Mutate<T::AccountId>>::mint_into(&who, new_total)?;
+				<Self as fungible::MutateHold<T::AccountId>>::hold(
+					&HoldReason::NativeIndex.into(),
+					&who,
+					new_reserved,
+				)?;
 
 				TokenType::Native
 			} else {
+				let old_reserved =
+					<T::Fungibles as fungibles::InspectHold<T::AccountId>>::balance_on_hold(
+						currency_id,
+						&(),
+						&who,
+					);
+				<T::Fungibles as fungibles::MutateHold<T::AccountId>>::release(
+					currency_id,
+					&(),
+					&who,
+					old_reserved,
+					Precision::Exact,
+				)?;
 				let to_burn =
 					<T::Fungibles as fungibles::Inspect<T::AccountId>>::balance(currency_id, &who);
 				<T::Fungibles as fungibles::Mutate<T::AccountId>>::burn_from(
@@ -529,7 +561,13 @@ pub mod pallet {
 				<T::Fungibles as fungibles::Mutate<T::AccountId>>::mint_into(
 					currency_id,
 					&who,
-					new_free,
+					new_total,
+				)?;
+				<T::Fungibles as fungibles::MutateHold<T::AccountId>>::hold(
+					currency_id,
+					&(),
+					&who,
+					new_reserved,
 				)?;
 
 				TokenType::Other
@@ -539,6 +577,7 @@ pub mod pallet {
 				currency_id,
 				who,
 				free: new_free,
+				reserved: new_reserved,
 			});
 
 			match token {

pallets/restricted-tokens/src/tests.rs

@@ -252,6 +252,7 @@ fn set_balance_native_works() {
 				1,
 				CurrencyId::Cfg,
 				200,
+				100,
 			));
 			assert_eq!(System::account(1).data.free, 200);
 		})
@@ -267,37 +268,51 @@ fn set_balance_foreign_works() {
 				1,
 				CurrencyId::AUSD,
 				200,
+				100,
 			));
 			assert!(orml_tokens::Pallet::<Runtime>::accounts(1, CurrencyId::AUSD).free == 200);
+			assert!(orml_tokens::Pallet::<Runtime>::accounts(1, CurrencyId::AUSD).reserved == 100);
 
 			assert_ok!(pallet_restricted_tokens::Pallet::<Runtime>::set_balance(
 				RuntimeOrigin::root(),
 				1,
 				CurrencyId::AUSD,
 				400,
+				200,
 			));
 			assert!(orml_tokens::Pallet::<Runtime>::accounts(1, CurrencyId::AUSD).free == 400);
+			assert!(orml_tokens::Pallet::<Runtime>::accounts(1, CurrencyId::AUSD).reserved == 200);
 
 			assert_ok!(pallet_restricted_tokens::Pallet::<Runtime>::set_balance(
 				RuntimeOrigin::root(),
 				1111,
 				CurrencyId::RestrictedCoin,
 				999,
+				80
 			));
 			assert!(
 				orml_tokens::Pallet::<Runtime>::accounts(1111, CurrencyId::RestrictedCoin).free
 					== 999
 			);
+			assert!(
+				orml_tokens::Pallet::<Runtime>::accounts(1111, CurrencyId::RestrictedCoin).reserved
+					== 80
+			);
 
 			assert_ok!(pallet_restricted_tokens::Pallet::<Runtime>::set_balance(
 				RuntimeOrigin::root(),
 				101,
 				CurrencyId::RestrictedCoin,
 				0,
+				100
 			));
 			assert!(
 				orml_tokens::Pallet::<Runtime>::accounts(101, CurrencyId::RestrictedCoin).free == 0
 			);
+			assert!(
+				orml_tokens::Pallet::<Runtime>::accounts(101, CurrencyId::RestrictedCoin).reserved
+					== 100
+			);
 		})
 }