support multi currency rewards (#601)

* support multi currency rewards

* update migration

* remove storage_version

* fixes

* add missing import

* update docs

* add missing import

* fix currency type

* make sure there're no empty entries

* import prelude::*

* update std import

* accumulate reward returns error if pool does not exist

* cleanup

* cleanup

* move PoolInfoV0 to migration

* test migration

* update migration test

* keep old map storage for migrate keys

* new migration

* fix migration

* update

* limit migration

* fix

Co-authored-by: wangjj9219 <[email protected]>

Author: ermalkaleci

rewards/Cargo.toml

@@ -13,13 +13,11 @@ codec = { package = "parity-scale-codec", version = "2.2.0", default-features =
 sp-runtime = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
 sp-io = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
 sp-std = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
+sp-core = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
 frame-support = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
 frame-system = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9", default-features = false }
 orml-traits = { path = "../traits", version = "0.4.1-dev", default-features = false }
 
-[dev-dependencies]
-sp-core = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.9" }
-
 [features]
 default = ["std"]
 std = [
@@ -28,6 +26,7 @@ std = [
 	"sp-runtime/std",
 	"sp-io/std",
 	"sp-std/std",
+	"sp-core/std",
 	"frame-support/std",
 	"frame-system/std",
 	"orml-traits/std",

rewards/src/lib.rs

@@ -1,34 +1,42 @@
 #![allow(clippy::unused_unit)]
 #![cfg_attr(not(feature = "std"), no_std)]
 
+pub mod migrations;
 mod mock;
 mod tests;
 
 use codec::{FullCodec, HasCompact, MaxEncodedLen};
-use frame_support::pallet_prelude::*;
+use frame_support::{pallet_prelude::*, weights::Weight};
+pub use migrations::PoolInfoV0;
+use orml_traits::RewardHandler;
+use sp_core::U256;
 use sp_runtime::{
-	traits::{AtLeast32BitUnsigned, Bounded, MaybeSerializeDeserialize, Member, Saturating, Zero},
-	FixedPointNumber, FixedPointOperand, FixedU128, RuntimeDebug,
-};
-use sp_std::{
-	cmp::{Eq, PartialEq},
-	fmt::Debug,
+	traits::{AtLeast32BitUnsigned, Convert, MaybeSerializeDeserialize, Member, Saturating, UniqueSaturatedInto, Zero},
+	FixedPointOperand, RuntimeDebug, SaturatedConversion,
 };
-
-use orml_traits::RewardHandler;
+use sp_std::{borrow::ToOwned, collections::btree_map::BTreeMap, fmt::Debug, prelude::*};
 
 /// The Reward Pool Info.
-#[derive(Clone, Encode, Decode, PartialEq, Eq, RuntimeDebug, Default, MaxEncodedLen)]
-pub struct PoolInfo<Share: HasCompact, Balance: HasCompact> {
+#[derive(Clone, Encode, Decode, PartialEq, Eq, RuntimeDebug)]
+pub struct PoolInfo<Share: HasCompact, Balance: HasCompact, CurrencyId: Ord> {
 	/// Total shares amount
-	#[codec(compact)]
 	pub total_shares: Share,
-	/// Total rewards amount
-	#[codec(compact)]
-	pub total_rewards: Balance,
-	/// Total withdrawn rewards amount
-	#[codec(compact)]
-	pub total_withdrawn_rewards: Balance,
+	/// Reward infos <reward_currency, (total_reward, total_withdrawn_reward)>
+	pub rewards: BTreeMap<CurrencyId, (Balance, Balance)>,
+}
+
+impl<Share, Balance, CurrencyId> Default for PoolInfo<Share, Balance, CurrencyId>
+where
+	Share: Default + HasCompact,
+	Balance: HasCompact,
+	CurrencyId: Ord,
+{
+	fn default() -> Self {
+		Self {
+			total_shares: Default::default(),
+			rewards: BTreeMap::new(),
+		}
+	}
 }
 
 pub use module::*;
@@ -59,24 +67,64 @@ pub mod module {
 			+ Debug
 			+ FixedPointOperand;
 
+		/// The old version of reward pool ID type.
+		/// NOTE: remove it after migration
+		type PoolIdV0: Parameter + Member + Clone + FullCodec;
+
+		/// The convertor to convert PoolIdV0 to PoolId
+		/// NOTE: remove it after migration
+		type PoolIdConvertor: Convert<Self::PoolIdV0, Option<Self::PoolId>>;
+
 		/// The reward pool ID type.
 		type PoolId: Parameter + Member + Clone + FullCodec;
 
+		type CurrencyId: Parameter + Member + Copy + MaybeSerializeDeserialize + Ord;
+
 		/// The `RewardHandler`
-		type Handler: RewardHandler<Self::AccountId, Balance = Self::Balance, PoolId = Self::PoolId>;
+		type Handler: RewardHandler<Self::AccountId, Self::CurrencyId, Balance = Self::Balance, PoolId = Self::PoolId>;
+	}
+
+	#[pallet::error]
+	pub enum Error<T> {
+		/// Pool does not exist
+		PoolDoesNotExist,
 	}
 
 	/// Stores reward pool info.
+	/// NOTE: remove it after migration
 	#[pallet::storage]
 	#[pallet::getter(fn pools)]
-	pub type Pools<T: Config> = StorageMap<_, Twox64Concat, T::PoolId, PoolInfo<T::Share, T::Balance>, ValueQuery>;
+	pub type Pools<T: Config> = StorageMap<_, Twox64Concat, T::PoolIdV0, PoolInfoV0<T::Share, T::Balance>, ValueQuery>;
+
+	/// Record reward pool info.
+	#[pallet::storage]
+	#[pallet::getter(fn pool_infos)]
+	pub type PoolInfos<T: Config> =
+		StorageMap<_, Twox64Concat, T::PoolId, PoolInfo<T::Share, T::Balance, T::CurrencyId>, ValueQuery>;
 
 	/// Record share amount and withdrawn reward amount for specific `AccountId`
 	/// under `PoolId`.
+	/// NOTE: remove it after migration
 	#[pallet::storage]
 	#[pallet::getter(fn share_and_withdrawn_reward)]
 	pub type ShareAndWithdrawnReward<T: Config> =
-		StorageDoubleMap<_, Twox64Concat, T::PoolId, Twox64Concat, T::AccountId, (T::Share, T::Balance), ValueQuery>;
+		StorageDoubleMap<_, Twox64Concat, T::PoolIdV0, Twox64Concat, T::AccountId, (T::Share, T::Balance), ValueQuery>;
+
+	/// Record share amount, reward currency and withdrawn reward amount for
+	/// specific `AccountId` under `PoolId`.
+	///
+	/// double_map (PoolId, AccountId) => (Share, BTreeMap<CurrencyId, Balance>)
+	#[pallet::storage]
+	#[pallet::getter(fn shares_and_withdrawn_rewards)]
+	pub type SharesAndWithdrawnRewards<T: Config> = StorageDoubleMap<
+		_,
+		Twox64Concat,
+		T::PoolId,
+		Twox64Concat,
+		T::AccountId,
+		(T::Share, BTreeMap<T::CurrencyId, T::Balance>),
+		ValueQuery,
+	>;
 
 	#[pallet::pallet]
 	pub struct Pallet<T>(_);
@@ -89,35 +137,73 @@ pub mod module {
 }
 
 impl<T: Config> Pallet<T> {
-	pub fn accumulate_reward(pool: &T::PoolId, reward_increment: T::Balance) {
-		if !reward_increment.is_zero() {
-			Pools::<T>::mutate(pool, |pool_info| {
-				pool_info.total_rewards = pool_info.total_rewards.saturating_add(reward_increment)
-			});
+	pub fn accumulate_reward(
+		pool: &T::PoolId,
+		reward_currency: T::CurrencyId,
+		reward_increment: T::Balance,
+	) -> DispatchResult {
+		if reward_increment.is_zero() {
+			return Ok(());
 		}
+		PoolInfos::<T>::mutate_exists(pool, |maybe_pool_info| -> DispatchResult {
+			let pool_info = maybe_pool_info.as_mut().ok_or(Error::<T>::PoolDoesNotExist)?;
+
+			pool_info
+				.rewards
+				.entry(reward_currency)
+				.and_modify(|(total_reward, _)| {
+					*total_reward = total_reward.saturating_add(reward_increment);
+				})
+				.or_insert((reward_increment, Zero::zero()));
+
+			Ok(())
+		})
 	}
 
 	pub fn add_share(who: &T::AccountId, pool: &T::PoolId, add_amount: T::Share) {
 		if add_amount.is_zero() {
 			return;
 		}
 
-		Pools::<T>::mutate(pool, |pool_info| {
-			let reward_inflation = if pool_info.total_shares.is_zero() {
-				Zero::zero()
-			} else {
-				let proportion = FixedU128::checked_from_rational(add_amount, pool_info.total_shares)
-					.unwrap_or_else(FixedU128::max_value);
-				proportion.saturating_mul_int(pool_info.total_rewards)
-			};
-
+		PoolInfos::<T>::mutate(pool, |pool_info| {
+			let initial_total_shares = pool_info.total_shares;
 			pool_info.total_shares = pool_info.total_shares.saturating_add(add_amount);
-			pool_info.total_rewards = pool_info.total_rewards.saturating_add(reward_inflation);
-			pool_info.total_withdrawn_rewards = pool_info.total_withdrawn_rewards.saturating_add(reward_inflation);
 
-			ShareAndWithdrawnReward::<T>::mutate(pool, who, |(share, withdrawn_rewards)| {
+			let mut withdrawn_inflation = Vec::<(T::CurrencyId, T::Balance)>::new();
+
+			pool_info
+				.rewards
+				.iter_mut()
+				.for_each(|(reward_currency, (total_reward, total_withdrawn_reward))| {
+					let reward_inflation = if initial_total_shares.is_zero() {
+						Zero::zero()
+					} else {
+						U256::from(add_amount.to_owned().saturated_into::<u128>())
+							.saturating_mul(total_reward.to_owned().saturated_into::<u128>().into())
+							.checked_div(initial_total_shares.to_owned().saturated_into::<u128>().into())
+							.unwrap_or_default()
+							.as_u128()
+							.saturated_into()
+					};
+					*total_reward = total_reward.saturating_add(reward_inflation);
+					*total_withdrawn_reward = total_withdrawn_reward.saturating_add(reward_inflation);
+
+					withdrawn_inflation.push((*reward_currency, reward_inflation));
+				});
+
+			SharesAndWithdrawnRewards::<T>::mutate(pool, who, |(share, withdrawn_rewards)| {
 				*share = share.saturating_add(add_amount);
-				*withdrawn_rewards = withdrawn_rewards.saturating_add(reward_inflation);
+				// update withdrawn inflation for each reward currency
+				withdrawn_inflation
+					.into_iter()
+					.for_each(|(reward_currency, reward_inflation)| {
+						withdrawn_rewards
+							.entry(reward_currency)
+							.and_modify(|withdrawn_reward| {
+								*withdrawn_reward = withdrawn_reward.saturating_add(reward_inflation);
+							})
+							.or_insert(reward_inflation);
+					});
 			});
 		});
 	}
@@ -130,26 +216,50 @@ impl<T: Config> Pallet<T> {
 		// claim rewards firstly
 		Self::claim_rewards(who, pool);
 
-		ShareAndWithdrawnReward::<T>::mutate_exists(pool, who, |share_info| {
+		SharesAndWithdrawnRewards::<T>::mutate_exists(pool, who, |share_info| {
 			if let Some((mut share, mut withdrawn_rewards)) = share_info.take() {
 				let remove_amount = remove_amount.min(share);
 
 				if remove_amount.is_zero() {
 					return;
 				}
 
-				Pools::<T>::mutate(pool, |pool_info| {
-					let proportion = FixedU128::checked_from_rational(remove_amount, share)
-						.expect("share is gte remove_amount and not zero which checked before; qed");
-					let withdrawn_rewards_to_remove = proportion.saturating_mul_int(withdrawn_rewards);
-
-					pool_info.total_shares = pool_info.total_shares.saturating_sub(remove_amount);
-					pool_info.total_rewards = pool_info.total_rewards.saturating_sub(withdrawn_rewards_to_remove);
-					pool_info.total_withdrawn_rewards = pool_info
-						.total_withdrawn_rewards
-						.saturating_sub(withdrawn_rewards_to_remove);
-
-					withdrawn_rewards = withdrawn_rewards.saturating_sub(withdrawn_rewards_to_remove);
+				PoolInfos::<T>::mutate_exists(pool, |maybe_pool_info| {
+					if let Some(mut pool_info) = maybe_pool_info.take() {
+						let removing_share = U256::from(remove_amount.saturated_into::<u128>());
+
+						pool_info.total_shares = pool_info.total_shares.saturating_sub(remove_amount);
+
+						// update withdrawn rewards for each reward currency
+						withdrawn_rewards
+							.iter_mut()
+							.for_each(|(reward_currency, withdrawn_reward)| {
+								let withdrawn_reward_to_remove: T::Balance = removing_share
+									.saturating_mul(withdrawn_reward.to_owned().saturated_into::<u128>().into())
+									.checked_div(share.saturated_into::<u128>().into())
+									.unwrap_or_default()
+									.as_u128()
+									.saturated_into();
+
+								if let Some((total_reward, total_withdrawn_reward)) =
+									pool_info.rewards.get_mut(reward_currency)
+								{
+									*total_reward = total_reward.saturating_sub(withdrawn_reward_to_remove);
+									*total_withdrawn_reward =
+										total_withdrawn_reward.saturating_sub(withdrawn_reward_to_remove);
+
+									// remove if all reward is withdrawn
+									if total_reward.is_zero() {
+										pool_info.rewards.remove(reward_currency);
+									}
+								}
+								*withdrawn_reward = withdrawn_reward.saturating_sub(withdrawn_reward_to_remove);
+							});
+
+						if !pool_info.total_shares.is_zero() {
+							*maybe_pool_info = Some(pool_info);
+						}
+					}
 				});
 
 				share = share.saturating_sub(remove_amount);
@@ -161,7 +271,7 @@ impl<T: Config> Pallet<T> {
 	}
 
 	pub fn set_share(who: &T::AccountId, pool: &T::PoolId, new_share: T::Share) {
-		let (share, _) = Self::share_and_withdrawn_reward(pool, who);
+		let (share, _) = Self::shares_and_withdrawn_rewards(pool, who);
 
 		if new_share > share {
 			Self::add_share(who, pool, new_share.saturating_sub(share));
@@ -171,33 +281,44 @@ impl<T: Config> Pallet<T> {
 	}
 
 	pub fn claim_rewards(who: &T::AccountId, pool: &T::PoolId) {
-		ShareAndWithdrawnReward::<T>::mutate(pool, who, |(share, withdrawn_rewards)| {
-			if share.is_zero() {
-				return;
-			}
-
-			Pools::<T>::mutate(pool, |pool_info| {
-				let proportion = FixedU128::checked_from_rational(*share, pool_info.total_shares).unwrap_or_default();
-				let reward_to_withdraw = proportion
-					.saturating_mul_int(pool_info.total_rewards)
-					.saturating_sub(*withdrawn_rewards)
-					.min(
-						pool_info
-							.total_rewards
-							.saturating_sub(pool_info.total_withdrawn_rewards),
-					);
-
-				if reward_to_withdraw.is_zero() {
+		SharesAndWithdrawnRewards::<T>::mutate_exists(pool, who, |maybe_share_withdrawn| {
+			if let Some((share, withdrawn_rewards)) = maybe_share_withdrawn {
+				if share.is_zero() {
 					return;
 				}
 
-				pool_info.total_withdrawn_rewards =
-					pool_info.total_withdrawn_rewards.saturating_add(reward_to_withdraw);
-				*withdrawn_rewards = withdrawn_rewards.saturating_add(reward_to_withdraw);
-
-				// pay reward to `who`
-				T::Handler::payout(who, pool, reward_to_withdraw);
-			});
+				PoolInfos::<T>::mutate(pool, |pool_info| {
+					let total_shares = U256::from(pool_info.total_shares.to_owned().saturated_into::<u128>());
+					pool_info.rewards.iter_mut().for_each(
+						|(reward_currency, (total_reward, total_withdrawn_reward))| {
+							let withdrawn_reward = withdrawn_rewards.get(reward_currency).copied().unwrap_or_default();
+
+							let total_reward_proportion: T::Balance =
+								U256::from(share.to_owned().saturated_into::<u128>())
+									.saturating_mul(U256::from(total_reward.to_owned().saturated_into::<u128>()))
+									.checked_div(total_shares)
+									.unwrap_or_default()
+									.as_u128()
+									.unique_saturated_into();
+
+							let reward_to_withdraw = total_reward_proportion
+								.saturating_sub(withdrawn_reward)
+								.min(total_reward.saturating_sub(*total_withdrawn_reward));
+
+							if reward_to_withdraw.is_zero() {
+								return;
+							}
+
+							*total_withdrawn_reward = total_withdrawn_reward.saturating_add(reward_to_withdraw);
+							withdrawn_rewards
+								.insert(*reward_currency, withdrawn_reward.saturating_add(reward_to_withdraw));
+
+							// pay reward to `who`
+							T::Handler::payout(who, pool, *reward_currency, reward_to_withdraw);
+						},
+					);
+				});
+			}
 		});
 	}
 }