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

[Simon Tatham via libc-commits]

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

>From b298f32fc662f03ce3f037b672fb4503b8f5bb83 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 1/3] [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..db221943d9efe4 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

>From 0f0a2838be9f9c91809aceef86a7bc51420780f0 Mon Sep 17 00:00:00 2001
From: Simon Tatham <simon.tatham at arm.com>
Date: Thu, 3 Oct 2024 09:16:59 +0100
Subject: [PATCH 2/3] remove rogue diagnostic parameter to quick_sort()

---
 libc/src/stdlib/quick_sort.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/libc/src/stdlib/quick_sort.h b/libc/src/stdlib/quick_sort.h
index db221943d9efe4..42816508c90f37 100644
--- a/libc/src/stdlib/quick_sort.h
+++ b/libc/src/stdlib/quick_sort.h
@@ -79,10 +79,10 @@ LIBC_INLINE void quick_sort(Array array) {
     // 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);
+      quick_sort(left);
       array.reset_bounds(right.get(0), right.size());
     } else {
-      quick_sort(right, depth + 1);
+      quick_sort(right);
       array.reset_bounds(left.get(0), left.size());
     }
   }

>From 4ade1bde5ae1272e52fb60845ebea9f66edcbcc6 Mon Sep 17 00:00:00 2001
From: Simon Tatham <simon.tatham at arm.com>
Date: Thu, 3 Oct 2024 09:18:59 +0100
Subject: [PATCH 3/3] remove braces that already contravened the style guide

---
 libc/src/stdlib/quick_sort.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/libc/src/stdlib/quick_sort.h b/libc/src/stdlib/quick_sort.h
index 42816508c90f37..2d1cb46edc6560 100644
--- a/libc/src/stdlib/quick_sort.h
+++ b/libc/src/stdlib/quick_sort.h
@@ -65,10 +65,9 @@ LIBC_INLINE void quick_sort(Array array) {
     if (array_size <= 1)
       return;
     size_t split_index = partition(array);
-    if (array_size <= 2) {
+    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);