Skip to content

Commit

Permalink
Implement internal accounting for stake accounting
Browse files Browse the repository at this point in the history
  • Loading branch information
@apbendi
apbendi committed Dec 12, 2023
1 parent a954b25 commit ca84d74
Show file tree
Hide file tree
Showing 2 changed files with 210 additions and 6 deletions.
31 changes: 28 additions & 3 deletions src/UniStaker.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,25 @@ import {SafeERC20} from "openzeppelin/token/ERC20/utils/SafeERC20.sol";
import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol";

contract UniStaker is ReentrancyGuard {
type DepositIdentifier is uint256;

struct Deposit {
uint256 balance;
address owner;
address delegatee;
}

IERC20 public immutable REWARDS_TOKEN;
IERC20Delegates public immutable STAKE_TOKEN;

DepositIdentifier private nextDepositId;

uint256 public totalSupply;

mapping(address depositor => uint256 amount) public totalDeposits;

mapping(DepositIdentifier depositId => Deposit deposit) public deposits;

mapping(address delegatee => DelegationSurrogate surrogate) public surrogates;

constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken) {
Expand All @@ -19,13 +35,17 @@ contract UniStaker is ReentrancyGuard {
}

function stake(uint256 _amount, address _delegatee)
public
external
nonReentrant
returns (uint256 _depositId)
returns (DepositIdentifier _depositId)
{
DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee);
_stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount);
_depositId = 1;
_depositId = _useDepositId();

totalSupply += _amount;
totalDeposits[msg.sender] += _amount;
deposits[_depositId] = Deposit({balance: _amount, owner: msg.sender, delegatee: _delegatee});
}

function _fetchOrDeploySurrogate(address _delegatee)
Expand All @@ -43,4 +63,9 @@ contract UniStaker is ReentrancyGuard {
function _stakeTokenSafeTransferFrom(address _from, address _to, uint256 _value) internal {
SafeERC20.safeTransferFrom(IERC20(address(STAKE_TOKEN)), _from, _to, _value);
}

function _useDepositId() internal returns (DepositIdentifier _depositId) {
_depositId = nextDepositId;
nextDepositId = DepositIdentifier.wrap(DepositIdentifier.unwrap(_depositId) + 1);
}
}
185 changes: 182 additions & 3 deletions test/UniStaker.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,22 @@ contract UniStakerTest is Test {

function _stake(address _depositor, uint256 _amount, address _delegatee)
internal
returns (uint256 _depositId)
returns (UniStaker.DepositIdentifier _depositId)
{
vm.startPrank(_depositor);
govToken.approve(address(uniStaker), _amount);
_depositId = uniStaker.stake(_amount, _delegatee);
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});
}
}

