Optimize Array.quicksort

.changeset/beige-cows-unite.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`Array`: Reduce reliance on recursion to prevent stack overflow and support larger arrays.

contracts/utils/Arrays.sol

@@ -114,25 +114,33 @@ library Arrays {
      */
     function _quickSort(uint256 begin, uint256 end, function(uint256, uint256) pure returns (bool) comp) private pure {
         unchecked {
-            if (end - begin < 0x40) return;
-
-            // Use first element as pivot
-            uint256 pivot = _mload(begin);
-            // Position where the pivot should be at the end of the loop
-            uint256 pos = begin;
-
-            for (uint256 it = begin + 0x20; it < end; it += 0x20) {
-                if (comp(_mload(it), pivot)) {
-                    // If the value stored at the iterator's position comes before the pivot, we increment the
-                    // position of the pivot and move the value there.
-                    pos += 0x20;
-                    _swap(pos, it);
+            while (end - begin > 0x20) {
+                // Use first element as pivot
+                uint256 pivot = _mload(begin);
+                // Position where the pivot should be at the end of the loop
+                uint256 pos = begin;
+
+                for (uint256 it = begin + 0x20; it < end; it += 0x20) {
+                    if (comp(_mload(it), pivot)) {
+                        // If the value stored at the iterator's position comes before the pivot, we increment the
+                        // position of the pivot and move the value there.
+                        pos += 0x20;
+                        _swap(pos, it);
+                    }
                 }
-            }
 
-            _swap(begin, pos); // Swap pivot into place
-            _quickSort(begin, pos, comp); // Sort the left side of the pivot
-            _quickSort(pos + 0x20, end, comp); // Sort the right side of the pivot
+                _swap(begin, pos); // Swap pivot into place
+
+                // Recurse on the smaller partition, iterate on the larger one.
+                uint256 middle = pos + 0x20;
+                if (pos - begin < end - middle) {
+                    _quickSort(begin, pos, comp);
+                    begin = middle;
+                } else {
+                    _quickSort(middle, end, comp);
+                    end = pos;
+                }
+            }
         }
     }
 

scripts/generate/templates/Arrays.js

@@ -62,25 +62,33 @@ const quickSort = `\
  */
 function _quickSort(uint256 begin, uint256 end, function(uint256, uint256) pure returns (bool) comp) private pure {
     unchecked {
-        if (end - begin < 0x40) return;
-
-        // Use first element as pivot
-        uint256 pivot = _mload(begin);
-        // Position where the pivot should be at the end of the loop
-        uint256 pos = begin;
-
-        for (uint256 it = begin + 0x20; it < end; it += 0x20) {
-            if (comp(_mload(it), pivot)) {
-                // If the value stored at the iterator's position comes before the pivot, we increment the
-                // position of the pivot and move the value there.
-                pos += 0x20;
-                _swap(pos, it);
+        while (end - begin > 0x20) {
+            // Use first element as pivot
+            uint256 pivot = _mload(begin);
+            // Position where the pivot should be at the end of the loop
+            uint256 pos = begin;
+
+            for (uint256 it = begin + 0x20; it < end; it += 0x20) {
+                if (comp(_mload(it), pivot)) {
+                    // If the value stored at the iterator's position comes before the pivot, we increment the
+                    // position of the pivot and move the value there.
+                    pos += 0x20;
+                    _swap(pos, it);
+                }
             }
-        }
 
-        _swap(begin, pos); // Swap pivot into place
-        _quickSort(begin, pos, comp); // Sort the left side of the pivot
-        _quickSort(pos + 0x20, end, comp); // Sort the right side of the pivot
+            _swap(begin, pos); // Swap pivot into place
+
+            // Recurse on the smaller partition, iterate on the larger one.
+            uint256 middle = pos + 0x20;
+            if (pos - begin < end - middle) {
+                _quickSort(begin, pos, comp);
+                begin = middle;
+            } else {
+                _quickSort(middle, end, comp);
+                end = pos;
+            }
+        }
     }
 }
 

test/utils/Arrays.test.js

@@ -133,7 +133,7 @@ describe('Arrays', function () {
 
       if (isValueType) {
         describe('sort', function () {
-          for (const length of [0, 1, 2, 8, 32, 128]) {
+          for (const length of [0, 1, 2, 8, 32, 128, 384]) {
             describe(`${name}[] of length ${length}`, function () {
               beforeEach(async function () {
                 this.array = Array.from({ length }, generators[name]);