Skip to content

Commit 4c481d6

Browse files
ernestognwAmxxcairoeth
authored
Implement 5.1 Full Audit Naming Suggestions (#5215)
Co-authored-by: Hadrien Croubois <[email protected]> Co-authored-by: cairo <[email protected]>
1 parent f6db286 commit 4c481d6

File tree

2 files changed

+42
-40
lines changed

2 files changed

+42
-40
lines changed

contracts/utils/structs/Heap.sol

Lines changed: 32 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,13 @@ import {StorageSlot} from "../StorageSlot.sol";
1313
* @dev Library for managing https://en.wikipedia.org/wiki/Binary_heap[binary heap] that can be used as
1414
* https://en.wikipedia.org/wiki/Priority_queue[priority queue].
1515
*
16-
* Heaps are represented as an tree of values where the first element (index 0) is the root, and where the node at
17-
* index i is the child of the node at index (i-1)/2 and the father of nodes at index 2*i+1 and 2*i+2. Each node
16+
* Heaps are represented as a tree of values where the first element (index 0) is the root, and where the node at
17+
* index i is the child of the node at index (i-1)/2 and the parent of nodes at index 2*i+1 and 2*i+2. Each node
1818
* stores an element of the heap.
1919
*
2020
* The structure is ordered so that each node is bigger than its parent. An immediate consequence is that the
2121
* highest priority value is the one at the root. This value can be looked up in constant time (O(1)) at
22-
* `heap.tree[0].value`
22+
* `heap.tree[0]`
2323
*
2424
* The structure is designed to perform the following operations with the corresponding complexities:
2525
*
@@ -42,7 +42,7 @@ library Heap {
4242
/**
4343
* @dev Binary heap that supports values of type uint256.
4444
*
45-
* Each element of that structure uses 2 storage slots.
45+
* Each element of that structure uses one storage slot.
4646
*/
4747
struct Uint256Heap {
4848
uint256[] tree;
@@ -184,7 +184,7 @@ library Heap {
184184
* @dev Perform heap maintenance on `self`, starting at `index` (with the `value`), using `comp` as a
185185
* comparator, and moving toward the leaves of the underlying tree.
186186
*
187-
* NOTE: This is a private function that is called in a trusted context with already cached parameters. `length`
187+
* NOTE: This is a private function that is called in a trusted context with already cached parameters. `size`
188188
* and `value` could be extracted from `self` and `index`, but that would require redundant storage read. These
189189
* parameters are not verified. It is the caller role to make sure the parameters are correct.
190190
*/
@@ -195,31 +195,33 @@ library Heap {
195195
uint256 value,
196196
function(uint256, uint256) view returns (bool) comp
197197
) private {
198-
// Check if there is a risk of overflow when computing the indices of the child nodes. If that is the case,
199-
// there cannot be child nodes in the tree, so sifting is done.
200-
if (index >= type(uint256).max / 2) return;
201-
202-
// Compute the indices of the potential child nodes
203-
uint256 lIndex = 2 * index + 1;
204-
uint256 rIndex = 2 * index + 2;
205-
206-
// Three cases:
207-
// 1. Both children exist: sifting may continue on one of the branch (selection required)
208-
// 2. Only left child exist: sifting may contineu on the left branch (no selection required)
209-
// 3. Neither child exist: sifting is done
210-
if (rIndex < size) {
211-
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
212-
uint256 rValue = self.tree.unsafeAccess(rIndex).value;
213-
if (comp(lValue, value) || comp(rValue, value)) {
214-
uint256 cIndex = comp(lValue, rValue).ternary(lIndex, rIndex);
215-
_swap(self, index, cIndex);
216-
_siftDown(self, size, cIndex, value, comp);
217-
}
218-
} else if (lIndex < size) {
219-
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
220-
if (comp(lValue, value)) {
221-
_swap(self, index, lIndex);
222-
_siftDown(self, size, lIndex, value, comp);
198+
unchecked {
199+
// Check if there is a risk of overflow when computing the indices of the child nodes. If that is the case,
200+
// there cannot be child nodes in the tree, so sifting is done.
201+
if (index >= type(uint256).max / 2) return;
202+
203+
// Compute the indices of the potential child nodes
204+
uint256 lIndex = 2 * index + 1;
205+
uint256 rIndex = 2 * index + 2;
206+
207+
// Three cases:
208+
// 1. Both children exist: sifting may continue on one of the branch (selection required)
209+
// 2. Only left child exist: sifting may continue on the left branch (no selection required)
210+
// 3. Neither child exist: sifting is done
211+
if (rIndex < size) {
212+
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
213+
uint256 rValue = self.tree.unsafeAccess(rIndex).value;
214+
if (comp(lValue, value) || comp(rValue, value)) {
215+
uint256 cIndex = comp(lValue, rValue).ternary(lIndex, rIndex);
216+
_swap(self, index, cIndex);
217+
_siftDown(self, size, cIndex, value, comp);
218+
}
219+
} else if (lIndex < size) {
220+
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
221+
if (comp(lValue, value)) {
222+
_swap(self, index, lIndex);
223+
_siftDown(self, size, lIndex, value, comp);
224+
}
223225
}
224226
}
225227
}

contracts/utils/structs/MerkleTree.sol

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ library MerkleTree {
4949

5050
/**
5151
* @dev Initialize a {Bytes32PushTree} using {Hashes-commutativeKeccak256} to hash internal nodes.
52-
* The capacity of the tree (i.e. number of leaves) is set to `2**levels`.
52+
* The capacity of the tree (i.e. number of leaves) is set to `2**treeDepth`.
5353
*
5454
* Calling this function on MerkleTree that was already setup and used will reset it to a blank state.
5555
*
@@ -59,8 +59,8 @@ library MerkleTree {
5959
* IMPORTANT: The zero value should be carefully chosen since it will be stored in the tree representing
6060
* empty leaves. It should be a value that is not expected to be part of the tree.
6161
*/
62-
function setup(Bytes32PushTree storage self, uint8 levels, bytes32 zero) internal returns (bytes32 initialRoot) {
63-
return setup(self, levels, zero, Hashes.commutativeKeccak256);
62+
function setup(Bytes32PushTree storage self, uint8 treeDepth, bytes32 zero) internal returns (bytes32 initialRoot) {
63+
return setup(self, treeDepth, zero, Hashes.commutativeKeccak256);
6464
}
6565

6666
/**
@@ -77,17 +77,17 @@ library MerkleTree {
7777
*/
7878
function setup(
7979
Bytes32PushTree storage self,
80-
uint8 levels,
80+
uint8 treeDepth,
8181
bytes32 zero,
8282
function(bytes32, bytes32) view returns (bytes32) fnHash
8383
) internal returns (bytes32 initialRoot) {
8484
// Store depth in the dynamic array
85-
Arrays.unsafeSetLength(self._sides, levels);
86-
Arrays.unsafeSetLength(self._zeros, levels);
85+
Arrays.unsafeSetLength(self._sides, treeDepth);
86+
Arrays.unsafeSetLength(self._zeros, treeDepth);
8787

8888
// Build each root of zero-filled subtrees
8989
bytes32 currentZero = zero;
90-
for (uint32 i = 0; i < levels; ++i) {
90+
for (uint32 i = 0; i < treeDepth; ++i) {
9191
Arrays.unsafeAccess(self._zeros, i).value = currentZero;
9292
currentZero = fnHash(currentZero, currentZero);
9393
}
@@ -129,20 +129,20 @@ library MerkleTree {
129129
function(bytes32, bytes32) view returns (bytes32) fnHash
130130
) internal returns (uint256 index, bytes32 newRoot) {
131131
// Cache read
132-
uint256 levels = self._zeros.length;
132+
uint256 treeDepth = depth(self);
133133

134134
// Get leaf index
135135
index = self._nextLeafIndex++;
136136

137137
// Check if tree is full.
138-
if (index >= 1 << levels) {
138+
if (index >= 1 << treeDepth) {
139139
Panic.panic(Panic.RESOURCE_ERROR);
140140
}
141141

142142
// Rebuild branch from leaf to root
143143
uint256 currentIndex = index;
144144
bytes32 currentLevelHash = leaf;
145-
for (uint32 i = 0; i < levels; i++) {
145+
for (uint32 i = 0; i < treeDepth; i++) {
146146
// Reaching the parent node, is currentLevelHash the left child?
147147
bool isLeft = currentIndex % 2 == 0;
148148

0 commit comments

Comments
 (0)