Skip to content

Commit

Permalink
update saturating calculation (#1016)
Browse files Browse the repository at this point in the history
  • Loading branch information
zjb0807 authored and xlc committed Jan 14, 2025
1 parent 9d09743 commit bf0909d
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 8 deletions.
4 changes: 2 additions & 2 deletions rewards/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ impl<T: Config> Pallet<T> {
.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::<u128>()
.saturated_into()
};
*total_reward = total_reward.saturating_add(reward_inflation);
Expand Down Expand Up @@ -239,7 +239,7 @@ impl<T: Config> Pallet<T> {
.saturating_mul(withdrawn_reward.to_owned().saturated_into::<u128>().into())
.checked_div(old_share.saturated_into::<u128>().into())
.unwrap_or_default()
.as_u128()
.saturated_into::<u128>()
.saturated_into();

if let Some((total_reward, total_withdrawn_reward)) =
Expand Down
4 changes: 2 additions & 2 deletions rewards/src/mock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@ use std::collections::HashMap;
use crate as rewards;

pub type AccountId = u128;
pub type Balance = u64;
pub type Share = u64;
pub type Balance = u128;
pub type Share = u128;
pub type PoolId = u32;
pub type CurrencyId = u32;

Expand Down
34 changes: 30 additions & 4 deletions rewards/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,18 +91,18 @@ fn add_share_should_work() {
);

// overflow occurs when saturating calculation
assert_ok!(RewardsModule::add_share(&ALICE, &DOT_POOL, u64::MAX));
assert_ok!(RewardsModule::add_share(&ALICE, &DOT_POOL, u128::MAX));

assert_eq!(
RewardsModule::pool_infos(DOT_POOL),
PoolInfo {
total_shares: u64::MAX,
rewards: vec![(NATIVE_COIN, (u64::MAX, u64::MAX))].into_iter().collect()
total_shares: u128::MAX,
rewards: vec![(NATIVE_COIN, (u128::MAX, u128::MAX))].into_iter().collect()
}
);
assert_eq!(
RewardsModule::shares_and_withdrawn_rewards(DOT_POOL, ALICE),
(u64::MAX, vec![(NATIVE_COIN, u64::MAX)].into_iter().collect())
(u128::MAX, vec![(NATIVE_COIN, u128::MAX)].into_iter().collect())
);
});
}
Expand Down Expand Up @@ -663,3 +663,29 @@ fn minimal_share_should_be_enforced() {
assert_ok!(RewardsModule::transfer_share_and_rewards(&ALICE, &DOT_POOL, 5, &BOB));
});
}

#[test]
fn overflow_should_work() {
ExtBuilder::default().build().execute_with(|| {
assert_ok!(RewardsModule::add_share(&ALICE, &DOT_POOL, 10));

// The DOT pool accumulate 21 rewards in the NATIVE COIN
assert_ok!(RewardsModule::accumulate_reward(&DOT_POOL, NATIVE_COIN, 21));
assert_eq!(
RewardsModule::pool_infos(DOT_POOL),
PoolInfo {
total_shares: 10,
rewards: vec![(NATIVE_COIN, (21, 0))].into_iter().collect()
}
);

assert_ok!(RewardsModule::add_share(&ALICE, &DOT_POOL, u128::MAX));
assert_eq!(
RewardsModule::pool_infos(DOT_POOL),
PoolInfo {
total_shares: u128::MAX,
rewards: vec![(NATIVE_COIN, (u128::MAX, u128::MAX))].into_iter().collect()
}
);
});
}

0 comments on commit bf0909d

Please sign in to comment.