[libc-commits] [libc] [libc] Bound the worst-case stack usage in qsort(). (PR #110849)

[Simon Tatham via libc-commits]

https://github.com/statham-arm created https://github.com/llvm/llvm-project/pull/110849

Previously, the Quicksort implementation was written in the obvious way: after each partitioning step, it explicitly recursed twice to sort the two sublists. Now it compares the two sublists' sizes, and recurses only to sort the smaller one. To handle the larger list it loops back round to the top of the function, so as to handle it within the existing stack frame.

This means that every recursive call is handling a list at most half that of its caller. So the maximum recursive call depth is O(lg N). Otherwise, in Quicksort's bad cases where each partition step peels off a small constant number of array elements, the stack usage could grow linearly with the array being sorted, i.e. it might be Θ(N).

I tested this code by manually constructing a List Of Doom that causes this particular quicksort implementation to hit its worst case, and confirming that it recursed very deeply in the old code and doesn't in the new code. But I haven't added that list to the test suite, because the List Of Doom has to be constructed in a way based on every detail of the quicksort algorithm (pivot choice and partitioning strategy), so it would silently stop being a useful regression test as soon as any detail changed.

>From 2a763fd6dea929b160d9f7f2048a3f3540de7214 Mon Sep 17 00:00:00 2001
From: Simon Tatham <simon.tatham at arm.com>
Date: Wed, 28 Aug 2024 13:34:45 +0100
Subject: [PATCH] [libc] Bound the worst-case stack usage in qsort().
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Previously, the Quicksort implementation was written in the obvious
way: after each partitioning step, it explicitly recursed twice to
sort the two sublists. Now it compares the two sublists' sizes, and
recurses only to sort the smaller one. To handle the larger list it
loops back round to the top of the function, so as to handle it within
the existing stack frame.

This means that every recursive call is handling a list at most half
that of its caller. So the maximum recursive call depth is O(lg N).
Otherwise, in Quicksort's bad cases where each partition step peels
off a small constant number of array elements, the stack usage could
grow linearly with the array being sorted, i.e. it might be Θ(N).

I tested this code by manually constructing a List Of Doom that causes
this particular quicksort implementation to hit its worst case, and
confirming that it recursed very deeply in the old code and doesn't in
the new code. But I haven't added that list to the test suite, because
the List Of Doom has to be constructed in a way based on every detail
of the quicksort algorithm (pivot choice and partitioning strategy),
so it would silently stop being a useful regression test as soon as
any detail changed.
---
 libc/src/stdlib/qsort_data.h |  6 ++++++
 libc/src/stdlib/quick_sort.h | 36 ++++++++++++++++++++++++++----------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/libc/src/stdlib/qsort_data.h b/libc/src/stdlib/qsort_data.h
index db045332708ae6..9d44d41f46580c 100644
--- a/libc/src/stdlib/qsort_data.h
+++ b/libc/src/stdlib/qsort_data.h
@@ -92,6 +92,12 @@ class Array {
   Array make_array(size_t i, size_t s) const {
     return Array(get(i), s, elem_size, compare);
   }
+
+  // Reset this Array to point at a different interval of the same items.
+  void reset_bounds(uint8_t *a, size_t s) {
+    array = a;
+    array_size = s;
+  }
 };
 
 using SortingRoutine = void(const Array &);
diff --git a/libc/src/stdlib/quick_sort.h b/libc/src/stdlib/quick_sort.h
index 89ec107161e3e5..a3fe1681591d26 100644
--- a/libc/src/stdlib/quick_sort.h
+++ b/libc/src/stdlib/quick_sort.h
@@ -59,17 +59,33 @@ static size_t partition(const Array &array) {
   }
 }
 
-LIBC_INLINE void quick_sort(const Array &array) {
-  const size_t array_size = array.size();
-  if (array_size <= 1)
-    return;
-  size_t split_index = partition(array);
-  if (array_size <= 2) {
-    // The partition operation sorts the two element array.
-    return;
+LIBC_INLINE void quick_sort(Array array) {
+  while (true) {
+    const size_t array_size = array.size();
+    if (array_size <= 1)
+      return;
+    size_t split_index = partition(array);
+    if (array_size <= 2) {
+      // The partition operation sorts the two element array.
+      return;
+    }
+
+    // Make Arrays describing the two sublists that still need sorting.
+    Array left = array.make_array(0, split_index);
+    Array right = array.make_array(split_index, array.size() - split_index);
+
+    // Recurse to sort the smaller of the two, and then loop round within this
+    // function to sort the larger. This way, recursive call depth is bounded
+    // by log2 of the total array size, because every recursive call is sorting
+    // a list at most half the length of the one in its caller.
+    if (left.size() < right.size()) {
+        quick_sort(left, depth+1);
+      array.reset_bounds(right.get(0), right.size());
+    } else {
+        quick_sort(right, depth+1);
+      array.reset_bounds(left.get(0), left.size());
+    }
   }
-  quick_sort(array.make_array(0, split_index));
-  quick_sort(array.make_array(split_index, array.size() - split_index));
 }
 
 } // namespace internal