From 46fd8090a73e1bf945f32ffaf87f958adabac8f2 Mon Sep 17 00:00:00 2001 From: Ben DiFrancesco Date: Fri, 15 Dec 2023 12:19:32 -0500 Subject: [PATCH] Add and track the concept of a beneficiary to a stake deposit Stakers will be able to assign the rewards earned by their stake to an arbitrary address of their choosing. This commit adds the machinery to track said "beneficiary" address. It is included in the Deposit struct. What's more, the total "earning power" of every address is tracked in a mapping independent of deposit balance. --- src/UniStaker.sol | 37 +++++- test/UniStaker.t.sol | 286 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 315 insertions(+), 8 deletions(-) diff --git a/src/UniStaker.sol b/src/UniStaker.sol index 0564c34..831f83a 100644 --- a/src/UniStaker.sol +++ b/src/UniStaker.sol @@ -16,6 +16,7 @@ contract UniStaker is ReentrancyGuard { uint256 balance; address owner; address delegatee; + address beneficiary; } IERC20 public immutable REWARDS_TOKEN; @@ -27,6 +28,8 @@ contract UniStaker is ReentrancyGuard { mapping(address depositor => uint256 amount) public totalDeposits; + mapping(address beneficiary => uint256 amount) public earningPower; + mapping(DepositIdentifier depositId => Deposit deposit) public deposits; mapping(address delegatee => DelegationSurrogate surrogate) public surrogates; @@ -41,13 +44,15 @@ contract UniStaker is ReentrancyGuard { nonReentrant returns (DepositIdentifier _depositId) { - DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee); - _stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount); - _depositId = _useDepositId(); + _depositId = _stake(_amount, _delegatee, msg.sender); + } - totalSupply += _amount; - totalDeposits[msg.sender] += _amount; - deposits[_depositId] = Deposit({balance: _amount, owner: msg.sender, delegatee: _delegatee}); + function stake(uint256 _amount, address _delegatee, address _beneficiary) + public + nonReentrant + returns (DepositIdentifier _depositId) + { + _depositId = _stake(_amount, _delegatee, _beneficiary); } function withdraw(DepositIdentifier _depositId, uint256 _amount) external nonReentrant { @@ -57,6 +62,7 @@ contract UniStaker is ReentrancyGuard { deposit.balance -= _amount; // overflow prevents withdrawing more than balance totalSupply -= _amount; totalDeposits[msg.sender] -= _amount; + earningPower[deposit.beneficiary] -= _amount; _stakeTokenSafeTransferFrom(address(surrogates[deposit.delegatee]), deposit.owner, _amount); } @@ -80,4 +86,23 @@ contract UniStaker is ReentrancyGuard { _depositId = nextDepositId; nextDepositId = DepositIdentifier.wrap(DepositIdentifier.unwrap(_depositId) + 1); } + + function _stake(uint256 _amount, address _delegatee, address _beneficiary) + internal + returns (DepositIdentifier _depositId) + { + DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee); + _stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount); + _depositId = _useDepositId(); + + totalSupply += _amount; + totalDeposits[msg.sender] += _amount; + earningPower[_beneficiary] += _amount; + deposits[_depositId] = Deposit({ + balance: _amount, + owner: msg.sender, + delegatee: _delegatee, + beneficiary: _beneficiary + }); + } } diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol index b4eca50..01dec25 100644 --- a/test/UniStaker.t.sol +++ b/test/UniStaker.t.sol @@ -41,13 +41,29 @@ contract UniStakerTest is Test { vm.stopPrank(); } + function _stake(address _depositor, uint256 _amount, address _delegatee, address _beneficiary) + internal + returns (UniStaker.DepositIdentifier _depositId) + { + vm.startPrank(_depositor); + govToken.approve(address(uniStaker), _amount); + _depositId = uniStaker.stake(_amount, _delegatee, _beneficiary); + vm.stopPrank(); + } + function _fetchDeposit(UniStaker.DepositIdentifier _depositId) internal view returns (UniStaker.Deposit memory) { - (uint256 _balance, address _owner, address _delegatee) = uniStaker.deposits(_depositId); - return UniStaker.Deposit({balance: _balance, owner: _owner, delegatee: _delegatee}); + (uint256 _balance, address _owner, address _delegatee, address _beneficiary) = + uniStaker.deposits(_depositId); + return UniStaker.Deposit({ + balance: _balance, + owner: _owner, + delegatee: _delegatee, + beneficiary: _beneficiary + }); } function _boundMintAndStake(address _depositor, uint256 _amount, address _delegatee) @@ -58,6 +74,17 @@ contract UniStakerTest is Test { _mintGovToken(_depositor, _boundedAmount); _depositId = _stake(_depositor, _boundedAmount, _delegatee); } + + function _boundMintAndStake( + address _depositor, + uint256 _amount, + address _delegatee, + address _beneficiary + ) internal returns (uint256 _boundedAmount, UniStaker.DepositIdentifier _depositId) { + _boundedAmount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _boundedAmount); + _depositId = _stake(_depositor, _boundedAmount, _delegatee, _beneficiary); + } } contract Constructor is UniStakerTest { @@ -298,6 +325,91 @@ contract Stake is UniStakerTest { assertEq(_deposit2.delegatee, _delegatee2); } + function testFuzz_AssignsEarningPowerToDepositorIfNoBeneficiaryIsSpecified( + address _depositor, + uint256 _amount, + address _delegatee + ) public { + _amount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _amount); + + UniStaker.DepositIdentifier _depositId = _stake(_depositor, _amount, _delegatee); + UniStaker.Deposit memory _deposit = _fetchDeposit(_depositId); + + assertEq(uniStaker.earningPower(_depositor), _amount); + assertEq(_deposit.beneficiary, _depositor); + } + + function testFuzz_AssignsEarningPowerToTheBeneficiaryProvided( + address _depositor, + uint256 _amount, + address _delegatee, + address _beneficiary + ) public { + _amount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _amount); + + UniStaker.DepositIdentifier _depositId = _stake(_depositor, _amount, _delegatee, _beneficiary); + UniStaker.Deposit memory _deposit = _fetchDeposit(_depositId); + + assertEq(uniStaker.earningPower(_beneficiary), _amount); + assertEq(_deposit.beneficiary, _beneficiary); + } + + function testFuzz_AssignsEarningPowerToDifferentBeneficiariesForDifferentDepositsFromTheSameDepositor( + address _depositor, + uint256 _amount1, + uint256 _amount2, + address _delegatee, + address _beneficiary1, + address _beneficiary2 + ) public { + vm.assume(_beneficiary1 != _beneficiary2); + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor, _amount1 + _amount2); + + // Perform both deposits and track their identifiers separately + UniStaker.DepositIdentifier _depositId1 = + _stake(_depositor, _amount1, _delegatee, _beneficiary1); + UniStaker.DepositIdentifier _depositId2 = + _stake(_depositor, _amount2, _delegatee, _beneficiary2); + UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1); + UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2); + + // Check that the earning power has been recorded independently + assertEq(_deposit1.beneficiary, _beneficiary1); + assertEq(uniStaker.earningPower(_beneficiary1), _amount1); + assertEq(_deposit2.beneficiary, _beneficiary2); + assertEq(uniStaker.earningPower(_beneficiary2), _amount2); + } + + function testFuzz_AssignsEarningPowerToTheSameBeneficiarySpecifiedByTwoDifferentDepositors( + address _depositor1, + address _depositor2, + uint256 _amount1, + uint256 _amount2, + address _delegatee, + address _beneficiary + ) public { + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor1, _amount1); + _mintGovToken(_depositor2, _amount2); + + // Perform both deposits and track their identifiers separately + UniStaker.DepositIdentifier _depositId1 = + _stake(_depositor1, _amount1, _delegatee, _beneficiary); + UniStaker.DepositIdentifier _depositId2 = + _stake(_depositor2, _amount2, _delegatee, _beneficiary); + UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1); + UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2); + + assertEq(_deposit1.beneficiary, _beneficiary); + assertEq(_deposit2.beneficiary, _beneficiary); + assertEq(uniStaker.earningPower(_beneficiary), _amount1 + _amount2); + } + mapping(UniStaker.DepositIdentifier depositId => bool isUsed) isIdUsed; function test_NeverReusesADepositIdentifier() public { @@ -454,6 +566,176 @@ contract Withdraw is UniStakerTest { assertEq(uniStaker.totalSupply(), _depositAmount1 + _depositAmount2 - _withdrawalAmount); } + function testFuzz_RemovesFullEarningPowerFromADepositorWhoHadSelfAssignedIt( + address _depositor, + uint256 _amount, + address _delegatee + ) public { + UniStaker.DepositIdentifier _depositId; + (_amount, _depositId) = _boundMintAndStake(_depositor, _amount, _delegatee); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _amount); + + assertEq(uniStaker.earningPower(_depositor), 0); + } + + function testFuzz_RemovesPartialEarningPowerFromADepositorWhoHadSelfAssignedIt( + address _depositor, + uint256 _depositAmount, + address _delegatee, + uint256 _withdrawalAmount + ) public { + UniStaker.DepositIdentifier _depositId; + (_depositAmount, _depositId) = _boundMintAndStake(_depositor, _depositAmount, _delegatee); + _withdrawalAmount = bound(_withdrawalAmount, 0, _depositAmount); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _withdrawalAmount); + + assertEq(uniStaker.earningPower(_depositor), _depositAmount - _withdrawalAmount); + } + + function testFuzz_RemovesFullEarningPowerFromABeneficiary( + address _depositor, + uint256 _amount, + address _delegatee, + address _beneficiary + ) public { + UniStaker.DepositIdentifier _depositId; + (_amount, _depositId) = _boundMintAndStake(_depositor, _amount, _delegatee, _beneficiary); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _amount); + + assertEq(uniStaker.earningPower(_beneficiary), 0); + } + + function testFuzz_RemovesPartialEarningPowerFromABeneficiary( + address _depositor, + uint256 _depositAmount, + address _delegatee, + address _beneficiary, + uint256 _withdrawalAmount + ) public { + UniStaker.DepositIdentifier _depositId; + (_depositAmount, _depositId) = + _boundMintAndStake(_depositor, _depositAmount, _delegatee, _beneficiary); + _withdrawalAmount = bound(_withdrawalAmount, 0, _depositAmount); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _withdrawalAmount); + + assertEq(uniStaker.earningPower(_beneficiary), _depositAmount - _withdrawalAmount); + } + + function testFuzz_RemovesPartialEarningPowerFromABeneficiaryAssignedByTwoDepositors( + address _depositor1, + address _depositor2, + uint256 _depositAmount1, + uint256 _depositAmount2, + address _delegatee, + address _beneficiary, + uint256 _withdrawalAmount1, + uint256 _withdrawalAmount2 + ) public { + UniStaker.DepositIdentifier _depositId1; + (_depositAmount1, _depositId1) = + _boundMintAndStake(_depositor1, _depositAmount1, _delegatee, _beneficiary); + _withdrawalAmount1 = bound(_withdrawalAmount1, 0, _depositAmount1); + + UniStaker.DepositIdentifier _depositId2; + (_depositAmount2, _depositId2) = + _boundMintAndStake(_depositor2, _depositAmount2, _delegatee, _beneficiary); + _withdrawalAmount2 = bound(_withdrawalAmount2, 0, _depositAmount2); + + vm.prank(_depositor1); + uniStaker.withdraw(_depositId1, _withdrawalAmount1); + + assertEq( + uniStaker.earningPower(_beneficiary), _depositAmount1 - _withdrawalAmount1 + _depositAmount2 + ); + + vm.prank(_depositor2); + uniStaker.withdraw(_depositId2, _withdrawalAmount2); + + assertEq( + uniStaker.earningPower(_beneficiary), + _depositAmount1 - _withdrawalAmount1 + _depositAmount2 - _withdrawalAmount2 + ); + } + + function testFuzz_RemovesPartialEarningPowerFromDifferentBeneficiariesOfTheSameDepositor( + address _depositor, + uint256 _depositAmount1, + uint256 _depositAmount2, + address _delegatee, + address _beneficiary1, + address _beneficiary2, + uint256 _withdrawalAmount1, + uint256 _withdrawalAmount2 + ) public { + vm.assume(_beneficiary1 != _beneficiary2); + + UniStaker.DepositIdentifier _depositId1; + (_depositAmount1, _depositId1) = + _boundMintAndStake(_depositor, _depositAmount1, _delegatee, _beneficiary1); + _withdrawalAmount1 = bound(_withdrawalAmount1, 0, _depositAmount1); + + UniStaker.DepositIdentifier _depositId2; + (_depositAmount2, _depositId2) = + _boundMintAndStake(_depositor, _depositAmount2, _delegatee, _beneficiary2); + _withdrawalAmount2 = bound(_withdrawalAmount2, 0, _depositAmount2); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId1, _withdrawalAmount1); + + assertEq(uniStaker.earningPower(_beneficiary1), _depositAmount1 - _withdrawalAmount1); + assertEq(uniStaker.earningPower(_beneficiary2), _depositAmount2); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId2, _withdrawalAmount2); + + assertEq(uniStaker.earningPower(_beneficiary1), _depositAmount1 - _withdrawalAmount1); + assertEq(uniStaker.earningPower(_beneficiary2), _depositAmount2 - _withdrawalAmount2); + } + + function testFuzz_RemovesPartialEarningPowerFromDifferentBeneficiariesAndDifferentDepositors( + address _depositor1, + address _depositor2, + uint256 _depositAmount1, + uint256 _depositAmount2, + address _delegatee, + address _beneficiary1, + address _beneficiary2, + uint256 _withdrawalAmount1, + uint256 _withdrawalAmount2 + ) public { + vm.assume(_beneficiary1 != _beneficiary2); + + UniStaker.DepositIdentifier _depositId1; + (_depositAmount1, _depositId1) = + _boundMintAndStake(_depositor1, _depositAmount1, _delegatee, _beneficiary1); + _withdrawalAmount1 = bound(_withdrawalAmount1, 0, _depositAmount1); + + UniStaker.DepositIdentifier _depositId2; + (_depositAmount2, _depositId2) = + _boundMintAndStake(_depositor2, _depositAmount2, _delegatee, _beneficiary2); + _withdrawalAmount2 = bound(_withdrawalAmount2, 0, _depositAmount2); + + vm.prank(_depositor1); + uniStaker.withdraw(_depositId1, _withdrawalAmount1); + + assertEq(uniStaker.earningPower(_beneficiary1), _depositAmount1 - _withdrawalAmount1); + assertEq(uniStaker.earningPower(_beneficiary2), _depositAmount2); + + vm.prank(_depositor2); + uniStaker.withdraw(_depositId2, _withdrawalAmount2); + + assertEq(uniStaker.earningPower(_beneficiary1), _depositAmount1 - _withdrawalAmount1); + assertEq(uniStaker.earningPower(_beneficiary2), _depositAmount2 - _withdrawalAmount2); + } + function testFuzz_RevertIf_TheWithdrawerIsNotTheDepositor( address _depositor, uint256 _amount,