Test votes migration for adding a new vote weight mechanism

.gitignore

@@ -114,6 +114,9 @@ contracts/governance/out
 packages/*/docs/README.md
 
 .DS_Store
+*~
+\#*\#
+.\#*
 
 # the snapshot that gets built for migrations sure does have a ton of files
 packages/migrations/0x_ganache_snapshot*

contracts/governance/src/ZeroExVotes.sol

@@ -31,7 +31,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
     address public immutable token;
     uint256 public immutable quadraticThreshold;
 
-    mapping(address => Checkpoint[]) private _checkpoints;
+    mapping(address => Checkpoint[]) internal _checkpoints;
     Checkpoint[] private _totalSupplyCheckpoints;
 
     constructor(address _token, uint256 _quadraticThreshold) {
@@ -48,7 +48,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
         _;
     }
 
-    function initialize() public onlyProxy initializer {
+    function initialize() public virtual onlyProxy initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
     }
@@ -146,7 +146,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
                     uint256 newWeight,
                     uint256 oldQuadraticWeight,
                     uint256 newQuadraticWeight
-                ) = _writeCheckpoint(_checkpoints[src], _subtract, srcBalance, amount);
+                ) = _writeCheckpoint(_checkpoints[src], _subtract, srcBalance, srcBalanceLastUpdated, amount);
 
                 emit DelegateVotesChanged(src, oldWeight, newWeight);
                 emit DelegateQuadraticVotesChanged(src, oldQuadraticWeight, newQuadraticWeight);
@@ -158,7 +158,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
                     uint256 newWeight,
                     uint256 oldQuadraticWeight,
                     uint256 newQuadraticWeight
-                ) = _writeCheckpoint(_checkpoints[dst], _add, dstBalance, amount);
+                ) = _writeCheckpoint(_checkpoints[dst], _add, dstBalance, dstBalanceLastUpdated, amount);
 
                 emit DelegateVotesChanged(dst, oldWeight, newWeight);
                 emit DelegateQuadraticVotesChanged(dst, oldQuadraticWeight, newQuadraticWeight);
@@ -178,6 +178,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
             _totalSupplyCheckpoints,
             _add,
             accountBalance,
+            0,
             amount
         );
 
@@ -196,6 +197,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
             _totalSupplyCheckpoints,
             _subtract,
             accountBalance,
+            0,
             amount
         );
 
@@ -211,7 +213,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
     function _checkpointsLookup(
         Checkpoint[] storage ckpts,
         uint256 blockNumber
-    ) private view returns (Checkpoint memory) {
+    ) internal view returns (Checkpoint memory) {
         // We run a binary search to look for the earliest checkpoint taken after `blockNumber`.
         //
         // Initially we check if the block is recent to narrow the search range.
@@ -258,8 +260,13 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
         Checkpoint[] storage ckpts,
         function(uint256, uint256) view returns (uint256) op,
         uint256 userBalance,
+        uint96 balanceLastUpdated,
         uint256 delta
-    ) private returns (uint256 oldWeight, uint256 newWeight, uint256 oldQuadraticWeight, uint256 newQuadraticWeight) {
+    )
+        internal
+        virtual
+        returns (uint256 oldWeight, uint256 newWeight, uint256 oldQuadraticWeight, uint256 newQuadraticWeight)
+    {
         uint256 pos = ckpts.length;
 
         Checkpoint memory oldCkpt = pos == 0 ? Checkpoint(0, 0, 0) : _unsafeAccess(ckpts, pos - 1);
@@ -313,7 +320,7 @@ contract ZeroExVotes is IZeroExVotes, Initializable, OwnableUpgradeable, UUPSUpg
      * Implementation from openzeppelin/token/ERC20/extensions/ERC20Votes.sol
      * https://github.com/ethereum/solidity/issues/9117
      */
