From 62115464c55d409a7290b4cd86d44bd0846d8e10 Mon Sep 17 00:00:00 2001 From: ferrell-code Date: Tue, 31 Aug 2021 22:03:56 -0400 Subject: [PATCH 1/3] migrate to mutate_exist --- rewards/src/lib.rs | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/rewards/src/lib.rs b/rewards/src/lib.rs index bbea342a2..710b212af 100644 --- a/rewards/src/lib.rs +++ b/rewards/src/lib.rs @@ -130,28 +130,33 @@ impl Pallet { // claim rewards firstly Self::claim_rewards(who, pool); - ShareAndWithdrawnReward::::mutate(pool, who, |(share, withdrawn_rewards)| { - let remove_amount = remove_amount.min(*share); + ShareAndWithdrawnReward::::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; - } + if remove_amount.is_zero() { + return; + } - Pools::::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); + Pools::::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); + 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); - }); + withdrawn_rewards = withdrawn_rewards.saturating_sub(withdrawn_rewards_to_remove); + }); - *share = share.saturating_sub(remove_amount); + share = share.saturating_sub(remove_amount); + if share != Zero::zero() { + *share_info = Some((share, withdrawn_rewards)); + } + } }); } From c1a21f652cfb7bdf58becd4cd5aaa1d65b218116 Mon Sep 17 00:00:00 2001 From: ferrell-code Date: Wed, 1 Sep 2021 01:46:51 -0400 Subject: [PATCH 2/3] test and add .is_zero() instead of == --- rewards/src/lib.rs | 2 +- rewards/src/tests.rs | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/rewards/src/lib.rs b/rewards/src/lib.rs index 710b212af..effbb069d 100644 --- a/rewards/src/lib.rs +++ b/rewards/src/lib.rs @@ -153,7 +153,7 @@ impl Pallet { }); share = share.saturating_sub(remove_amount); - if share != Zero::zero() { + if !share.is_zero() { *share_info = Some((share, withdrawn_rewards)); } } diff --git a/rewards/src/tests.rs b/rewards/src/tests.rs index 18f30f770..360e037f6 100644 --- a/rewards/src/tests.rs +++ b/rewards/src/tests.rs @@ -341,3 +341,35 @@ fn accumulate_reward_should_work() { ); }); } + +#[test] +fn share_to_zero_removes_storage() { + ExtBuilder::default().build().execute_with(|| { + assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, ALICE), false); + RewardsModule::add_share(&ALICE, &DOT_POOL, 100); + RewardsModule::add_share(&BOB, &DOT_POOL, 100); + Pools::::mutate(DOT_POOL, |pool_info| { + pool_info.total_rewards += 10000; + }); + + assert_eq!( + RewardsModule::pools(DOT_POOL), + PoolInfo { + total_shares: 200, + total_rewards: 10000, + total_withdrawn_rewards: 0, + } + ); + + // chescks if key is removed + assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, ALICE), true); + RewardsModule::remove_share(&ALICE, &DOT_POOL, 100); + assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, ALICE), false); + + RewardsModule::remove_share(&BOB, &DOT_POOL, 50); + assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, BOB), true); + + RewardsModule::remove_share(&BOB, &DOT_POOL, 100); + assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, BOB), false); + }); +} From 7e6a19a75848642da26852193ef2f37fc564746f Mon Sep 17 00:00:00 2001 From: ferrell-code Date: Wed, 1 Sep 2021 01:52:17 -0400 Subject: [PATCH 3/3] spelling --- rewards/src/tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rewards/src/tests.rs b/rewards/src/tests.rs index 360e037f6..57cea136f 100644 --- a/rewards/src/tests.rs +++ b/rewards/src/tests.rs @@ -361,7 +361,7 @@ fn share_to_zero_removes_storage() { } ); - // chescks if key is removed + // checks if key is removed assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, ALICE), true); RewardsModule::remove_share(&ALICE, &DOT_POOL, 100); assert_eq!(ShareAndWithdrawnReward::::contains_key(DOT_POOL, ALICE), false);