diff --git a/lib/std/math/big/int.zig b/lib/std/math/big/int.zig index f52495b739..b2997536a6 100644 --- a/lib/std/math/big/int.zig +++ b/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); } diff --git a/lib/std/math/big/int_test.zig b/lib/std/math/big/int_test.zig index d1c8b3ae8a..f4b294245a 100644 --- a/lib/std/math/big/int_test.zig +++ b/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);