Commit efa4f76c8b

@@ -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;

@@ -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();