contract Constructor is UniStakerTest {
Expand All @@ -50,7 +59,7 @@ contract Constructor is UniStakerTest {
}

contract Stake is UniStakerTest {
function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAUserStakes(
function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAnAccountStakes(
address _depositor,
uint256 _amount,
address _delegatee
Expand Down Expand Up @@ -91,7 +100,7 @@ contract Stake is UniStakerTest {
assertEq(govToken.balanceOf(address(_surrogate)), _amount1 + _amount2);
}

function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenUsersStakesToDifferentDelegatees(
function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenAccountsStakesToDifferentDelegatees(
address _depositor1,
uint256 _amount1,
address _depositor2,
Expand Down Expand Up @@ -122,4 +131,174 @@ contract Stake is UniStakerTest {
assertEq(govToken.delegates(address(_surrogate2)), _delegatee2);
assertEq(govToken.balanceOf(address(_surrogate2)), _amount2);
}

function testFuzz_UpdatesTheTotalSupplyWhenAnAccountStakes(
address _depositor,
uint256 _amount,
address _delegatee
) public {
_amount = _boundMintAmount(_amount);
_mintGovToken(_depositor, _amount);

_stake(_depositor, _amount, _delegatee);

assertEq(uniStaker.totalSupply(), _amount);
}

function testFuzz_UpdatesTheTotalSupplyWhenTwoAccountsStake(
address _depositor1,
uint256 _amount1,
address _depositor2,
uint256 _amount2,
address _delegatee1,
address _delegatee2
) public {
_amount1 = _boundMintAmount(_amount1);
_amount2 = _boundMintAmount(_amount2);
_mintGovToken(_depositor1, _amount1);
_mintGovToken(_depositor2, _amount2);

_stake(_depositor1, _amount1, _delegatee1);
assertEq(uniStaker.totalSupply(), _amount1);

_stake(_depositor2, _amount2, _delegatee2);
assertEq(uniStaker.totalSupply(), _amount1 + _amount2);
}

function testFuzz_UpdatesAnAccountsTotalDepositsWhenItStakes(
address _depositor,
uint256 _amount1,
uint256 _amount2,
address _delegatee1,
address _delegatee2
) public {
_amount1 = _boundMintAmount(_amount1);
_amount2 = _boundMintAmount(_amount2);
_mintGovToken(_depositor, _amount1 + _amount2);

// First stake + check total
_stake(_depositor, _amount1, _delegatee1);
assertEq(uniStaker.totalDeposits(_depositor), _amount1);

// Second stake + check total
_stake(_depositor, _amount2, _delegatee2);
assertEq(uniStaker.totalDeposits(_depositor), _amount1 + _amount2);
}

function testFuzz_UpdatesDifferentAccountsTotalDepositsIndependently(
address _depositor1,
uint256 _amount1,
address _depositor2,
uint256 _amount2,
address _delegatee1,
address _delegatee2
) public {
vm.assume(_depositor1 != _depositor2);
_amount1 = _boundMintAmount(_amount1);
_amount2 = _boundMintAmount(_amount2);
_mintGovToken(_depositor1, _amount1);
_mintGovToken(_depositor2, _amount2);

_stake(_depositor1, _amount1, _delegatee1);
assertEq(uniStaker.totalDeposits(_depositor1), _amount1);

_stake(_depositor2, _amount2, _delegatee2);
assertEq(uniStaker.totalDeposits(_depositor2), _amount2);
}

function testFuzz_TracksTheBalanceForASpecificDeposit(
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(_deposit.balance, _amount);
assertEq(_deposit.owner, _depositor);
assertEq(_deposit.delegatee, _delegatee);
}

function testFuzz_TracksTheBalanceForDifferentDepositsFromTheSameAccountIndependently(
address _depositor,
uint256 _amount1,
uint256 _amount2,
address _delegatee1,
address _delegatee2
) public {
_amount1 = _boundMintAmount(_amount1);
_amount2 = _boundMintAmount(_amount2);
_mintGovToken(_depositor, _amount1 + _amount2);

// Perform both deposits and track their identifiers separately
UniStaker.DepositIdentifier _depositId1 = _stake(_depositor, _amount1, _delegatee1);
UniStaker.DepositIdentifier _depositId2 = _stake(_depositor, _amount2, _delegatee2);
UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1);
UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2);

// Check that the deposits have been recorded independently
assertEq(_deposit1.balance, _amount1);
assertEq(_deposit1.owner, _depositor);
assertEq(_deposit1.delegatee, _delegatee1);
assertEq(_deposit2.balance, _amount2);
assertEq(_deposit2.owner, _depositor);
assertEq(_deposit2.delegatee, _delegatee2);
}

function testFuzz_TracksTheBalanceForDepositsFromDifferentAccountsIndependently(
address _depositor1,
address _depositor2,
uint256 _amount1,
uint256 _amount2,
address _delegatee1,
address _delegatee2
) 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, _delegatee1);
UniStaker.DepositIdentifier _depositId2 = _stake(_depositor2, _amount2, _delegatee2);
UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1);
UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2);

// Check that the deposits have been recorded independently
assertEq(_deposit1.balance, _amount1);
assertEq(_deposit1.owner, _depositor1);
assertEq(_deposit1.delegatee, _delegatee1);
assertEq(_deposit2.balance, _amount2);
assertEq(_deposit2.owner, _depositor2);
assertEq(_deposit2.delegatee, _delegatee2);
}

mapping(UniStaker.DepositIdentifier depositId => bool isUsed) isIdUsed;

function test_NeverReusesADepositIdentifier() public {
address _depositor = address(0xdeadbeef);
uint256 _amount = 116;
address _delegatee = address(0xaceface);

UniStaker.DepositIdentifier _depositId;

for (uint256 _i; _i < 10_000; _i++) {
// Perform the stake and save the deposit identifier
_amount = _bound(_amount, 0, 100_000_000_000e18);
_mintGovToken(_depositor, _amount);
_depositId = _stake(_depositor, _amount, _delegatee);

// Ensure the identifier hasn't yet been used
assertFalse(isIdUsed[_depositId]);
// Record the fact this deposit Id has been used
isIdUsed[_depositId] = true;

// Reset all the inputs for the next deposit by hashing the last inputs
_depositor = address(uint160(uint256(keccak256(abi.encode(_depositor)))));
_amount = uint256(keccak256(abi.encode(_amount)));
_delegatee = address(uint160(uint256(keccak256(abi.encode(_delegatee)))));
}
}
}

0 comments on commit ca84d74

Please sign in to comment.