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,46 +133,49 @@ describe('Arrays', function () {
 
       if (isValueType) {
         describe('sort', function () {
-          for (const length of [0, 1, 2, 8, 32, 128]) {
-            describe(`${name}[] of length ${length}`, function () {
-              beforeEach(async function () {
-                this.array = Array.from({ length }, generators[name]);
-              });
-
-              afterEach(async function () {
-                const expected = Array.from(this.array).sort(comparator);
-                const reversed = Array.from(expected).reverse();
-                await expect(this.instance.sort(this.array)).to.eventually.deep.equal(expected);
-                await expect(this.instance.sortReverse(this.array)).to.eventually.deep.equal(reversed);
-              });
-
-              it('sort array', async function () {
-                // nothing to do here, beforeEach and afterEach already take care of everything.
-              });
-
-              if (length > 1) {
-                it('sort array for identical elements', async function () {
-                  // duplicate the first value to all elements
-                  this.array.fill(this.array.at(0));
+          for (const length of [0, 1, 2, 8, 32, 128, 384]) {
+            describe(
+              [length > 32 && '[skip-on-coverage]', `${name}[] of length ${length}`].filter(Boolean).join(' '),
+              function () {
+                beforeEach(async function () {
+                  this.array = Array.from({ length }, generators[name]);
                 });
 
-                it('sort already sorted array', async function () {
-                  // pre-sort the elements
-                  this.array.sort(comparator);
+                afterEach(async function () {
+                  const expected = Array.from(this.array).sort(comparator);
+                  const reversed = Array.from(expected).reverse();
+                  await expect(this.instance.sort(this.array)).to.eventually.deep.equal(expected);
+                  await expect(this.instance.sortReverse(this.array)).to.eventually.deep.equal(reversed);
                 });
 
-                it('sort reversed array', async function () {
-                  // pre-sort in reverse order
-                  this.array.sort(comparator).reverse();
+                it('sort array', async function () {
+                  // nothing to do here, beforeEach and afterEach already take care of everything.
                 });
 
-                it('sort almost sorted array', async function () {
-                  // pre-sort + rotate (move the last element to the front) for an almost sorted effect
-                  this.array.sort(comparator);
-                  this.array.unshift(this.array.pop());
-                });
-              }
-            });
+                if (length > 1) {
+                  it('sort array for identical elements', async function () {
+                    // duplicate the first value to all elements
+                    this.array.fill(this.array.at(0));
+                  });
+
+                  it('sort already sorted array', async function () {
+                    // pre-sort the elements
+                    this.array.sort(comparator);
+                  });
+
+                  it('sort reversed array', async function () {
+                    // pre-sort in reverse order
+                    this.array.sort(comparator).reverse();
+                  });
+
+                  it('sort almost sorted array', async function () {
+                    // pre-sort + rotate (move the last element to the front) for an almost sorted effect
+                    this.array.sort(comparator);
+                    this.array.unshift(this.array.pop());
+                  });
+                }
+              },
+            );
           }
         });