From cfeda3c18220937e6eb2a5ad1445011a87b13785 Mon Sep 17 00:00:00 2001 From: Aaron Ang <67321817+aaron-ang@users.noreply.github.com> Date: Tue, 25 Nov 2025 11:45:02 -0800 Subject: [PATCH] std.math.big: require sufficient capacity for aliased params --- lib/std/math/big/int.zig | 117 +++++++++++++++++++++++++--------- lib/std/math/big/int_test.zig | 13 ++-- 2 files changed, 95 insertions(+), 35 deletions(-) diff --git a/lib/std/math/big/int.zig b/lib/std/math/big/int.zig index 5a00947ac365..892c29835a18 100644 --- a/lib/std/math/big/int.zig +++ b/lib/std/math/big/int.zig @@ -2979,7 +2979,9 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn addScalar(r: *Managed, a: *const Managed, scalar: anytype) Allocator.Error!void { - try r.ensureAddScalarCapacity(a.toConst(), scalar); + const needed = @max(a.len(), calcLimbLen(scalar)) + 1; + const aliased = limbsAliasDistinct(r, a); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.addScalar(a.toConst(), scalar); r.setMetadata(m.positive, m.len); @@ -2991,7 +2993,9 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn add(r: *Managed, a: *const Managed, b: *const Managed) Allocator.Error!void { - try r.ensureAddCapacity(a.toConst(), b.toConst()); + const needed = @max(a.len(), b.len()) + 1; + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.add(a.toConst(), b.toConst()); r.setMetadata(m.positive, m.len); @@ -3009,7 +3013,9 @@ pub const Managed = struct { signedness: Signedness, bit_count: usize, ) Allocator.Error!bool { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); const wrapped = m.addWrap(a.toConst(), b.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3022,7 +3028,9 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn addSat(r: *Managed, a: *const Managed, b: *const Managed, signedness: Signedness, bit_count: usize) Allocator.Error!void { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.addSat(a.toConst(), b.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3034,7 +3042,9 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn sub(r: *Managed, a: *const Managed, b: *const Managed) !void { - try r.ensureCapacity(@max(a.len(), b.len()) + 1); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = @max(a.len(), b.len()) + 1; + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.sub(a.toConst(), b.toConst()); r.setMetadata(m.positive, m.len); @@ -3052,7 +3062,9 @@ pub const Managed = struct { signedness: Signedness, bit_count: usize, ) Allocator.Error!bool { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); const wrapped = m.subWrap(a.toConst(), b.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3071,7 +3083,9 @@ pub const Managed = struct { signedness: Signedness, bit_count: usize, ) Allocator.Error!void { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.subSat(a.toConst(), b.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3090,7 +3104,9 @@ pub const Managed = struct { alias_count += 1; if (rma.limbs.ptr == b.limbs.ptr) alias_count += 1; - try rma.ensureMulCapacity(a.toConst(), b.toConst()); + const needed = a.len() + b.len() + 1; + const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b); + try rma.ensureAliasAwareCapacity(needed, capacity_alias); var m = rma.toMutable(); if (alias_count == 0) { m.mulNoAlias(a.toConst(), b.toConst(), rma.allocator); @@ -3122,8 +3138,9 @@ pub const Managed = struct { alias_count += 1; if (rma.limbs.ptr == b.limbs.ptr) alias_count += 1; - - try rma.ensureTwosCompCapacity(bit_count); + const needed = calcTwosCompLimbCount(bit_count); + const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b); + try rma.ensureAliasAwareCapacity(needed, capacity_alias); var m = rma.toMutable(); if (alias_count == 0) { m.mulWrapNoAlias(a.toConst(), b.toConst(), signedness, bit_count, rma.allocator); @@ -3140,16 +3157,30 @@ pub const Managed = struct { try r.ensureCapacity(calcTwosCompLimbCount(bit_count)); } - pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void { - try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1); + /// True if both `Managed` instances are distinct but share the same limbs buffer. + fn limbsAliasDistinct(a: *const Managed, b: *const Managed) bool { + return @intFromPtr(a) != @intFromPtr(b) and a.limbs.ptr == b.limbs.ptr; + } + + /// Ensures capacity if not aliased. Requires enough capacity if aliased. + fn ensureAliasAwareCapacity(r: *Managed, needed: usize, aliased: bool) !void { + if (aliased) { + assert(needed <= r.limbs.len); + } else { + try r.ensureCapacity(needed); + } + } + + pub fn ensureAddScalarCapacity(r: *Managed, a: *const Managed, scalar: anytype) !void { + try r.ensureCapacity(@max(a.len(), calcLimbLen(scalar)) + 1); } - pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void { - try r.ensureCapacity(@max(a.limbs.len, b.limbs.len) + 1); + pub fn ensureAddCapacity(r: *Managed, a: *const Managed, b: *const Managed) !void { + try r.ensureCapacity(@max(a.len(), b.len()) + 1); } - pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void { - try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1); + pub fn ensureMulCapacity(rma: *Managed, a: *const Managed, b: *const Managed) !void { + try rma.ensureCapacity(a.len() + b.len() + 1); } /// q = a / b (rem r) @@ -3158,8 +3189,10 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn divFloor(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void { - try q.ensureCapacity(a.len()); - try r.ensureCapacity(b.len()); + const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b); + const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + try q.ensureAliasAwareCapacity(a.len(), q_alias); + try r.ensureAliasAwareCapacity(b.len(), r_alias); var mq = q.toMutable(); var mr = r.toMutable(); const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len())); @@ -3175,8 +3208,10 @@ pub const Managed = struct { /// /// Returns an error if memory could not be allocated. pub fn divTrunc(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void { - try q.ensureCapacity(a.len()); - try r.ensureCapacity(b.len()); + const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b); + const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + try q.ensureAliasAwareCapacity(a.len(), q_alias); + try r.ensureAliasAwareCapacity(b.len(), r_alias); var mq = q.toMutable(); var mr = r.toMutable(); const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len())); @@ -3189,7 +3224,9 @@ pub const Managed = struct { /// r = a << shift, in other words, r = a * 2^shift /// r and a may alias. pub fn shiftLeft(r: *Managed, a: *const Managed, shift: usize) !void { - try r.ensureCapacity(a.len() + (shift / limb_bits) + 1); + const aliased = limbsAliasDistinct(r, a); + const needed = a.len() + (shift / limb_bits) + 1; + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.shiftLeft(a.toConst(), shift); r.setMetadata(m.positive, m.len); @@ -3198,7 +3235,9 @@ pub const Managed = struct { /// r = a <<| shift with 2s-complement saturating semantics. /// r and a may alias. pub fn shiftLeftSat(r: *Managed, a: *const Managed, shift: usize, signedness: Signedness, bit_count: usize) !void { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.shiftLeftSat(a.toConst(), shift, signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3220,7 +3259,9 @@ pub const Managed = struct { return; } - try r.ensureCapacity(a.len() - (shift / limb_bits)); + const aliased = limbsAliasDistinct(r, a); + const needed = a.len() - (shift / limb_bits); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.shiftRight(a.toConst(), shift); r.setMetadata(m.positive, m.len); @@ -3229,7 +3270,9 @@ pub const Managed = struct { /// r = ~a under 2s-complement wrapping semantics. /// r and a may alias. pub fn bitNotWrap(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void { - try r.ensureTwosCompCapacity(bit_count); + const aliased = limbsAliasDistinct(r, a); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.bitNotWrap(a.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3239,7 +3282,9 @@ pub const Managed = struct { /// /// a and b are zero-extended to the longer of a or b. pub fn bitOr(r: *Managed, a: *const Managed, b: *const Managed) !void { - try r.ensureCapacity(@max(a.len(), b.len())); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + const needed = @max(a.len(), b.len()); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.bitOr(a.toConst(), b.toConst()); r.setMetadata(m.positive, m.len); @@ -3251,7 +3296,8 @@ pub const Managed = struct { if (b.isPositive()) b.len() else if (a.isPositive()) a.len() else a.len() + 1 else if (a.isPositive()) a.len() else if (b.isPositive()) b.len() else b.len() + 1; - try r.ensureCapacity(cap); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + try r.ensureAliasAwareCapacity(cap, aliased); var m = r.toMutable(); m.bitAnd(a.toConst(), b.toConst()); r.setMetadata(m.positive, m.len); @@ -3260,7 +3306,8 @@ pub const Managed = struct { /// r = a ^ b pub fn bitXor(r: *Managed, a: *const Managed, b: *const Managed) !void { const cap = @max(a.len(), b.len()) + @intFromBool(a.isPositive() != b.isPositive()); - try r.ensureCapacity(cap); + const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b); + try r.ensureAliasAwareCapacity(cap, aliased); var m = r.toMutable(); m.bitXor(a.toConst(), b.toConst()); @@ -3272,7 +3319,9 @@ pub const Managed = struct { /// /// rma's allocator is used for temporary storage to boost multiplication performance. pub fn gcd(rma: *Managed, x: *const Managed, y: *const Managed) !void { - try rma.ensureCapacity(@min(x.len(), y.len())); + const aliased = limbsAliasDistinct(rma, x) or limbsAliasDistinct(rma, y); + const needed = @min(x.len(), y.len()); + try rma.ensureAliasAwareCapacity(needed, aliased); var m = rma.toMutable(); var limbs_buffer = std.array_list.Managed(Limb).init(rma.allocator); defer limbs_buffer.deinit(); @@ -3350,7 +3399,9 @@ pub const Managed = struct { /// r = truncate(Int(signedness, bit_count), a) pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void { - try r.ensureCapacity(calcTwosCompLimbCount(bit_count)); + const aliased = limbsAliasDistinct(r, a); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.truncate(a.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3358,7 +3409,9 @@ pub const Managed = struct { /// r = saturate(Int(signedness, bit_count), a) pub fn saturate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void { - try r.ensureCapacity(calcTwosCompLimbCount(bit_count)); + const aliased = limbsAliasDistinct(r, a); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.saturate(a.toConst(), signedness, bit_count); r.setMetadata(m.positive, m.len); @@ -3367,7 +3420,9 @@ pub const Managed = struct { /// r = @popCount(a) with 2s-complement semantics. /// r and a may be aliases. pub fn popCount(r: *Managed, a: *const Managed, bit_count: usize) !void { - try r.ensureCapacity(calcTwosCompLimbCount(bit_count)); + const aliased = limbsAliasDistinct(r, a); + const needed = calcTwosCompLimbCount(bit_count); + try r.ensureAliasAwareCapacity(needed, aliased); var m = r.toMutable(); m.popCount(a.toConst(), bit_count); r.setMetadata(m.positive, m.len); diff --git a/lib/std/math/big/int_test.zig b/lib/std/math/big/int_test.zig index eee6a0c7cee2..2162eed106bd 100644 --- a/lib/std/math/big/int_test.zig +++ b/lib/std/math/big/int_test.zig @@ -583,7 +583,9 @@ test "bitcount + sizeInBaseUpperBound" { try testing.expect(a.sizeInBaseUpperBound(2) >= 32); try testing.expect(a.sizeInBaseUpperBound(10) >= 10); - try a.shiftLeft(&a, 5000); + const shift = 5000; + try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1); + try a.shiftLeft(&a, shift); try testing.expectEqual(5032, a.bitCountAbs()); try testing.expect(a.sizeInBaseUpperBound(2) >= 5032); a.setSign(false); @@ -2380,7 +2382,9 @@ test "truncate negative multi to single" { test "truncate multi unsigned many" { var a = try Managed.initSet(testing.allocator, 1); defer a.deinit(); - try a.shiftLeft(&a, 1023); + const shift = 1023; + try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1); + try a.shiftLeft(&a, shift); var b = try Managed.init(testing.allocator); defer b.deinit(); @@ -3263,7 +3267,7 @@ test "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.ensureAddCapacity(&a, &b); try a.add(&a, &b); try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(19740274219868223167)); @@ -3277,7 +3281,7 @@ test "regression test for realloc with alias" { var b = try Managed.initSet(testing.allocator, 9079598147510263717870894449029933369491131786514446266146); defer b.deinit(); - try a.ensureAddCapacity(a.toConst(), b.toConst()); + try a.ensureAddCapacity(&a, &b); try a.add(&a, &b); try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(14691098406862188148944207245954912110548093601382197697835)); @@ -3692,6 +3696,7 @@ test "mul multi-multi alias r with a and b" { var a = try Managed.initSet(testing.allocator, 2 * maxInt(Limb)); defer a.deinit(); + try a.ensureMulCapacity(&a, &a); try a.mul(&a, &a); var want = try Managed.initSet(testing.allocator, 4 * maxInt(Limb) * maxInt(Limb));