-    function _unsafeAccess(Checkpoint[] storage ckpts, uint256 pos) private pure returns (Checkpoint storage result) {
+    function _unsafeAccess(Checkpoint[] storage ckpts, uint256 pos) internal pure returns (Checkpoint storage result) {
         assembly ("memory-safe") {
             mstore(0, ckpts.slot)
             result.slot := add(keccak256(0, 0x20), pos)

contracts/governance/test/CubeRoot.sol

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+library CubeRoot {
+    /// @dev Returns the cube root of `x`.
+    /// Credit to pleasemarkdarkly under MIT license
+    // Originaly from https://github.com/pleasemarkdarkly/fei-protocol-core-hh/blob/main/contracts/utils/Roots.sol
+    function cbrt(uint y) internal pure returns (uint z) {
+        // Newton's method https://en.wikipedia.org/wiki/Cube_root#Numerical_methods
+        if (y > 7) {
+            z = y;
+            uint x = y / 3 + 1;
+            while (x < z) {
+                z = x;
+                x = (y / (x * x) + (2 * x)) / 3;
+            }
+        } else if (y != 0) {
+            z = 1;
+        }
+    }
+}

contracts/governance/test/ZeroExVotesMigration.sol

@@ -0,0 +1,215 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+
+  Copyright 2023 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+pragma solidity ^0.8.19;
+
+import {ZeroExVotes} from "../src/ZeroExVotes.sol";
+import {SafeCast} from "@openzeppelin/utils/math/SafeCast.sol";
+import {Math} from "@openzeppelin/utils/math/Math.sol";
+import {CubeRoot} from "./CubeRoot.sol";
+
+contract ZeroExVotesMigration is ZeroExVotes {
+    uint32 public migrationBlock;
+
+    constructor(address _token, uint256 _quadraticThreshold) ZeroExVotes(_token, _quadraticThreshold) {}
+
+    function initialize() public virtual override onlyProxy reinitializer(2) {
+        migrationBlock = uint32(block.number);
+    }
+
+    struct CheckpointMigration {
+        uint32 fromBlock;
+        uint96 votes;
+        uint96 quadraticVotes;
+        uint32 migratedVotes;
+    }
+
+    function _toMigration(Checkpoint storage ckpt) internal pure returns (CheckpointMigration storage result) {
+        assembly {
+            result.slot := ckpt.slot
+        }
+    }
+
+    function _toMigration(Checkpoint[] storage ckpt) internal pure returns (CheckpointMigration[] storage result) {
+        assembly {
+            result.slot := ckpt.slot
+        }
+    }
+
+    function getMigratedVotes(address account) public view returns (uint256) {
+        uint256 pos = _checkpoints[account].length;
+        if (pos == 0) {
+            return 0;
+        }
+        Checkpoint storage ckpt = _unsafeAccess(_checkpoints[account], pos - 1);
+        if (ckpt.fromBlock <= migrationBlock) {
+            return 0;
+        }
+        return _toMigration(ckpt).migratedVotes;
+    }
+
+    function getPastMigratedVotes(address account, uint256 blockNumber) public view returns (uint256) {
+        require(blockNumber < block.number, "ZeroExVotesMigration: block not yet mined");
+        if (blockNumber <= migrationBlock) {
+            return 0;
+        }
+
+        Checkpoint storage checkpoint = _checkpointsLookupStorage(_checkpoints[account], blockNumber);
+        if (checkpoint.fromBlock <= migrationBlock) {
+            return 0;
+        }
+        return _toMigration(checkpoint).migratedVotes;
+    }
+
+    function _checkpointsLookupStorage(
+        Checkpoint[] storage ckpts,
+        uint256 blockNumber
+    ) internal view returns (Checkpoint storage result) {
+        // We run a binary search to look for the earliest checkpoint taken after `blockNumber`.
+        //
+        // Initially we check if the block is recent to narrow the search range.
+        // During the loop, the index of the wanted checkpoint remains in the range [low-1, high).
+        // With each iteration, either `low` or `high` is moved towards the middle of the range to maintain the
+        // invariant.
+        // - If the middle checkpoint is after `blockNumber`, we look in [low, mid)
+        // - If the middle checkpoint is before or equal to `blockNumber`, we look in [mid+1, high)
+        // Once we reach a single value (when low == high), we've found the right checkpoint at the index high-1, if not
+        // out of bounds (in which case we're looking too far in the past and the result is 0).
+        // Note that if the latest checkpoint available is exactly for `blockNumber`, we end up with an index that is
+        // past the end of the array, so we technically don't find a checkpoint after `blockNumber`, but it works out
+        // the same.
+        uint256 length = ckpts.length;
+
+        uint256 low = 0;
+        uint256 high = length;
+
+        if (length > 5) {
+            uint256 mid = length - Math.sqrt(length);
+            if (_unsafeAccess(ckpts, mid).fromBlock > blockNumber) {
+                high = mid;
+            } else {
+                low = mid + 1;
+            }
+        }
+
+        while (low < high) {
+            uint256 mid = Math.average(low, high);
+            if (_unsafeAccess(ckpts, mid).fromBlock > blockNumber) {
+                high = mid;
+            } else {
+                low = mid + 1;
+            }
+        }
+
+        // Leaving here for posterity this is the original OZ implementation which we've replaced
+        // return high == 0 ? 0 : _unsafeAccess(ckpts, high - 1).votes;
+        // Checkpoint memory checkpoint = high == 0 ? Checkpoint(0, 0, 0) : _unsafeAccess(ckpts, high - 1);
+        // return checkpoint;
+        // TODO: bad. very bad. only works on accident
+        if (high > 0) {
+            result = _unsafeAccess(ckpts, high - 1);
+        } else {
+            // suppress compiler warning, which really shouldn't be suppressed
+            assembly {
+                result.slot := 0xDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEF
+            }
+        }
+    }
+
+    // TODO: we're not handling totalSupply
+
+    // TODO: need to return the migrated weight
+    function _writeCheckpoint(
+        Checkpoint[] storage ckpts,
+        function(uint256, uint256) view returns (uint256) op,
+        uint256 userBalance,
+        uint96 balanceLastUpdated,
+        uint256 delta
+    )
+        internal
+        virtual
+        override
+        returns (uint256 oldWeight, uint256 newWeight, uint256 oldQuadraticWeight, uint256 newQuadraticWeight)
+    {
+        uint256 pos = ckpts.length;
+
+        CheckpointMigration memory oldCkpt = pos == 0
+            ? CheckpointMigration(0, 0, 0, 0)
+            : _toMigration(_unsafeAccess(ckpts, pos - 1));
+
+        oldWeight = oldCkpt.votes;
+        newWeight = op(oldWeight, delta);
+
+        oldQuadraticWeight = oldCkpt.quadraticVotes;
+
+        if (pos > 0) {
+            deductOldWeightFromCheckpoint(oldCkpt, userBalance, balanceLastUpdated);
+        }
+
+        // if wallet > threshold, calculate quadratic power over the treshold only, below threshold is linear
+        uint256 newBalance = op(userBalance, delta);
+        uint256 newQuadraticBalance = newBalance <= quadraticThreshold
+            ? newBalance
+            : quadraticThreshold + Math.sqrt((newBalance - quadraticThreshold) * 1e18);
+        newQuadraticWeight = oldCkpt.quadraticVotes + newQuadraticBalance;
+        uint256 newMigratedWeight = oldCkpt.migratedVotes + CubeRoot.cbrt(newBalance);
+
+        if (pos > 0 && oldCkpt.fromBlock == block.number) {
+            addCheckpoint(ckpts, pos, newWeight, newQuadraticWeight, newMigratedWeight);
+        } else {
+            _toMigration(ckpts).push(
+                CheckpointMigration({
+                    fromBlock: SafeCast.toUint32(block.number),
+                    votes: SafeCast.toUint96(newWeight),
+                    quadraticVotes: SafeCast.toUint96(newQuadraticWeight),
+                    migratedVotes: SafeCast.toUint32(newMigratedWeight)
+                })
+            );
+        }
+    }
+
+    function deductOldWeightFromCheckpoint(
+        CheckpointMigration memory oldCkpt,
+        uint256 userBalance,
+        uint96 balanceLastUpdated
+    ) internal {
+        // Remove the entire sqrt userBalance from quadratic voting power.
+        // Note that `userBalance` is value _after_ transfer.
+        uint256 oldQuadraticVotingPower = userBalance <= quadraticThreshold
+            ? userBalance
+            : quadraticThreshold + Math.sqrt((userBalance - quadraticThreshold) * 1e18);
+        oldCkpt.quadraticVotes -= SafeCast.toUint96(oldQuadraticVotingPower);
+
+        if (oldCkpt.fromBlock > migrationBlock && balanceLastUpdated > migrationBlock) {
+            oldCkpt.migratedVotes -= SafeCast.toUint32(CubeRoot.cbrt(userBalance));
+        }
+    }
+
+    function addCheckpoint(
+        Checkpoint[] storage ckpts,
+        uint256 pos,
+        uint256 newWeight,
+        uint256 newQuadraticWeight,
+        uint256 newMigratedWeight
+    ) internal {
+        CheckpointMigration storage chpt = _toMigration(_unsafeAccess(ckpts, pos - 1));
+        chpt.votes = SafeCast.toUint96(newWeight);
+        chpt.quadraticVotes = SafeCast.toUint96(newQuadraticWeight);
+        chpt.migratedVotes = SafeCast.toUint32(newMigratedWeight);
+    }
+}