authorgravatar for thatlemon@gmail.comLemonBoy <thatlemon@gmail.com> 2020-10-10 00:46:53+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-09 22:16:48-04:00
loga31b70c4b8d0bed67463b2f54e74198baa93329f
tree7be5aefedf537b92c8860b924617e55855e9daab
parentfbc6a00b0a939f8752bbd571284c46ab58a8fcc4

std: Add/Fix/Change parts of big.int

* Add an optimized squaring routine under the `sqr` name. Algorithms for squaring bigger numbers efficiently will come in a PR later. * Fix a bug where a multiplication was done twice if the threshold for the use of Karatsuba algorithm was crossed. Add a test to make sure this won't happen again. * Streamline `pow` method, take a `Const` parameter. * Minor tweaks to `pow`, avoid bit-reversing the exponent.

2 files changed, 137 insertions(+), 37 deletions(-)

lib/std/math/big/int.zig+103-27
......@@ -446,6 +446,26 @@ pub const Mutable = struct {
446446 rma.positive = (a.positive == b.positive);
447447 }
448448
449 /// rma = a * a
450 ///
451 /// `rma` may not alias with `a`.
452 ///
453 /// Asserts the result fits in `rma`. An upper bound on the number of limbs needed by
454 /// rma is given by `2 * a.limbs.len + 1`.
455 ///
456 /// If `allocator` is provided, it will be used for temporary storage to improve
457 /// multiplication performance. `error.OutOfMemory` is handled with a fallback algorithm.
458 pub fn sqrNoAlias(rma: *Mutable, a: Const, opt_allocator: ?*Allocator) void {
459 assert(rma.limbs.ptr != a.limbs.ptr); // illegal aliasing
460
461 mem.set(Limb, rma.limbs, 0);
462
463 llsquare_basecase(rma.limbs, a.limbs);
464
465 rma.normalize(2 * a.limbs.len + 1);
466 rma.positive = true;
467 }
468
449469 /// q = a / b (rem r)
450470 ///
451471 /// a / b are floored (rounded towards 0).
......@@ -1827,7 +1847,28 @@ pub const Managed = struct {
18271847 rma.setMetadata(m.positive, m.len);
18281848 }
18291849
1830 pub fn pow(rma: *Managed, a: Managed, b: u32) !void {
1850 /// r = a * a
1851 pub fn sqr(rma: *Managed, a: Const) !void {
1852 const needed_limbs = 2 * a.limbs.len + 1;
1853
1854 if (rma.limbs.ptr == a.limbs.ptr) {
1855 var m = try Managed.initCapacity(rma.allocator, needed_limbs);
1856 errdefer m.deinit();
1857 var m_mut = m.toMutable();
1858 m_mut.sqrNoAlias(a, rma.allocator);
1859 m.setMetadata(m_mut.positive, m_mut.len);
1860
1861 rma.deinit();
1862 rma.swap(&m);
1863 } else {
1864 try rma.ensureCapacity(needed_limbs);
1865 var rma_mut = rma.toMutable();
1866 rma_mut.sqrNoAlias(a, rma.allocator);
1867 rma.setMetadata(rma_mut.positive, rma_mut.len);
1868 }
1869 }
1870
1871 pub fn pow(rma: *Managed, a: Const, b: u32) !void {
18311872 const needed_limbs = calcPowLimbsBufferLen(a.bitCountAbs(), b);
18321873
18331874 const limbs_buffer = try rma.allocator.alloc(Limb, needed_limbs);
......@@ -1837,7 +1878,7 @@ pub const Managed = struct {
18371878 var m = try Managed.initCapacity(rma.allocator, needed_limbs);
18381879 errdefer m.deinit();
18391880 var m_mut = m.toMutable();
1840 try m_mut.pow(a.toConst(), b, limbs_buffer);
1881 try m_mut.pow(a, b, limbs_buffer);
18411882 m.setMetadata(m_mut.positive, m_mut.len);
18421883
18431884 rma.deinit();
......@@ -1845,7 +1886,7 @@ pub const Managed = struct {
18451886 } else {
18461887 try rma.ensureCapacity(needed_limbs);
18471888 var rma_mut = rma.toMutable();
1848 try rma_mut.pow(a.toConst(), b, limbs_buffer);
1889 try rma_mut.pow(a, b, limbs_buffer);
18491890 rma.setMetadata(rma_mut.positive, rma_mut.len);
18501891 }
18511892 }
......@@ -1869,11 +1910,14 @@ fn llmulacc(opt_allocator: ?*Allocator, r: []Limb, a: []const Limb, b: []const L
18691910 assert(r.len >= x.len + y.len + 1);
18701911
18711912 // 48 is a pretty abitrary size chosen based on performance of a factorial program.
1872 if (x.len > 48) {
1873 if (opt_allocator) |allocator| {
1874 llmulacc_karatsuba(allocator, r, x, y) catch |err| switch (err) {
1875 error.OutOfMemory => {}, // handled below
1876 };
1913 k_mul: {
1914 if (x.len > 48) {
1915 if (opt_allocator) |allocator| {
1916 llmulacc_karatsuba(allocator, r, x, y) catch |err| switch (err) {
1917 error.OutOfMemory => break :k_mul, // handled below
1918 };
1919 return;
1920 }
18771921 }
18781922 }
18791923
......@@ -2203,6 +2247,42 @@ fn llxor(r: []Limb, a: []const Limb, b: []const Limb) void {
22032247 }
22042248}
22052249
2250/// r MUST NOT alias x.
2251fn llsquare_basecase(r: []Limb, x: []const Limb) void {
2252 @setRuntimeSafety(debug_safety);
2253
2254 const x_norm = x;
2255 assert(r.len >= 2 * x_norm.len + 1);
2256
2257 // Compute the square of a N-limb bigint with only (N^2 + N)/2
2258 // multiplications by exploting the symmetry of the coefficients around the
2259 // diagonal:
2260 //
2261 // a b c *
2262 // a b c =
2263 // -------------------
2264 // ca cb cc +
2265 // ba bb bc +
2266 // aa ab ac
2267 //
2268 // Note that:
2269 // - Each mixed-product term appears twice for each column,
2270 // - Squares are always in the 2k (0 <= k < N) column
2271
2272 for (x_norm) |v, i| {
2273 // Accumulate all the x[i]*x[j] (with x!=j) products
2274 llmulDigit(r[2 * i + 1 ..], x_norm[i + 1 ..], v);
2275 }
2276
2277 // Each product appears twice, multiply by 2
2278 llshl(r, r[0 .. 2 * x_norm.len], 1);
2279
2280 for (x_norm) |v, i| {
2281 // Compute and add the squares
2282 llmulDigit(r[2 * i ..], x[i .. i + 1], v);
2283 }
2284}
2285
22062286/// Knuth 4.6.3
22072287fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
22082288 var tmp1: []Limb = undefined;
......@@ -2212,9 +2292,9 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
22122292 // variable, use the output limbs and another temporary set to overcome this
22132293 // limitation.
22142294 // The initial assignment makes the result end in `r` so an extra memory
2215 // copy is saved, each 1 flips the index twice so it's a no-op so count the
2216 // 0.
2217 const b_leading_zeros = @intCast(u5, @clz(u32, b));
2295 // copy is saved, each 1 flips the index twice so it's only the zeros that
2296 // matter.
2297 const b_leading_zeros = @clz(u32, b);
22182298 const exp_zeros = @popCount(u32, ~b) - b_leading_zeros;
22192299 if (exp_zeros & 1 != 0) {
22202300 tmp1 = tmp_limbs;
......@@ -2224,32 +2304,28 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
22242304 tmp2 = tmp_limbs;
22252305 }
22262306
2227 const a_norm = a[0..llnormalize(a)];
2228
2229 mem.copy(Limb, tmp1, a_norm);
2230 mem.set(Limb, tmp1[a_norm.len..], 0);
2307 mem.copy(Limb, tmp1, a);
2308 mem.set(Limb, tmp1[a.len..], 0);
22312309
22322310 // Scan the exponent as a binary number, from left to right, dropping the
22332311 // most significant bit set.
2234 const exp_bits = @intCast(u5, 31 - b_leading_zeros);
2235 var exp = @bitReverse(u32, b) >> 1 + b_leading_zeros;
2312 // Square the result if the current bit is zero, square and multiply by a if
2313 // it is one.
2314 var exp_bits = 32 - 1 - b_leading_zeros;
2315 var exp = b << @intCast(u5, 1 + b_leading_zeros);
22362316
2237 var i: u5 = 0;
2317 var i: usize = 0;
22382318 while (i < exp_bits) : (i += 1) {
22392319 // Square
2240 {
2241 mem.set(Limb, tmp2, 0);
2242 const op = tmp1[0..llnormalize(tmp1)];
2243 llmulacc(null, tmp2, op, op);
2244 mem.swap([]Limb, &tmp1, &tmp2);
2245 }
2320 mem.set(Limb, tmp2, 0);
2321 llsquare_basecase(tmp2, tmp1[0..llnormalize(tmp1)]);
2322 mem.swap([]Limb, &tmp1, &tmp2);
22462323 // Multiply by a
2247 if (exp & 1 != 0) {
2324 if (@shlWithOverflow(u32, exp, 1, &exp)) {
22482325 mem.set(Limb, tmp2, 0);
2249 llmulacc(null, tmp2, tmp1[0..llnormalize(tmp1)], a_norm);
2326 llmulacc(null, tmp2, tmp1[0..llnormalize(tmp1)], a);
22502327 mem.swap([]Limb, &tmp1, &tmp2);
22512328 }
2252 exp >>= 1;
22532329 }
22542330}
22552331
lib/std/math/big/int_test.zig+34-10
......@@ -720,6 +720,27 @@ test "big.int mul 0*0" {
720720 testing.expect((try c.to(u32)) == 0);
721721}
722722
723test "big.int mul large" {
724 var a = try Managed.initCapacity(testing.allocator, 50);
725 defer a.deinit();
726 var b = try Managed.initCapacity(testing.allocator, 100);
727 defer b.deinit();
728 var c = try Managed.initCapacity(testing.allocator, 100);
729 defer c.deinit();
730
731 // Generate a number that's large enough to cross the thresholds for the use
732 // of subquadratic algorithms
733 for (a.limbs) |*p| {
734 p.* = std.math.maxInt(Limb);
735 }
736 a.setMetadata(true, 50);
737
738 try b.mul(a.toConst(), a.toConst());
739 try c.sqr(a.toConst());
740
741 testing.expect(b.eq(c));
742}
743
723744test "big.int div single-single no rem" {
724745 var a = try Managed.initSet(testing.allocator, 50);
725746 defer a.deinit();
......@@ -1483,11 +1504,14 @@ test "big.int const to managed" {
14831504
14841505test "big.int pow" {
14851506 {
1486 var a = try Managed.initSet(testing.allocator, 10);
1507 var a = try Managed.initSet(testing.allocator, -3);
14871508 defer a.deinit();
14881509
1489 try a.pow(a, 8);
1490 testing.expectEqual(@as(u32, 100000000), try a.to(u32));
1510 try a.pow(a.toConst(), 3);
1511 testing.expectEqual(@as(i32, -27), try a.to(i32));
1512
1513 try a.pow(a.toConst(), 4);
1514 testing.expectEqual(@as(i32, 531441), try a.to(i32));
14911515 }
14921516 {
14931517 var a = try Managed.initSet(testing.allocator, 10);
......@@ -1497,9 +1521,9 @@ test "big.int pow" {
14971521 defer y.deinit();
14981522
14991523 // y and a are not aliased
1500 try y.pow(a, 123);
1524 try y.pow(a.toConst(), 123);
15011525 // y and a are aliased
1502 try a.pow(a, 123);
1526 try a.pow(a.toConst(), 123);
15031527
15041528 testing.expect(a.eq(y));
15051529
......@@ -1517,18 +1541,18 @@ test "big.int pow" {
15171541 var a = try Managed.initSet(testing.allocator, 0);
15181542 defer a.deinit();
15191543
1520 try a.pow(a, 100);
1544 try a.pow(a.toConst(), 100);
15211545 testing.expectEqual(@as(i32, 0), try a.to(i32));
15221546
15231547 try a.set(1);
1524 try a.pow(a, 0);
1548 try a.pow(a.toConst(), 0);
15251549 testing.expectEqual(@as(i32, 1), try a.to(i32));
1526 try a.pow(a, 100);
1550 try a.pow(a.toConst(), 100);
15271551 testing.expectEqual(@as(i32, 1), try a.to(i32));
15281552 try a.set(-1);
1529 try a.pow(a, 15);
1553 try a.pow(a.toConst(), 15);
15301554 testing.expectEqual(@as(i32, -1), try a.to(i32));
1531 try a.pow(a, 16);
1555 try a.pow(a.toConst(), 16);
15321556 testing.expectEqual(@as(i32, 1), try a.to(i32));
15331557 }
15341558}