Add timestamp to delegator balance updates

contracts/governance/src/IZeroExVotes.sol

@@ -112,13 +112,17 @@ 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(
         address src,
         address dst,
         uint256 srcBalance,
         uint256 dstBalance,
+        uint96 srcBalanceLastUpdated,
+        uint96 dstBalanceLastUpdated,
         uint256 amount
     ) external returns (bool);
 

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
+        );
     }
 }

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) {

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);
+    }
 }