authorgravatar for Validark@pm.meNiles Salter <Validark@pm.me> 2023-06-13 14:55:58-06:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-06-13 16:55:58-04:00
log700ea694b293565ececb7571e4d9613d2c143ca6
tree309bcb58d4892e8b38fdee6a972ffc1788a2e4ac
parent129afba460bd654230c8d678a53d7a1b5fec9132
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

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

2 files changed, 71 insertions(+), 23 deletions(-)

lib/std/sort.zig+63-13
......@@ -74,29 +74,29 @@ pub fn heap(
7474/// Sorts in ascending order with respect to the given `lessThan` function.
7575pub fn heapContext(a: usize, b: usize, context: anytype) void {
7676 // build the heap in linear time.
77 var i = b / 2;
78 while (i > a) : (i -= 1) {
79 siftDown(i - 1, b, context);
77 var i = a + (b - a) / 2;
78 while (i > a) {
79 i -= 1;
80 siftDown(a, i, b, context);
8081 }
8182
8283 // pop maximal elements from the heap.
8384 i = b;
84 while (i > a) : (i -= 1) {
85 context.swap(a, i - 1);
86 siftDown(a, i - 1, context);
85 while (i > a) {
86 i -= 1;
87 context.swap(a, i);
88 siftDown(a, a, i, context);
8789 }
8890}
8991
90fn siftDown(root: usize, n: usize, context: anytype) void {
92fn siftDown(a: usize, root: usize, n: usize, context: anytype) void {
9193 var node = root;
9294 while (true) {
93 var child = 2 * node + 1;
95 var child = a + 2 * (node - a) + 1;
9496 if (child >= n) break;
9597
9698 // choose the greater child.
97 if (child + 1 < n and context.lessThan(child, child + 1)) {
98 child += 1;
99 }
99 child += @boolToInt(child + 1 < n and context.lessThan(child, child + 1));
100100
101101 // stop if the invariant holds at `node`.
102102 if (!context.lessThan(node, child)) break;
......@@ -138,6 +138,13 @@ const sort_funcs = &[_]fn (comptime type, anytype, anytype, comptime anytype) vo
138138 heap,
139139};
140140
141const context_sort_funcs = &[_]fn (usize, usize, anytype) void{
142 // blockContext,
143 pdqContext,
144 insertionContext,
145 heapContext,
146};
147
141148const IdAndValue = struct {
142149 id: usize,
143150 value: i32,
......@@ -248,11 +255,15 @@ test "sort" {
248255 &[_]i32{ 2, 1, 3 },
249256 &[_]i32{ 1, 2, 3 },
250257 },
258 &[_][]const i32{
259 &[_]i32{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 55, 32, 39, 58, 21, 88, 43, 22, 59 },
260 &[_]i32{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 21, 22, 32, 39, 43, 55, 58, 59, 88 },
261 },
251262 };
252263
253264 inline for (sort_funcs) |sortFn| {
254265 for (u8cases) |case| {
255 var buf: [8]u8 = undefined;
266 var buf: [20]u8 = undefined;
256267 const slice = buf[0..case[0].len];
257268 @memcpy(slice, case[0]);
258269 sortFn(u8, slice, {}, asc_u8);
......@@ -260,7 +271,7 @@ test "sort" {
260271 }
261272
262273 for (i32cases) |case| {
263 var buf: [8]i32 = undefined;
274 var buf: [20]i32 = undefined;
264275 const slice = buf[0..case[0].len];
265276 @memcpy(slice, case[0]);
266277 sortFn(i32, slice, {}, asc_i32);
......@@ -308,6 +319,45 @@ test "sort descending" {
308319 }
309320}
310321
322test "sort with context in the middle of a slice" {
323 const Context = struct {
324 items: []i32,
325
326 pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
327 return ctx.items[a] < ctx.items[b];
328 }
329
330 pub fn swap(ctx: @This(), a: usize, b: usize) void {
331 return mem.swap(i32, &ctx.items[a], &ctx.items[b]);
332 }
333 };
334
335 const i32cases = [_][]const []const i32{
336 &[_][]const i32{
337 &[_]i32{ 0, 1, 8, 3, 6, 5, 4, 2, 9, 7, 10, 55, 32, 39, 58, 21, 88, 43, 22, 59 },
338 &[_]i32{ 50, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 21, 22, 32, 39, 43, 55, 58, 59, 88 },
339 },
340 };
341
342 const ranges = [_]struct { start: usize, end: usize }{
343 .{ .start = 10, .end = 20 },
344 .{ .start = 1, .end = 11 },
345 .{ .start = 3, .end = 7 },
346 };
347
348 inline for (context_sort_funcs) |sortFn| {
349 for (i32cases) |case| {
350 for (ranges) |range| {
351 var buf: [20]i32 = undefined;
352 const slice = buf[0..case[0].len];
353 @memcpy(slice, case[0]);
354 sortFn(range.start, range.end, Context{ .items = slice });
355 try testing.expectEqualSlices(i32, slice[range.start..range.end], case[1][range.start..range.end]);
356 }
357 }
358 }
359}
360
311361test "sort fuzz testing" {
312362 var prng = std.rand.DefaultPrng.init(0x12345678);
313363 const random = prng.random();
lib/std/sort/pdq.zig+8-10
......@@ -43,7 +43,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
4343 // slices of up to this length get sorted using insertion sort.
4444 const max_insertion = 24;
4545 // number of allowed imbalanced partitions before switching to heap sort.
46 const max_limit = std.math.floorPowerOfTwo(usize, b) + 1;
46 const max_limit = std.math.floorPowerOfTwo(usize, b - a) + 1;
4747
4848 // set upper bound on stack memory usage.
4949 const Range = struct { a: usize, b: usize, limit: usize };
......@@ -100,7 +100,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
100100 // if the chosen pivot is equal to the predecessor, then it's the smallest element in the
101101 // slice. Partition the slice into elements equal to and elements greater than the pivot.
102102 // This case is usually hit when the slice contains many duplicate elements.
103 if (range.a > 0 and !context.lessThan(range.a - 1, pivot)) {
103 if (range.a > a and !context.lessThan(range.a - 1, pivot)) {
104104 range.a = partitionEqual(range.a, range.b, pivot, context);
105105 continue;
106106 }
......@@ -284,13 +284,13 @@ fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
284284 if (len >= 8) {
285285 if (len >= shortest_ninther) {
286286 // find medians in the neighborhoods of `i`, `j` and `k`
287 i = sort3(i - 1, i, i + 1, &swaps, context);
288 j = sort3(j - 1, j, j + 1, &swaps, context);
289 k = sort3(k - 1, k, k + 1, &swaps, context);
287 sort3(i - 1, i, i + 1, &swaps, context);
288 sort3(j - 1, j, j + 1, &swaps, context);
289 sort3(k - 1, k, k + 1, &swaps, context);
290290 }
291291
292 // find the median among `i`, `j` and `k`
293 j = sort3(i, j, k, &swaps, context);
292 // find the median among `i`, `j` and `k` and stores it in `j`
293 sort3(i, j, k, &swaps, context);
294294 }
295295
296296 pivot.* = j;
......@@ -301,7 +301,7 @@ fn chosePivot(a: usize, b: usize, pivot: *usize, context: anytype) Hint {
301301 };
302302}
303303
304fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) usize {
304fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) void {
305305 if (context.lessThan(b, a)) {
306306 swaps.* += 1;
307307 context.swap(b, a);
......@@ -316,8 +316,6 @@ fn sort3(a: usize, b: usize, c: usize, swaps: *usize, context: anytype) usize {
316316 swaps.* += 1;
317317 context.swap(b, a);
318318 }
319
320 return b;
321319}
322320
323321fn reverseRange(a: usize, b: usize, context: anytype) void {