⚡️ Optimize RedBlackTreeLib values function (#784)

Vectorized · Jan 9, 2024 · b2a5a8f · b2a5a8f
1 parent ac6d2cc
commit b2a5a8f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 21 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -914,9 +914,9 @@ ReceiverTest:test__codesize() (gas: 3310)
 RedBlackTreeLibTest:testRedBlackTreeBenchUint160() (gas: 3438464)
 RedBlackTreeLibTest:testRedBlackTreeBenchUint256() (gas: 5850774)
 RedBlackTreeLibTest:testRedBlackTreeClear() (gas: 57543)
-RedBlackTreeLibTest:testRedBlackTreeClear(uint256) (runs: 256, μ: 296953, ~: 214791)
-RedBlackTreeLibTest:testRedBlackTreeInsertAndRemove(uint256) (runs: 256, μ: 647105, ~: 525248)
-RedBlackTreeLibTest:testRedBlackTreeInsertAndRemove2(uint256) (runs: 256, μ: 416355, ~: 388854)
+RedBlackTreeLibTest:testRedBlackTreeClear(uint256) (runs: 256, μ: 293796, ~: 216591)
+RedBlackTreeLibTest:testRedBlackTreeInsertAndRemove(uint256) (runs: 256, μ: 688219, ~: 521466)
+RedBlackTreeLibTest:testRedBlackTreeInsertAndRemove2(uint256) (runs: 256, μ: 430879, ~: 388854)
 RedBlackTreeLibTest:testRedBlackTreeInsertAndRemove3() (gas: 21537965)
 RedBlackTreeLibTest:testRedBlackTreeInsertBenchStep() (gas: 3711341)
 RedBlackTreeLibTest:testRedBlackTreeInsertBenchUint160() (gas: 3476959)
@@ -925,17 +925,17 @@ RedBlackTreeLibTest:testRedBlackTreeInsertOneGas() (gas: 45574)
 RedBlackTreeLibTest:testRedBlackTreeInsertTenGas() (gas: 283398)
 RedBlackTreeLibTest:testRedBlackTreeInsertThreeGas() (gas: 96286)
 RedBlackTreeLibTest:testRedBlackTreeInsertTwoGas() (gas: 69823)
-RedBlackTreeLibTest:testRedBlackTreeNearest(uint256) (runs: 256, μ: 232208, ~: 224742)
-RedBlackTreeLibTest:testRedBlackTreeNearestAfter(uint256) (runs: 256, μ: 247456, ~: 246502)
-RedBlackTreeLibTest:testRedBlackTreeNearestBefore(uint256) (runs: 256, μ: 230754, ~: 192926)
+RedBlackTreeLibTest:testRedBlackTreeNearest(uint256) (runs: 256, μ: 227782, ~: 223313)
+RedBlackTreeLibTest:testRedBlackTreeNearestAfter(uint256) (runs: 256, μ: 247391, ~: 245995)
+RedBlackTreeLibTest:testRedBlackTreeNearestBefore(uint256) (runs: 256, μ: 235439, ~: 234078)
 RedBlackTreeLibTest:testRedBlackTreePointers() (gas: 91923)
 RedBlackTreeLibTest:testRedBlackTreeRejectsEmptyValue() (gas: 3238)
 RedBlackTreeLibTest:testRedBlackTreeRemoveViaPointer() (gas: 58216)
 RedBlackTreeLibTest:testRedBlackTreeTreeFullReverts() (gas: 50359)
 RedBlackTreeLibTest:testRedBlackTreeTryInsertAndRemove() (gas: 56162)
-RedBlackTreeLibTest:testRedBlackTreeValues() (gas: 177033)
-RedBlackTreeLibTest:testRedBlackTreeValues(uint256) (runs: 256, μ: 277183, ~: 218609)
-RedBlackTreeLibTest:test__codesize() (gas: 14806)
+RedBlackTreeLibTest:testRedBlackTreeValues() (gas: 175206)
+RedBlackTreeLibTest:testRedBlackTreeValues(uint256) (runs: 256, μ: 266948, ~: 218626)
+RedBlackTreeLibTest:test__codesize() (gas: 14795)
 ReentrancyGuardTest:testRecursiveDirectUnguardedCall() (gas: 34256)
 ReentrancyGuardTest:testRecursiveIndirectUnguardedCall() (gas: 47773)
 ReentrancyGuardTest:testRevertGuardLocked() (gas: 53925)

diff --git a/src/utils/RedBlackTreeLib.sol b/src/utils/RedBlackTreeLib.sol
@@ -101,24 +101,23 @@ library RedBlackTreeLib {
         /// @solidity memory-safe-assembly
         assembly {
             function visit(current_) {
-                let cursor_ := or(mload(0x00), current_)
-                let packed_ := sload(cursor_)
-                let left_ := and(packed_, _BITMASK_KEY)
-                if left_ { visit(left_) }
-                let value_ := shr(_BITPOS_PACKED_VALUE, packed_)
-                if iszero(value_) { value_ := sload(or(cursor_, _BIT_FULL_VALUE_SLOT)) }
-                mstore(mload(0x20), value_)
-                mstore(0x20, add(0x20, mload(0x20)))
-                let right_ := and(shr(_BITPOS_RIGHT, packed_), _BITMASK_KEY)
-                if right_ { visit(right_) }
+                if iszero(current_) { leave } // If the current node is null, leave.
+                current_ := or(mload(0x00), current_) // Current node's storage slot.
+                let packed_ := sload(current_)
+                visit(and(packed_, _BITMASK_KEY)) // Visit left child.
+                let value_ := shr(_BITPOS_PACKED_VALUE, packed_) // Current value.
+                if iszero(value_) { value_ := sload(or(current_, _BIT_FULL_VALUE_SLOT)) }
+                mstore(mload(0x20), value_) // Append the value to `results`.
+                mstore(0x20, add(0x20, mload(0x20))) // Advance the offset into `results`.
+                visit(and(shr(_BITPOS_RIGHT, packed_), _BITMASK_KEY)) // Visit right child.
             }
             result := mload(0x40)
             let rootPacked := sload(nodes)
             mstore(result, and(rootPacked, _BITMASK_KEY)) // Length of `result`.
             mstore(0x00, nodes) // Cache the nodes pointer in scratch space.
-            mstore(0x20, add(result, 0x20)) // Cache the results offset in scratch space.
+            mstore(0x20, add(result, 0x20)) // Cache the offset into `results` in scratch space.
             mstore(0x40, add(mload(0x20), shl(5, mload(result)))) // Allocate memory.
-            visit(shr(128, rootPacked))
+            visit(shr(128, rootPacked)) // Start the tree traversal from the root node.
         }
     }