From bf0909dd6458924781fb36faca59923f67ce4f28 Mon Sep 17 00:00:00 2001 From: zjb0807 Date: Tue, 14 Jan 2025 09:36:13 +0800 Subject: [PATCH] update saturating calculation (#1016) --- rewards/src/lib.rs | 4 ++-- rewards/src/mock.rs | 4 ++-- rewards/src/tests.rs | 34 ++++++++++++++++++++++++++++++---- 3 files changed, 34 insertions(+), 8 deletions(-) diff --git a/rewards/src/lib.rs b/rewards/src/lib.rs index 85e127d1b..ab33f8683 100644 --- a/rewards/src/lib.rs +++ b/rewards/src/lib.rs @@ -171,7 +171,7 @@ impl Pallet { .saturating_mul(total_reward.to_owned().saturated_into::().into()) .checked_div(initial_total_shares.to_owned().saturated_into::().into()) .unwrap_or_default() - .as_u128() + .saturated_into::() .saturated_into() }; *total_reward = total_reward.saturating_add(reward_inflation); @@ -239,7 +239,7 @@ impl Pallet { .saturating_mul(withdrawn_reward.to_owned().saturated_into::().into()) .checked_div(old_share.saturated_into::().into()) .unwrap_or_default() - .as_u128() + .saturated_into::() .saturated_into(); if let Some((total_reward, total_withdrawn_reward)) = diff --git a/rewards/src/mock.rs b/rewards/src/mock.rs index 6fb3618e4..a238605dd 100644 --- a/rewards/src/mock.rs +++ b/rewards/src/mock.rs @@ -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; diff --git a/rewards/src/tests.rs b/rewards/src/tests.rs index 3235d5c6e..6cf00451e 100644 --- a/rewards/src/tests.rs +++ b/rewards/src/tests.rs @@ -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()) ); }); } @@ -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() + } + ); + }); +}