Fix pdqSort+heapSort for ranges besides 0..len (#15982)

1 files changed, 8 insertions, 10 deletions

diff --git a/lib/std/sort/pdq.zig b/lib/std/sort/pdq.zig
index e7042b0c76..23678a79c6 100644
--- a/lib/std/sort/pdq.zig
+++ b/lib/std/sort/pdq.zig
@@ -43,7 +43,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
     // slices of up to this length get sorted using insertion sort.
     const max_insertion = 24;
     // number of allowed imbalanced partitions before switching to heap sort.
-    const max_limit = std.math.floorPowerOfTwo(usize, b) + 1;
+    const max_limit = std.math.floorPowerOfTwo(usize, b - a) + 1;
 
     // set upper bound on stack memory usage.
     const Range = struct { a: usize, b: usize, limit: usize };
@@ -100,7 +100,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
             // if the chosen pivot is equal to the predecessor, then it's the smallest element in the
             // slice. Partition the slice into elements equal to and elements greater than the pivot.
             // This case is usually hit when the slice contains many duplicate elements.
-            if (range.a > 0 and !context.lessThan(range.a - 1, pivot)) {
+            if (range.a > a and !context.lessThan(range.a - 1, pivot)) {
                 range.a = partitionEqual(range.a, range.b, pivot, context);
                 continue;
             }
@@ -284,13 +284,13 @@ fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
     if (len >= 8) {
         if (len >= shortest_ninther) {
             // find medians in the neighborhoods of `i`, `j` and `k`
-            i = sort3(i - 1, i, i + 1, &swaps, context);
-            j = sort3(j - 1, j, j + 1, &swaps, context);
-            k = sort3(k - 1, k, k + 1, &swaps, context);
+            sort3(i - 1, i, i + 1, &swaps, context);
+            sort3(j - 1, j, j + 1, &swaps, context);
+            sort3(k - 1, k, k + 1, &swaps, context);
         }
 
-        // find the median among `i`, `j` and `k`
-        j = sort3(i, j, k, &swaps, context);
+        // find the median among `i`, `j` and `k` and stores it in `j`
+        sort3(i, j, k, &swaps, context);
     }
 
     pivot.* = j;
@@ -301,7 +301,7 @@ fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
     };
 }
 
-fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) usize {
+fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) void {
     if (context.lessThan(b, a)) {
         swaps.* += 1;
         context.swap(b, a);
@@ -316,8 +316,6 @@ fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) usize {
         swaps.* += 1;
         context.swap(b, a);
     }
-
-    return b;
 }
 
 fn reverseRange(a: usize, b: usize, context: anytype) void {