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

Signed-off-by: tison <wander4096@gmail.com>
This commit is contained in:
tison 2023-11-08 14:39:08 +08:00 committed by GitHub
parent 77bc8e7b67
commit ee47643b6e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 2 deletions

View File

@ -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;

View File

@ -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));
}