support multi currency rewards

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,25 +1,25 @@
 #![allow(clippy::unused_unit)]
 #![cfg_attr(not(feature = "std"), no_std)]
 
+mod migrations;
 mod mock;
 mod tests;
 
+pub use migrations::migrate_to_multi_currency_reward;
+
 use codec::{FullCodec, HasCompact, MaxEncodedLen};
-use frame_support::pallet_prelude::*;
+use frame_support::{pallet_prelude::*, weights::Weight};
+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, 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> {
+pub struct PoolInfoV0<Share: HasCompact, Balance: HasCompact> {
 	/// Total shares amount
 	#[codec(compact)]
 	pub total_shares: Share,
@@ -31,6 +31,29 @@ pub struct PoolInfo<Share: HasCompact, Balance: HasCompact> {
 	pub total_withdrawn_rewards: Balance,
 }
 
+/// The Reward Pool Info.
+#[derive(Clone, Encode, Decode, PartialEq, Eq, RuntimeDebug)]
+pub struct PoolInfo<Share: HasCompact, Balance: HasCompact, CurrencyId: Ord> {
+	/// Total shares amount
+	pub total_shares: Share,
+	/// 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::*;
 
 #[frame_support::pallet]
@@ -62,21 +85,39 @@ pub mod module {
 		/// 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.
 	#[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::PoolId, PoolInfo<T::Share, T::Balance, T::CurrencyId>, ValueQuery>;
 
-	/// Record share amount and withdrawn reward amount for specific `AccountId`
-	/// under `PoolId`.
+	/// 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 share_and_withdrawn_reward)]
-	pub type ShareAndWithdrawnReward<T: Config> =
-		StorageDoubleMap<_, Twox64Concat, T::PoolId, Twox64Concat, T::AccountId, (T::Share, T::Balance), ValueQuery>;
+	pub type ShareAndWithdrawnReward<T: Config> = StorageDoubleMap<
+		_,
+		Twox64Concat,
+		T::PoolId,
+		Twox64Concat,
+		T::AccountId,
+		(T::Share, BTreeMap<T::CurrencyId, T::Balance>),
+		ValueQuery,
+	>;
 
 	#[pallet::pallet]
 	pub struct Pallet<T>(_);
@@ -89,12 +130,27 @@ 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(());
 		}
+		Pools::<T>::mutate_exists(pool, |maybe_pool_info| -> DispatchResult {
+			ensure!(maybe_pool_info.is_some(), Error::<T>::PoolDoesNotExist);
+			if let Some(pool_info) = maybe_pool_info {
+				if let Some((total_reward, _)) = pool_info.rewards.get_mut(&reward_currency) {
+					*total_reward = total_reward.saturating_add(reward_increment);
+				} else {
+					pool_info
+						.rewards
+						.insert(reward_currency, (reward_increment, Zero::zero()));
+				}
+			}
+			Ok(())
+		})
 	}
 
 	pub fn add_share(who: &T::AccountId, pool: &T::PoolId, add_amount: T::Share) {
@@ -103,21 +159,45 @@ impl<T: Config> Pallet<T> {
 		}
 
 		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)
-			};
-
+			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);
+
+			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));
+				});
 
 			ShareAndWithdrawnReward::<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)| {
+						let withdrawn_reward = withdrawn_rewards
+							.get(&reward_currency)
+							.copied()
+							.unwrap_or_default()
+							.saturating_add(reward_inflation);
+
+						withdrawn_rewards.insert(reward_currency, withdrawn_reward);
+					});
 			});
 		});
 	}
@@ -138,18 +218,42 @@ impl<T: Config> Pallet<T> {
 					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);
+				Pools::<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);
@@ -171,33 +275,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() {
+		ShareAndWithdrawnReward::<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);
-			});
+				Pools::<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);
+						},
+					);
+				});
+			}
 		});
 	}
 }

rewards/src/migrations.rs

@@ -0,0 +1,37 @@
+use super::*;
+
+pub fn migrate_to_multi_currency_reward<T: Config>(
+	get_reward_currency: impl Fn(&T::PoolId) -> T::CurrencyId,
+) -> Weight {
+	let mut reads_writes: Weight = 0;
+	Pools::<T>::translate::<PoolInfoV0<T::Share, T::Balance>, _>(|pool_id, old_pool_info| {
+		reads_writes += 1;
+		let currency_id = get_reward_currency(&pool_id);
+
+		let mut rewards = BTreeMap::new();
+		rewards.insert(
+			currency_id,
+			(old_pool_info.total_rewards, old_pool_info.total_withdrawn_rewards),
+		);
+
+		Some(PoolInfo {
+			total_shares: old_pool_info.total_shares,
+			rewards,
+		})
+	});
+
+	ShareAndWithdrawnReward::<T>::translate::<(T::Share, T::Balance), _>(
+		|pool_id, _who, (shares, withdrawn_rewards)| {
+			reads_writes += 1;
+			let currency_id = get_reward_currency(&pool_id);
+
+			let mut withdrawn = BTreeMap::new();
+			withdrawn.insert(currency_id, withdrawn_rewards);
+
+			Some((shares, withdrawn))
+		},
+	);
+
+	// Return the weight consumed by the migration.
+	T::DbWeight::get().reads_writes(reads_writes, reads_writes)
+}