Commit a31b70c4b8d0bed67463b2f54e74198baa93329f | zig

author	LemonBoy <thatlemon@gmail.com> 2020-10-10 00:46:53+02:00
committer	Andrew Kelley <andrew@ziglang.org> 2020-10-09 22:16:48-04:00
log	a31b70c4b8d0bed67463b2f54e74198baa93329f
tree	7be5aefedf537b92c8860b924617e55855e9daab
parent	fbc6a00b0a939f8752bbd571284c46ab58a8fcc4

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
    int_test.zig

lib/std/math/big/int.zig+103-27

...	@@ -446,6 +446,26 @@ pub const Mutable = struct {	...	@@ -446,6 +446,26 @@ pub const Mutable = struct {
446	rma.positive = (a.positive == b.positive);	446	rma.positive = (a.positive == b.positive);
447	}	447	}
448		448
		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
449	/// q = a / b (rem r)	469	/// q = a / b (rem r)
450	///	470	///
451	/// a / b are floored (rounded towards 0).	471	/// a / b are floored (rounded towards 0).
...	@@ -1827,7 +1847,28 @@ pub const Managed = struct {	...	@@ -1827,7 +1847,28 @@ pub const Managed = struct {
1827	rma.setMetadata(m.positive, m.len);	1847	rma.setMetadata(m.positive, m.len);
1828	}	1848	}
1829		1849
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 {
1831	const needed_limbs = calcPowLimbsBufferLen(a.bitCountAbs(), b);	1872	const needed_limbs = calcPowLimbsBufferLen(a.bitCountAbs(), b);
1832		1873
1833	const limbs_buffer = try rma.allocator.alloc(Limb, needed_limbs);	1874	const limbs_buffer = try rma.allocator.alloc(Limb, needed_limbs);
...	@@ -1837,7 +1878,7 @@ pub const Managed = struct {	...	@@ -1837,7 +1878,7 @@ pub const Managed = struct {
1837	var m = try Managed.initCapacity(rma.allocator, needed_limbs);	1878	var m = try Managed.initCapacity(rma.allocator, needed_limbs);
1838	errdefer m.deinit();	1879	errdefer m.deinit();
1839	var m_mut = m.toMutable();	1880	var m_mut = m.toMutable();
1840	try m_mut.pow(a.toConst(), b, limbs_buffer);	1881	try m_mut.pow(a, b, limbs_buffer);
1841	m.setMetadata(m_mut.positive, m_mut.len);	1882	m.setMetadata(m_mut.positive, m_mut.len);
1842		1883
1843	rma.deinit();	1884	rma.deinit();
...	@@ -1845,7 +1886,7 @@ pub const Managed = struct {	...	@@ -1845,7 +1886,7 @@ pub const Managed = struct {
1845	} else {	1886	} else {
1846	try rma.ensureCapacity(needed_limbs);	1887	try rma.ensureCapacity(needed_limbs);
1847	var rma_mut = rma.toMutable();	1888	var rma_mut = rma.toMutable();
1848	try rma_mut.pow(a.toConst(), b, limbs_buffer);	1889	try rma_mut.pow(a, b, limbs_buffer);
1849	rma.setMetadata(rma_mut.positive, rma_mut.len);	1890	rma.setMetadata(rma_mut.positive, rma_mut.len);
1850	}	1891	}
1851	}	1892	}
...	@@ -1869,11 +1910,14 @@ fn llmulacc(opt_allocator: ?*Allocator, r: []Limb, a: []const Limb, b: []const L	...	@@ -1869,11 +1910,14 @@ fn llmulacc(opt_allocator: ?*Allocator, r: []Limb, a: []const Limb, b: []const L
1869	assert(r.len >= x.len + y.len + 1);	1910	assert(r.len >= x.len + y.len + 1);
1870		1911
1871	// 48 is a pretty abitrary size chosen based on performance of a factorial program.	1912	// 48 is a pretty abitrary size chosen based on performance of a factorial program.
1872	if (x.len > 48) {	1913	k_mul: {
1873	if (opt_allocator) \|allocator\| {	1914	if (x.len > 48) {
1874	llmulacc_karatsuba(allocator, r, x, y) catch \|err\| switch (err) {	1915	if (opt_allocator) \|allocator\| {
1875	error.OutOfMemory => {}, // handled below	1916	llmulacc_karatsuba(allocator, r, x, y) catch \|err\| switch (err) {
1876	};	1917	error.OutOfMemory => break :k_mul, // handled below
		1918	};
		1919	return;
		1920	}
1877	}	1921	}
1878	}	1922	}
1879		1923
...	@@ -2203,6 +2247,42 @@ fn llxor(r: []Limb, a: []const Limb, b: []const Limb) void {	...	@@ -2203,6 +2247,42 @@ fn llxor(r: []Limb, a: []const Limb, b: []const Limb) void {
2203	}	2247	}
2204	}	2248	}
2205		2249
		2250	/// r MUST NOT alias x.
		2251	fn 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
