bigint: add ensureAdd(Scalar)Capacity, note aliasing requirements

lib/std/math/big/int.zig

@@ -335,10 +335,10 @@ pub const Mutable = struct {
             }
         } else {
             if (a.limbs.len >= b.limbs.len) {
-                lladd(r.limbs[0..], a.limbs[0..a.limbs.len], b.limbs[0..b.limbs.len]);
+                lladd(r.limbs[0..], a.limbs, b.limbs);
                 r.normalize(a.limbs.len + 1);
             } else {
-                lladd(r.limbs[0..], b.limbs[0..b.limbs.len], a.limbs[0..a.limbs.len]);
+                lladd(r.limbs[0..], b.limbs, a.limbs);
                 r.normalize(b.limbs.len + 1);
             }
 
@@ -1685,12 +1685,14 @@ pub const Managed = struct {
 
     /// r = a + scalar
     ///
-    /// r and a may be aliases.
+    /// r and a may be aliases. If r aliases a, then caller must call
+    /// `r.ensureAddScalarCapacity` prior to calling `add`.
     /// scalar is a primitive integer type.
     ///
     /// Returns an error if memory could not be allocated.
     pub fn addScalar(r: *Managed, a: Const, scalar: anytype) Allocator.Error!void {
-        try r.ensureCapacity(math.max(a.limbs.len, calcLimbLen(scalar)) + 1);
+        assert((r.limbs.ptr != a.limbs.ptr) or r.limbs.len >= math.max(a.limbs.len, calcLimbLen(scalar)) + 1);
+        try r.ensureAddScalarCapacity(a, scalar);
         var m = r.toMutable();
         m.addScalar(a, scalar);
         r.setMetadata(m.positive, m.len);
@@ -1698,11 +1700,13 @@ pub const Managed = struct {
 
     /// r = a + b
     ///
-    /// r, a and b may be aliases.
+    /// r, a and b may be aliases. If r aliases a or b, then caller must call
+    /// `r.ensureAddCapacity` prior to calling `add`.
     ///
     /// Returns an error if memory could not be allocated.
     pub fn add(r: *Managed, a: Const, b: Const) Allocator.Error!void {
-        try r.ensureCapacity(math.max(a.limbs.len, b.limbs.len) + 1);
+        assert((r.limbs.ptr != a.limbs.ptr and r.limbs.ptr != b.limbs.ptr) or r.limbs.len >= math.max(a.limbs.len, b.limbs.len) + 1);
+        try r.ensureAddCapacity(a, b);
         var m = r.toMutable();
         m.add(a, b);
         r.setMetadata(m.positive, m.len);
@@ -1748,6 +1752,14 @@ pub const Managed = struct {
         rma.setMetadata(m.positive, m.len);
     }
 
+    pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void {
+        try r.ensureCapacity(math.max(a.limbs.len, calcLimbLen(scalar)) + 1);
+    }
+
+    pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void {
+        try r.ensureCapacity(math.max(a.limbs.len, b.limbs.len) + 1);
+    }
+
     pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void {
         try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1);
     }

lib/std/math/big/int_test.zig

@@ -538,6 +538,17 @@ test "big.int add sign" {
     try testing.expect((try a.to(i32)) == -3);
 }
 
+test "big.int add scalar" {
+    var a = try Managed.initSet(testing.allocator, 50);
+    defer a.deinit();
+
+    var b = try Managed.init(testing.allocator);
+    defer b.deinit();
+    try b.addScalar(a.toConst(), 5);
+
+    try testing.expect((try b.to(u32)) == 55);
+}
+
 test "big.int sub single-single" {
     var a = try Managed.initSet(testing.allocator, 50);
     defer a.deinit();
@@ -1563,7 +1574,7 @@ test "big.int pow" {
     }
 }
 
-test "bigint regression test for 1 limb overflow with alias" {
+test "big.int regression test for 1 limb overflow with alias" {
     // Note these happen to be two consecutive Fibonacci sequence numbers, the
     // first two whose sum exceeds 2**64.
     var a = try Managed.initSet(testing.allocator, 7540113804746346429);
@@ -1571,12 +1582,13 @@ test "bigint regression test for 1 limb overflow with alias" {
     var b = try Managed.initSet(testing.allocator, 12200160415121876738);
     defer b.deinit();
 
+    try a.ensureAddCapacity(a.toConst(), b.toConst());
     try a.add(a.toConst(), b.toConst());
 
     try testing.expect(a.toConst().orderAgainstScalar(19740274219868223167) == .eq);
 }
 
-test "bigint regression test for something else with alias" {
+test "big.int regression test for realloc with alias" {
     // Note these happen to be two consecutive Fibonacci sequence numbers, the
     // second of which is the first such number to exceed 2**192.
     var a = try Managed.initSet(testing.allocator, 5611500259351924431073312796924978741056961814867751431689);
@@ -1584,6 +1596,7 @@ test "bigint regression test for something else with alias" {
     var b = try Managed.initSet(testing.allocator, 9079598147510263717870894449029933369491131786514446266146);
     defer b.deinit();
 
+    try a.ensureAddCapacity(a.toConst(), b.toConst());
     try a.add(a.toConst(), b.toConst());
 
     try testing.expect(a.toConst().orderAgainstScalar(14691098406862188148944207245954912110548093601382197697835) == .eq);