Extend Checkpoints with new sizes and lookup mechanisms

contracts/governance/extensions/GovernorVotesQuorumFraction.sol

@@ -50,8 +50,8 @@ abstract contract GovernorVotesQuorumFraction is GovernorVotes {
         }
 
         // Optimistic search, check the latest checkpoint
-        Checkpoints.Checkpoint memory latest = _quorumNumeratorHistory._checkpoints[length - 1];
-        if (latest._blockNumber <= blockNumber) {
+        Checkpoints.Checkpoint224 memory latest = _quorumNumeratorHistory._checkpoints[length - 1];
+        if (latest._key <= blockNumber) {
             return latest._value;
         }
 
@@ -107,7 +107,7 @@ abstract contract GovernorVotesQuorumFraction is GovernorVotes {
         // Make sure we keep track of the original numerator in contracts upgraded from a version without checkpoints.
         if (oldQuorumNumerator != 0 && _quorumNumeratorHistory._checkpoints.length == 0) {
             _quorumNumeratorHistory._checkpoints.push(
-                Checkpoints.Checkpoint({_blockNumber: 0, _value: SafeCast.toUint224(oldQuorumNumerator)})
+                Checkpoints.Checkpoint224({_key: 0, _value: SafeCast.toUint224(oldQuorumNumerator)})
             );
         }
 

contracts/mocks/CheckpointsMock.sol

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../utils/Checkpoints.sol";
+
+contract CheckpointsMock {
+    using Checkpoints for Checkpoints.History;
+
+    Checkpoints.History private _totalCheckpoints;
+
+    function latest() public view returns (uint256) {
+        return _totalCheckpoints.latest();
+    }
+
+    function push(uint256 value) public returns (uint256, uint256) {
+        return _totalCheckpoints.push(value);
+    }
+
+    function getAtBlock(uint256 blockNumber) public view returns (uint256) {
+        return _totalCheckpoints.getAtBlock(blockNumber);
+    }
+
+    function length() public view returns (uint256) {
+        return _totalCheckpoints._checkpoints.length;
+    }
+}
+
+contract Checkpoints224Mock {
+    using Checkpoints for Checkpoints.Checkpoint224[];
+
+    Checkpoints.Checkpoint224[] private _totalCheckpoints;
+
+    function latest() public view returns (uint224) {
+        return _totalCheckpoints.latest();
+    }
+
+    function push(uint32 key, uint224 value) public returns (uint224, uint224) {
+        return _totalCheckpoints.push(key, value);
+    }
+
+    function lowerLookup(uint32 key) public view returns (uint224) {
+        return _totalCheckpoints.lowerLookup(key);
+    }
+
+    function upperLookup(uint32 key) public view returns (uint224) {
+        return _totalCheckpoints.upperLookup(key);
+    }
+
+    function upperLookupExpEnd(uint32 key) public view returns (uint224) {
+        return _totalCheckpoints.upperLookupExpEnd(key);
+    }
+
+    function length() public view returns (uint256) {
+        return _totalCheckpoints.length;
+    }
+}
+
+contract Checkpoints160Mock {
+    using Checkpoints for Checkpoints.Checkpoint160[];
+
+    Checkpoints.Checkpoint160[] private _totalCheckpoints;
+
+    function latest() public view returns (uint160) {
+        return _totalCheckpoints.latest();
+    }
+
+    function push(uint96 key, uint160 value) public returns (uint160, uint160) {
+        return _totalCheckpoints.push(key, value);
+    }
+
+    function lowerLookup(uint96 key) public view returns (uint160) {
+        return _totalCheckpoints.lowerLookup(key);
+    }
+
+    function upperLookup(uint96 key) public view returns (uint160) {
+        return _totalCheckpoints.upperLookup(key);
+    }
+
+    function upperLookupExpEnd(uint96 key) public view returns (uint224) {
+        return _totalCheckpoints.upperLookupExpEnd(key);
+    }
+
+    function length() public view returns (uint256) {
+        return _totalCheckpoints.length;
+    }
+}

contracts/utils/Checkpoints.sol

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: MIT
 // OpenZeppelin Contracts (last updated v4.5.0) (utils/Checkpoints.sol)
+
 pragma solidity ^0.8.0;
 
 import "./math/Math.sol";
@@ -15,21 +16,187 @@ import "./math/SafeCast.sol";
  * _Available since v4.5._
  */
 library Checkpoints {
-    struct Checkpoint {
-        uint32 _blockNumber;
+    struct Checkpoint224 {
+        uint32 _key;
         uint224 _value;
     }
 
+    function latest(Checkpoint224[] storage self) internal view returns (uint224) {
+        uint256 pos = self.length;
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function push(
+        Checkpoint224[] storage self,
+        uint32 key,
+        uint224 value
+    ) internal returns (uint224, uint224) {
+        uint256 pos = self.length;
+        uint224 old = latest(self);
+        if (pos > 0 && self[pos - 1]._key == key) {
+            self[pos - 1]._value = value;
+        } else {
+            self.push(Checkpoint224({_key: key, _value: value}));
+        }
+        return (old, value);
+    }
+
+    function lowerLookup(Checkpoint224[] storage self, uint32 key) internal view returns (uint224) {
+        uint256 length = self.length;
+        uint256 pos = _lowerDichotomicLookup(self, key, 0, length);
+        return pos == length ? 0 : self[pos]._value;
+    }
+
+    function upperLookup(Checkpoint224[] storage self, uint32 key) internal view returns (uint224) {
+        uint256 length = self.length;
+        uint256 pos = _upperDichotomicLookup(self, key, 0, length);
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function upperLookupExpEnd(Checkpoint224[] storage self, uint32 key) internal view returns (uint224) {
+        uint256 length = self.length;
+        uint256 offset = 1;
+
+        while (offset <= length && self[length - offset]._key > key) {
+            offset <<= 1;
+        }
+
+        uint256 low = offset < length ? length - offset : 0;
+        uint256 high = length - (offset >> 1);
+        uint256 pos = _upperDichotomicLookup(self, key, low, high);
+
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function _upperDichotomicLookup(
+        Checkpoint224[] storage self,
+        uint32 key,
+        uint256 low,
+        uint256 high
+    ) private view returns (uint256) {
+        while (low < high) {
+            uint256 mid = Math.average(low, high);
+            if (self[mid]._key > key) {
+                high = mid;
+            } else {
+                low = mid + 1;
+            }
+        }
+        return high;
+    }
+
+    function _lowerDichotomicLookup(
+        Checkpoint224[] storage self,
+        uint32 key,
+        uint256 low,
+        uint256 high
+    ) private view returns (uint256) {
+        while (low < high) {
+            uint256 mid = Math.average(low, high);
+            if (self[mid]._key < key) {
+                low = mid + 1;
+            } else {
+                high = mid;
+            }
+        }
+        return high;
+    }
+
+    struct Checkpoint160 {
+        uint96 _key;
+        uint160 _value;
+    }
+
+    function latest(Checkpoint160[] storage self) internal view returns (uint160) {
+        uint256 pos = self.length;
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function push(
+        Checkpoint160[] storage self,
+        uint96 key,
+        uint160 value
+    ) internal returns (uint160, uint160) {
+        uint256 pos = self.length;
+        uint160 old = latest(self);
+        if (pos > 0 && self[pos - 1]._key == key) {
+            self[pos - 1]._value = value;
+        } else {
+            self.push(Checkpoint160({_key: key, _value: value}));
+        }
+        return (old, value);
+    }
+
+    function lowerLookup(Checkpoint160[] storage self, uint96 key) internal view returns (uint160) {
+        uint256 length = self.length;
+        uint256 pos = _lowerDichotomicLookup(self, key, 0, length);
+        return pos == length ? 0 : self[pos]._value;
+    }
+
+    function upperLookup(Checkpoint160[] storage self, uint96 key) internal view returns (uint160) {
+        uint256 length = self.length;
+        uint256 pos = _upperDichotomicLookup(self, key, 0, length);
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function upperLookupExpEnd(Checkpoint160[] storage self, uint96 key) internal view returns (uint224) {
+        uint256 length = self.length;
+        uint256 offset = 1;
+
+        while (offset <= length && self[length - offset]._key > key) {
+            offset <<= 1;
+        }
+
+        uint256 low = offset < length ? length - offset : 0;
+        uint256 high = length - (offset >> 1);
+        uint256 pos = _upperDichotomicLookup(self, key, low, high);
+
+        return pos == 0 ? 0 : self[pos - 1]._value;
+    }
+
+    function _upperDichotomicLookup(
+        Checkpoint160[] storage self,
+        uint96 key,
+        uint256 low,
+        uint256 high
+    ) private view returns (uint256) {
+        while (low < high) {
+            uint256 mid = Math.average(low, high);
+            if (self[mid]._key > key) {
+                high = mid;
+            } else {
+                low = mid + 1;
+            }
+        }
+        return high;
+    }
+
+    function _lowerDichotomicLookup(
+        Checkpoint160[] storage self,
+        uint96 key,
+        uint256 low,
+        uint256 high
+    ) private view returns (uint256) {
+        while (low < high) {
+            uint256 mid = Math.average(low, high);
+            if (self[mid]._key < key) {
+                low = mid + 1;
+            } else {
+                high = mid;
+            }
+        }
+        return high;
+    }
+
     struct History {
-        Checkpoint[] _checkpoints;
+        Checkpoint224[] _checkpoints;
     }
 
     /**
      * @dev Returns the value in the latest checkpoint, or zero if there are no checkpoints.
      */
     function latest(History storage self) internal view returns (uint256) {
-        uint256 pos = self._checkpoints.length;
-        return pos == 0 ? 0 : self._checkpoints[pos - 1]._value;
+        return latest(self._checkpoints);
     }
 
     /**
@@ -39,17 +206,7 @@ library Checkpoints {
     function getAtBlock(History storage self, uint256 blockNumber) internal view returns (uint256) {
         require(blockNumber < block.number, "Checkpoints: block not yet mined");
 
-        uint256 high = self._checkpoints.length;
-        uint256 low = 0;
-        while (low < high) {
-            uint256 mid = Math.average(low, high);
-            if (self._checkpoints[mid]._blockNumber > blockNumber) {
-                high = mid;
-            } else {
-                low = mid + 1;
-            }
-        }
-        return high == 0 ? 0 : self._checkpoints[high - 1]._value;
+        return upperLookupExpEnd(self._checkpoints, SafeCast.toUint32(blockNumber));
     }
 
     /**
@@ -58,16 +215,7 @@ library Checkpoints {
      * Returns previous value and new value.
      */
     function push(History storage self, uint256 value) internal returns (uint256, uint256) {
-        uint256 pos = self._checkpoints.length;
-        uint256 old = latest(self);
-        if (pos > 0 && self._checkpoints[pos - 1]._blockNumber == block.number) {
-            self._checkpoints[pos - 1]._value = SafeCast.toUint224(value);
-        } else {
-            self._checkpoints.push(
-                Checkpoint({_blockNumber: SafeCast.toUint32(block.number), _value: SafeCast.toUint224(value)})
-            );
-        }
-        return (old, value);
+        return push(self._checkpoints, SafeCast.toUint32(block.number), SafeCast.toUint224(value));
     }
 
     /**

scripts/generate/run.js

@@ -14,6 +14,8 @@ function getVersion (path) {
 }
 
 for (const [ file, template ] of Object.entries({
+  'utils/Checkpoints.sol': './templates/Checkpoints',
+  'mocks/CheckpointsMock.sol': './templates/CheckpointsMock',
   'utils/math/SafeCast.sol': './templates/SafeCast',
   'mocks/SafeCastMock.sol': './templates/SafeCastMock',
 })) {