feat: allow rewards to be paid in any currency

crates/reward/src/lib.rs

@@ -26,7 +26,7 @@ use sp_runtime::{
     traits::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, MaybeSerializeDeserialize, Saturating, Zero},
     ArithmeticError,
 };
-use sp_std::{cmp::PartialOrd, convert::TryInto, fmt::Debug};
+use sp_std::{cmp::PartialOrd, collections::btree_set::BTreeSet, convert::TryInto, fmt::Debug};
 
 pub(crate) type SignedFixedPoint<T, I = ()> = <T as Config<I>>::SignedFixedPoint;
 
@@ -100,7 +100,18 @@ pub mod pallet {
     }
 
     #[pallet::hooks]
-    impl<T: Config<I>, I: 'static> Hooks<T::BlockNumber> for Pallet<T, I> {}
+    impl<T: Config<I>, I: 'static> Hooks<T::BlockNumber> for Pallet<T, I> {
+        fn on_runtime_upgrade() -> Weight {
+            RewardPerToken::<T, I>::iter_keys().for_each(|(_currency_id, pool_id)| {
+                RewardCurrencies::<T, I>::mutate(pool_id, |reward_currencies| {
+                    reward_currencies.insert(T::GetNativeCurrencyId::get());
+                    reward_currencies.insert(T::GetWrappedCurrencyId::get());
+                });
+            });
+
+            Weight::zero()
+        }
+    }
 
     /// The total stake deposited to this reward pool.
     #[pallet::storage]
