Skip to content

Commit cfeda3c

Browse files
aaron-angaaron-ang
committed
std.math.big: require sufficient capacity for aliased params
1 parent 0f1a6ae commit cfeda3c

File tree

2 files changed

+95
-35
lines changed

2 files changed

+95
-35
lines changed

‎lib/std/math/big/int.zig‎

Lines changed: 86 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -2979,7 +2979,9 @@ pub const Managed = struct {
29792979
///
29802980
/// Returns an error if memory could not be allocated.
29812981
pub fn addScalar(r: *Managed, a: *const Managed, scalar: anytype) Allocator.Error!void {
2982-
try r.ensureAddScalarCapacity(a.toConst(), scalar);
2982+
const needed = @max(a.len(), calcLimbLen(scalar)) + 1;
2983+
const aliased = limbsAliasDistinct(r, a);
2984+
try r.ensureAliasAwareCapacity(needed, aliased);
29832985
var m = r.toMutable();
29842986
m.addScalar(a.toConst(), scalar);
29852987
r.setMetadata(m.positive, m.len);
@@ -2991,7 +2993,9 @@ pub const Managed = struct {
29912993
///
29922994
/// Returns an error if memory could not be allocated.
29932995
pub fn add(r: *Managed, a: *const Managed, b: *const Managed) Allocator.Error!void {
2994-
try r.ensureAddCapacity(a.toConst(), b.toConst());
2996+
const needed = @max(a.len(), b.len()) + 1;
2997+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
2998+
try r.ensureAliasAwareCapacity(needed, aliased);
29952999
var m = r.toMutable();
29963000
m.add(a.toConst(), b.toConst());
29973001
r.setMetadata(m.positive, m.len);
@@ -3009,7 +3013,9 @@ pub const Managed = struct {
30093013
signedness: Signedness,
30103014
bit_count: usize,
30113015
) Allocator.Error!bool {
3012-
try r.ensureTwosCompCapacity(bit_count);
3016+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3017+
const needed = calcTwosCompLimbCount(bit_count);
3018+
try r.ensureAliasAwareCapacity(needed, aliased);
30133019
var m = r.toMutable();
30143020
const wrapped = m.addWrap(a.toConst(), b.toConst(), signedness, bit_count);
30153021
r.setMetadata(m.positive, m.len);
@@ -3022,7 +3028,9 @@ pub const Managed = struct {
30223028
///
30233029
/// Returns an error if memory could not be allocated.
30243030
pub fn addSat(r: *Managed, a: *const Managed, b: *const Managed, signedness: Signedness, bit_count: usize) Allocator.Error!void {
3025-
try r.ensureTwosCompCapacity(bit_count);
3031+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3032+
const needed = calcTwosCompLimbCount(bit_count);
3033+
try r.ensureAliasAwareCapacity(needed, aliased);
30263034
var m = r.toMutable();
30273035
m.addSat(a.toConst(), b.toConst(), signedness, bit_count);
30283036
r.setMetadata(m.positive, m.len);
@@ -3034,7 +3042,9 @@ pub const Managed = struct {
30343042
///
30353043
/// Returns an error if memory could not be allocated.
30363044
pub fn sub(r: *Managed, a: *const Managed, b: *const Managed) !void {
3037-
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3045+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3046+
const needed = @max(a.len(), b.len()) + 1;
3047+
try r.ensureAliasAwareCapacity(needed, aliased);
30383048
var m = r.toMutable();
30393049
m.sub(a.toConst(), b.toConst());
30403050
r.setMetadata(m.positive, m.len);
@@ -3052,7 +3062,9 @@ pub const Managed = struct {
30523062
signedness: Signedness,
30533063
bit_count: usize,
30543064
) Allocator.Error!bool {
3055-
try r.ensureTwosCompCapacity(bit_count);
3065+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3066+
const needed = calcTwosCompLimbCount(bit_count);
3067+
try r.ensureAliasAwareCapacity(needed, aliased);
30563068
var m = r.toMutable();
30573069
const wrapped = m.subWrap(a.toConst(), b.toConst(), signedness, bit_count);
30583070
r.setMetadata(m.positive, m.len);
@@ -3071,7 +3083,9 @@ pub const Managed = struct {
30713083
signedness: Signedness,
30723084
bit_count: usize,
30733085
) Allocator.Error!void {
3074-
try r.ensureTwosCompCapacity(bit_count);
3086+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3087+
const needed = calcTwosCompLimbCount(bit_count);
3088+
try r.ensureAliasAwareCapacity(needed, aliased);
30753089
var m = r.toMutable();
30763090
m.subSat(a.toConst(), b.toConst(), signedness, bit_count);
30773091
r.setMetadata(m.positive, m.len);
@@ -3090,7 +3104,9 @@ pub const Managed = struct {
30903104
alias_count += 1;
30913105
if (rma.limbs.ptr == b.limbs.ptr)
30923106
alias_count += 1;
3093-
try rma.ensureMulCapacity(a.toConst(), b.toConst());
3107+
const needed = a.len() + b.len() + 1;
3108+
const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b);
3109+
try rma.ensureAliasAwareCapacity(needed, capacity_alias);
30943110
var m = rma.toMutable();
30953111
if (alias_count == 0) {
30963112
m.mulNoAlias(a.toConst(), b.toConst(), rma.allocator);
@@ -3122,8 +3138,9 @@ pub const Managed = struct {
31223138
alias_count += 1;
31233139
if (rma.limbs.ptr == b.limbs.ptr)
31243140
alias_count += 1;
3125-
3126-
try rma.ensureTwosCompCapacity(bit_count);
3141+
const needed = calcTwosCompLimbCount(bit_count);
3142+
const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b);
3143+
try rma.ensureAliasAwareCapacity(needed, capacity_alias);
31273144
var m = rma.toMutable();
31283145
if (alias_count == 0) {
31293146
m.mulWrapNoAlias(a.toConst(), b.toConst(), signedness, bit_count, rma.allocator);
@@ -3140,16 +3157,30 @@ pub const Managed = struct {
31403157
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
31413158
}
31423159

3143-
pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void {
3144-
try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1);
3160+
/// True if both `Managed` instances are distinct but share the same limbs buffer.
3161+
fn limbsAliasDistinct(a: *const Managed, b: *const Managed) bool {
3162+
return @intFromPtr(a) != @intFromPtr(b) and a.limbs.ptr == b.limbs.ptr;
3163+
}
3164+
3165+
/// Ensures capacity if not aliased. Requires enough capacity if aliased.
3166+
fn ensureAliasAwareCapacity(r: *Managed, needed: usize, aliased: bool) !void {
3167+
if (aliased) {
3168+
assert(needed <= r.limbs.len);
3169+
} else {
3170+
try r.ensureCapacity(needed);
3171+
}
3172+
}
3173+
3174+
pub fn ensureAddScalarCapacity(r: *Managed, a: *const Managed, scalar: anytype) !void {
3175+
try r.ensureCapacity(@max(a.len(), calcLimbLen(scalar)) + 1);
31453176
}
31463177

3147-
pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void {
3148-
try r.ensureCapacity(@max(a.limbs.len, b.limbs.len) + 1);
3178+
pub fn ensureAddCapacity(r: *Managed, a: *const Managed, b: *const Managed) !void {
3179+
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
31493180
}
31503181

3151-
pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void {
3152-
try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1);
3182+
pub fn ensureMulCapacity(rma: *Managed, a: *const Managed, b: *const Managed) !void {
3183+
try rma.ensureCapacity(a.len() + b.len() + 1);
31533184
}
31543185

31553186
/// q = a / b (rem r)
@@ -3158,8 +3189,10 @@ pub const Managed = struct {
31583189
///
31593190
/// Returns an error if memory could not be allocated.
31603191
pub fn divFloor(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3161-
try q.ensureCapacity(a.len());
3162-
try r.ensureCapacity(b.len());
3192+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3193+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3194+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3195+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31633196
var mq = q.toMutable();
31643197
var mr = r.toMutable();
31653198
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3175,8 +3208,10 @@ pub const Managed = struct {
31753208
///
31763209
/// Returns an error if memory could not be allocated.
31773210
pub fn divTrunc(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3178-
try q.ensureCapacity(a.len());
3179-
try r.ensureCapacity(b.len());
3211+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3212+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3213+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3214+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31803215
var mq = q.toMutable();
31813216
var mr = r.toMutable();
31823217
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3189,7 +3224,9 @@ pub const Managed = struct {
31893224
/// r = a << shift, in other words, r = a * 2^shift
31903225
/// r and a may alias.
31913226
pub fn shiftLeft(r: *Managed, a: *const Managed, shift: usize) !void {
3192-
try r.ensureCapacity(a.len() + (shift / limb_bits) + 1);
3227+
const aliased = limbsAliasDistinct(r, a);
3228+
const needed = a.len() + (shift / limb_bits) + 1;
3229+
try r.ensureAliasAwareCapacity(needed, aliased);
31933230
var m = r.toMutable();
31943231
m.shiftLeft(a.toConst(), shift);
31953232
r.setMetadata(m.positive, m.len);
@@ -3198,7 +3235,9 @@ pub const Managed = struct {
31983235
/// r = a <<| shift with 2s-complement saturating semantics.
31993236
/// r and a may alias.
32003237
pub fn shiftLeftSat(r: *Managed, a: *const Managed, shift: usize, signedness: Signedness, bit_count: usize) !void {
3201-
try r.ensureTwosCompCapacity(bit_count);
3238+
const aliased = limbsAliasDistinct(r, a);
3239+
const needed = calcTwosCompLimbCount(bit_count);
3240+
try r.ensureAliasAwareCapacity(needed, aliased);
32023241
var m = r.toMutable();
32033242
m.shiftLeftSat(a.toConst(), shift, signedness, bit_count);
32043243
r.setMetadata(m.positive, m.len);
@@ -3220,7 +3259,9 @@ pub const Managed = struct {
32203259
return;
32213260
}
32223261

3223-
try r.ensureCapacity(a.len() - (shift / limb_bits));
3262+
const aliased = limbsAliasDistinct(r, a);
3263+
const needed = a.len() - (shift / limb_bits);
3264+
try r.ensureAliasAwareCapacity(needed, aliased);
32243265
var m = r.toMutable();
32253266
m.shiftRight(a.toConst(), shift);
32263267
r.setMetadata(m.positive, m.len);
@@ -3229,7 +3270,9 @@ pub const Managed = struct {
32293270
/// r = ~a under 2s-complement wrapping semantics.
32303271
/// r and a may alias.
32313272
pub fn bitNotWrap(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3232-
try r.ensureTwosCompCapacity(bit_count);
3273+
const aliased = limbsAliasDistinct(r, a);
3274+
const needed = calcTwosCompLimbCount(bit_count);
3275+
try r.ensureAliasAwareCapacity(needed, aliased);
32333276
var m = r.toMutable();
32343277
m.bitNotWrap(a.toConst(), signedness, bit_count);
32353278
r.setMetadata(m.positive, m.len);
@@ -3239,7 +3282,9 @@ pub const Managed = struct {
32393282
///
32403283
/// a and b are zero-extended to the longer of a or b.
32413284
pub fn bitOr(r: *Managed, a: *const Managed, b: *const Managed) !void {
3242-
try r.ensureCapacity(@max(a.len(), b.len()));
3285+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3286+
const needed = @max(a.len(), b.len());
3287+
try r.ensureAliasAwareCapacity(needed, aliased);
32433288
var m = r.toMutable();
32443289
m.bitOr(a.toConst(), b.toConst());
32453290
r.setMetadata(m.positive, m.len);
@@ -3251,7 +3296,8 @@ pub const Managed = struct {
32513296
if (b.isPositive()) b.len() else if (a.isPositive()) a.len() else a.len() + 1
32523297
else if (a.isPositive()) a.len() else if (b.isPositive()) b.len() else b.len() + 1;
32533298

3254-
try r.ensureCapacity(cap);
3299+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3300+
try r.ensureAliasAwareCapacity(cap, aliased);
32553301
var m = r.toMutable();
32563302
m.bitAnd(a.toConst(), b.toConst());
32573303
r.setMetadata(m.positive, m.len);
@@ -3260,7 +3306,8 @@ pub const Managed = struct {
32603306
/// r = a ^ b
32613307
pub fn bitXor(r: *Managed, a: *const Managed, b: *const Managed) !void {
32623308
const cap = @max(a.len(), b.len()) + @intFromBool(a.isPositive() != b.isPositive());
3263-
try r.ensureCapacity(cap);
3309+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3310+
try r.ensureAliasAwareCapacity(cap, aliased);
32643311

32653312
var m = r.toMutable();
32663313
m.bitXor(a.toConst(), b.toConst());
@@ -3272,7 +3319,9 @@ pub const Managed = struct {
32723319
///
32733320
/// rma's allocator is used for temporary storage to boost multiplication performance.
32743321
pub fn gcd(rma: *Managed, x: *const Managed, y: *const Managed) !void {
3275-
try rma.ensureCapacity(@min(x.len(), y.len()));
3322+
const aliased = limbsAliasDistinct(rma, x) or limbsAliasDistinct(rma, y);
3323+
const needed = @min(x.len(), y.len());
3324+
try rma.ensureAliasAwareCapacity(needed, aliased);
32763325
var m = rma.toMutable();
32773326
var limbs_buffer = std.array_list.Managed(Limb).init(rma.allocator);
32783327
defer limbs_buffer.deinit();
@@ -3350,15 +3399,19 @@ pub const Managed = struct {
33503399

33513400
/// r = truncate(Int(signedness, bit_count), a)
33523401
pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3353-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3402+
const aliased = limbsAliasDistinct(r, a);
3403+
const needed = calcTwosCompLimbCount(bit_count);
3404+
try r.ensureAliasAwareCapacity(needed, aliased);
33543405
var m = r.toMutable();
33553406
m.truncate(a.toConst(), signedness, bit_count);
33563407
r.setMetadata(m.positive, m.len);
33573408
}
33583409

33593410
/// r = saturate(Int(signedness, bit_count), a)
33603411
pub fn saturate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3361-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3412+
const aliased = limbsAliasDistinct(r, a);
3413+
const needed = calcTwosCompLimbCount(bit_count);
3414+
try r.ensureAliasAwareCapacity(needed, aliased);
33623415
var m = r.toMutable();
33633416
m.saturate(a.toConst(), signedness, bit_count);
33643417
r.setMetadata(m.positive, m.len);
@@ -3367,7 +3420,9 @@ pub const Managed = struct {
33673420
/// r = @popCount(a) with 2s-complement semantics.
33683421
/// r and a may be aliases.
33693422
pub fn popCount(r: *Managed, a: *const Managed, bit_count: usize) !void {
3370-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3423+
const aliased = limbsAliasDistinct(r, a);
3424+
const needed = calcTwosCompLimbCount(bit_count);
3425+
try r.ensureAliasAwareCapacity(needed, aliased);
33713426
var m = r.toMutable();
33723427
m.popCount(a.toConst(), bit_count);
33733428
r.setMetadata(m.positive, m.len);

‎lib/std/math/big/int_test.zig‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -583,7 +583,9 @@ test "bitcount + sizeInBaseUpperBound" {
583583
try testing.expect(a.sizeInBaseUpperBound(2) >= 32);
584584
try testing.expect(a.sizeInBaseUpperBound(10) >= 10);
585585

586-
try a.shiftLeft(&a, 5000);
586+
const shift = 5000;
587+
try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1);
588+
try a.shiftLeft(&a, shift);
587589
try testing.expectEqual(5032, a.bitCountAbs());
588590
try testing.expect(a.sizeInBaseUpperBound(2) >= 5032);
589591
a.setSign(false);
@@ -2380,7 +2382,9 @@ test "truncate negative multi to single" {
23802382
test "truncate multi unsigned many" {
23812383
var a = try Managed.initSet(testing.allocator, 1);
23822384
defer a.deinit();
2383-
try a.shiftLeft(&a, 1023);
2385+
const shift = 1023;
2386+
try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1);
2387+
try a.shiftLeft(&a, shift);
23842388

23852389
var b = try Managed.init(testing.allocator);
23862390
defer b.deinit();
@@ -3263,7 +3267,7 @@ test "regression test for 1 limb overflow with alias" {
32633267
var b = try Managed.initSet(testing.allocator, 12200160415121876738);
32643268
defer b.deinit();
32653269

3266-
try a.ensureAddCapacity(a.toConst(), b.toConst());
3270+
try a.ensureAddCapacity(&a, &b);
32673271
try a.add(&a, &b);
32683272

32693273
try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(19740274219868223167));
@@ -3277,7 +3281,7 @@ test "regression test for realloc with alias" {
32773281
var b = try Managed.initSet(testing.allocator, 9079598147510263717870894449029933369491131786514446266146);
32783282
defer b.deinit();
32793283

3280-
try a.ensureAddCapacity(a.toConst(), b.toConst());
3284+
try a.ensureAddCapacity(&a, &b);
32813285
try a.add(&a, &b);
32823286

32833287
try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(14691098406862188148944207245954912110548093601382197697835));
@@ -3692,6 +3696,7 @@ test "mul multi-multi alias r with a and b" {
36923696
var a = try Managed.initSet(testing.allocator, 2 * maxInt(Limb));
36933697
defer a.deinit();
36943698

3699+
try a.ensureMulCapacity(&a, &a);
36953700
try a.mul(&a, &a);
36963701

36973702
var want = try Managed.initSet(testing.allocator, 4 * maxInt(Limb) * maxInt(Limb));

0 commit comments

Comments
 (0)