wip: Begin invariant test suite

Co-authored-by: Alexander Keating <[email protected]>

foundry.toml

@@ -34,3 +34,12 @@
   # with assume statements because we don't have the surrogate address until it's deployed later in
   # the test.
   include_storage = false
+
+[invariant]
+  call_override = false
+  depth = 15
+  dictionary_weight = 80
+  fail_on_revert = false
+  include_push_bytes = true
+  include_storage = true
+  runs = 100

test/UniStaker.invariants.t.sol

@@ -0,0 +1,67 @@
+pragma solidity ^0.8.23;
+
+import {Test} from "forge-std/Test.sol";
+import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
+
+import {UniStaker} from "src/UniStaker.sol";
+import {UniStakerHandler} from "test/helpers/UniStaker.handler.sol";
+import {ERC20VotesMock} from "test/mocks/MockERC20Votes.sol";
+import {ERC20Fake} from "test/fakes/ERC20Fake.sol";
+
+contract UniStakerInvariants is Test {
+  UniStakerHandler public handler;
+  UniStaker public uniStaker;
+  ERC20Fake rewardToken;
+  ERC20VotesMock govToken;
+  address rewardsNotifier;
+
+  function setUp() public {
+    // deploy UniStaker
+    rewardToken = new ERC20Fake();
+    vm.label(address(rewardToken), "Rewards Token");
+
+    govToken = new ERC20VotesMock();
+    vm.label(address(govToken), "Governance Token");
+
+    rewardsNotifier = address(0xaffab1ebeef);
+    vm.label(rewardsNotifier, "Rewards Notifier");
+    uniStaker = new UniStaker(rewardToken, govToken, rewardsNotifier);
+    handler = new UniStakerHandler(uniStaker);
+
+    bytes4[] memory selectors = new bytes4[](1);
+    selectors[0] = UniStakerHandler.stake.selector;
+
+    targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
+
+    targetContract(address(handler));
+  }
+
+  function invariant_Sum_of_beneficiary_earning_power_equals_total_stake() public {
+    assertEq(uniStaker.totalSupply(), handler.reduceBeneficiaries(0, this.accumulateEarningPower));
+  }
+
+  function invariant_Sum_of_surrogate_balance_equals_total_stake() public {
+    assertEq(uniStaker.totalSupply(), handler.reduceDelegates(0, this.accumulateSurrogateBalance));
+  }
+
+  function accumulateEarningPower(uint256 earningPower, address caller)
+    external
+    view
+    returns (uint256)
+  {
+    return earningPower + uniStaker.earningPower(caller);
+  }
+
+  function accumulateSurrogateBalance(uint256 balance, address delegate)
+    external
+    view
+    returns (uint256)
+  {
+    address surrogateAddr = address(uniStaker.surrogates(delegate));
+    return balance + IERC20(address(uniStaker.STAKE_TOKEN())).balanceOf(surrogateAddr);
+  }
+
+  function invariant_callSummary() public view {
+    handler.callSummary();
+  }
+}

test/helpers/AddressSet.sol

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.23;
+
+struct AddressSet {
+  address[] addrs;
+  mapping(address => bool) saved;
+}
+
+library LibAddressSet {
+  function add(AddressSet storage s, address addr) internal {
+    if (!s.saved[addr]) {
+      s.addrs.push(addr);
+      s.saved[addr] = true;
+    }
+  }
+
+  function contains(AddressSet storage s, address addr) internal view returns (bool) {
+    return s.saved[addr];
+  }
+
+  function count(AddressSet storage s) internal view returns (uint256) {
+    return s.addrs.length;
+  }
+
+  function rand(AddressSet storage s, uint256 seed) internal view returns (address) {
+    if (s.addrs.length > 0) return s.addrs[seed % s.addrs.length];
+    else return address(0);
+  }
+
+  function forEach(AddressSet storage s, function(address) external func) internal {
+    for (uint256 i; i < s.addrs.length; ++i) {
+      func(s.addrs[i]);
+    }
+  }
+
+  function reduce(
+    AddressSet storage s,
+    uint256 acc,
+    function(uint256,address) external returns (uint256) func
+  ) internal returns (uint256) {
+    for (uint256 i; i < s.addrs.length; ++i) {
+      acc = func(acc, s.addrs[i]);
+    }
+    return acc;
+  }
+}

test/helpers/UniStaker.handler.sol

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.13;
+
+import {CommonBase} from "forge-std/Base.sol";
+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 {UniStaker} from "src/UniStaker.sol";
+import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
+
+contract UniStakerHandler is CommonBase, StdCheats, StdUtils {
+  using LibAddressSet for AddressSet;
+
+  UniStaker public uniStaker;
+  IERC20 public stakeToken;
+
+  uint256 public ghost_stakeSum;
+
+  mapping(bytes32 => uint256) public calls;
+
+  AddressSet internal _actors;
+  address internal currentActor;
+
+  AddressSet internal _delegates;
+  AddressSet internal _beneficiaries;
+
+  modifier createActor() {
+    currentActor = msg.sender;
+    _actors.add(msg.sender);
+    _;
+  }
+
+  modifier useActor(uint256 actorIndexSeed) {
+    currentActor = _actors.rand(actorIndexSeed);
+    _;
+  }
+
+  modifier countCall(bytes32 key) {
+    calls[key]++;
+    _;
+  }
+
+  constructor(UniStaker _uniStaker) {
+    uniStaker = _uniStaker;
+    stakeToken = IERC20(address(_uniStaker.STAKE_TOKEN()));
+  }
+
+  function _mintStakeToken(address _to, uint256 _amount) internal {
+    vm.assume(_to != address(0));
+    deal(address(stakeToken), _to, _amount);
+  }
+
+  function stake(uint256 _amount, address _delegatee, address _beneficiary)
+    public
+    createActor
+    countCall("stake")
+  {
+    // 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
+    _amount = bound(_amount, 0, 100_000_000e18);
+
+    // assumes user has stake amount
+    _mintStakeToken(currentActor, _amount);
+
+    vm.startPrank(currentActor);
+    stakeToken.approve(address(uniStaker), _amount);
+    uniStaker.stake(_amount, _delegatee, _beneficiary);
+    vm.stopPrank();
+
+    ghost_stakeSum += _amount;
+  }
+
+  function _getBeneficiaryEarningPower(address _beneficiary) internal view returns (uint256) {
+    return uniStaker.earningPower(_beneficiary);
+  }
+
+  function forEachActor(function(address) external func) public {
+    return _actors.forEach(func);
+  }
+
+  function reduceActors(uint256 acc, function(uint256,address) external returns (uint256) func)
+    public
+    returns (uint256)
+  {
+    return _actors.reduce(acc, func);
+  }
+
+  function reduceBeneficiaries(
+    uint256 acc,
+    function(uint256,address) external returns (uint256) func
+  ) public returns (uint256) {
+    return _beneficiaries.reduce(acc, func);
+  }
+
+  function reduceDelegates(uint256 acc, function(uint256,address) external returns (uint256) func)
+    public
+    returns (uint256)
+  {
+    return _delegates.reduce(acc, func);
+  }
+
+  function actors() external view returns (address[] memory) {
+    return _actors.addrs;
+  }
+
+  function callSummary() external view {
+    console.log("Call summary:");
+    console.log("-------------------");
+    console.log("stake", calls["stake"]);
+    console.log("-------------------");
+  }
+
+  receive() external payable {}
+}