Implement the surrogate deployment & transfer mechanics of staking

src/UniStaker.sol

@@ -1,15 +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 {
+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

@@ -2,7 +2,7 @@
 pragma solidity 0.8.23;
 
 import {Test, console2} from "forge-std/Test.sol";
-import {UniStaker} from "src/UniStaker.sol";
+import {UniStaker, DelegationSurrogate} from "src/UniStaker.sol";
 import {ERC20VotesMock} from "test/mocks/MockERC20Votes.sol";
 import {ERC20Fake} from "test/fakes/ERC20Fake.sol";
 
@@ -19,6 +19,26 @@ contract UniStakerTest is Test {
     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();
   }
 }
 
@@ -28,3 +48,78 @@ contract Constructor is UniStakerTest {
     assertEq(address(uniStaker.STAKE_TOKEN()), address(govToken));
   }
 }
+
+contract Stake is UniStakerTest {
+  function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAUserStakes(
+    address _depositor,
+    uint256 _amount,
+    address _delegatee
+  ) public {
+    _amount = _boundMintAmount(_amount);
+    _mintGovToken(_depositor, _amount);
+    _stake(_depositor, _amount, _delegatee);
+
+    DelegationSurrogate _surrogate = uniStaker.surrogates(_delegatee);
+
+    assertEq(govToken.balanceOf(address(_surrogate)), _amount);
+    assertEq(govToken.delegates(address(_surrogate)), _delegatee);
+  }
+
+  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 2x the balance
+    assertEq(address(uniStaker.surrogates(_delegatee)), address(_surrogate));
+    assertEq(govToken.delegates(address(_surrogate)), _delegatee);
+    assertEq(govToken.balanceOf(address(_surrogate)), _amount1 + _amount2);
+  }
+
+  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);
+  }
+}