std.math.big: fix sqrt with bits > limb_bits

Signed-off-by: tison <wander4096@gmail.com>

2 files changed, 21 insertions, 2 deletions

diff --git a/lib/std/math/big/int.zig b/lib/std/math/big/int.zig
index b182993885..346da746d0 100644
--- a/lib/std/math/big/int.zig
+++ b/lib/std/math/big/int.zig
@@ -69,7 +69,7 @@ pub fn calcPowLimbsBufferLen(a_bit_count: usize, y: usize) usize {
 pub fn calcSqrtLimbsBufferLen(a_bit_count: usize) usize {
     const a_limb_count = (a_bit_count - 1) / limb_bits + 1;
     const shift = (a_bit_count + 1) / 2;
-    const u_s_rem_limb_count = 1 + ((shift - 1) / limb_bits + 1);
+    const u_s_rem_limb_count = 1 + ((shift / limb_bits) + 1);
     return a_limb_count + 3 * u_s_rem_limb_count + calcDivLimbsBufferLen(a_limb_count, u_s_rem_limb_count);
 }
 
@@ -1380,7 +1380,7 @@ pub const Mutable = struct {
         var u = b: {
             const start = buf_index;
             const shift = (a.bitCountAbs() + 1) / 2;
-            buf_index += 1 + ((shift - 1) / limb_bits + 1);
+            buf_index += 1 + ((shift / limb_bits) + 1);
             var m = Mutable.init(limbs_buffer[start..buf_index], 1);
             m.shiftLeft(m.toConst(), shift); // u must be >= ⌊√a⌋, and should be as small as possible for efficiency
             break :b m;
diff --git a/lib/std/math/big/int_test.zig b/lib/std/math/big/int_test.zig
index ac4326ec4e..a925f2c4b2 100644
--- a/lib/std/math/big/int_test.zig
+++ b/lib/std/math/big/int_test.zig
@@ -3161,6 +3161,7 @@ test "big.int.Managed sqrt(0) = 0" {
     try a.setString(10, "0");
 
     try res.sqrt(&a);
+    try testing.expectEqual(@as(i32, 0), try res.to(i32));
 }
 
 test "big.int.Managed sqrt(-1) = error" {
@@ -3175,3 +3176,21 @@ test "big.int.Managed sqrt(-1) = error" {
 
     try testing.expectError(error.SqrtOfNegativeNumber, res.sqrt(&a));
 }
+
+test "big.int.Managed sqrt(n) succeed with res.bitCountAbs() >= usize bits" {
+    const allocator = testing.allocator;
+    var a = try Managed.initSet(allocator, 1);
+    defer a.deinit();
+
+    var res = try Managed.initSet(allocator, 1);
+    defer res.deinit();
+
+    // a.bitCountAbs() = 127 so the first attempt has 64 bits >= usize bits
+    try a.setString(10, "136036462105870278006290938611834481486");
+    try res.sqrt(&a);
+
+    var expected = try Managed.initSet(allocator, 1);
+    defer expected.deinit();
+    try expected.setString(10, "11663466984815033033");
+    try std.testing.expectEqual(std.math.Order.eq, expected.order(res));
+}