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