Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+95 −35
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 86 additions & 31 deletions lib/std/math/big/int.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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)
Expand All @@ -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()));
Expand All @@ -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()));
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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());
Expand All @@ -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();
Expand Down Expand Up @@ -3350,15 +3399,19 @@ 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);
}

/// 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);
Expand All @@ -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);
Expand Down
13 changes: 9 additions & 4 deletions lib/std/math/big/int_test.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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));
Expand All @@ -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));
Expand Down Expand Up @@ -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));
Expand Down