diff --git a/contracts/utils/structs/Heap.sol b/contracts/utils/structs/Heap.sol index c8e75f240..17a778f3a 100644 --- a/contracts/utils/structs/Heap.sol +++ b/contracts/utils/structs/Heap.sol @@ -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: * @@ -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; @@ -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. */ @@ -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; + 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; + // 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); + // 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); + } } } } diff --git a/contracts/utils/structs/MerkleTree.sol b/contracts/utils/structs/MerkleTree.sol index c3768ed8c..6a78c4b3c 100644 --- a/contracts/utils/structs/MerkleTree.sol +++ b/contracts/utils/structs/MerkleTree.sol @@ -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. * @@ -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); } /** @@ -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); } @@ -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;