Skip to content

Commit c8011c6

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

File tree

2 files changed

+101
-29
lines changed

2 files changed

+101
-29
lines changed

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

Lines changed: 92 additions & 25 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,26 +3157,54 @@ pub const Managed = struct {
31403157
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
31413158
}
31423159

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+
31433174
pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void {
31443175
try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1);
31453176
}
31463177

3178+
pub fn ensureAddScalarCapacityManaged(r: *Managed, a: *const Managed, scalar: anytype) !void {
3179+
try r.ensureCapacity(@max(a.len(), calcLimbLen(scalar)) + 1);
3180+
}
3181+
31473182
pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void {
31483183
try r.ensureCapacity(@max(a.limbs.len, b.limbs.len) + 1);
31493184
}
31503185

3186+
pub fn ensureAddCapacityManaged(r: *Managed, a: *const Managed, b: *const Managed) !void {
3187+
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3188+
}
3189+
31513190
pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void {
31523191
try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1);
31533192
}
31543193

3194+
pub fn ensureMulCapacityManaged(rma: *Managed, a: *const Managed, b: *const Managed) !void {
3195+
try rma.ensureCapacity(a.len() + b.len() + 1);
3196+
}
3197+
31553198
/// q = a / b (rem r)
31563199
///
31573200
/// a / b are floored (rounded towards 0).
31583201
///
31593202
/// Returns an error if memory could not be allocated.
31603203
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());
3204+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3205+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3206+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3207+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31633208
var mq = q.toMutable();
31643209
var mr = r.toMutable();
31653210
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3175,8 +3220,10 @@ pub const Managed = struct {
31753220
///
31763221
/// Returns an error if memory could not be allocated.
31773222
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());
3223+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3224+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3225+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3226+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31803227
var mq = q.toMutable();
31813228
var mr = r.toMutable();
31823229
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3189,7 +3236,9 @@ pub const Managed = struct {
31893236
/// r = a << shift, in other words, r = a * 2^shift
31903237
/// r and a may alias.
31913238
pub fn shiftLeft(r: *Managed, a: *const Managed, shift: usize) !void {
3192-
try r.ensureCapacity(a.len() + (shift / limb_bits) + 1);
3239+
const aliased = limbsAliasDistinct(r, a);
3240+
const needed = a.len() + (shift / limb_bits) + 1;
3241+
try r.ensureAliasAwareCapacity(needed, aliased);
31933242
var m = r.toMutable();
31943243
m.shiftLeft(a.toConst(), shift);
31953244
r.setMetadata(m.positive, m.len);
@@ -3198,7 +3247,9 @@ pub const Managed = struct {
31983247
/// r = a <<| shift with 2s-complement saturating semantics.
31993248
/// r and a may alias.
32003249
pub fn shiftLeftSat(r: *Managed, a: *const Managed, shift: usize, signedness: Signedness, bit_count: usize) !void {
3201-
try r.ensureTwosCompCapacity(bit_count);
3250+
const aliased = limbsAliasDistinct(r, a);
3251+
const needed = calcTwosCompLimbCount(bit_count);
3252+
try r.ensureAliasAwareCapacity(needed, aliased);
32023253
var m = r.toMutable();
32033254
m.shiftLeftSat(a.toConst(), shift, signedness, bit_count);
32043255
r.setMetadata(m.positive, m.len);
@@ -3220,7 +3271,9 @@ pub const Managed = struct {
32203271
return;
32213272
}
32223273

3223-
try r.ensureCapacity(a.len() - (shift / limb_bits));
3274+
const aliased = limbsAliasDistinct(r, a);
3275+
const needed = a.len() - (shift / limb_bits);
3276+
try r.ensureAliasAwareCapacity(needed, aliased);
32243277
var m = r.toMutable();
32253278
m.shiftRight(a.toConst(), shift);
32263279
r.setMetadata(m.positive, m.len);
@@ -3229,7 +3282,9 @@ pub const Managed = struct {
32293282
/// r = ~a under 2s-complement wrapping semantics.
32303283
/// r and a may alias.
32313284
pub fn bitNotWrap(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3232-
try r.ensureTwosCompCapacity(bit_count);
3285+
const aliased = limbsAliasDistinct(r, a);
3286+
const needed = calcTwosCompLimbCount(bit_count);
3287+
try r.ensureAliasAwareCapacity(needed, aliased);
32333288
var m = r.toMutable();
32343289
m.bitNotWrap(a.toConst(), signedness, bit_count);
32353290
r.setMetadata(m.positive, m.len);
@@ -3239,7 +3294,9 @@ pub const Managed = struct {
32393294
///
32403295
/// a and b are zero-extended to the longer of a or b.
32413296
pub fn bitOr(r: *Managed, a: *const Managed, b: *const Managed) !void {
3242-
try r.ensureCapacity(@max(a.len(), b.len()));
3297+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3298+
const needed = @max(a.len(), b.len());
3299+
try r.ensureAliasAwareCapacity(needed, aliased);
32433300
var m = r.toMutable();
32443301
m.bitOr(a.toConst(), b.toConst());
32453302
r.setMetadata(m.positive, m.len);
@@ -3251,7 +3308,8 @@ pub const Managed = struct {
32513308
if (b.isPositive()) b.len() else if (a.isPositive()) a.len() else a.len() + 1
32523309
else if (a.isPositive()) a.len() else if (b.isPositive()) b.len() else b.len() + 1;
32533310

3254-
try r.ensureCapacity(cap);
3311+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3312+
try r.ensureAliasAwareCapacity(cap, aliased);
32553313
var m = r.toMutable();
32563314
m.bitAnd(a.toConst(), b.toConst());
32573315
r.setMetadata(m.positive, m.len);
@@ -3260,7 +3318,8 @@ pub const Managed = struct {
32603318
/// r = a ^ b
32613319
pub fn bitXor(r: *Managed, a: *const Managed, b: *const Managed) !void {
32623320
const cap = @max(a.len(), b.len()) + @intFromBool(a.isPositive() != b.isPositive());
3263-
try r.ensureCapacity(cap);
3321+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3322+
try r.ensureAliasAwareCapacity(cap, aliased);
32643323

32653324
var m = r.toMutable();
32663325
m.bitXor(a.toConst(), b.toConst());
@@ -3272,7 +3331,9 @@ pub const Managed = struct {
32723331
///
32733332
/// rma's allocator is used for temporary storage to boost multiplication performance.
32743333
pub fn gcd(rma: *Managed, x: *const Managed, y: *const Managed) !void {
3275-
try rma.ensureCapacity(@min(x.len(), y.len()));
3334+
const aliased = limbsAliasDistinct(rma, x) or limbsAliasDistinct(rma, y);
3335+
const needed = @min(x.len(), y.len());
3336+
try rma.ensureAliasAwareCapacity(needed, aliased);
32763337
var m = rma.toMutable();
32773338
var limbs_buffer = std.array_list.Managed(Limb).init(rma.allocator);
32783339
defer limbs_buffer.deinit();
@@ -3350,15 +3411,19 @@ pub const Managed = struct {
33503411

33513412
/// r = truncate(Int(signedness, bit_count), a)
33523413
pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3353-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3414+
const aliased = limbsAliasDistinct(r, a);
3415+
const needed = calcTwosCompLimbCount(bit_count);
3416+
try r.ensureAliasAwareCapacity(needed, aliased);
33543417
var m = r.toMutable();
33553418
m.truncate(a.toConst(), signedness, bit_count);
33563419
r.setMetadata(m.positive, m.len);
33573420
}
33583421

33593422
/// r = saturate(Int(signedness, bit_count), a)
33603423
pub fn saturate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3361-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3424+
const aliased = limbsAliasDistinct(r, a);
3425+
const needed = calcTwosCompLimbCount(bit_count);
3426+
try r.ensureAliasAwareCapacity(needed, aliased);
33623427
var m = r.toMutable();
33633428
m.saturate(a.toConst(), signedness, bit_count);
33643429
r.setMetadata(m.positive, m.len);
@@ -3367,7 +3432,9 @@ pub const Managed = struct {
33673432
/// r = @popCount(a) with 2s-complement semantics.
33683433
/// r and a may be aliases.
33693434
pub fn popCount(r: *Managed, a: *const Managed, bit_count: usize) !void {
3370-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3435+
const aliased = limbsAliasDistinct(r, a);
3436+
const needed = calcTwosCompLimbCount(bit_count);
3437+
try r.ensureAliasAwareCapacity(needed, aliased);
33713438
var m = r.toMutable();
33723439
m.popCount(a.toConst(), bit_count);
33733440
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.ensureAddCapacityManaged(&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.ensureAddCapacityManaged(&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.ensureMulCapacityManaged(&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)