authorgravatar for 2762638a@student.gla.ac.ukPaul Anderson <2762638a@student.gla.ac.uk> 2026-09-01 00:07:30+01:00
committergravatar for alichraghi@noreply.codeberg.orgAli Cheraghi <alichraghi@noreply.codeberg.org> 2026-09-05 17:50:56+02:00
logaf24fd11a5c37ef5531ddb86356cd6c55ed71d77
treef6468929ad369fa6822e8f50c14d242220daac0f
parent1dd182d24b1b0da00ef0803898024734204dca5c

std.sort.pdq: push larger partition to stack to avoid stack overflow

1 files changed, 57 insertions(+), 23 deletions(-)

lib/std/sort/pdq.zig+57-23
......@@ -35,6 +35,15 @@ const Hint = enum {
3535 unknown,
3636};
3737
38const Range = struct {
39 a: usize,
40 b: usize,
41 limit: usize,
42 leftmost: bool,
43 balanced: bool,
44 partitioned: bool,
45};
46
3847/// Unstable in-place sort. O(n) best case, O(n*log(n)) worst case and average case.
3948/// O(log(n)) memory (no allocator required).
4049/// `context` must have methods `swap` and `lessThan`,
......@@ -46,16 +55,19 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
4655 // number of allowed imbalanced partitions before switching to heap sort.
4756 const max_limit = if (b > a) math.log2_int(usize, b - a) else 0;
4857
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 };
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 };
5368 var top: usize = 0;
5469
5570 while (true) {
56 var was_balanced = true;
57 var was_partitioned = true;
58
5971 while (true) {
6072 const len = range.b - range.a;
6173
......@@ -76,7 +88,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
7688
7789 // if the last partitioning was imbalanced, try breaking patterns in the slice by shuffling
7890 // some elements around. Hopefully we'll choose a better pivot this time.
79 if (!was_balanced) {
91 if (!range.balanced) {
8092 breakPatterns(range.a, range.b, context);
8193 range.limit -= 1;
8294 }
......@@ -95,7 +107,7 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
95107
96108 // if the last partitioning was decently balanced and didn't shuffle elements, and if pivot
97109 // selection predicts the slice is likely already sorted...
98 if (was_balanced and was_partitioned and hint == .increasing) {
110 if (range.balanced and range.partitioned and hint == .increasing) {
99111 // try identifying several out-of-order elements and shifting them to correct
100112 // positions. If the slice ends up being completely sorted, we're done.
101113 if (partialInsertionSort(range.a, range.b, context)) break;
......@@ -111,23 +123,45 @@ pub fn pdqContext(a: usize, b: usize, context: anytype) void {
111123
112124 // partition the slice.
113125 var mid = pivot;
114 was_partitioned = partition(range.a, range.b, &mid, context);
126 const was_partitioned = partition(range.a, range.b, &mid, context);
115127
116128 const left_len = mid - range.a;
117 const right_len = range.b - mid;
129 const right_len = range.b - (mid + 1);
118130 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 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;
131165 }
132166
133167 top = math.sub(usize, top, 1) catch break;