Add checkpoint variant with uint256 keys and values (#5748)

Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

test/utils/structs/Checkpoints.t.sol

@@ -7,6 +7,114 @@ import {Test} from "forge-std/Test.sol";
 import {SafeCast} from "@openzeppelin/contracts/utils/math/SafeCast.sol";
 import {Checkpoints} from "@openzeppelin/contracts/utils/structs/Checkpoints.sol";
 
+contract CheckpointsTrace256Test is Test {
+    using Checkpoints for Checkpoints.Trace256;
+
+    // Maximum gap between keys used during the fuzzing tests: the `_prepareKeys` function will make sure that
+    // key#n+1 is in the [key#n, key#n + _KEY_MAX_GAP] range.
+    uint8 internal constant _KEY_MAX_GAP = 64;
+
+    Checkpoints.Trace256 internal _ckpts;
+
+    // helpers
+    function _boundUint256(uint256 x, uint256 min, uint256 max) internal pure returns (uint256) {
+        return bound(x, min, max);
+    }
+
+    function _prepareKeys(uint256[] memory keys, uint256 maxSpread) internal pure {
+        uint256 lastKey = 0;
+        for (uint256 i = 0; i < keys.length; ++i) {
+            uint256 key = _boundUint256(keys[i], lastKey, lastKey + maxSpread);
+            keys[i] = key;
+            lastKey = key;
+        }
+    }
+
+    function _assertLatestCheckpoint(bool exist, uint256 key, uint256 value) internal view {
+        (bool _exist, uint256 _key, uint256 _value) = _ckpts.latestCheckpoint();
+        assertEq(_exist, exist);
+        assertEq(_key, key);
+        assertEq(_value, value);
+    }
+
+    // tests
+    function testPush(uint256[] memory keys, uint256[] memory values, uint256 pastKey) public {
+        vm.assume(values.length > 0 && values.length <= keys.length);
+        _prepareKeys(keys, _KEY_MAX_GAP);
+
+        // initial state
+        assertEq(_ckpts.length(), 0);
+        assertEq(_ckpts.latest(), 0);
+        _assertLatestCheckpoint(false, 0, 0);
+
+        uint256 duplicates = 0;
+        for (uint256 i = 0; i < keys.length; ++i) {
+            uint256 key = keys[i];
+            uint256 value = values[i % values.length];
+            if (i > 0 && key == keys[i - 1]) ++duplicates;
+
+            // push
+            _ckpts.push(key, value);
+
+            // check length & latest
+            assertEq(_ckpts.length(), i + 1 - duplicates);
+            assertEq(_ckpts.latest(), value);
+            _assertLatestCheckpoint(true, key, value);
+        }
+
+        if (keys.length > 0) {
+            uint256 lastKey = keys[keys.length - 1];
+            if (lastKey > 0) {
+                pastKey = _boundUint256(pastKey, 0, lastKey - 1);
+
+                vm.expectRevert();
+                this.push(pastKey, values[keys.length % values.length]);
+            }
+        }
+    }
+
+    // used to test reverts
+    function push(uint256 key, uint256 value) external {
+        _ckpts.push(key, value);
+    }
+
+    function testLookup(uint256[] memory keys, uint256[] memory values, uint256 lookup) public {
+        vm.assume(values.length > 0 && values.length <= keys.length);
+        _prepareKeys(keys, _KEY_MAX_GAP);
+
+        uint256 lastKey = keys.length == 0 ? 0 : keys[keys.length - 1];
+        lookup = _boundUint256(lookup, 0, lastKey + _KEY_MAX_GAP);
+
+        uint256 upper = 0;
+        uint256 lower = 0;
+        uint256 lowerKey = type(uint256).max;
+        for (uint256 i = 0; i < keys.length; ++i) {
+            uint256 key = keys[i];
+            uint256 value = values[i % values.length];
+
+            // push
+            _ckpts.push(key, value);
+
+            // track expected result of lookups
+            if (key <= lookup) {
+                upper = value;
+            }
+            // find the first key that is not smaller than the lookup key
+            if (key >= lookup && (i == 0 || keys[i - 1] < lookup)) {
+                lowerKey = key;
+            }
+            if (key == lowerKey) {
+                lower = value;
+            }
+        }
+
+        // check lookup
+        assertEq(_ckpts.lowerLookup(lookup), lower);
+        assertEq(_ckpts.upperLookup(lookup), upper);
+        assertEq(_ckpts.upperLookupRecent(lookup), upper);
+    }
+}
+
 contract CheckpointsTrace224Test is Test {
     using Checkpoints for Checkpoints.Trace224;
 

test/utils/structs/Checkpoints.test.js

@@ -2,23 +2,24 @@ const { ethers } = require('hardhat');
 const { expect } = require('chai');
 const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
 
-const { VALUE_SIZES } = require('../../../scripts/generate/templates/Checkpoints.opts');
+const { OPTS } = require('../../../scripts/generate/templates/Checkpoints.opts');
 
 describe('Checkpoints', function () {
-  for (const length of VALUE_SIZES) {
-    describe(`Trace${length}`, function () {
+  for (const opt of OPTS) {
+    describe(opt.historyTypeName, function () {
       const fixture = async () => {
         const mock = await ethers.deployContract('$Checkpoints');
         const methods = {
-          at: (...args) => mock.getFunction(`$at_Checkpoints_Trace${length}`)(0, ...args),
-          latest: (...args) => mock.getFunction(`$latest_Checkpoints_Trace${length}`)(0, ...args),
-          latestCheckpoint: (...args) => mock.getFunction(`$latestCheckpoint_Checkpoints_Trace${length}`)(0, ...args),
-          length: (...args) => mock.getFunction(`$length_Checkpoints_Trace${length}`)(0, ...args),
-          push: (...args) => mock.getFunction(`$push(uint256,uint${256 - length},uint${length})`)(0, ...args),
-          lowerLookup: (...args) => mock.getFunction(`$lowerLookup(uint256,uint${256 - length})`)(0, ...args),
-          upperLookup: (...args) => mock.getFunction(`$upperLookup(uint256,uint${256 - length})`)(0, ...args),
+          at: (...args) => mock.getFunction(`$at_Checkpoints_${opt.historyTypeName}`)(0, ...args),
+          latest: (...args) => mock.getFunction(`$latest_Checkpoints_${opt.historyTypeName}`)(0, ...args),
+          latestCheckpoint: (...args) =>
+            mock.getFunction(`$latestCheckpoint_Checkpoints_${opt.historyTypeName}`)(0, ...args),
+          length: (...args) => mock.getFunction(`$length_Checkpoints_${opt.historyTypeName}`)(0, ...args),
+          push: (...args) => mock.getFunction(`$push(uint256,${opt.keyTypeName},${opt.valueTypeName})`)(0, ...args),
+          lowerLookup: (...args) => mock.getFunction(`$lowerLookup(uint256,${opt.keyTypeName})`)(0, ...args),
+          upperLookup: (...args) => mock.getFunction(`$upperLookup(uint256,${opt.keyTypeName})`)(0, ...args),
           upperLookupRecent: (...args) =>
-            mock.getFunction(`$upperLookupRecent(uint256,uint${256 - length})`)(0, ...args),
+            mock.getFunction(`$upperLookupRecent(uint256,${opt.keyTypeName})`)(0, ...args),
         };
 
         return { mock, methods };