authorgravatar for keitam913@yahoo.co.jpMizuochi Keita <keitam913@yahoo.co.jp> 2023-05-24 23:54:51+09:00
committergravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2023-05-29 13:04:32+03:00
log4422af8be96d243db7840b840737069c38be8afb
treebc7105f165cb8877b9365498a440b81ea1f7f8de
parent6e6a61a3847092be8a754f70f19ad3c779648ba3

std.math.big.int: Add Sqrt

Implemented with reference to Modern Computer Arithmetic, Algorithm 1.13. https://members.loria.fr/PZimmermann/mca/pub226.html The below optimization ideas are derived from Go's big package. * Minimize initial loop value * Reuse loop values math/big/int.go: https://cs.opensource.google/go/go/+/refs/tags/go1.20.4:src/math/big/int.go;l=1286

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

lib/std/math/big/int.zig+78
......@@ -66,6 +66,13 @@ pub fn calcPowLimbsBufferLen(a_bit_count: usize, y: usize) usize {
6666 return 2 + (a_bit_count * y + (limb_bits - 1)) / limb_bits;
6767}
6868
69pub fn calcSqrtLimbsBufferLen(a_bit_count: usize) usize {
70 const a_limb_count = (a_bit_count - 1) / limb_bits + 1;
71 const shift = (a_bit_count + 1) / 2;
72 const u_s_rem_limb_count = 1 + ((shift - 1) / limb_bits + 1);
73 return a_limb_count + 3 * u_s_rem_limb_count + calcDivLimbsBufferLen(a_limb_count, u_s_rem_limb_count);
74}
75
6976// Compute the number of limbs required to store a 2s-complement number of `bit_count` bits.
7077pub fn calcTwosCompLimbCount(bit_count: usize) usize {
7178 return std.math.divCeil(usize, bit_count, @bitSizeOf(Limb)) catch unreachable;
......@@ -1344,6 +1351,64 @@ pub const Mutable = struct {
13441351 r.positive = a.positive or (b & 1) == 0;
13451352 }
13461353
1354 /// r = ⌊√a⌋
1355 ///
1356 /// r may alias a.
1357 ///
1358 /// Asserts that `r` has enough limbs to store the result. Upper bound is
1359 /// `(a.limbs.len - 1) / 2 + 1`.
1360 ///
1361 /// `limbs_buffer` is used for temporary storage.
1362 /// The amount required is given by `calcSqrtLimbsBufferLen`.
1363 pub fn sqrt(
1364 r: *Mutable,
1365 a: Const,
1366 limbs_buffer: []Limb,
1367 ) void {
1368 // Brent and Zimmermann, Modern Computer Arithmetic, Algorithm 1.13 SqrtInt
1369 // https://members.loria.fr/PZimmermann/mca/pub226.html
1370 var buf_index: usize = 0;
1371 var t = b: {
1372 const start = buf_index;
1373 buf_index += a.limbs.len;
1374 break :b Mutable.init(limbs_buffer[start..buf_index], 0);
1375 };
1376 var u = b: {
1377 const start = buf_index;
1378 const shift = (a.bitCountAbs() + 1) / 2;
1379 buf_index += 1 + ((shift - 1) / limb_bits + 1);
1380 var m = Mutable.init(limbs_buffer[start..buf_index], 1);
1381 m.shiftLeft(m.toConst(), shift); // u must be >= ⌊√a⌋, and should be as small as possible for efficiency
1382 break :b m;
1383 };
1384 var s = b: {
1385 const start = buf_index;
1386 buf_index += u.limbs.len;
1387 break :b u.toConst().toMutable(limbs_buffer[start..buf_index]);
1388 };
1389 var rem = b: {
1390 const start = buf_index;
1391 buf_index += s.limbs.len;
1392 break :b Mutable.init(limbs_buffer[start..buf_index], 0);
1393 };
1394
1395 while (true) {
1396 t.divFloor(&rem, a, s.toConst(), limbs_buffer[buf_index..]);
1397 t.add(t.toConst(), s.toConst());
1398 u.shiftRight(t.toConst(), 1);
1399
1400 if (u.toConst().order(s.toConst()).compare(.gte)) {
1401 r.copy(s.toConst());
1402 return;
1403 }
1404
1405 // Avoid copying u to s by swapping u and s
1406 var tmp_s = s;
1407 s = u;
1408 u = tmp_s;
1409 }
1410 }
1411
13471412 /// rma may not alias x or y.
13481413 /// x and y may alias each other.
13491414 /// Asserts that `rma` has enough limbs to store the result. Upper bound is given by `calcGcdNoAliasLimbLen`.
......@@ -3140,6 +3205,19 @@ pub const Managed = struct {
31403205 }
31413206 }
31423207
3208 /// r = ⌊√a⌋
3209 pub fn sqrt(rma: *Managed, a: *const Managed) !void {
3210 const needed_limbs = calcSqrtLimbsBufferLen(a.bitCountAbs());
3211
3212 const limbs_buffer = try rma.allocator.alloc(Limb, needed_limbs);
3213 defer rma.allocator.free(limbs_buffer);
3214
3215 try rma.ensureCapacity((a.len() - 1) / 2 + 1);
3216 var m = rma.toMutable();
3217 m.sqrt(a.toConst(), limbs_buffer);
3218 rma.setMetadata(m.positive, m.len);
3219 }
3220
31433221 /// r = truncate(Int(signedness, bit_count), a)
31443222 pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
31453223 try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
lib/std/math/big/int_test.zig+30
......@@ -2622,6 +2622,36 @@ test "big.int pow" {
26222622 }
26232623}
26242624
2625test "big.int sqrt" {
2626 var r = try Managed.init(testing.allocator);
2627 defer r.deinit();
2628 var a = try Managed.init(testing.allocator);
2629 defer a.deinit();
2630
2631 // not aliased
2632 try r.set(0);
2633 try a.set(25);
2634 try r.sqrt(&a);
2635 try testing.expectEqual(@as(i32, 5), try r.to(i32));
2636
2637 // aliased
2638 try a.set(25);
2639 try a.sqrt(&a);
2640 try testing.expectEqual(@as(i32, 5), try a.to(i32));
2641
2642 // bottom
2643 try r.set(0);
2644 try a.set(24);
2645 try r.sqrt(&a);
2646 try testing.expectEqual(@as(i32, 4), try r.to(i32));
2647
2648 // large number
2649 try r.set(0);
2650 try a.set(0x1_0000_0000_0000);
2651 try r.sqrt(&a);
2652 try testing.expectEqual(@as(i32, 0x100_0000), try r.to(i32));
2653}
2654
26252655test "big.int regression test for 1 limb overflow with alias" {
26262656 // Note these happen to be two consecutive Fibonacci sequence numbers, the
26272657 // first two whose sum exceeds 2**64.