mirrors/zig
1
0
Fork 0
mirror of https://codeberg.org/ziglang/zig.git synced 2025-12-06 13:54:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

std.math.big: require sufficient capacity for aliased params

Browse source
This commit is contained in:
Aaron Ang 2025-11-22 19:55:12 -08:00
parent d828115dab
commit d59b0ae44a
2 changed files with 103 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

119
lib/std/math/big/int.zig
View file

@ -2971,7 +2971,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 = r.limbs.ptr == a.limbs.ptr;
try r.ensureAliasAwareCapacity(needed, aliased);
var m = r.toMutable();
m.addScalar(a.toConst(), scalar);
r.setMetadata(m.positive, m.len);
@ -2983,7 +2985,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
try r.ensureAliasAwareCapacity(needed, aliased);
var m = r.toMutable();
m.add(a.toConst(), b.toConst());
r.setMetadata(m.positive, m.len);
@ -3001,7 +3005,9 @@ pub const Managed = struct {
signedness: Signedness,
bit_count: usize,
) Allocator.Error!bool {
try r.ensureTwosCompCapacity(bit_count);
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3014,7 +3020,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3026,7 +3034,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3044,7 +3054,9 @@ pub const Managed = struct {
signedness: Signedness,
bit_count: usize,
) Allocator.Error!bool {
try r.ensureTwosCompCapacity(bit_count);
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3063,7 +3075,9 @@ pub const Managed = struct {
signedness: Signedness,
bit_count: usize,
) Allocator.Error!void {
try r.ensureTwosCompCapacity(bit_count);
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3082,7 +3096,8 @@ 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;
try rma.ensureAliasAwareCapacity(needed, alias_count > 0);
var m = rma.toMutable();
if (alias_count == 0) {
m.mulNoAlias(a.toConst(), b.toConst(), rma.allocator);
@ -3114,8 +3129,8 @@ 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);
try rma.ensureAliasAwareCapacity(needed, alias_count > 0);
var m = rma.toMutable();
if (alias_count == 0) {
m.mulWrapNoAlias(a.toConst(), b.toConst(), signedness, bit_count, rma.allocator);
@ -3132,26 +3147,58 @@ pub const Managed = struct {
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
}
/// Ensures capacity only when parameters do not alias the result.
///
/// When aliasing is detected, this function requires the caller to have already
/// ensured sufficient capacity. This prevents use-after-free bugs that occur when
/// reallocating memory while const parameter pointers reference that same memory.
///
/// Callers using aliasing must pre-allocate capacity using the appropriate
/// `ensure*Capacity` helper before calling the operation.
///
/// See: https://github.com/ziglang/zig/issues/6167
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, scalar: anytype) !void {
try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1);
}
pub fn ensureAddScalarCapacityManaged(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 ensureAddCapacityManaged(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 ensureMulCapacityManaged(rma: *Managed, a: *const Managed, b: *const Managed) !void {
try rma.ensureCapacity(a.len() + b.len() + 1);
}
/// q = a / b (rem r)
///
/// a / b are floored (rounded towards 0).
///
/// 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 = q.limbs.ptr == a.limbs.ptr or q.limbs.ptr == b.limbs.ptr;
const r_alias = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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()));
@ -3167,8 +3214,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 = q.limbs.ptr == a.limbs.ptr or q.limbs.ptr == b.limbs.ptr;
const r_alias = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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()));
@ -3181,7 +3230,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 = r.limbs.ptr == a.limbs.ptr;
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);
@ -3190,7 +3241,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 = r.limbs.ptr == a.limbs.ptr;
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);
@ -3212,7 +3265,9 @@ pub const Managed = struct {
return;
}
try r.ensureCapacity(a.len() - (shift / limb_bits));
const aliased = r.limbs.ptr == a.limbs.ptr;
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);
@ -3221,7 +3276,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 = r.limbs.ptr == a.limbs.ptr;
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);
@ -3231,7 +3288,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
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);
@ -3243,7 +3302,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
try r.ensureAliasAwareCapacity(cap, aliased);
var m = r.toMutable();
m.bitAnd(a.toConst(), b.toConst());
r.setMetadata(m.positive, m.len);
@ -3252,7 +3312,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 = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
try r.ensureAliasAwareCapacity(cap, aliased);
var m = r.toMutable();
m.bitXor(a.toConst(), b.toConst());
@ -3264,7 +3325,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 = rma.limbs.ptr == x.limbs.ptr or rma.limbs.ptr == y.limbs.ptr;
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();
@ -3342,7 +3405,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 = r.limbs.ptr == a.limbs.ptr;
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);
@ -3350,7 +3415,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 = r.limbs.ptr == a.limbs.ptr;
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);
@ -3359,7 +3426,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 = r.limbs.ptr == a.limbs.ptr;
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);

13
lib/std/math/big/int_test.zig
View file

@ -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.ensureAddCapacityManaged(&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.ensureAddCapacityManaged(&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.ensureMulCapacityManaged(&a, &a);
try a.mul(&a, &a);
var want = try Managed.initSet(testing.allocator, 4 * maxInt(Limb) * maxInt(Limb));