Stake Deposit w/ Delegation Implementation

foundry.toml

@@ -26,3 +26,11 @@
   single_line_statement_blocks = "single"
   tab_width = 2
   wrap_comments = true
+
+[fuzz]
+  # We turn on this setting to prevent the fuzzer from picking DelegationSurrogate contracts,
+  # including before they're actually even deployed, as some other entity in the test, for example
+  # depositor. This makes no sense and breaks test assertions, but is extremely difficult to handle
+  # with assume statements because we don't have the surrogate address until it's deployed later in
+  # the test.
+  include_storage = false

src/UniStaker.sol

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.23;
+
+import {DelegationSurrogate} from "src/DelegationSurrogate.sol";
+import {IERC20Delegates} from "src/interfaces/IERC20Delegates.sol";
+import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
+import {SafeERC20} from "openzeppelin/token/ERC20/utils/SafeERC20.sol";
+import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol";
+
+contract UniStaker is ReentrancyGuard {
+  IERC20 public immutable REWARDS_TOKEN;
+  IERC20Delegates public immutable STAKE_TOKEN;
+
+  mapping(address delegatee => DelegationSurrogate surrogate) public surrogates;
+
+  constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken) {
+    REWARDS_TOKEN = _rewardsToken;
+    STAKE_TOKEN = _stakeToken;
+  }
+
+  function stake(uint256 _amount, address _delegatee)
+    public
+    nonReentrant
+    returns (uint256 _depositId)
+  {
+    DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee);
+    _stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount);
+    _depositId = 1;
+  }
+
+  function _fetchOrDeploySurrogate(address _delegatee)
+    internal
+    returns (DelegationSurrogate _surrogate)
+  {
+    _surrogate = surrogates[_delegatee];
+
+    if (address(_surrogate) == address(0)) {
+      _surrogate = new DelegationSurrogate(STAKE_TOKEN, _delegatee);
+      surrogates[_delegatee] = _surrogate;
+    }
+  }
+
+  function _stakeTokenSafeTransferFrom(address _from, address _to, uint256 _value) internal {
+    SafeERC20.safeTransferFrom(IERC20(address(STAKE_TOKEN)), _from, _to, _value);
+  }
+}

test/UniStaker.t.sol

