This commit is contained in:
Aaron Ang 2025-11-25 16:16:36 +01:00 committed by GitHub
commit 454340e532
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 103 additions and 29 deletions

View file

@ -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 = 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);
@ -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 = 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);
@ -3009,7 +3013,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);
@ -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 = 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);
@ -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 = 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);
@ -3052,7 +3062,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);
@ -3071,7 +3083,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);
@ -3090,7 +3104,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);
@ -3122,8 +3137,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);
@ -3140,26 +3155,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()));
@ -3175,8 +3222,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()));
@ -3189,7 +3238,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);
@ -3198,7 +3249,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);
@ -3220,7 +3273,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);
@ -3229,7 +3284,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);
@ -3239,7 +3296,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);
@ -3251,7 +3310,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);
@ -3260,7 +3320,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());
@ -3272,7 +3333,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();
@ -3350,7 +3413,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);
@ -3358,7 +3423,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);
@ -3367,7 +3434,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);

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