Browse Source

Implement 5.1 Full Audit Naming Suggestions (#5215)

Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>
Co-authored-by: cairo <cairoeth@protonmail.com>
Ernesto García 1 year ago
parent
commit
4c481d6584
2 changed files with 42 additions and 40 deletions
  1. 32 30
      contracts/utils/structs/Heap.sol
  2. 10 10
      contracts/utils/structs/MerkleTree.sol

+ 32 - 30
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;
-
-        // 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);
+                }
             }
         }
     }

+ 10 - 10
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;