From e2a2e890b4fd9937639b6d38c6f30bbc9f61ca1b Mon Sep 17 00:00:00 2001
From: elenadimitrova <elena@arenabg.com>
Date: Mon, 13 Feb 2023 19:15:02 +0200
Subject: [PATCH] Remove keeping track of total balances for voting

---
 contracts/governance/src/IZeroExVotes.sol     | 11 ++++++
 contracts/governance/src/ZeroExVotes.sol      | 19 ++++++++--
 .../governance/test/ZRXWrappedTokenTest.t.sol | 36 +++++++++++++++++++
 3 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/contracts/governance/src/IZeroExVotes.sol b/contracts/governance/src/IZeroExVotes.sol
index 4dc9847d79..146ab821bf 100644
--- a/contracts/governance/src/IZeroExVotes.sol
+++ b/contracts/governance/src/IZeroExVotes.sol
@@ -85,6 +85,17 @@ interface IZeroExVotes {
      */
     function getPastTotalSupply(uint256 blockNumber) external view returns (uint256);
 
+    /**
+     * @dev Retrieve the sqrt of `totalSupply` at the end of `blockNumber`. Note, this value is the square root of the
+     * sum of all balances.
+     * It is but NOT the sum of all the sqrt of the delegated votes!
+     *
+     * Requirements:
+     *
+     * - `blockNumber` must have been already mined
+     */
+    function getPastQuadraticTotalSupply(uint256 blockNumber) external view returns (uint256);
+
     /**
      * @dev Moves the voting power for an account with balance `delegateBalance` from `srcDelegatee` to `dstDelegatee`.
      * Note that if the delegator isn't delegating to anyone before the function call `srcDelegatee` = address(0)
diff --git a/contracts/governance/src/ZeroExVotes.sol b/contracts/governance/src/ZeroExVotes.sol
index 1126c461c1..cf2c77b9ed 100644
--- a/contracts/governance/src/ZeroExVotes.sol
+++ b/contracts/governance/src/ZeroExVotes.sol
@@ -93,14 +93,27 @@ contract ZeroExVotes is IZeroExVotes {
     /**
      * @inheritdoc IZeroExVotes
      */
-    // TODO we probably don't need to keep track of total supply checkpoints as all governance values are fixed numbers
     function getPastTotalSupply(uint256 blockNumber) public view override returns (uint256) {
         require(blockNumber < block.number, "ZeroExVotes: block not yet mined");
 
+        // Note that due to the disabled updates of `_totalSupplyCheckpoints` in `writeCheckpointTotalSupply` function
+        // this always returns 0.
         Checkpoint memory checkpoint = _checkpointsLookup(_totalSupplyCheckpoints, blockNumber);
         return checkpoint.votes;
     }
 
+    /**
+     * @inheritdoc IZeroExVotes
+     */
+    function getPastQuadraticTotalSupply(uint256 blockNumber) public view override returns (uint256) {
+        require(blockNumber < block.number, "ZeroExVotes: block not yet mined");
+
+        // Note that due to the disabled updates of `_totalSupplyCheckpoints` in `writeCheckpointTotalSupply` function
+        // this always returns 0.
+        Checkpoint memory checkpoint = _checkpointsLookup(_totalSupplyCheckpoints, blockNumber);
+        return checkpoint.quadraticVotes;
+    }
+
     /**
      * @inheritdoc IZeroExVotes
      */
@@ -219,7 +232,9 @@ contract ZeroExVotes is IZeroExVotes {
      * @inheritdoc IZeroExVotes
      */
     function writeCheckpointTotalSupply(uint256 totalSupply) public override onlyToken {
-        _writeCheckpoint(_totalSupplyCheckpoints, totalSupply, Math.sqrt(totalSupply));
+        // Currently we don't keep track of total supply checkpoints as all governance values are fixed numbers
+        // i.e. governance quorum is not a percentage of total
+        // _writeCheckpoint(_totalSupplyCheckpoints, totalSupply, Math.sqrt(totalSupply));
     }
 
     /**
diff --git a/contracts/governance/test/ZRXWrappedTokenTest.t.sol b/contracts/governance/test/ZRXWrappedTokenTest.t.sol
index 392ba03700..5ab9eb8cca 100644
--- a/contracts/governance/test/ZRXWrappedTokenTest.t.sol
+++ b/contracts/governance/test/ZRXWrappedTokenTest.t.sol
@@ -103,6 +103,42 @@ contract ZRXWrappedTokenTest is BaseTest {
         assertEq(tokenBalance2, 100e18 - wTokenBalance2 - tokenBalance4);
     }
 
+    function testWrappedZRXTotalsAreCorrect() public {
+        // Wrap 1e18 and check total supply is correct
+        vm.startPrank(account2);
+        token.approve(address(wToken), 1e18);
+        wToken.depositFor(account2, 1e18);
+        vm.stopPrank();
+        uint256 wTokenBalance = wToken.totalSupply();
+        assertEq(wTokenBalance, 1e18);
+
+        // Wrap 2e18 more and check total supply is correct
+        vm.startPrank(account3);
+        token.approve(address(wToken), 2e18);
+        wToken.depositFor(account3, 2e18);
+        vm.stopPrank();
+        wTokenBalance = wToken.totalSupply();
+        assertEq(wTokenBalance, 1e18 + 2e18);
+
+        // Unwrap 1e7 and check total supply is correct
+        vm.startPrank(account2);
+        wToken.withdrawTo(account2, 1e7);
+        vm.stopPrank();
+        wTokenBalance = wToken.totalSupply();
+        assertEq(wTokenBalance, 3e18 - 1e7);
+
+        // Unwrap 8e17 and check total supply is correct
+        vm.startPrank(account2);
+        wToken.withdrawTo(account2, 8e17);
+        vm.stopPrank();
+        wTokenBalance = wToken.totalSupply();
+        assertEq(wTokenBalance, 3e18 - 1e7 - 8e17);
+
+        // We are not keeping record of total balances so check they are zero
+        assertEq(votes.getPastTotalSupply(0), 0);
+        assertEq(votes.getPastQuadraticTotalSupply(0), 0);
+    }
+
     function testShouldNotBeAbleToReinitialiseTheZeroExVotes() public {
         vm.expectRevert("ZeroExVotes: Already initialized");
         votes.initialize(account2);