Skip to content

Commit

Permalink
Implement 5.1 Full Audit Naming Suggestions (#5215)
Browse files Browse the repository at this point in the history 
Co-authored-by: Hadrien Croubois <[email protected]>
Co-authored-by: cairo <[email protected]>
  • Loading branch information
@ernestognw @Amxx @cairoeth
3 people authored Sep 25, 2024
1 parent f6db286 commit 4c481d6
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 40 deletions.
62 changes: 32 additions & 30 deletions contracts/utils/structs/Heap.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,13 @@ import {StorageSlot} from "../StorageSlot.sol";
* @dev Library for managing https://en.wikipedia.org/wiki/Binary_heap[binary heap] that can be used as
* https://en.wikipedia.org/wiki/Priority_queue[priority queue].
*
* Heaps are represented as an tree of values where the first element (index 0) is the root, and where the node at
* 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
* Heaps are represented as a tree of values where the first element (index 0) is the root, and where the node at
* 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
* stores an element of the heap.
*
* The structure is ordered so that each node is bigger than its parent. An immediate consequence is that the
* highest priority value is the one at the root. This value can be looked up in constant time (O(1)) at
* `heap.tree[0].value`
* `heap.tree[0]`
*
* The structure is designed to perform the following operations with the corresponding complexities:
*
Expand All @@ -42,7 +42,7 @@ library Heap {
/**
* @dev Binary heap that supports values of type uint256.
*
* Each element of that structure uses 2 storage slots.
* Each element of that structure uses one storage slot.
*/
struct Uint256Heap {
uint256[] tree;
Expand Down Expand Up @@ -184,7 +184,7 @@ library Heap {
* @dev Perform heap maintenance on `self`, starting at `index` (with the `value`), using `comp` as a
* comparator, and moving toward the leaves of the underlying tree.
*
* NOTE: This is a private function that is called in a trusted context with already cached parameters. `length`
* NOTE: This is a private function that is called in a trusted context with already cached parameters. `size`
* and `value` could be extracted from `self` and `index`, but that would require redundant storage read. These
* parameters are not verified. It is the caller role to make sure the parameters are correct.
*/
Expand All @@ -195,31 +195,33 @@ library Heap {
uint256 value,
function(uint256, uint256) view returns (bool) comp
) private {
// Check if there is a risk of overflow when computing the indices of the child nodes. If that is the case,
// there cannot be child nodes in the tree, so sifting is done.
if (index >= type(uint256).max / 2) return;

// Compute the indices of the potential child nodes
uint256 lIndex = 2 * index + 1;
uint256 rIndex = 2 * index + 2;

// Three cases:
// 1. Both children exist: sifting may continue on one of the branch (selection required)
// 2. Only left child exist: sifting may contineu on the left branch (no selection required)
// 3. Neither child exist: sifting is done
if (rIndex < size) {
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
uint256 rValue = self.tree.unsafeAccess(rIndex).value;
if (comp(lValue, value) || comp(rValue, value)) {
uint256 cIndex = comp(lValue, rValue).ternary(lIndex, rIndex);
_swap(self, index, cIndex);
_siftDown(self, size, cIndex, value, comp);
}
} else if (lIndex < size) {
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
if (comp(lValue, value)) {
_swap(self, index, lIndex);
_siftDown(self, size, lIndex, value, comp);
unchecked {
// Check if there is a risk of overflow when computing the indices of the child nodes. If that is the case,
// there cannot be child nodes in the tree, so sifting is done.
if (index >= type(uint256).max / 2) return;

// Compute the indices of the potential child nodes
uint256 lIndex = 2 * index + 1;
uint256 rIndex = 2 * index + 2;

// Three cases:
// 1. Both children exist: sifting may continue on one of the branch (selection required)
// 2. Only left child exist: sifting may continue on the left branch (no selection required)
// 3. Neither child exist: sifting is done
if (rIndex < size) {
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
uint256 rValue = self.tree.unsafeAccess(rIndex).value;
if (comp(lValue, value) || comp(rValue, value)) {
uint256 cIndex = comp(lValue, rValue).ternary(lIndex, rIndex);
_swap(self, index, cIndex);
_siftDown(self, size, cIndex, value, comp);
}
} else if (lIndex < size) {
uint256 lValue = self.tree.unsafeAccess(lIndex).value;
if (comp(lValue, value)) {
_swap(self, index, lIndex);
_siftDown(self, size, lIndex, value, comp);
}
}
}
}
Expand Down
20 changes: 10 additions & 10 deletions contracts/utils/structs/MerkleTree.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ library MerkleTree {

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

/**
Expand All @@ -77,17 +77,17 @@ library MerkleTree {
*/
function setup(
Bytes32PushTree storage self,
uint8 levels,
uint8 treeDepth,
bytes32 zero,
function(bytes32, bytes32) view returns (bytes32) fnHash
) internal returns (bytes32 initialRoot) {
// Store depth in the dynamic array
Arrays.unsafeSetLength(self._sides, levels);
Arrays.unsafeSetLength(self._zeros, levels);
Arrays.unsafeSetLength(self._sides, treeDepth);
Arrays.unsafeSetLength(self._zeros, treeDepth);

// Build each root of zero-filled subtrees
bytes32 currentZero = zero;
for (uint32 i = 0; i < levels; ++i) {
for (uint32 i = 0; i < treeDepth; ++i) {
Arrays.unsafeAccess(self._zeros, i).value = currentZero;
currentZero = fnHash(currentZero, currentZero);
}
Expand Down Expand Up @@ -129,20 +129,20 @@ library MerkleTree {
function(bytes32, bytes32) view returns (bytes32) fnHash
) internal returns (uint256 index, bytes32 newRoot) {
// Cache read
uint256 levels = self._zeros.length;
uint256 treeDepth = depth(self);

// Get leaf index
index = self._nextLeafIndex++;

// Check if tree is full.
if (index >= 1 << levels) {
if (index >= 1 << treeDepth) {
Panic.panic(Panic.RESOURCE_ERROR);
}

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

Expand Down

0 comments on commit 4c481d6

Please sign in to comment.