2206	/// Knuth 4.6.3	2286	/// Knuth 4.6.3
2207	fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {	2287	fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
2208	var tmp1: []Limb = undefined;	2288	var tmp1: []Limb = undefined;
...	@@ -2212,9 +2292,9 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {	...	@@ -2212,9 +2292,9 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
2212	// variable, use the output limbs and another temporary set to overcome this	2292	// variable, use the output limbs and another temporary set to overcome this
2213	// limitation.	2293	// limitation.
2214	// The initial assignment makes the result end in `r` so an extra memory	2294	// 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	2295	// copy is saved, each 1 flips the index twice so it's only the zeros that
2216	// 0.	2296	// matter.
2217	const b_leading_zeros = @intCast(u5, @clz(u32, b));	2297	const b_leading_zeros = @clz(u32, b);
2218	const exp_zeros = @popCount(u32, ~b) - b_leading_zeros;	2298	const exp_zeros = @popCount(u32, ~b) - b_leading_zeros;
2219	if (exp_zeros & 1 != 0) {	2299	if (exp_zeros & 1 != 0) {
2220	tmp1 = tmp_limbs;	2300	tmp1 = tmp_limbs;
...	@@ -2224,32 +2304,28 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {	...	@@ -2224,32 +2304,28 @@ fn llpow(r: []Limb, a: []const Limb, b: u32, tmp_limbs: []Limb) void {
2224	tmp2 = tmp_limbs;	2304	tmp2 = tmp_limbs;
2225	}	2305	}
2226		2306
2227	const a_norm = a[0..llnormalize(a)];	2307	mem.copy(Limb, tmp1, a);
2228		2308	mem.set(Limb, tmp1[a.len..], 0);
2229	mem.copy(Limb, tmp1, a_norm);
2230	mem.set(Limb, tmp1[a_norm.len..], 0);
2231		2309
2232	// Scan the exponent as a binary number, from left to right, dropping the	2310	// Scan the exponent as a binary number, from left to right, dropping the
2233	// most significant bit set.	2311	// most significant bit set.
2234	const exp_bits = @intCast(u5, 31 - b_leading_zeros);	2312	// Square the result if the current bit is zero, square and multiply by a if
2235	var exp = @bitReverse(u32, b) >> 1 + b_leading_zeros;	2313	// it is one.
		2314	var exp_bits = 32 - 1 - b_leading_zeros;
		2315	var exp = b << @intCast(u5, 1 + b_leading_zeros);
2236		2316
2237	var i: u5 = 0;	2317	var i: usize = 0;
2238	while (i < exp_bits) : (i += 1) {	2318	while (i < exp_bits) : (i += 1) {
2239	// Square	2319	// Square
2240	{	2320	mem.set(Limb, tmp2, 0);
2241	mem.set(Limb, tmp2, 0);	2321	llsquare_basecase(tmp2, tmp1[0..llnormalize(tmp1)]);
2242	const op = tmp1[0..llnormalize(tmp1)];	2322	mem.swap([]Limb, &tmp1, &tmp2);
2243	llmulacc(null, tmp2, op, op);
2244	mem.swap([]Limb, &tmp1, &tmp2);
2245	}
2246	// Multiply by a	2323	// Multiply by a
2247	if (exp & 1 != 0) {	2324	if (@shlWithOverflow(u32, exp, 1, &exp)) {
2248	mem.set(Limb, tmp2, 0);	2325	mem.set(Limb, tmp2, 0);
2249	llmulacc(null, tmp2, tmp1[0..llnormalize(tmp1)], a_norm);	2326	llmulacc(null, tmp2, tmp1[0..llnormalize(tmp1)], a);
2250	mem.swap([]Limb, &tmp1, &tmp2);	2327	mem.swap([]Limb, &tmp1, &tmp2);
2251	}	2328	}
2252	exp >>= 1;
2253	}	2329	}
2254	}	2330	}
2255		2331

lib/std/math/big/int_test.zig+34-10

 @@ -720,6 +720,27 @@ test "big.int mul 0*0" {
     testing.expect((try c.to(u32)) == 0);
 }
+test "big.int mul large" {
+    var a = try Managed.initCapacity(testing.allocator, 50);
+    defer a.deinit();
+    var b = try Managed.initCapacity(testing.allocator, 100);
+    defer b.deinit();
+    var c = try Managed.initCapacity(testing.allocator, 100);
+    defer c.deinit();
+    // Generate a number that's large enough to cross the thresholds for the use
+    // of subquadratic algorithms
+    for (a.limbs) |*p| {
+        p.* = std.math.maxInt(Limb);
+    }
+    a.setMetadata(true, 50);
+    try b.mul(a.toConst(), a.toConst());
+    try c.sqr(a.toConst());
+    testing.expect(b.eq(c));
+}
 test "big.int div single-single no rem" {
     var a = try Managed.initSet(testing.allocator, 50);
     defer a.deinit();
 @@ -1483,11 +1504,14 @@ test "big.int const to managed" {
 test "big.int pow" {
     {
-        var a = try Managed.initSet(testing.allocator, 10);
+        var a = try Managed.initSet(testing.allocator, -3);
         defer a.deinit();
-        try a.pow(a, 8);
+        try a.pow(a.toConst(), 3);
-        testing.expectEqual(@as(u32, 100000000), try a.to(u32));
+        testing.expectEqual(@as(i32, -27), try a.to(i32));
+        try a.pow(a.toConst(), 4);
+        testing.expectEqual(@as(i32, 531441), try a.to(i32));
     }
     {
         var a = try Managed.initSet(testing.allocator, 10);
 @@ -1497,9 +1521,9 @@ test "big.int pow" {
         defer y.deinit();
         // y and a are not aliased
-        try y.pow(a, 123);
+        try y.pow(a.toConst(), 123);
         // y and a are aliased
-        try a.pow(a, 123);
+        try a.pow(a.toConst(), 123);
         testing.expect(a.eq(y));
 @@ -1517,18 +1541,18 @@ test "big.int pow" {
         var a = try Managed.initSet(testing.allocator, 0);
         defer a.deinit();
-        try a.pow(a, 100);
+        try a.pow(a.toConst(), 100);
         testing.expectEqual(@as(i32, 0), try a.to(i32));
         try a.set(1);
-        try a.pow(a, 0);
+        try a.pow(a.toConst(), 0);
         testing.expectEqual(@as(i32, 1), try a.to(i32));
-        try a.pow(a, 100);
+        try a.pow(a.toConst(), 100);
         testing.expectEqual(@as(i32, 1), try a.to(i32));
         try a.set(-1);
-        try a.pow(a, 15);
+        try a.pow(a.toConst(), 15);
         testing.expectEqual(@as(i32, -1), try a.to(i32));
-        try a.pow(a, 16);
+        try a.pow(a.toConst(), 16);
         testing.expectEqual(@as(i32, 1), try a.to(i32));
     }
 }