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
38const Range = struct {
39 a: usize,
40 b: usize,
41 limit: usize,
42 leftmost: bool,
43 balanced: bool,
44 partitioned: bool,
45};
46
47/// Unstable in-place sort. O(n) best case, O(n*log(n)) worst case and average case.
48/// O(log(n)) memory (no allocator required).
49/// `context` must have methods `swap` and `lessThan`,
50/// which each take 2 `usize` parameters indicating the index of an item.
51/// Sorts in ascending order with respect to `lessThan`.
52pub fn pdqContext(a: usize, b: usize, context: anytype) void {
53 // slices of up to this length get sorted using insertion sort.
54 const max_insertion = 24;
55 // number of allowed imbalanced partitions before switching to heap sort.
56 const max_limit = if (b > a) math.log2_int(usize, b - a) else 0;
57
58 // stack usage is bounded by log_2(n) due to placing longer partition onto stack each iteration.
59 var stack: [@bitSizeOf(usize)]Range = undefined;
60 var range = Range{
61 .a = a,
62 .b = b,
63 .limit = max_limit,
64 .leftmost = true,
65 .balanced = true,
66 .partitioned = true,
67 };
68 var top: usize = 0;
69
70 while (true) {
71 while (true) {
72 const len = range.b - range.a;
73
74 // very short slices get sorted using insertion sort.
75 if (len <= max_insertion) {
76 if (range.leftmost) {
77 break sort.insertionContext(range.a, range.b, context);
78 } else {
79 break unguardedInsertionContext(range.a, range.b, context);
80 }
81 }
82
83 // if too many bad pivot choices were made, simply fall back to heapsort in order to
84 // guarantee O(n*log(n)) worst-case.
85 if (range.limit == 0) {
86 break sort.heapContext(range.a, range.b, context);
87 }
88
89 // if the last partitioning was imbalanced, try breaking patterns in the slice by shuffling
90 // some elements around. Hopefully we'll choose a better pivot this time.
91 if (!range.balanced) {
92 breakPatterns(range.a, range.b, context);
93 range.limit -= 1;
94 }
95
96 // choose a pivot and try guessing whether the slice is already sorted.
97 var pivot: usize = 0;
98 var hint = chosePivot(range.a, range.b, &pivot, context);
99
100 if (hint == .decreasing) {
101 // The maximum number of swaps was performed, so items are likely
102 // in reverse order. Reverse it to make sorting faster.
103 reverseRange(range.a, range.b, context);
104 pivot = (range.b - 1) - (pivot - range.a);
105 hint = .increasing;
106 }
107
108 // if the last partitioning was decently balanced and didn't shuffle elements, and if pivot
109 // selection predicts the slice is likely already sorted...
110 if (range.balanced and range.partitioned and hint == .increasing) {
111 // try identifying several out-of-order elements and shifting them to correct
112 // positions. If the slice ends up being completely sorted, we're done.
113 if (partialInsertionSort(range.a, range.b, context)) break;
114 }
115
116 // if the chosen pivot is equal to the predecessor, then it's the smallest element in the
117 // slice. Partition the slice into elements equal to and elements greater than the pivot.
118 // This case is usually hit when the slice contains many duplicate elements.
119 if (range.a > a and !context.lessThan(range.a - 1, pivot)) {
120 range.a = partitionEqual(range.a, range.b, pivot, context);
121 continue;
122 }
123
124 // partition the slice.
125 var mid = pivot;
126 const was_partitioned = partition(range.a, range.b, &mid, context);
127
128 const left_len = mid - range.a;
129 const right_len = range.b - (mid + 1);
130 const balanced_threshold = len / 8;
131
132 const left_is_smaller = left_len < right_len;
133
134 const smaller_len = if (left_is_smaller) left_len else right_len;
135 const was_balanced = smaller_len >= balanced_threshold;
136
137 const smaller_start_offset = if (left_is_smaller) range.a else mid + 1;
138 const larger_start_offset = if (left_is_smaller) mid + 1 else range.a;
139 const smaller_end_exclusive_offset = if (left_is_smaller) mid else range.b;
140 const larger_end_exclusive_offset = if (left_is_smaller) range.b else mid;
141
142 const smaller_is_leftmost = if (left_is_smaller) range.leftmost else false;
143 const larger_is_leftmost = if (left_is_smaller) false else range.leftmost;
144
145 // defer sorting the larger range until later to ensure stack usage is always less than log_2(n):
146 // as if we always push more than half the range to the stack then each time we push to the stack
147 // we reduce the amount of items we can push to it in later iterations by at least n/2
148 // therefore the count of items on the stack can never be more than log_2(n)
149 stack[top] = .{
150 .a = larger_start_offset,
151 .b = larger_end_exclusive_offset,
152 .limit = range.limit,
153 .leftmost = larger_is_leftmost,
154 .balanced = was_balanced,
155 .partitioned = was_partitioned,
156 };
157 top += 1;
158
159 // sort the smaller range immediately
160 range.a = smaller_start_offset;
161 range.b = smaller_end_exclusive_offset;
162 range.leftmost = smaller_is_leftmost;
163 range.balanced = true; // this either already true, or the range is small so we don't care
164 range.partitioned = was_partitioned;
165 }
166
167 top = math.sub(usize, top, 1) catch break;
168 range = stack[top];
169 }
170}
171
172/// Insertion sort that assumes `items[a-1]` exists and is <= all elements in `[a, b)`,
173/// allowing the inner loop to skip the bounds check.
174fn unguardedInsertionContext(a: usize, b: usize, context: anytype) void {
175 var i = a + 1;
176 while (i < b) : (i += 1) {
177 var j = i;
178 while (context.lessThan(j, j - 1)) : (j -= 1) {
179 context.swap(j, j - 1);
180 }
181 }
182}
183
184/// partitions `items[a..b]` into elements smaller than `items[pivot]`,
185/// followed by elements greater than or equal to `items[pivot]`.
186///
187/// sets the new pivot.
188/// returns `true` if already partitioned.
189fn partition(a: usize, b: usize, pivot: *usize, context: anytype) bool {
190 // move pivot to the first place
191 context.swap(a, pivot.*);
192
193 var i = a + 1;
194 var j = b - 1;
195
196 while (i <= j and context.lessThan(i, a)) i += 1;
197 while (i <= j and !context.lessThan(j, a)) j -= 1;
198
199 // check if items are already partitioned (no item to swap)
200 if (i > j) {
201 // put pivot back to the middle
202 context.swap(j, a);
203 pivot.* = j;
204 return true;
205 }
206
207 context.swap(i, j);
208 i += 1;
209 j -= 1;
210
211 const block_size = 64;
212 var offsets_l: [block_size]u8 align(std.atomic.cache_line) = undefined;
213 var offsets_r: [block_size]u8 align(std.atomic.cache_line) = undefined;
214
215 var offsets_l_base = i;
216 var offsets_r_base = j;
217 var num_l: usize = 0;
218 var num_r: usize = 0;
219 var start_l: usize = 0;
220 var start_r: usize = 0;
221
222 while (i <= j) {
223 const num_unknown = j + 1 - i;
224 const left_split = if (num_l == 0)
225 @min(block_size, if (num_r == 0) num_unknown / 2 else num_unknown)
226 else
227 0;
228 const right_split = if (num_r == 0)
229 @min(block_size, num_unknown - left_split)
230 else
231 0;
232
233 for (0..left_split) |k| {
234 offsets_l[num_l] = @intCast(k);
235 num_l += @intFromBool(!context.lessThan(i + k, a));
236 }
237 i += left_split;
238
239 for (0..right_split) |k| {
240 offsets_r[num_r] = @intCast(k);
241 num_r += @intFromBool(context.lessThan(j - k, a));
242 }
243 j -= right_split;
244
245 const num = @min(num_l, num_r);
246 for (0..num) |m| {
247 context.swap(
248 offsets_l_base + offsets_l[start_l + m],
249 offsets_r_base - offsets_r[start_r + m],
250 );
251 }
252 num_l -= num;
253 num_r -= num;
254 start_l += num;
255 start_r += num;
256
257 if (num_l == 0) {
258 start_l = 0;
259 offsets_l_base = i;
260 }
261 if (num_r == 0) {
262 start_r = 0;
263 offsets_r_base = j;
264 }
265 }
266
267 if (num_l > 0) {
268 while (num_l > 0) {
269 num_l -= 1;
270 context.swap(offsets_l_base + offsets_l[start_l + num_l], j);
271 j -= 1;
272 }
273 i = j + 1;
274 }
275 if (num_r > 0) {
276 while (num_r > 0) {
277 num_r -= 1;
278 context.swap(offsets_r_base - offsets_r[start_r + num_r], i);
279 i += 1;
280 }
281 j = i - 1;
282 }
283
284 context.swap(j, a);
285 pivot.* = j;
286 return false;
287}
288
289/// partitions items into elements equal to `items[pivot]`
290/// followed by elements greater than `items[pivot]`.
291///
292/// it assumed that `items[a..b]` does not contain elements smaller than the `items[pivot]`.
293fn partitionEqual(a: usize, b: usize, pivot: usize, context: anytype) usize {
294 // move pivot to the first place
295 context.swap(a, pivot);
296
297 var i = a + 1;
298 var j = b - 1;
299
300 while (true) {
301 while (i <= j and !context.lessThan(a, i)) i += 1;
302 while (i <= j and context.lessThan(a, j)) j -= 1;
303 if (i > j) break;
304
305 context.swap(i, j);
306 i += 1;
307 j -= 1;
308 }
309
310 return i;
311}
312
313/// partially sorts a slice by shifting several out-of-order elements around.
314///
315/// returns `true` if the slice is sorted at the end. This function is `O(n)` worst-case.
316fn partialInsertionSort(a: usize, b: usize, context: anytype) bool {
317 @branchHint(.cold);
318
319 // maximum number of adjacent out-of-order pairs that will get shifted
320 const max_steps = 5;
321 // if the slice is shorter than this, don't shift any elements
322 const shortest_shifting = 50;
323
324 var i = a + 1;
325 for (0..max_steps) |_| {
326 // find the next pair of adjacent out-of-order elements.
327 while (i < b and !context.lessThan(i, i - 1)) i += 1;
328
329 // are we done?
330 if (i == b) return true;
331
332 // don't shift elements on short arrays, that has a performance cost.
333 if (b - a < shortest_shifting) return false;
334
335 // swap the found pair of elements. This puts them in correct order.
336 context.swap(i, i - 1);
337
338 // shift the smaller element to the left.
339 if (i - a >= 2) {
340 var j = i - 1;
341 while (j > a) : (j -= 1) {
342 if (!context.lessThan(j, j - 1)) break;
343 context.swap(j, j - 1);
344 }
345 }
346
347 // shift the greater element to the right.
348 if (b - i >= 2) {
349 var j = i + 1;
350 while (j < b) : (j += 1) {
351 if (!context.lessThan(j, j - 1)) break;
352 context.swap(j, j - 1);
353 }
354 }
355 }
356
357 return false;
358}
359
360fn breakPatterns(a: usize, b: usize, context: anytype) void {
361 @branchHint(.cold);
362
363 const len = b - a;
364 if (len < 8) return;
365
366 var rand = @as(u64, @intCast(len));
367 const modulus = math.ceilPowerOfTwoAssert(u64, len);
368
369 var i = a + (len / 4) * 2 - 1;
370 while (i <= a + (len / 4) * 2 + 1) : (i += 1) {
371 // xorshift64
372 rand ^= rand << 13;
373 rand ^= rand >> 7;
374 rand ^= rand << 17;
375
376 var other = @as(usize, @intCast(rand & (modulus - 1)));
377 if (other >= len) other -= len;
378 context.swap(i, a + other);
379 }
380}
381
382/// chooses a pivot in `items[a..b]`.
383/// swaps likely_sorted when `items[a..b]` seems to be already sorted.
384fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
385 // minimum length for using the Tukey's ninther method
386 const shortest_ninther = 50;
387 // max_swaps is the maximum number of swaps allowed in this function
388 const max_swaps = 4 * 3;
389
390 const len = b - a;
391 const i = a + len / 4 * 1;
392 const j = a + len / 4 * 2;
393 const k = a + len / 4 * 3;
394 var swaps: usize = 0;
395
396 if (len >= 8) {
397 if (len >= shortest_ninther) {
398 // find medians in the neighborhoods of `i`, `j` and `k`
399 sort3(i - 1, i, i + 1, &swaps, context);
400 sort3(j - 1, j, j + 1, &swaps, context);
401 sort3(k - 1, k, k + 1, &swaps, context);
402 }
403
404 // find the median among `i`, `j` and `k` and stores it in `j`
405 sort3(i, j, k, &swaps, context);
406 }
407
408 pivot.* = j;
409 return switch (swaps) {
410 0 => .increasing,
411 max_swaps => .decreasing,
412 else => .unknown,
413 };
414}
415
416fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) void {
417 if (context.lessThan(b, a)) {
418 swaps.* += 1;
419 context.swap(b, a);
420 }
421
422 if (context.lessThan(c, b)) {
423 swaps.* += 1;
424 context.swap(c, b);
425 }
426
427 if (context.lessThan(b, a)) {
428 swaps.* += 1;
429 context.swap(b, a);
430 }
431}
432
433fn reverseRange(a: usize, b: usize, context: anytype) void {
434 var i = a;
435 var j = b - 1;
436 while (i < j) {
437 context.swap(i, j);
438 i += 1;
439 j -= 1;
440 }
441}
442
443test "pdqContext respects arbitrary range boundaries" {
444 // Regression test for issue #25250
445 // pdqsort should never access indices outside the specified [a, b) range
446 var data: [2000]i32 = @splat(0);
447
448 // Fill with data that triggers the partialInsertionSort path
449 for (0..data.len) |i| {
450 data[i] = @intCast(@mod(@as(i32, @intCast(i)) * 7, 100));
451 }
452
453 const TestContext = struct {
454 items: []i32,
455 range_start: usize,
456 range_end: usize,
457
458 pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
459 // Assert indices are within the expected range
460 testing.expect(a >= ctx.range_start and a < ctx.range_end) catch @panic("index a out of range");
461 testing.expect(b >= ctx.range_start and b < ctx.range_end) catch @panic("index b out of range");
462 return ctx.items[a] < ctx.items[b];
463 }
464
465 pub fn swap(ctx: @This(), a: usize, b: usize) void {
466 // Assert indices are within the expected range
467 testing.expect(a >= ctx.range_start and a < ctx.range_end) catch @panic("index a out of range");
468 testing.expect(b >= ctx.range_start and b < ctx.range_end) catch @panic("index b out of range");
469 mem.swap(i32, &ctx.items[a], &ctx.items[b]);
470 }
471 };
472
473 // Test sorting a sub-range that doesn't start at 0
474 const start = 1118;
475 const end = 1764;
476 const ctx = TestContext{
477 .items = &data,
478 .range_start = start,
479 .range_end = end,
480 };
481
482 pdqContext(start, end, ctx);
483
484 // Verify the range is sorted
485 for ((start + 1)..end) |i| {
486 try testing.expect(data[i - 1] <= data[i]);
487 }
488}