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(...@@ -36,6 +36,8 @@ pub fn insertion(
36/// O(1) memory (no allocator required).36/// O(1) memory (no allocator required).
37/// Sorts in ascending order with respect to the given `lessThan` function.37/// Sorts in ascending order with respect to the given `lessThan` function.
38pub fn insertionContext(a: usize, b: usize, context: anytype) void {38pub fn insertionContext(a: usize, b: usize, context: anytype) void {
39 assert(a <= b);
40
39 var i = a + 1;41 var i = a + 1;
40 while (i < b) : (i += 1) {42 while (i < b) : (i += 1) {
41 var j = i;43 var j = i;
...@@ -73,6 +75,7 @@ pub fn heap(...@@ -73,6 +75,7 @@ pub fn heap(
73/// O(1) memory (no allocator required).75/// O(1) memory (no allocator required).
74/// Sorts in ascending order with respect to the given `lessThan` function.76/// Sorts in ascending order with respect to the given `lessThan` function.
75pub fn heapContext(a: usize, b: usize, context: anytype) void {77pub fn heapContext(a: usize, b: usize, context: anytype) void {
78 assert(a <= b);
76 // build the heap in linear time.79 // build the heap in linear time.
77 var i = a + (b - a) / 2;80 var i = a + (b - a) / 2;
78 while (i > a) {81 while (i > a) {
...@@ -89,22 +92,33 @@ pub fn heapContext(a: usize, b: usize, context: anytype) void {...@@ -89,22 +92,33 @@ pub fn heapContext(a: usize, b: usize, context: anytype) void {
89 }92 }
90}93}
9194
92fn siftDown(a: usize, root: usize, n: usize, context: anytype) void {95fn siftDown(a: usize, target: usize, b: usize, context: anytype) void {
93 var node = root;96 var cur = target;
94 while (true) {97 while (true) {
95 var child = a + 2 * (node - a) + 1;98 // When we don't overflow from the multiply below, the following expression equals (2*cur) - (2*a) + a + 1
96 if (child >= n) break;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.107 // `next_child` is at most `b`, therefore no overflow is possible
99 child += @intFromBool(child + 1 < n and context.lessThan(child, child + 1));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`.115 // stop if the Heap invariant holds at `cur`.
102 if (!context.lessThan(node, child)) break;116 if (context.lessThan(child, cur)) break;
103117
104 // swap `node` with the greater child,118 // swap `cur` with the greater child,
105 // move one step down, and continue sifting.119 // move one step down, and continue sifting.
106 context.swap(node, child);120 context.swap(child, cur);
107 node = child;121 cur = child;
108 }122 }
109}123}
110124