Skip to content

Commit b84db20

Browse files
arr00Amxx
arr00
and
Amxx
authored
Add checkpoint variant with uint256 keys and values (#5748)
Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>
1 parent 6079eb3 commit b84db20

File tree

7 files changed

+337
-15
lines changed

7 files changed

+337
-15
lines changed

‎.changeset/eight-radios-check.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'openzeppelin-solidity': minor
3+
---
4+
5+
`Checkpoints`: Add a new checkpoint variant `Checkpoint256` using `uint256` type for the value and key.

‎contracts/utils/structs/Checkpoints.sol

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,209 @@ library Checkpoints {
1919
*/
2020
error CheckpointUnorderedInsertion();
2121

22+
struct Trace256 {
23+
Checkpoint256[] _checkpoints;
24+
}
25+
26+
struct Checkpoint256 {
27+
uint256 _key;
28+
uint256 _value;
29+
}
30+
31+
/**
32+
* @dev Pushes a (`key`, `value`) pair into a Trace256 so that it is stored as the checkpoint.
33+
*
34+
* Returns previous value and new value.
35+
*
36+
* IMPORTANT: Never accept `key` as a user input, since an arbitrary `type(uint256).max` key set will disable the
37+
* library.
38+
*/
39+
function push(
40+
Trace256 storage self,
41+
uint256 key,
42+
uint256 value
43+
) internal returns (uint256 oldValue, uint256 newValue) {
44+
return _insert(self._checkpoints, key, value);
45+
}
46+
47+
/**
48+
* @dev Returns the value in the first (oldest) checkpoint with key greater or equal than the search key, or zero if
49+
* there is none.
50+
*/
51+
function lowerLookup(Trace256 storage self, uint256 key) internal view returns (uint256) {
52+
uint256 len = self._checkpoints.length;
53+
uint256 pos = _lowerBinaryLookup(self._checkpoints, key, 0, len);
54+
return pos == len ? 0 : _unsafeAccess(self._checkpoints, pos)._value;
55+
}
56+
57+
/**
58+
* @dev Returns the value in the last (most recent) checkpoint with key lower or equal than the search key, or zero
59+
* if there is none.
60+
*/
61+
function upperLookup(Trace256 storage self, uint256 key) internal view returns (uint256) {
62+
uint256 len = self._checkpoints.length;
63+
uint256 pos = _upperBinaryLookup(self._checkpoints, key, 0, len);
64+
return pos == 0 ? 0 : _unsafeAccess(self._checkpoints, pos - 1)._value;
65+
}
66+
67+
/**
68+
* @dev Returns the value in the last (most recent) checkpoint with key lower or equal than the search key, or zero
69+
* if there is none.
70+
*
71+
* NOTE: This is a variant of {upperLookup} that is optimized to find "recent" checkpoint (checkpoints with high
72+
* keys).
73+
*/
74+
function upperLookupRecent(Trace256 storage self, uint256 key) internal view returns (uint256) {
75+
uint256 len = self._checkpoints.length;
76+
77+
uint256 low = 0;
78+
uint256 high = len;
79+
80+
if (len > 5) {
81+
uint256 mid = len - Math.sqrt(len);
82+
if (key < _unsafeAccess(self._checkpoints, mid)._key) {
83+
high = mid;
84+
} else {
85+
low = mid + 1;
86+
}
87+
}
88+
89+
uint256 pos = _upperBinaryLookup(self._checkpoints, key, low, high);
90+
91+
return pos == 0 ? 0 : _unsafeAccess(self._checkpoints, pos - 1)._value;
92+
}
93+
94+
/**
95+
* @dev Returns the value in the most recent checkpoint, or zero if there are no checkpoints.
96+
*/
97+
function latest(Trace256 storage self) internal view returns (uint256) {
98+
uint256 pos = self._checkpoints.length;
99+
return pos == 0 ? 0 : _unsafeAccess(self._checkpoints, pos - 1)._value;
100+
}
101+
102+
/**
103+
* @dev Returns whether there is a checkpoint in the structure (i.e. it is not empty), and if so the key and value
104+
* in the most recent checkpoint.
105+
*/
106+
function latestCheckpoint(Trace256 storage self) internal view returns (bool exists, uint256 _key, uint256 _value) {
107+
uint256 pos = self._checkpoints.length;
108+
if (pos == 0) {
109+
return (false, 0, 0);
110+
} else {
111+
Checkpoint256 storage ckpt = _unsafeAccess(self._checkpoints, pos - 1);
112+
return (true, ckpt._key, ckpt._value);
113+
}
114+
}
115+
116+
/**
117+
* @dev Returns the number of checkpoints.
118+
*/
119+
function length(Trace256 storage self) internal view returns (uint256) {
120+
return self._checkpoints.length;
121+
}
122+
123+
/**
124+
* @dev Returns checkpoint at given position.
125+
*/
126+
function at(Trace256 storage self, uint32 pos) internal view returns (Checkpoint256 memory) {
127+
return self._checkpoints[pos];
128+
}
129+
130+
/**
131+
* @dev Pushes a (`key`, `value`) pair into an ordered list of checkpoints, either by inserting a new checkpoint,
132+
* or by updating the last one.
133+
*/
134+
function _insert(
135+
Checkpoint256[] storage self,
136+
uint256 key,
137+
uint256 value
138+
) private returns (uint256 oldValue, uint256 newValue) {
139+
uint256 pos = self.length;
140+
141+
if (pos > 0) {
142+
Checkpoint256 storage last = _unsafeAccess(self, pos - 1);
143+
uint256 lastKey = last._key;
144+
uint256 lastValue = last._value;
145+
146+
// Checkpoint keys must be non-decreasing.
147+
if (lastKey > key) {
148+
revert CheckpointUnorderedInsertion();
149+
}
150+
151+
// Update or push new checkpoint
152+
if (lastKey == key) {
153+
last._value = value;
154+
} else {
155+
self.push(Checkpoint256({_key: key, _value: value}));
156+
}
157+
return (lastValue, value);
158+
} else {
159+
self.push(Checkpoint256({_key: key, _value: value}));
160+
return (0, value);
161+
}
162+
}
163+
164+
/**
165+
* @dev Return the index of the first (oldest) checkpoint with key strictly bigger than the search key, or `high`
166+
* if there is none. `low` and `high` define a section where to do the search, with inclusive `low` and exclusive
167+
* `high`.
168+
*
169+
* WARNING: `high` should not be greater than the array's length.
170+
*/
171+
function _upperBinaryLookup(
172+
Checkpoint256[] storage self,
173+
uint256 key,
174+
uint256 low,
175+
uint256 high
176+
) private view returns (uint256) {
177+
while (low < high) {
178+
uint256 mid = Math.average(low, high);
179+
if (_unsafeAccess(self, mid)._key > key) {
180+
high = mid;
181+
} else {
182+
low = mid + 1;
183+
}
184+
}
185+
return high;
186+
}
187+
188+
/**
189+
* @dev Return the index of the first (oldest) checkpoint with key greater or equal than the search key, or `high`
190+
* if there is none. `low` and `high` define a section where to do the search, with inclusive `low` and exclusive
191+
* `high`.
192+
*
193+
* WARNING: `high` should not be greater than the array's length.
194+
*/
195+
function _lowerBinaryLookup(
196+
Checkpoint256[] storage self,
197+
uint256 key,
198+
uint256 low,
199+
uint256 high
200+
) private view returns (uint256) {
201+
while (low < high) {
202+
uint256 mid = Math.average(low, high);
203+
if (_unsafeAccess(self, mid)._key < key) {
204+
low = mid + 1;
205+
} else {
206+
high = mid;
207+
}
208+
}
209+
return high;
210+
}
211+
212+
/**
213+
* @dev Access an element of the array without performing bounds check. The position is assumed to be within bounds.
214+
*/
215+
function _unsafeAccess(
216+
Checkpoint256[] storage self,
217+
uint256 pos
218+
) private pure returns (Checkpoint256 storage result) {
219+
assembly {
220+
mstore(0, self.slot)
221+
result.slot := add(keccak256(0, 0x20), mul(pos, 2))
222+
}
223+
}
224+
22225
struct Trace224 {
23226
Checkpoint224[] _checkpoints;
24227
}

‎scripts/generate/templates/Checkpoints.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ function _unsafeAccess(
223223
) private pure returns (${opts.checkpointTypeName} storage result) {
224224
assembly {
225225
mstore(0, self.slot)
226-
result.slot := add(keccak256(0, 0x20), pos)
226+
result.slot := add(keccak256(0, 0x20), ${opts.checkpointSize === 1 ? 'pos' : `mul(pos, ${opts.checkpointSize})`})
227227
}
228228
}
229229
`;

‎scripts/generate/templates/Checkpoints.opts.js

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
// OPTIONS
2-
const VALUE_SIZES = [224, 208, 160];
2+
const VALUE_SIZES = [256, 224, 208, 160];
33

44
const defaultOpts = size => ({
55
historyTypeName: `Trace${size}`,
66
checkpointTypeName: `Checkpoint${size}`,
77
checkpointFieldName: '_checkpoints',
8-
keyTypeName: `uint${256 - size}`,
8+
checkpointSize: size < 256 ? 1 : 2,
9+
keyTypeName: size < 256 ? `uint${256 - size}` : 'uint256',
910
keyFieldName: '_key',
1011
valueTypeName: `uint${size}`,
1112
valueFieldName: '_value',

‎scripts/generate/templates/Checkpoints.t.js

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,11 @@ Checkpoints.${opts.historyTypeName} internal _ckpts;
2424
function _bound${capitalize(opts.keyTypeName)}(${opts.keyTypeName} x, ${opts.keyTypeName} min, ${
2525
opts.keyTypeName
2626
} max) internal pure returns (${opts.keyTypeName}) {
27-
return SafeCast.to${capitalize(opts.keyTypeName)}(bound(uint256(x), uint256(min), uint256(max)));
27+
return ${
28+
opts.keyTypeName === 'uint256'
29+
? 'bound(x, min, max)'
30+
: `SafeCast.to${capitalize(opts.keyTypeName)}(bound(uint256(x), uint256(min), uint256(max)))`
31+
};
2832
}
2933
3034
function _prepareKeys(${opts.keyTypeName}[] memory keys, ${opts.keyTypeName} maxSpread) internal pure {

‎test/utils/structs/Checkpoints.t.sol

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,114 @@ import {Test} from "forge-std/Test.sol";
77
import {SafeCast} from "@openzeppelin/contracts/utils/math/SafeCast.sol";
88
import {Checkpoints} from "@openzeppelin/contracts/utils/structs/Checkpoints.sol";
99

10+
contract CheckpointsTrace256Test is Test {
11+
using Checkpoints for Checkpoints.Trace256;
12+
13+
// Maximum gap between keys used during the fuzzing tests: the `_prepareKeys` function will make sure that
14+
// key#n+1 is in the [key#n, key#n + _KEY_MAX_GAP] range.
15+
uint8 internal constant _KEY_MAX_GAP = 64;
16+
17+
Checkpoints.Trace256 internal _ckpts;
18+
19+
// helpers
20+
function _boundUint256(uint256 x, uint256 min, uint256 max) internal pure returns (uint256) {
21+
return bound(x, min, max);
22+
}
23+
24+
function _prepareKeys(uint256[] memory keys, uint256 maxSpread) internal pure {
25+
uint256 lastKey = 0;
26+
for (uint256 i = 0; i < keys.length; ++i) {
27+
uint256 key = _boundUint256(keys[i], lastKey, lastKey + maxSpread);
28+
keys[i] = key;
29+
lastKey = key;
30+
}
31+
}
32+
33+
function _assertLatestCheckpoint(bool exist, uint256 key, uint256 value) internal view {
34+
(bool _exist, uint256 _key, uint256 _value) = _ckpts.latestCheckpoint();
35+
assertEq(_exist, exist);
36+
assertEq(_key, key);
37+
assertEq(_value, value);
38+
}
39+
40+
// tests
41+
function testPush(uint256[] memory keys, uint256[] memory values, uint256 pastKey) public {
42+
vm.assume(values.length > 0 && values.length <= keys.length);
43+
_prepareKeys(keys, _KEY_MAX_GAP);
44+
45+
// initial state
46+
assertEq(_ckpts.length(), 0);
47+
assertEq(_ckpts.latest(), 0);
48+
_assertLatestCheckpoint(false, 0, 0);
49+
50+
uint256 duplicates = 0;
51+
for (uint256 i = 0; i < keys.length; ++i) {
52+
uint256 key = keys[i];
53+
uint256 value = values[i % values.length];
54+
if (i > 0 && key == keys[i - 1]) ++duplicates;
55+
56+
// push
57+
_ckpts.push(key, value);
58+
59+
// check length & latest
60+
assertEq(_ckpts.length(), i + 1 - duplicates);
61+
assertEq(_ckpts.latest(), value);
62+
_assertLatestCheckpoint(true, key, value);
63+
}
64+
65+
if (keys.length > 0) {
66+
uint256 lastKey = keys[keys.length - 1];
67+
if (lastKey > 0) {
68+
pastKey = _boundUint256(pastKey, 0, lastKey - 1);
69+
70+
vm.expectRevert();
71+
this.push(pastKey, values[keys.length % values.length]);
72+
}
73+
}
74+
}
75+
76+
// used to test reverts
77+
function push(uint256 key, uint256 value) external {
78+
_ckpts.push(key, value);
79+
}
80+
81+
function testLookup(uint256[] memory keys, uint256[] memory values, uint256 lookup) public {
82+
vm.assume(values.length > 0 && values.length <= keys.length);
83+
_prepareKeys(keys, _KEY_MAX_GAP);
84+
85+
uint256 lastKey = keys.length == 0 ? 0 : keys[keys.length - 1];
86+
lookup = _boundUint256(lookup, 0, lastKey + _KEY_MAX_GAP);
87+
88+
uint256 upper = 0;
89+
uint256 lower = 0;
90+
uint256 lowerKey = type(uint256).max;
91+
for (uint256 i = 0; i < keys.length; ++i) {
92+
uint256 key = keys[i];
93+
uint256 value = values[i % values.length];
94+
95+
// push
96+
_ckpts.push(key, value);
97+
98+
// track expected result of lookups
99+
if (key <= lookup) {
100+
upper = value;
101+
}
102+
// find the first key that is not smaller than the lookup key
103+
if (key >= lookup && (i == 0 || keys[i - 1] < lookup)) {
104+
lowerKey = key;
105+
}
106+
if (key == lowerKey) {
107+
lower = value;
108+
}
109+
}
110+
111+
// check lookup
112+
assertEq(_ckpts.lowerLookup(lookup), lower);
113+
assertEq(_ckpts.upperLookup(lookup), upper);
114+
assertEq(_ckpts.upperLookupRecent(lookup), upper);
115+
}
116+
}
117+
10118
contract CheckpointsTrace224Test is Test {
11119
using Checkpoints for Checkpoints.Trace224;
12120

0 commit comments

Comments
 (0)