diff --git a/contracts/governance/src/IZeroExVotes.sol b/contracts/governance/src/IZeroExVotes.sol index 4d2ed4d7fc..9bc7df5e06 100644 --- a/contracts/governance/src/IZeroExVotes.sol +++ b/contracts/governance/src/IZeroExVotes.sol @@ -112,6 +112,8 @@ interface IZeroExVotes { * @param dst the delegatee we are moving voting power to * @param srcBalance balance of the delegator whose delegatee is `src`. This is value _after_ the transfer. * @param dstBalance balance of the delegator whose delegatee is `dst`. This is value _after_ the transfer. + * @param srcBalanceLastUpdated timestamp when balance of `src` was last updated. + * @param dstBalanceLastUpdated timestamp when balance of `dst` was last updated. * @param amount The amount of tokens transferred from the source delegate to destination delegate. */ function moveVotingPower( @@ -119,6 +121,8 @@ interface IZeroExVotes { address dst, uint256 srcBalance, uint256 dstBalance, + uint96 srcBalanceLastUpdated, + uint96 dstBalanceLastUpdated, uint256 amount ) external returns (bool); diff --git a/contracts/governance/src/ZRXWrappedToken.sol b/contracts/governance/src/ZRXWrappedToken.sol index 8a572585a9..d57b468454 100644 --- a/contracts/governance/src/ZRXWrappedToken.sol +++ b/contracts/governance/src/ZRXWrappedToken.sol @@ -22,12 +22,18 @@ import "@openzeppelin/token/ERC20/ERC20.sol"; import "@openzeppelin/token/ERC20/extensions/draft-ERC20Permit.sol"; import "@openzeppelin/token/ERC20/extensions/ERC20Wrapper.sol"; import "@openzeppelin/governance/utils/IVotes.sol"; +import "@openzeppelin/utils/math/SafeCast.sol"; import "./IZeroExVotes.sol"; import "./CallWithGas.sol"; contract ZRXWrappedToken is ERC20, ERC20Permit, ERC20Wrapper { using CallWithGas for address; + struct DelegateInfo { + address delegate; + uint96 balanceLastUpdated; + } + constructor( IERC20 wrappedToken, IZeroExVotes _zeroExVotes @@ -36,7 +42,7 @@ contract ZRXWrappedToken is ERC20, ERC20Permit, ERC20Wrapper { } IZeroExVotes public immutable zeroExVotes; - mapping(address => address) private _delegates; + mapping(address => DelegateInfo) private _delegates; bytes32 private constant _DELEGATION_TYPEHASH = keccak256("Delegation(address delegatee,uint256 nonce,uint256 expiry)"); @@ -55,10 +61,26 @@ contract ZRXWrappedToken is ERC20, ERC20Permit, ERC20Wrapper { function _afterTokenTransfer(address from, address to, uint256 amount) internal override(ERC20) { super._afterTokenTransfer(from, to, amount); - uint256 fromBalance = delegates(from) == address(0) ? 0 : balanceOf(from) + amount; - uint256 toBalance = delegates(to) == address(0) ? 0 : balanceOf(to) - amount; + DelegateInfo memory fromDelegate = delegateInfo(from); + DelegateInfo memory toDelegate = delegateInfo(to); + + uint256 fromBalance = fromDelegate.delegate == address(0) ? 0 : balanceOf(from) + amount; + uint256 toBalance = toDelegate.delegate == address(0) ? 0 : balanceOf(to) - amount; - zeroExVotes.moveVotingPower(delegates(from), delegates(to), fromBalance, toBalance, amount); + if (fromDelegate.delegate != address(0)) + _delegates[from].balanceLastUpdated = SafeCast.toUint96(block.timestamp); + + if (toDelegate.delegate != address(0)) _delegates[to].balanceLastUpdated = SafeCast.toUint96(block.timestamp); + + zeroExVotes.moveVotingPower( + fromDelegate.delegate, + toDelegate.delegate, + fromBalance, + toBalance, + fromDelegate.balanceLastUpdated, + toDelegate.balanceLastUpdated, + amount + ); } function _mint(address account, uint256 amount) internal override(ERC20) { @@ -81,6 +103,14 @@ contract ZRXWrappedToken is ERC20, ERC20Permit, ERC20Wrapper { * @dev Get the address `account` is currently delegating to. */ function delegates(address account) public view returns (address) { + return _delegates[account].delegate; + } + + function delegatorBalanceLastUpdated(address account) public view returns (uint96) { + return _delegates[account].balanceLastUpdated; + } + + function delegateInfo(address account) public view returns (DelegateInfo memory) { return _delegates[account]; } @@ -112,12 +142,21 @@ contract ZRXWrappedToken is ERC20, ERC20Permit, ERC20Wrapper { * Emits events {DelegateChanged} and {IZeroExVotes-DelegateVotesChanged}. */ function _delegate(address delegator, address delegatee) internal virtual { - address currentDelegate = delegates(delegator); + DelegateInfo memory delegateInfo = delegateInfo(delegator); uint256 delegatorBalance = balanceOf(delegator); - _delegates[delegator] = delegatee; - emit DelegateChanged(delegator, currentDelegate, delegatee); + _delegates[delegator] = DelegateInfo(delegatee, SafeCast.toUint96(block.timestamp)); + + emit DelegateChanged(delegator, delegateInfo.delegate, delegatee); - zeroExVotes.moveVotingPower(currentDelegate, delegatee, delegatorBalance, 0, delegatorBalance); + zeroExVotes.moveVotingPower( + delegateInfo.delegate, + delegatee, + delegatorBalance, + 0, + delegateInfo.balanceLastUpdated, + 0, + delegatorBalance + ); } } diff --git a/contracts/governance/src/ZeroExVotes.sol b/contracts/governance/src/ZeroExVotes.sol index 5e01e0c90f..1dfec72790 100644 --- a/contracts/governance/src/ZeroExVotes.sol +++ b/contracts/governance/src/ZeroExVotes.sol @@ -136,6 +136,8 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg address dst, uint256 srcBalance, uint256 dstBalance, + uint96 srcBalanceLastUpdated, + uint96 dstBalanceLastUpdated, uint256 amount ) public virtual onlyToken returns (bool) { if (src != dst) { diff --git a/contracts/governance/test/ZRXWrappedTokenTest.t.sol b/contracts/governance/test/ZRXWrappedTokenTest.t.sol index 101c873239..1145c662f6 100644 --- a/contracts/governance/test/ZRXWrappedTokenTest.t.sol +++ b/contracts/governance/test/ZRXWrappedTokenTest.t.sol @@ -258,4 +258,57 @@ contract ZRXWrappedTokenTest is BaseTest { assertEq(votes.getVotes(account3), 0); assertEq(votes.getQuadraticVotes(account3), 0); } + + function testShouldUpdateVotingPowerWhenDepositing() public { + // Account 2 wraps ZRX and delegates voting power to itself + vm.startPrank(account2); + token.approve(address(wToken), 10e18); + wToken.depositFor(account2, 7e18); + wToken.delegate(account2); + + assertEq(votes.getVotes(account2), 7e18); + assertEq(votes.getQuadraticVotes(account2), 7e18); + + wToken.depositFor(account2, 2e18); + assertEq(votes.getVotes(account2), 9e18); + assertEq(votes.getQuadraticVotes(account2), 9e18); + } + + function testShouldUpdateVotingPowerWhenWithdrawing() public { + // Account 2 wraps ZRX and delegates voting power to itself + vm.startPrank(account2); + token.approve(address(wToken), 10e18); + wToken.depositFor(account2, 10e18); + wToken.delegate(account2); + + assertEq(votes.getVotes(account2), 10e18); + assertEq(votes.getQuadraticVotes(account2), 10e18); + + wToken.withdrawTo(account2, 2e18); + assertEq(votes.getVotes(account2), 8e18); + assertEq(votes.getQuadraticVotes(account2), 8e18); + } + + function testShouldSetDelegateBalanceLastUpdatedOnTransfer() public { + ZRXWrappedToken.DelegateInfo memory account2DelegateInfo = wToken.delegateInfo(account2); + assertEq(account2DelegateInfo.delegate, address(0)); + assertEq(account2DelegateInfo.balanceLastUpdated, 0); + + // Account 2 wraps ZRX and delegates voting power to account3 + vm.startPrank(account2); + token.approve(address(wToken), 10e18); + wToken.depositFor(account2, 10e18); + wToken.delegate(account3); + + account2DelegateInfo = wToken.delegateInfo(account2); + assertEq(account2DelegateInfo.delegate, account3); + assertEq(account2DelegateInfo.balanceLastUpdated, block.timestamp); + + vm.warp(101); + wToken.transfer(account3, 3e18); + + account2DelegateInfo = wToken.delegateInfo(account2); + assertEq(account2DelegateInfo.delegate, account3); + assertEq(account2DelegateInfo.balanceLastUpdated, 101); + } }