authorgravatar for Validark@pm.meNiles Salter <Validark@pm.me> 2023-06-22 11:32:28-06:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-06-22 17:32:28+00:00
log7d511d642845c860ec212b9ee39e371d3f70a68b
treee4e58d563eaff0c35b1b56a734914d0761b5e1d2
parentc60896743d3b0776a44ac63582f51b6e64892528
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

[heapsort] Protect against integer overflow

(Firstly, I changed `n` to `b`, as that is less confusing. It's not a length, it's a right boundary.) The invariant maintained is `cur < b`. In the worst case `2*cur + 1` results in a maximum of `2b`. Since `2b` is not guaranteed to be lower than `maxInt`, we have to add one overflow check to `siftDown` to make sure we avoid undefined behavior. LLVM also seems to have a nicer time compiling this version of the function. It is about 2x faster in my tests (I think LLVM was stumped by the `child += @intFromBool` line), and adding/removing the overflow check has a negligible performance difference on my machine. Of course, we could check `2b <= maxInt` in the parent function, and dispatch to a version of the function without the overflow check in the common case, but that probably is not worth the code size just to eliminate a single instruction.

1 files changed, 25 insertions(+), 11 deletions(-)

lib/std/sort.zig+25-11
......@@ -36,6 +36,8 @@ pub fn insertion(
3636/// O(1) memory (no allocator required).
3737/// Sorts in ascending order with respect to the given `lessThan` function.
3838pub fn insertionContext(a: usize, b: usize, context: anytype) void {
39 assert(a <= b);
40
3941 var i = a + 1;
4042 while (i < b) : (i += 1) {
4143 var j = i;
......@@ -73,6 +75,7 @@ pub fn heap(
7375/// O(1) memory (no allocator required).
7476/// Sorts in ascending order with respect to the given `lessThan` function.
7577pub fn heapContext(a: usize, b: usize, context: anytype) void {
78 assert(a <= b);
7679 // build the heap in linear time.
7780 var i = a + (b - a) / 2;
7881 while (i > a) {
......@@ -89,22 +92,33 @@ pub fn heapContext(a: usize, b: usize, context: anytype) void {
8992 }
9093}
9194
92fn siftDown(a: usize, root: usize, n: usize, context: anytype) void {
93 var node = root;
95fn siftDown(a: usize, target: usize, b: usize, context: anytype) void {
96 var cur = target;
9497 while (true) {
95 var child = a + 2 * (node - a) + 1;
96 if (child >= n) break;
98 // When we don't overflow from the multiply below, the following expression equals (2*cur) - (2*a) + a + 1
99 // The `+ a + 1` is safe because:
100 // for `a > 0` then `2a >= a + 1`.
101 // for `a = 0`, the expression equals `2*cur+1`. `2*cur` is an even number, therefore adding 1 is safe.
102 var child = (math.mul(usize, cur - a, 2) catch break) + a + 1;
103
104 // stop if we overshot the boundary
105 if (!(child < b)) break;
97106
98 // choose the greater child.
99 child += @intFromBool(child + 1 < n and context.lessThan(child, child + 1));
107 // `next_child` is at most `b`, therefore no overflow is possible
108 const next_child = child + 1;
109
110 // store the greater child in `child`
111 if (next_child < b and context.lessThan(child, next_child)) {
112 child = next_child;
113 }
100114
101 // stop if the invariant holds at `node`.
102 if (!context.lessThan(node, child)) break;
115 // stop if the Heap invariant holds at `cur`.
116 if (context.lessThan(child, cur)) break;
103117
104 // swap `node` with the greater child,
118 // swap `cur` with the greater child,
105119 // move one step down, and continue sifting.
106 context.swap(node, child);
107 node = child;
120 context.swap(child, cur);
121 cur = child;
108122 }
109123}
110124