authorgravatar for thatlemon@gmail.comLemonBoy <thatlemon@gmail.com> 2020-09-30 15:35:05+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-04 02:24:40-04:00
log538d485782629a358c2dd0f8e34d74813e3c9526
treed67182c6085d69a63d3f495d0c180b2b149d52c7
parent2de53592a1d84a1476f662e20d7339d25d4716fe

std: Add pow(a,b) for big ints

Implemented following Knuth's "Evaluation of Powers" chapter in TAOCP, some extra complexity is needed to make sure there's no aliasing and avoid allocating too many limbs. A brief example to illustrate why the last point is important: consider 10^123, since 10 is well within the limits of a single limb we can safely say that the result will surely fit in: ⌈log2(10)⌉ bit * 123 = 492 bits = 7 limbs A naive calculation using only the number of limbs yields: 1 limb * 123 = 123 limbs The space savings are noticeable.

2 files changed, 165 insertions(+), 0 deletions(-)

lib/std/math/big/int.zig+120
......@@ -58,6 +58,11 @@ pub fn calcSetStringLimbCount(base: u8, string_len: usize) usize {
5858 return (string_len + (limb_bits / base - 1)) / (limb_bits / base);
5959}
6060
61pub fn calcPowLimbsBufferLen(a_bit_count: usize, y: usize) usize {
62 // The 1 accounts for the multiplication carry
63 return 1 + (a_bit_count * y + (limb_bits - 1)) / limb_bits;
64}
65
6166/// a + b * c + *carry, sets carry to the overflow bits
6267pub fn addMulLimbWithCarry(a: Limb, b: Limb, c: Limb, carry: *Limb) Limb {
6368 @setRuntimeSafety(debug_safety);
......@@ -597,6 +602,52 @@ pub const Mutable = struct {
597602 return gcdLehmer(rma, x_copy, y_copy, limbs_buffer);
598603 }
599604
605 /// q = a ^ b
606 ///
607 /// r may not alias a.
608 ///
609 /// Asserts that `r` has enough limbs to store the result. Upper bound is
610 /// `calcPowLimbsBufferLen(a.bitCountAbs(), b)`.
611 ///
612 /// `limbs_buffer` is used for temporary storage.
613 /// The amount required is given by `calcPowLimbsBufferLen`.
614 pub fn pow(r: *Mutable, a: Const, b: u32, limbs_buffer: []Limb) !void {
615 assert(r.limbs.ptr != a.limbs.ptr); // illegal aliasing
616
617 // Handle all the trivial cases first
618 switch (b) {
619 0 => {
620 // a^0 = 1
621 return r.set(1);
622 },
623 1 => {
624 // a^1 = a
625 return r.copy(a);
626 },
627 else => {},
628 }
629
630 if (a.eqZero()) {
631 // 0^b = 0
632 return r.set(0);
633 } else if (a.limbs.len == 1 and a.limbs[0] == 1) {
634 // 1^b = 1 and -1^b = ±1
635 r.set(1);
636 r.positive = a.positive or (b & 1) == 0;
637 return;
638 }
639
640 // Here a>1 and b>1
641 const needed_limbs = calcPowLimbsBufferLen(a.bitCountAbs(), b);
642 assert(r.limbs.len >= needed_limbs);
643 assert(limbs_buffer.len >= needed_limbs);
644
645 llpow(r.limbs, a.limbs, b, limbs_buffer);
646
647 r.normalize(needed_limbs);
648 r.positive = a.positive or (b & 1) == 0;
649 }
650
600651 /// rma may not alias x or y.
601652 /// x and y may alias each other.
602653 /// Asserts that `rma` has enough limbs to store the result. Upper bound is given by `calcGcdNoAliasLimbLen`.
......@@ -1775,6 +1826,29 @@ pub const Managed = struct {
17751826 try m.gcd(x.toConst(), y.toConst(), &limbs_buffer);
17761827 rma.setMetadata(m.positive, m.len);
17771828 }
1829
1830 pub fn pow(rma: *Managed, a: Managed, b: u32) !void {
1831 const needed_limbs = calcPowLimbsBufferLen(a.bitCountAbs(), b);
1832
1833 const limbs_buffer = try rma.allocator.alloc(Limb, needed_limbs);
1834 defer rma.allocator.free(limbs_buffer);
1835
1836 if (rma.limbs.ptr == a.limbs.ptr) {
1837 var m = try Managed.initCapacity(rma.allocator, needed_limbs);
1838 errdefer m.deinit();
1839 var m_mut = m.toMutable();
1840 try m_mut.pow(a.toConst(), b, limbs_buffer);
1841 m.setMetadata(m_mut.positive, m_mut.len);
1842
1843 rma.deinit();
1844 rma.swap(&m);
1845 } else {
1846 try rma.ensureCapacity(needed_limbs);
1847 var rma_mut = rma.toMutable();
1848 try rma_mut.pow(a.toConst(), b, limbs_buffer);
1849 rma.setMetadata(rma_mut.positive, rma_mut.len);
1850 }
1851 }
17781852};
17791853
17801854/// Knuth 4.3.1, Algorithm M.
......@@ -2129,6 +2203,52 @@ fn llxor(r: []Limb, a: []const Limb, b: []const Limb) void {
21292203 }
21302204}
21312205
2206/// Knuth 4.6.3
2207fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
2208 mem.copy(Limb, r, a);
2209 mem.set(Limb, r[a.len..], 0);
2210
2211 // Multiplication requires no aliasing between the operand and the result
2212 // variable, use the output limbs and another temporary set to overcome this
2213 // limit.
2214 // Note that the order is important in the code below.
2215 var list = [_][]Limb{ r, tmp_limbs };
2216 var index: usize = 0;
2217
2218 // Scan the exponent as a binary number, from left to right, dropping the
2219 // most significant bit set
2220 var exp = @bitReverse(u32, b) >> (1 + @intCast(u5, @clz(u32, b)));
2221 while (exp != 0) : (exp >>= 1) {
2222 // Square
2223 {
2224 const cur_buf = list[index];
2225 const cur_buf_len = llnormalize(cur_buf);
2226 const cur_buf_out = list[index ^ 1];
2227
2228 mem.set(Limb, cur_buf_out, 0);
2229 llmulacc(null, cur_buf_out, cur_buf[0..cur_buf_len], cur_buf[0..cur_buf_len]);
2230
2231 index ^= 1;
2232 }
2233
2234 if ((exp & 1) != 0) {
2235 // Multiply
2236 const cur_buf = list[index];
2237 const cur_buf_len = llnormalize(cur_buf);
2238 const cur_buf_out = list[index ^ 1];
2239
2240 mem.set(Limb, cur_buf_out, 0);
2241 llmulacc(null, cur_buf_out, cur_buf, a);
2242
2243 index ^= 1;
2244 }
2245 }
2246
2247 if (index != 0) {
2248 mem.copy(Limb, r, tmp_limbs);
2249 }
2250}
2251
21322252// Storage must live for the lifetime of the returned value
21332253fn fixedIntFromSignedDoubleLimb(A: SignedDoubleLimb, storage: []Limb) Mutable {
21342254 assert(storage.len >= 2);
lib/std/math/big/int_test.zig+45
......@@ -1480,3 +1480,48 @@ test "big.int const to managed" {
14801480
14811481 testing.expect(a.toConst().eq(b.toConst()));
14821482}
1483
1484test "big.int pow" {
1485 {
1486 var a = try Managed.initSet(testing.allocator, 10);
1487 defer a.deinit();
1488
1489 var y = try Managed.init(testing.allocator);
1490 defer y.deinit();
1491
1492 // y and a are not aliased
1493 try y.pow(a, 123);
1494 // y and a are aliased
1495 try a.pow(a, 123);
1496
1497 testing.expect(a.eq(y));
1498
1499 const ys = try y.toString(testing.allocator, 16, false);
1500 defer testing.allocator.free(ys);
1501 testing.expectEqualSlices(
1502 u8,
1503 "183425a5f872f126e00a5ad62c839075cd6846c6fb0230887c7ad7a9dc530fcb" ++
1504 "4933f60e8000000000000000000000000000000",
1505 ys,
1506 );
1507 }
1508 // Special cases
1509 {
1510 var a = try Managed.initSet(testing.allocator, 0);
1511 defer a.deinit();
1512
1513 try a.pow(a, 100);
1514 testing.expectEqual(@as(i32, 0), try a.to(i32));
1515
1516 try a.set(1);
1517 try a.pow(a, 0);
1518 testing.expectEqual(@as(i32, 1), try a.to(i32));
1519 try a.pow(a, 100);
1520 testing.expectEqual(@as(i32, 1), try a.to(i32));
1521 try a.set(-1);
1522 try a.pow(a, 15);
1523 testing.expectEqual(@as(i32, -1), try a.to(i32));
1524 try a.pow(a, 16);
1525 testing.expectEqual(@as(i32, 1), try a.to(i32));
1526 }
1527}