From 53e05581deec55c9ab128826309d14f2292b6638 Mon Sep 17 00:00:00 2001
From: wildmolasses <changes@gmail.com>
Date: Wed, 21 Feb 2024 12:13:27 -0500
Subject: [PATCH] Invariant tests

---
 test/UniStaker.invariants.t.sol    | 31 ++++++++++-------
 test/helpers/UniStaker.handler.sol | 56 +++++++++++-------------------
 2 files changed, 38 insertions(+), 49 deletions(-)

diff --git a/test/UniStaker.invariants.t.sol b/test/UniStaker.invariants.t.sol
index 1b05c03..a09ed34 100644
--- a/test/UniStaker.invariants.t.sol
+++ b/test/UniStaker.invariants.t.sol
@@ -43,6 +43,12 @@ contract UniStakerInvariants is Test {
     targetContract(address(handler));
   }
 
+  // Invariants
+
+  function invariant_Sum_of_all_depositor_balances_equals_total_stake() public {
+    assertEq(uniStaker.totalStaked(), handler.reduceDepositors(0, this.accumulateDeposits));
+  }
+
   function invariant_Sum_of_beneficiary_earning_power_equals_total_stake() public {
     assertEq(uniStaker.totalStaked(), handler.reduceBeneficiaries(0, this.accumulateEarningPower));
   }
@@ -51,11 +57,7 @@ contract UniStakerInvariants is Test {
     assertEq(uniStaker.totalStaked(), handler.reduceDelegates(0, this.accumulateSurrogateBalance));
   }
 
-  function invariant_Sum_of_all_depositor_balances_equals_total_stake() public {
-    assertEq(uniStaker.totalStaked(), handler.reduceDepositors(0, this.accumulateDeposits));
-  }
-
-  function invariant_Total_staked_minus_withdrawals_equals_total_stake() public {
+  function invariant_Cumulative_staked_minus_withdrawals_equals_total_stake() public {
     assertEq(uniStaker.totalStaked(), handler.ghost_stakeSum() - handler.ghost_stakeWithdrawn());
   }
 
@@ -66,13 +68,20 @@ contract UniStakerInvariants is Test {
     );
   }
 
-  function invariant_Sum_of_beneficiary_unclaimed_rewards_equals_rewards_left() public {
-    assertEq(
-      rewardToken.balanceOf(address(uniStaker)),
-      handler.reduceBeneficiaries(0, this.accumulateUnclaimedReward)
+  function invariant_Unclaimed_reward_LTE_total_rewards() public {
+    assertLe(
+      handler.reduceBeneficiaries(0, this.accumulateUnclaimedReward),
+      rewardToken.balanceOf(address(uniStaker))
     );
   }
 
+  // Used to see distribution of non-reverting calls
+  function invariant_callSummary() public view {
+    handler.callSummary();
+  }
+
+  // Helpers
+
   function accumulateDeposits(uint256 balance, address depositor) external view returns (uint256) {
     return balance + uniStaker.depositorTotalStaked(depositor);
   }
@@ -101,8 +110,4 @@ contract UniStakerInvariants is Test {
     address surrogateAddr = address(uniStaker.surrogates(delegate));
     return balance + IERC20(address(uniStaker.STAKE_TOKEN())).balanceOf(surrogateAddr);
   }
-
-  function invariant_callSummary() public view {
-    handler.callSummary();
-  }
 }
diff --git a/test/helpers/UniStaker.handler.sol b/test/helpers/UniStaker.handler.sol
index df25435..bc405f5 100644
--- a/test/helpers/UniStaker.handler.sol
+++ b/test/helpers/UniStaker.handler.sol
@@ -12,43 +12,28 @@ import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
 contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
   using LibAddressSet for AddressSet;
 
+  // system setup
   UniStaker public uniStaker;
   IERC20 public stakeToken;
   IERC20 public rewardToken;
   address public admin;
 
-  uint256 public ghost_stakeSum;
-  uint256 public ghost_stakeWithdrawn;
-  uint256 public ghost_depositCount;
-  uint256 public ghost_rewardsClaimed;
-  uint256 public ghost_rewardsNotified;
-
-  mapping(bytes32 => uint256) public calls;
-
+  // actors, deposit state
   address internal currentActor;
-
   AddressSet internal _depositors;
   AddressSet internal _delegates;
   AddressSet internal _beneficiaries;
   AddressSet internal _surrogates;
   AddressSet internal _rewardNotifiers;
-
   mapping(address => uint256[]) internal _depositIds;
+  mapping(bytes32 => uint256) public calls;
 
-  function _getActorRandDepositId(uint256 _randomDepositSeed) internal view returns (uint256) {
-    return _depositIds[currentActor][_randomDepositSeed % _depositIds[currentActor].length];
-  }
-
-  function _createDepositor() internal {
-    currentActor = msg.sender;
-    // Surrogates can't stake. We won't include them as potential depositors.
-    vm.assume(!_surrogates.contains(currentActor));
-    _depositors.add(msg.sender);
-  }
-
-  function _useActor(AddressSet storage _set, uint256 _randomActorSeed) internal {
-    currentActor = _set.rand(_randomActorSeed);
-  }
+  // ghost vars
+  uint256 public ghost_stakeSum;
+  uint256 public ghost_stakeWithdrawn;
+  uint256 public ghost_depositCount;
+  uint256 public ghost_rewardsClaimed;
+  uint256 public ghost_rewardsNotified;
 
   modifier countCall(bytes32 key) {
     calls[key]++;
@@ -101,10 +86,6 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
   {
     _createDepositor();
 
-    // TODO: decide if we want reverts in stake
-    //_beneficiary = address(uint160(bound(uint160(_beneficiary), 1, type(uint160).max)));
-    //_delegatee = address(uint160(bound(uint160(_delegatee), 1, type(uint160).max)));
-
     _beneficiaries.add(_beneficiary);
     _delegates.add(_delegatee);
     // todo: adjust upper bound
@@ -175,12 +156,19 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
     skip(_seconds);
   }
 
-  function _getBeneficiaryEarningPower(address _beneficiary) internal view returns (uint256) {
-    return uniStaker.earningPower(_beneficiary);
+  function _getActorRandDepositId(uint256 _randomDepositSeed) internal view returns (uint256) {
+    return _depositIds[currentActor][_randomDepositSeed % _depositIds[currentActor].length];
+  }
+
+  function _createDepositor() internal {
+    currentActor = msg.sender;
+    // Surrogates can't stake. We won't include them as potential depositors.
+    vm.assume(!_surrogates.contains(currentActor));
+    _depositors.add(msg.sender);
   }
 
-  function forEachActor(function(address) external func) public {
-    return _depositors.forEach(func);
+  function _useActor(AddressSet storage _set, uint256 _randomActorSeed) internal {
+    currentActor = _set.rand(_randomActorSeed);
   }
 
   function reduceDepositors(uint256 acc, function(uint256,address) external returns (uint256) func)
@@ -204,10 +192,6 @@ contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
     return _delegates.reduce(acc, func);
   }
 
-  function actors() external view returns (address[] memory) {
-    return _depositors.addrs;
-  }
-
   function callSummary() external view {
     console.log("\nCall summary:");
     console.log("-------------------");