Commit 2b99b04285

Tau <jonathan.haehne@hotmail.com>
2024-07-11 13:57:37
Fix right shift on negative BigInts
Closes #17662.
1 parent 3ad81c4
Changed files (2)
lib
std
lib/std/math/big/int.zig
@@ -1173,7 +1173,9 @@ pub const Mutable = struct {
     /// Asserts there is enough memory to fit the result. The upper bound Limb count is
     /// `a.limbs.len - (shift / (@sizeOf(Limb) * 8))`.
     pub fn shiftRight(r: *Mutable, a: Const, shift: usize) void {
-        if (a.limbs.len <= shift / limb_bits) {
+        const full_limbs_shifted_out = shift / limb_bits;
+        const remaining_bits_shifted_out = shift % limb_bits;
+        if (a.limbs.len <= full_limbs_shifted_out) {
             // Shifting negative numbers converges to -1 instead of 0
             if (a.positive) {
                 r.len = 1;
@@ -1186,14 +1188,29 @@ pub const Mutable = struct {
             }
             return;
         }
+        const nonzero_negative_shiftout = if (a.positive) false else nonzero: {
+            for (a.limbs[0..full_limbs_shifted_out]) |x| {
+                if (x != 0)
+                    break :nonzero true;
+            }
+            if (remaining_bits_shifted_out == 0)
+                break :nonzero false;
+            const not_covered: Log2Limb = @intCast(limb_bits - remaining_bits_shifted_out);
+            break :nonzero a.limbs[full_limbs_shifted_out] << not_covered != 0;
+        };
 
         llshr(r.limbs[0..], a.limbs[0..a.limbs.len], shift);
-        r.normalize(a.limbs.len - (shift / limb_bits));
-        r.positive = a.positive;
-        // Shifting negative numbers converges to -1 instead of 0
-        if (!r.positive and r.len == 1 and r.limbs[0] == 0) {
-            r.limbs[0] = 1;
+
+        r.len = a.limbs.len - full_limbs_shifted_out;
+        if (nonzero_negative_shiftout) {
+            if (full_limbs_shifted_out > 0) {
+                r.limbs[a.limbs.len - full_limbs_shifted_out] = 0;
+                r.len += 1;
+            }
+            r.addScalar(r.toConst(), -1);
         }
+        r.normalize(r.len);
+        r.positive = a.positive;
     }
 
     /// r = ~a under 2s complement wrapping semantics.
lib/std/math/big/int_test.zig
@@ -2066,6 +2066,23 @@ test "shift-right negative" {
     defer arg3.deinit();
     try a.shiftRight(&arg3, 1232);
     try testing.expect((try a.to(i32)) == -1); // -10 >> 1232 == -1
+
+    var arg4 = try Managed.initSet(testing.allocator, -5);
+    defer arg4.deinit();
+    try a.shiftRight(&arg4, 2);
+    try testing.expect(try a.to(i32) == -2); // -5 >> 2 == -2
+
+    var arg5 = try Managed.initSet(testing.allocator, -0xffff0000eeee1111dddd2222cccc3333);
+    defer arg5.deinit();
+    try a.shiftRight(&arg5, 67);
+    try testing.expect(try a.to(i64) == -0x1fffe0001dddc223);
+
+    var arg6 = try Managed.initSet(testing.allocator, -0x1ffffffffffffffff);
+    defer arg6.deinit();
+    try a.shiftRight(&arg6, 1);
+    try a.shiftRight(&a, 1);
+    a.setSign(true);
+    try testing.expect(try a.to(u64) == 0x8000000000000000);
 }
 
 test "sat shift-left simple unsigned" {