big.int.depositBits: fix handling of different limb counts

Author: bnuuydev

lib/std/math/big/int.zig

@@ -1763,6 +1763,8 @@ pub const Mutable = struct {
     /// result = @depositBits(source, mask)
     ///
     /// Asserts that `source` and `mask` are positive
+    /// The value in `result` may use the same number of or less limbs than `mask`.
+    /// `result` is assumed to have sufficient length to store the result.
     pub fn depositBits(result: *Mutable, source: Const, mask: Const) void {
         assert(source.positive);
         assert(mask.positive);
@@ -1771,10 +1773,14 @@ pub const Mutable = struct {
         @memset(result.limbs, 0);
 
         var shift: usize = 0;
-        for (mask.limbs, result.limbs) |mask_limb, *result_limb| {
+        for (mask.limbs, 0..) |mask_limb, i| {
             const shift_bits: Log2Limb = @intCast(shift % limb_bits);
             const shift_limbs = shift / limb_bits;
 
+            if (shift_limbs >= source.limbs.len) break;
+
+            const result_limb = &result.limbs[i];
+
             var source_limb = source.limbs[shift_limbs] >> shift_bits;
             if (shift_bits != 0 and shift_limbs + 1 < source.limbs.len) {
                 source_limb += source.limbs[shift_limbs + 1] << @intCast(limb_bits - shift_bits);
@@ -1793,15 +1799,19 @@ pub const Mutable = struct {
     /// result = @extractBits(source, mask)
     ///
     /// Asserts that `source` and `mask` are positive
+    /// The value in `result` may use the same number of or less limbs than `mask`.
+    /// `result` is assumed to have sufficient length to store the result.
     pub fn extractBits(result: *Mutable, source: Const, mask: Const) void {
         assert(source.positive);
         assert(mask.positive);
 
         result.positive = true;
         @memset(result.limbs, 0);
 
+        const len = @min(source.limbs.len, mask.limbs.len);
+
         var shift: usize = 0;
-        for (source.limbs, mask.limbs) |source_limb, mask_limb| {
+        for (source.limbs[0..len], mask.limbs[0..len]) |source_limb, mask_limb| {
             const pext_limb = @extractBits(source_limb, mask_limb);
             const shift_bits: Log2Limb = @intCast(shift % limb_bits);
             const shift_limbs = shift / limb_bits;

lib/std/math/big/int_test.zig

@@ -2968,6 +2968,9 @@ test "big int extractBits" {
 
     try extractBitsTest(0x12345678_90123456_78901234_56789012, 0xff << 64, 0x56);
     try extractBitsTest(0x12345678_90123456_78901234_56789012, (0xff << 64) | 0xff00f, 0x56892);
+
+    try extractBitsTest(0x12345678_90123456_78901234_56789012, 0xf0f0, 0x91);
+    try extractBitsTest(0x12345678_90123456, 0xffffffff_ffffffff, 0x12345678_90123456);
 }
 
 fn extractBitsTest(comptime source: comptime_int, comptime mask: comptime_int, comptime expected: comptime_int) !void {
@@ -2992,6 +2995,8 @@ test "big int depositBits" {
 
     try depositBitsTest(0x1234, 0xff << 64, 0x34_00000000_00000000);
     try depositBitsTest(0x12345678, (0xff << 64) | 0xff00f, 0x45_00000000_00067008);
+
+    try depositBitsTest(0x0, 0xff_ffffffff_ffffffff, 0x0);
 }
 
 fn depositBitsTest(comptime source: comptime_int, comptime mask: comptime_int, comptime expected: comptime_int) !void {