1const std = @import("../std.zig");
  2const sort = std.sort;
  3const mem = std.mem;
  4const math = std.math;
  5const testing = std.testing;
  6
  7/// Unstable in-place sort. n best case, n*log(n) worst case and average case.
  8/// log(n) memory (no allocator required).
  9///
 10/// Sorts in ascending order with respect to the given `lessThan` function.
 11pub fn pdq(
 12    comptime T: type,
 13    items: []T,
 14    context: anytype,
 15    comptime lessThanFn: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
 16) void {
 17    const Context = struct {
 18        items: []T,
 19        sub_ctx: @TypeOf(context),
 20
 21        pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
 22            return lessThanFn(ctx.sub_ctx, ctx.items[a], ctx.items[b]);
 23        }
 24
 25        pub fn swap(ctx: @This(), a: usize, b: usize) void {
 26            return mem.swap(T, &ctx.items[a], &ctx.items[b]);
 27        }
 28    };
 29    pdqContext(0, items.len, Context{ .items = items, .sub_ctx = context });
 30}
 31
 32const Hint = enum {
 33    increasing,
 34    decreasing,
 35    unknown,
 36};
 37
 38/// Unstable in-place sort. O(n) best case, O(n*log(n)) worst case and average case.
 39/// O(log(n)) memory (no allocator required).
 40/// `context` must have methods `swap` and `lessThan`,
 41/// which each take 2 `usize` parameters indicating the index of an item.
 42/// Sorts in ascending order with respect to `lessThan`.
 43pub fn pdqContext(a: usize, b: usize, context: anytype) void {
 44    // slices of up to this length get sorted using insertion sort.
 45    const max_insertion = 24;
 46    // number of allowed imbalanced partitions before switching to heap sort.
 47    const max_limit = std.math.floorPowerOfTwo(usize, b - a) + 1;
 48
 49    // set upper bound on stack memory usage.
 50    const Range = struct { a: usize, b: usize, limit: usize };
 51    const stack_size = math.log2(math.maxInt(usize) + 1);
 52    var stack: [stack_size]Range = undefined;
 53    var range = Range{ .a = a, .b = b, .limit = max_limit };
 54    var top: usize = 0;
 55
 56    while (true) {
 57        var was_balanced = true;
 58        var was_partitioned = true;
 59
 60        while (true) {
 61            const len = range.b - range.a;
 62
 63            // very short slices get sorted using insertion sort.
 64            if (len <= max_insertion) {
 65                break sort.insertionContext(range.a, range.b, context);
 66            }
 67
 68            // if too many bad pivot choices were made, simply fall back to heapsort in order to
 69            // guarantee O(n*log(n)) worst-case.
 70            if (range.limit == 0) {
 71                break sort.heapContext(range.a, range.b, context);
 72            }
 73
 74            // if the last partitioning was imbalanced, try breaking patterns in the slice by shuffling
 75            // some elements around. Hopefully we'll choose a better pivot this time.
 76            if (!was_balanced) {
 77                breakPatterns(range.a, range.b, context);
 78                range.limit -= 1;
 79            }
 80
 81            // choose a pivot and try guessing whether the slice is already sorted.
 82            var pivot: usize = 0;
 83            var hint = chosePivot(range.a, range.b, &pivot, context);
 84
 85            if (hint == .decreasing) {
 86                // The maximum number of swaps was performed, so items are likely
 87                // in reverse order. Reverse it to make sorting faster.
 88                reverseRange(range.a, range.b, context);
 89                pivot = (range.b - 1) - (pivot - range.a);
 90                hint = .increasing;
 91            }
 92
 93            // if the last partitioning was decently balanced and didn't shuffle elements, and if pivot
 94            // selection predicts the slice is likely already sorted...
 95            if (was_balanced and was_partitioned and hint == .increasing) {
 96                // try identifying several out-of-order elements and shifting them to correct
 97                // positions. If the slice ends up being completely sorted, we're done.
 98                if (partialInsertionSort(range.a, range.b, context)) break;
 99            }
100
101            // if the chosen pivot is equal to the predecessor, then it's the smallest element in the
102            // slice. Partition the slice into elements equal to and elements greater than the pivot.
103            // This case is usually hit when the slice contains many duplicate elements.
104            if (range.a > a and !context.lessThan(range.a - 1, pivot)) {
105                range.a = partitionEqual(range.a, range.b, pivot, context);
106                continue;
107            }
108
109            // partition the slice.
110            var mid = pivot;
111            was_partitioned = partition(range.a, range.b, &mid, context);
112
113            const left_len = mid - range.a;
114            const right_len = range.b - mid;
115            const balanced_threshold = len / 8;
116            if (left_len < right_len) {
117                was_balanced = left_len >= balanced_threshold;
118                stack[top] = .{ .a = range.a, .b = mid, .limit = range.limit };
119                top += 1;
120                range.a = mid + 1;
121            } else {
122                was_balanced = right_len >= balanced_threshold;
123                stack[top] = .{ .a = mid + 1, .b = range.b, .limit = range.limit };
124                top += 1;
125                range.b = mid;
126            }
127        }
128
129        top = math.sub(usize, top, 1) catch break;
130        range = stack[top];
131    }
132}
133
134/// partitions `items[a..b]` into elements smaller than `items[pivot]`,
135/// followed by elements greater than or equal to `items[pivot]`.
136///
137/// sets the new pivot.
138/// returns `true` if already partitioned.
139fn partition(a: usize, b: usize, pivot: *usize, context: anytype) bool {
140    // move pivot to the first place
141    context.swap(a, pivot.*);
142
143    var i = a + 1;
144    var j = b - 1;
145
146    while (i <= j and context.lessThan(i, a)) i += 1;
147    while (i <= j and !context.lessThan(j, a)) j -= 1;
148
149    // check if items are already partitioned (no item to swap)
150    if (i > j) {
151        // put pivot back to the middle
152        context.swap(j, a);
153        pivot.* = j;
154        return true;
155    }
156
157    context.swap(i, j);
158    i += 1;
159    j -= 1;
160
161    while (true) {
162        while (i <= j and context.lessThan(i, a)) i += 1;
163        while (i <= j and !context.lessThan(j, a)) j -= 1;
164        if (i > j) break;
165
166        context.swap(i, j);
167        i += 1;
168        j -= 1;
169    }
170
171    // TODO: Enable the BlockQuicksort optimization
172
173    context.swap(j, a);
174    pivot.* = j;
175    return false;
176}
177
178/// partitions items into elements equal to `items[pivot]`
179/// followed by elements greater than `items[pivot]`.
180///
181/// it assumed that `items[a..b]` does not contain elements smaller than the `items[pivot]`.
182fn partitionEqual(a: usize, b: usize, pivot: usize, context: anytype) usize {
183    // move pivot to the first place
184    context.swap(a, pivot);
185
186    var i = a + 1;
187    var j = b - 1;
188
189    while (true) {
190        while (i <= j and !context.lessThan(a, i)) i += 1;
191        while (i <= j and context.lessThan(a, j)) j -= 1;
192        if (i > j) break;
193
194        context.swap(i, j);
195        i += 1;
196        j -= 1;
197    }
198
199    return i;
200}
201
202/// partially sorts a slice by shifting several out-of-order elements around.
203///
204/// returns `true` if the slice is sorted at the end. This function is `O(n)` worst-case.
205fn partialInsertionSort(a: usize, b: usize, context: anytype) bool {
206    @branchHint(.cold);
207
208    // maximum number of adjacent out-of-order pairs that will get shifted
209    const max_steps = 5;
210    // if the slice is shorter than this, don't shift any elements
211    const shortest_shifting = 50;
212
213    var i = a + 1;
214    for (0..max_steps) |_| {
215        // find the next pair of adjacent out-of-order elements.
216        while (i < b and !context.lessThan(i, i - 1)) i += 1;
217
218        // are we done?
219        if (i == b) return true;
220
221        // don't shift elements on short arrays, that has a performance cost.
222        if (b - a < shortest_shifting) return false;
223
224        // swap the found pair of elements. This puts them in correct order.
225        context.swap(i, i - 1);
226
227        // shift the smaller element to the left.
228        if (i - a >= 2) {
229            var j = i - 1;
230            while (j > a) : (j -= 1) {
231                if (!context.lessThan(j, j - 1)) break;
232                context.swap(j, j - 1);
233            }
234        }
235
236        // shift the greater element to the right.
237        if (b - i >= 2) {
238            var j = i + 1;
239            while (j < b) : (j += 1) {
240                if (!context.lessThan(j, j - 1)) break;
241                context.swap(j, j - 1);
242            }
243        }
244    }
245
246    return false;
247}
248
249fn breakPatterns(a: usize, b: usize, context: anytype) void {
250    @branchHint(.cold);
251
252    const len = b - a;
253    if (len < 8) return;
254
255    var rand = @as(u64, @intCast(len));
256    const modulus = math.ceilPowerOfTwoAssert(u64, len);
257
258    var i = a + (len / 4) * 2 - 1;
259    while (i <= a + (len / 4) * 2 + 1) : (i += 1) {
260        // xorshift64
261        rand ^= rand << 13;
262        rand ^= rand >> 7;
263        rand ^= rand << 17;
264
265        var other = @as(usize, @intCast(rand & (modulus - 1)));
266        if (other >= len) other -= len;
267        context.swap(i, a + other);
268    }
269}
270
271/// chooses a pivot in `items[a..b]`.
272/// swaps likely_sorted when `items[a..b]` seems to be already sorted.
273fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
274    // minimum length for using the Tukey's ninther method
275    const shortest_ninther = 50;
276    // max_swaps is the maximum number of swaps allowed in this function
277    const max_swaps = 4 * 3;
278
279    const len = b - a;
280    const i = a + len / 4 * 1;
281    const j = a + len / 4 * 2;
282    const k = a + len / 4 * 3;
283    var swaps: usize = 0;
284
285    if (len >= 8) {
286        if (len >= shortest_ninther) {
287            // find medians in the neighborhoods of `i`, `j` and `k`
288            sort3(i - 1, i, i + 1, &swaps, context);
289            sort3(j - 1, j, j + 1, &swaps, context);
290            sort3(k - 1, k, k + 1, &swaps, context);
291        }
292
293        // find the median among `i`, `j` and `k` and stores it in `j`
294        sort3(i, j, k, &swaps, context);
295    }
296
297    pivot.* = j;
298    return switch (swaps) {
299        0 => .increasing,
300        max_swaps => .decreasing,
301        else => .unknown,
302    };
303}
304
305fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) void {
306    if (context.lessThan(b, a)) {
307        swaps.* += 1;
308        context.swap(b, a);
309    }
310
311    if (context.lessThan(c, b)) {
312        swaps.* += 1;
313        context.swap(c, b);
314    }
315
316    if (context.lessThan(b, a)) {
317        swaps.* += 1;
318        context.swap(b, a);
319    }
320}
321
322fn reverseRange(a: usize, b: usize, context: anytype) void {
323    var i = a;
324    var j = b - 1;
325    while (i < j) {
326        context.swap(i, j);
327        i += 1;
328        j -= 1;
329    }
330}
331
332test "pdqContext respects arbitrary range boundaries" {
333    // Regression test for issue #25250
334    // pdqsort should never access indices outside the specified [a, b) range
335    var data: [2000]i32 = @splat(0);
336
337    // Fill with data that triggers the partialInsertionSort path
338    for (0..data.len) |i| {
339        data[i] = @intCast(@mod(@as(i32, @intCast(i)) * 7, 100));
340    }
341
342    const TestContext = struct {
343        items: []i32,
344        range_start: usize,
345        range_end: usize,
346
347        pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
348            // Assert indices are within the expected range
349            testing.expect(a >= ctx.range_start and a < ctx.range_end) catch @panic("index a out of range");
350            testing.expect(b >= ctx.range_start and b < ctx.range_end) catch @panic("index b out of range");
351            return ctx.items[a] < ctx.items[b];
352        }
353
354        pub fn swap(ctx: @This(), a: usize, b: usize) void {
355            // Assert indices are within the expected range
356            testing.expect(a >= ctx.range_start and a < ctx.range_end) catch @panic("index a out of range");
357            testing.expect(b >= ctx.range_start and b < ctx.range_end) catch @panic("index b out of range");
358            mem.swap(i32, &ctx.items[a], &ctx.items[b]);
359        }
360    };
361
362    // Test sorting a sub-range that doesn't start at 0
363    const start = 1118;
364    const end = 1764;
365    const ctx = TestContext{
366        .items = &data,
367        .range_start = start,
368        .range_end = end,
369    };
370
371    pdqContext(start, end, ctx);
372
373    // Verify the range is sorted
374    for ((start + 1)..end) |i| {
375        try testing.expect(data[i - 1] <= data[i]);
376    }
377}