diff --git a/test/Staker.invariants.t.sol b/test/Staker.invariants.t.sol index f6fbd20..28973ac 100644 --- a/test/Staker.invariants.t.sol +++ b/test/Staker.invariants.t.sol @@ -77,6 +77,15 @@ contract StakerInvariants is Test { ); } + function invariant_Sum_of_unclaimed_reward_should_be_less_than_or_equal_to_total_rewards_balance() + public + { + assertLe( + handler.reduceDeposits(0, this.accumulateUnclaimedReward), + rewardToken.balanceOf(address(govStaker)) + ); + } + function invariant_RewardPerTokenAccumulatedCheckpoint_should_be_greater_or_equal_to_the_last_rewardPerTokenAccumulatedCheckpoint( ) public view { assertGe( @@ -96,6 +105,13 @@ contract StakerInvariants is Test { return balance + govStaker.depositorTotalStaked(depositor); } + function accumulateUnclaimedReward( + uint256 unclaimedReward, + StakerHarness.DepositIdentifier depositId + ) external view returns (uint256) { + return unclaimedReward + govStaker.unclaimedReward(depositId); + } + function accumulateSurrogateBalance(uint256 balance, address delegate) external view diff --git a/test/helpers/DepositIdSet.sol b/test/helpers/DepositIdSet.sol new file mode 100644 index 0000000..c29efa4 --- /dev/null +++ b/test/helpers/DepositIdSet.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later +pragma solidity ^0.8.23; + +import {Staker} from "src/Staker.sol"; + +struct DepositIdSet { + Staker.DepositIdentifier[] ids; + mapping(Staker.DepositIdentifier => bool) saved; +} + +library LibDepositIdSet { + function reduce( + DepositIdSet storage s, + uint256 acc, + function(uint256,Staker.DepositIdentifier) external returns (uint256) func + ) internal returns (uint256) { + for (uint256 i; i < s.ids.length; ++i) { + acc = func(acc, s.ids[i]); + } + return acc; + } + + function add(DepositIdSet storage s, Staker.DepositIdentifier id) internal { + if (!s.saved[id]) { + s.ids.push(id); + s.saved[id] = true; + } + } +} diff --git a/test/helpers/Staker.handler.sol b/test/helpers/Staker.handler.sol index cd3497e..c422a58 100644 --- a/test/helpers/Staker.handler.sol +++ b/test/helpers/Staker.handler.sol @@ -6,11 +6,13 @@ import {StdCheats} from "forge-std/StdCheats.sol"; import {StdUtils} from "forge-std/StdUtils.sol"; import {console} from "forge-std/console.sol"; import {AddressSet, LibAddressSet} from "../helpers/AddressSet.sol"; +import {DepositIdSet, LibDepositIdSet} from "../helpers/DepositIdSet.sol"; import {Staker} from "src/Staker.sol"; import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol"; contract StakerHandler is CommonBase, StdCheats, StdUtils { using LibAddressSet for AddressSet; + using LibDepositIdSet for DepositIdSet; // system setup Staker public govStaker; @@ -23,6 +25,7 @@ contract StakerHandler is CommonBase, StdCheats, StdUtils { AddressSet internal _depositors; AddressSet internal _delegates; AddressSet internal _claimers; + DepositIdSet internal _depositIdSet; AddressSet internal _surrogates; AddressSet internal _rewardNotifiers; mapping(address => uint256[]) internal _depositIds; @@ -113,6 +116,7 @@ contract StakerHandler is CommonBase, StdCheats, StdUtils { // update handler state _depositIds[_currentActor].push(ghost_depositCount); + _depositIdSet.add(Staker.DepositIdentifier.wrap(ghost_depositCount)); ghost_depositCount++; _surrogates.add(address(govStaker.surrogates(_delegatee))); ghost_stakeSum += _amount; @@ -207,6 +211,13 @@ contract StakerHandler is CommonBase, StdCheats, StdUtils { return _claimers.reduce(acc, func); } + function reduceDeposits( + uint256 acc, + function(uint256,Staker.DepositIdentifier) external returns (uint256) func + ) public returns (uint256) { + return _depositIdSet.reduce(acc, func); + } + function reduceDelegates(uint256 acc, function(uint256,address) external returns (uint256) func) public returns (uint256)