@@ -0,0 +1,139 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.23;
+
+import {Test, console2} from "forge-std/Test.sol";
+import {UniStaker, DelegationSurrogate, IERC20, IERC20Delegates} from "src/UniStaker.sol";
+import {ERC20VotesMock} from "test/mocks/MockERC20Votes.sol";
+import {ERC20Fake} from "test/fakes/ERC20Fake.sol";
+
+contract UniStakerTest is Test {
+  ERC20Fake rewardToken;
+  ERC20VotesMock govToken;
+  UniStaker uniStaker;
+
+  function setUp() public {
+    rewardToken = new ERC20Fake();
+    vm.label(address(rewardToken), "Reward Token");
+
+    govToken = new ERC20VotesMock();
+    vm.label(address(govToken), "Governance Token");
+
+    uniStaker = new UniStaker(rewardToken, govToken);
+    vm.label(address(uniStaker), "UniStaker");
+  }
+
+  function _boundMintAmount(uint256 _amount) internal view returns (uint256) {
+    return bound(_amount, 0, 100_000_000_000e18);
+  }
+
+  function _mintGovToken(address _to, uint256 _amount) internal {
+    vm.assume(_to != address(0));
+    govToken.mint(_to, _amount);
+  }
+
+  function _stake(address _depositor, uint256 _amount, address _delegatee)
+    internal
+    returns (uint256 _depositId)
+  {
+    vm.startPrank(_depositor);
+    govToken.approve(address(uniStaker), _amount);
+    _depositId = uniStaker.stake(_amount, _delegatee);
+    vm.stopPrank();
+  }
+}
+
+contract Constructor is UniStakerTest {
+  function test_SetsTheRewardTokenAndStakeToken() public {
+    assertEq(address(uniStaker.REWARDS_TOKEN()), address(rewardToken));
+    assertEq(address(uniStaker.STAKE_TOKEN()), address(govToken));
+  }
+
+  function testFuzz_SetsTheRewardTokenAndStakeTokenToArbitraryAddresses(
+    address _rewardToken,
+    address _stakeToken
+  ) public {
+    UniStaker _uniStaker = new UniStaker(IERC20(_rewardToken), IERC20Delegates(_stakeToken));
+    assertEq(address(_uniStaker.REWARDS_TOKEN()), address(_rewardToken));
+    assertEq(address(_uniStaker.STAKE_TOKEN()), address(_stakeToken));
+  }
+}
+
+contract Stake is UniStakerTest {
+  function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAUserStakes(
+    address _depositor,
+    uint256 _amount,
+    address _delegatee
+  ) public {
+    _amount = bound(_amount, 1, type(uint224).max);
+    _mintGovToken(_depositor, _amount);
+    _stake(_depositor, _amount, _delegatee);
+
+    DelegationSurrogate _surrogate = uniStaker.surrogates(_delegatee);
+
+    assertEq(govToken.balanceOf(address(_surrogate)), _amount);
+    assertEq(govToken.delegates(address(_surrogate)), _delegatee);
+    assertEq(govToken.balanceOf(_depositor), 0);
+  }
+
+  function testFuzz_TransfersToAnExistingSurrogateWhenStakedToTheSameDelegatee(
+    address _depositor1,
+    uint256 _amount1,
+    address _depositor2,
+    uint256 _amount2,
+    address _delegatee
+  ) public {
+    _amount1 = _boundMintAmount(_amount1);
+    _amount2 = _boundMintAmount(_amount2);
+    _mintGovToken(_depositor1, _amount1);
+    _mintGovToken(_depositor2, _amount2);
+
+    // Perform first stake with this delegatee
+    _stake(_depositor1, _amount1, _delegatee);
+    // Remember the surrogate which was deployed for this delegatee
+    DelegationSurrogate _surrogate = uniStaker.surrogates(_delegatee);
+
+    // Perform the second stake with this delegatee
+    _stake(_depositor2, _amount2, _delegatee);
+
+    // Ensure surrogate for this delegatee hasn't changed and has summed stake balance
+    assertEq(address(uniStaker.surrogates(_delegatee)), address(_surrogate));
+    assertEq(govToken.delegates(address(_surrogate)), _delegatee);
+    assertEq(govToken.balanceOf(address(_surrogate)), _amount1 + _amount2);
+    assertEq(govToken.balanceOf(_depositor1), 0);
+    assertEq(govToken.balanceOf(_depositor2), 0);
+  }
+
+  function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenUsersStakesToDifferentDelegatees(
+    address _depositor1,
+    uint256 _amount1,
+    address _depositor2,
+    uint256 _amount2,
+    address _delegatee1,
+    address _delegatee2
+  ) public {
+    vm.assume(_delegatee1 != _delegatee2);
+    _amount1 = _boundMintAmount(_amount1);
+    _amount2 = _boundMintAmount(_amount2);
+    _mintGovToken(_depositor1, _amount1);
+    _mintGovToken(_depositor2, _amount2);
+
+    // Perform first stake with first delegatee
+    _stake(_depositor1, _amount1, _delegatee1);
+    // Remember the surrogate which was deployed for first delegatee
+    DelegationSurrogate _surrogate1 = uniStaker.surrogates(_delegatee1);
+
+    // Perform second stake with second delegatee
+    _stake(_depositor2, _amount2, _delegatee2);
+    // Remember the surrogate which was deployed for first delegatee
+    DelegationSurrogate _surrogate2 = uniStaker.surrogates(_delegatee2);
+
+    // Ensure surrogates are different with discreet delegation & balances
+    assertTrue(_surrogate1 != _surrogate2);
+    assertEq(govToken.delegates(address(_surrogate1)), _delegatee1);
+    assertEq(govToken.balanceOf(address(_surrogate1)), _amount1);
+    assertEq(govToken.delegates(address(_surrogate2)), _delegatee2);
+    assertEq(govToken.balanceOf(address(_surrogate2)), _amount2);
+    assertEq(govToken.balanceOf(_depositor1), 0);
+    assertEq(govToken.balanceOf(_depositor2), 0);
+  }
+}

test/fakes/ERC20Fake.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.23;
+
+import {ERC20} from "openzeppelin/token/ERC20/ERC20.sol";
+
+/// @dev An ERC20 token that allows for public minting for use in tests.
+contract ERC20Fake is ERC20 {
+  constructor() ERC20("Fake Token", "FAKE") {}
+
+  /// @dev Public mint function useful for testing
+  function mint(address _account, uint256 _value) public {
+    _mint(_account, _value);
+  }
+}