@@ -145,6 +156,11 @@ pub mod pallet {
         ValueQuery,
     >;
 
+    /// Track the currencies used for rewards.
+    #[pallet::storage]
+    pub type RewardCurrencies<T: Config<I>, I: 'static = ()> =
+        StorageMap<_, Blake2_128Concat, T::PoolId, BTreeSet<T::CurrencyId>, ValueQuery>;
+
     #[pallet::pallet]
     #[pallet::without_storage_info]
     pub struct Pallet<T, I = ()>(_);
@@ -199,6 +215,23 @@ macro_rules! checked_sub_mut {
 
 // "Internal" functions, callable by code.
 impl<T: Config<I>, I: 'static> Pallet<T, I> {
+    // TODO: remove this after the migration, I have added
+    // this because it's not clear whether the capacity migration
+    // will run before or after the pallet hook and we need to make
+    // sure these are included for any deposits / withdrawals
+    fn add_default_currencies(pool_id: &T::PoolId) {
+        let reward_currencies = RewardCurrencies::<T, I>::get(pool_id);
+        if reward_currencies.contains(&T::GetNativeCurrencyId::get())
+            && reward_currencies.contains(&T::GetWrappedCurrencyId::get())
+        {
+            return;
+        }
+        RewardCurrencies::<T, I>::mutate(pool_id, |reward_currencies| {
+            reward_currencies.insert(T::GetNativeCurrencyId::get());
+            reward_currencies.insert(T::GetWrappedCurrencyId::get());
+        });
+    }
+
     pub fn stake(pool_id: &T::PoolId, stake_id: &T::StakeId) -> SignedFixedPoint<T, I> {
         Stake::<T, I>::get((pool_id, stake_id))
     }
@@ -219,7 +252,8 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
         checked_add_mut!(Stake<T, I>, (pool_id, stake_id), &amount);
         checked_add_mut!(TotalStake<T, I>, pool_id, &amount);
 
-        for currency_id in [T::GetNativeCurrencyId::get(), T::GetWrappedCurrencyId::get()] {
+        Self::add_default_currencies(pool_id);
+        for currency_id in RewardCurrencies::<T, I>::get(pool_id) {
             <RewardTally<T, I>>::mutate(currency_id, (pool_id, stake_id), |reward_tally| {
                 let reward_per_token = Self::reward_per_token(currency_id, pool_id);
                 let reward_per_token_mul_amount =
@@ -245,12 +279,17 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
         currency_id: T::CurrencyId,
         reward: SignedFixedPoint<T, I>,
     ) -> DispatchResult {
-        let total_stake = Self::total_stake(pool_id);
         if reward.is_zero() {
             return Ok(());
         }
+        let total_stake = Self::total_stake(pool_id);
         ensure!(!total_stake.is_zero(), Error::<T, I>::ZeroTotalStake);
 
+        // track currency for future deposits / withdrawals
+        RewardCurrencies::<T, I>::mutate(pool_id, |reward_currencies| {
+            reward_currencies.insert(currency_id);
+        });
+
         let reward_div_total_stake = reward.checked_div(&total_stake).ok_or(ArithmeticError::Underflow)?;
         checked_add_mut!(RewardPerToken<T, I>, currency_id, pool_id, &reward_div_total_stake);
         checked_add_mut!(TotalRewards<T, I>, currency_id, &reward);
@@ -293,7 +332,8 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
         checked_sub_mut!(Stake<T, I>, (pool_id, stake_id), &amount);
         checked_sub_mut!(TotalStake<T, I>, pool_id, &amount);
 
-        for currency_id in [T::GetNativeCurrencyId::get(), T::GetWrappedCurrencyId::get()] {
+        Self::add_default_currencies(pool_id);
+        for currency_id in RewardCurrencies::<T, I>::get(pool_id) {
             <RewardTally<T, I>>::mutate(currency_id, (pool_id, stake_id), |reward_tally| {
                 let reward_per_token = Self::reward_per_token(currency_id, pool_id);
                 let reward_per_token_mul_amount =

crates/reward/src/mock.rs

@@ -1,7 +1,12 @@
 use crate as reward;
 use crate::{Config, Error};
 use frame_support::{parameter_types, traits::Everything};
-pub use primitives::{CurrencyId, CurrencyId::Token, TokenSymbol::*, VaultCurrencyPair, VaultId};
+pub use primitives::{
+    CurrencyId,
+    CurrencyId::{ForeignAsset, Token},
+    TokenSymbol::*,
+    VaultCurrencyPair, VaultId,
+};
 use sp_arithmetic::FixedI128;
 use sp_core::H256;
 use sp_runtime::{

crates/reward/src/tests.rs

@@ -147,3 +147,30 @@ fn should_distribute_with_many_rewards() {
         assert_ok!(Reward::withdraw_reward(&(), &BOB, Token(IBTC)), bob_reward);
     })
 }
+
+macro_rules! assert_approx_eq {
+    ($left:expr, $right:expr, $delta:expr) => {
+        assert!(if $left > $right { $left - $right } else { $right - $left } <= $delta)
+    };
+}
+
+#[test]
+fn should_distribute_with_different_rewards() {
+    run_test(|| {
+        assert_ok!(Reward::deposit_stake(&(), &ALICE, fixed!(100)));
+        assert_ok!(Reward::distribute_reward(&(), Token(IBTC), fixed!(1000)));
+        assert_ok!(Reward::deposit_stake(&(), &ALICE, fixed!(100)));
+        assert_ok!(Reward::distribute_reward(&(), Token(INTR), fixed!(1000)));
+        assert_ok!(Reward::deposit_stake(&(), &ALICE, fixed!(100)));
+        assert_ok!(Reward::distribute_reward(&(), Token(DOT), fixed!(1000)));
+        assert_ok!(Reward::deposit_stake(&(), &ALICE, fixed!(100)));
+        assert_ok!(Reward::distribute_reward(&(), ForeignAsset(0), fixed!(1000)));
+
+        assert_ok!(Reward::withdraw_stake(&(), &ALICE, fixed!(300)));
+
+        assert_approx_eq!(Reward::compute_reward(&(), &ALICE, Token(IBTC)).unwrap(), 1000, 1);
+        assert_approx_eq!(Reward::compute_reward(&(), &ALICE, Token(INTR)).unwrap(), 1000, 1);
+        assert_approx_eq!(Reward::compute_reward(&(), &ALICE, Token(DOT)).unwrap(), 1000, 1);
+        assert_approx_eq!(Reward::compute_reward(&(), &ALICE, ForeignAsset(0)).unwrap(), 1000, 1);
+    })
+}