diff --git a/lib/std/math/big/int.zig b/lib/std/math/big/int.zig index 1f472acd58..4e814f06a9 100644 --- a/lib/std/math/big/int.zig +++ b/lib/std/math/big/int.zig @@ -835,6 +835,75 @@ pub const Mutable = struct { r.positive = a.positive; } + /// r = a <<| shift with 2s-complement saturating semantics. + /// + /// r and a may alias. + /// + /// Asserts there is enough memory to fit the result. The upper bound Limb count is + /// r is `calcTwosCompLimbCount(bit_count)`. + pub fn shiftLeftSat(r: *Mutable, a: Const, shift: usize, signedness: std.builtin.Signedness, bit_count: usize) void { + // Special case: When the argument is negative, but the result is supposed to be unsigned, + // return 0 in all cases. + if (!a.positive and signedness == .unsigned) { + r.set(0); + return; + } + + // Check whether the shift is going to overflow. This is the case + // when (in 2s complement) any bit above `bit_count - shift` is set in the unshifted value. + // Note, the sign bit is not counted here. + + // Handle shifts larger than the target type. This also deals with + // 0-bit integers. + if (bit_count <= shift) { + // In this case, there is only no overflow if `a` is zero. + if (a.eqZero()) { + r.set(0); + } else { + r.setTwosCompIntLimit(if (a.positive) .max else .min, signedness, bit_count); + } + return; + } + + const checkbit = bit_count - shift - @boolToInt(signedness == .signed); + // If `checkbit` and more significant bits are zero, no overflow will take place. + + if (checkbit >= a.limbs.len * limb_bits) { + // `checkbit` is outside the range of a, so definitely no overflow will take place. We + // can defer to a normal shift. + // Note that if `a` is normalized (which we assume), this checks for set bits in the upper limbs. + + // Note, in this case r should already have enough limbs required to perform the normal shift. + // In this case the shift of the most significant limb may still overflow. + r.shiftLeft(a, shift); + return; + } else if (checkbit < (a.limbs.len - 1) * limb_bits) { + // `checkbit` is not in the most significant limb. If `a` is normalized the most significant + // limb will not be zero, so in this case we need to saturate. Note that `a.limbs.len` must be + // at least one according to normalization rules. + + r.setTwosCompIntLimit(if (a.positive) .max else .min, signedness, bit_count); + return; + } + + // Generate a mask with the bits to check in the most signficant limb. We'll need to check + // all bits with equal or more significance than checkbit. + // const msb = @truncate(Log2Limb, checkbit); + // const checkmask = (@as(Limb, 1) << msb) -% 1; + + if (a.limbs[a.limbs.len - 1] >> @truncate(Log2Limb, checkbit) != 0) { + // Need to saturate. + r.setTwosCompIntLimit(if (a.positive) .max else .min, signedness, bit_count); + return; + } + + // This shift should not be able to overflow, so invoke llshl and normalize manually + // to avoid the extra required limb. + llshl(r.limbs[0..], a.limbs[0..a.limbs.len], shift); + r.normalize(a.limbs.len + (shift / limb_bits)); + r.positive = a.positive; + } + /// r = a >> shift /// r and a may alias. /// @@ -2401,6 +2470,14 @@ pub const Managed = struct { r.setMetadata(m.positive, m.len); } + /// r = a <<| shift with 2s-complement saturating semantics. + pub fn shiftLeftSat(r: *Managed, a: Managed, shift: usize, signedness: std.builtin.Signedness, bit_count: usize) !void { + try r.ensureTwosCompCapacity(bit_count); + var m = r.toMutable(); + m.shiftLeftSat(a.toConst(), shift, signedness, bit_count); + r.setMetadata(m.positive, m.len); + } + /// r = a >> shift pub fn shiftRight(r: *Managed, a: Managed, shift: usize) !void { if (a.len() <= shift / limb_bits) { @@ -2949,10 +3026,18 @@ fn lldiv1(quo: []Limb, rem: *Limb, a: []const Limb, b: Limb) void { fn llshl(r: []Limb, a: []const Limb, shift: usize) void { @setRuntimeSafety(debug_safety); assert(a.len >= 1); - assert(r.len >= a.len + (shift / limb_bits) + 1); + + const interior_limb_shift = @truncate(Log2Limb, shift); + + // We only need the extra limb if the shift of the last element overflows. + // This is useful for the implementation of `shiftLeftSat`. + if (a[a.len - 1] << interior_limb_shift >> interior_limb_shift != a[a.len - 1]) { + assert(r.len >= a.len + (shift / limb_bits) + 1); + } else { + assert(r.len >= a.len + (shift / limb_bits)); + } const limb_shift = shift / limb_bits + 1; - const interior_limb_shift = @intCast(Log2Limb, shift % limb_bits); var carry: Limb = 0; var i: usize = 0; @@ -2979,7 +3064,7 @@ fn llshr(r: []Limb, a: []const Limb, shift: usize) void { assert(r.len >= a.len - (shift / limb_bits)); const limb_shift = shift / limb_bits; - const interior_limb_shift = @intCast(Log2Limb, shift % limb_bits); + const interior_limb_shift = @truncate(Log2Limb, shift); var carry: Limb = 0; var i: usize = 0; diff --git a/lib/std/math/big/int_test.zig b/lib/std/math/big/int_test.zig index ef3d61e616..fb7b4bd4ad 100644 --- a/lib/std/math/big/int_test.zig +++ b/lib/std/math/big/int_test.zig @@ -1773,6 +1773,92 @@ test "big.int shift-left negative" { try testing.expect((try a.to(i32)) == -10 >> 1232); } +test "big.int sat shift-left simple unsigned" { + var a = try Managed.initSet(testing.allocator, 0xffff); + defer a.deinit(); + try a.shiftLeftSat(a, 16, .unsigned, 21); + + try testing.expect((try a.to(u64)) == 0x1fffff); +} + +test "big.int sat shift-left simple unsigned no sat" { + var a = try Managed.initSet(testing.allocator, 1); + defer a.deinit(); + try a.shiftLeftSat(a, 16, .unsigned, 21); + + try testing.expect((try a.to(u64)) == 0x10000); +} + +test "big.int sat shift-left multi unsigned" { + var a = try Managed.initSet(testing.allocator, 16); + defer a.deinit(); + try a.shiftLeftSat(a, @bitSizeOf(DoubleLimb) - 3, .unsigned, @bitSizeOf(DoubleLimb) - 1); + + try testing.expect((try a.to(DoubleLimb)) == maxInt(DoubleLimb) >> 1); +} + +test "big.int sat shift-left unsigned shift > bitcount" { + var a = try Managed.initSet(testing.allocator, 1); + defer a.deinit(); + try a.shiftLeftSat(a, 10, .unsigned, 10); + + try testing.expect((try a.to(u10)) == maxInt(u10)); +} + +test "big.int sat shift-left unsigned zero" { + var a = try Managed.initSet(testing.allocator, 0); + defer a.deinit(); + try a.shiftLeftSat(a, 1, .unsigned, 0); + + try testing.expect((try a.to(u64)) == 0); +} + +test "big.int sat shift-left unsigned negative" { + var a = try Managed.initSet(testing.allocator, -100); + defer a.deinit(); + try a.shiftLeftSat(a, 0, .unsigned, 0); + + try testing.expect((try a.to(u64)) == 0); +} + +test "big.int sat shift-left signed simple negative" { + var a = try Managed.initSet(testing.allocator, -100); + defer a.deinit(); + try a.shiftLeftSat(a, 3, .signed, 10); + + try testing.expect((try a.to(i10)) == minInt(i10)); +} + +test "big.int sat shift-left signed simple positive" { + var a = try Managed.initSet(testing.allocator, 100); + defer a.deinit(); + try a.shiftLeftSat(a, 3, .signed, 10); + + try testing.expect((try a.to(i10)) == maxInt(i10)); +} + +test "big.int sat shift-left signed multi positive" { + const x = 1; + const shift = @bitSizeOf(SignedDoubleLimb) - 1; + + var a = try Managed.initSet(testing.allocator, x); + defer a.deinit(); + try a.shiftLeftSat(a, shift, .signed, @bitSizeOf(SignedDoubleLimb)); + + try testing.expect((try a.to(SignedDoubleLimb)) == @as(SignedDoubleLimb, x) <<| shift); +} + +test "big.int sat shift-left signed multi negative" { + const x = -1; + const shift = @bitSizeOf(SignedDoubleLimb) - 1; + + var a = try Managed.initSet(testing.allocator, x); + defer a.deinit(); + try a.shiftLeftSat(a, shift, .signed, @bitSizeOf(SignedDoubleLimb)); + + try testing.expect((try a.to(SignedDoubleLimb)) == @as(SignedDoubleLimb, x) <<| shift); +} + test "big.int bitwise and simple" { var a = try Managed.initSet(testing.allocator, 0xffffffff11111111); defer a.deinit();