std.math.sqrt_int: fixed odd size integers types

This commit is contained in:
Filippo Casarin 2021-03-24 20:59:15 +01:00 committed by Veikka Tuominen
parent 13dc34f779
commit 8b0dbca32d

View File

@ -39,15 +39,12 @@ pub fn sqrt(x: anytype) Sqrt(@TypeOf(x)) {
} }
fn sqrt_int(comptime T: type, value: T) Sqrt(T) { fn sqrt_int(comptime T: type, value: T) Sqrt(T) {
switch (T) { if (@typeInfo(T).Int.bits <= 2) {
u0 => return 0, return if (value == 0) 0 else 1; // shortcut for small number of bits to simplify general case
u1 => return value, } else {
else => {},
}
var op = value; var op = value;
var res: T = 0; var res: T = 0;
var one: T = 1 << (@typeInfo(T).Int.bits - 2); var one: T = 1 << ((@typeInfo(T).Int.bits - 1) & -2); // highest power of four that fits into T
// "one" starts at the highest power of four <= than the argument. // "one" starts at the highest power of four <= than the argument.
while (one > op) { while (one > op) {
@ -55,39 +52,37 @@ fn sqrt_int(comptime T: type, value: T) Sqrt(T) {
} }
while (one != 0) { while (one != 0) {
if (op >= res + one) { var c = op >= res + one;
op -= res + one; if (c) op -= res + one;
res += 2 * one;
}
res >>= 1; res >>= 1;
if (c) res += one;
one >>= 2; one >>= 2;
} }
const ResultType = Sqrt(T); return @intCast(Sqrt(T), res);
return @intCast(ResultType, res); }
} }
test "math.sqrt_int" { test "math.sqrt_int" {
try expect(sqrt_int(u0, 0) == 0);
try expect(sqrt_int(u1, 1) == 1);
try expect(sqrt_int(u32, 3) == 1); try expect(sqrt_int(u32, 3) == 1);
try expect(sqrt_int(u32, 4) == 2); try expect(sqrt_int(u32, 4) == 2);
try expect(sqrt_int(u32, 5) == 2); try expect(sqrt_int(u32, 5) == 2);
try expect(sqrt_int(u32, 8) == 2); try expect(sqrt_int(u32, 8) == 2);
try expect(sqrt_int(u32, 9) == 3); try expect(sqrt_int(u32, 9) == 3);
try expect(sqrt_int(u32, 10) == 3); try expect(sqrt_int(u32, 10) == 3);
try expect(sqrt_int(u0, 0) == 0);
try expect(sqrt_int(u1, 1) == 1);
try expect(sqrt_int(u2, 3) == 1);
try expect(sqrt_int(u3, 4) == 2);
try expect(sqrt_int(u4, 8) == 2);
try expect(sqrt_int(u4, 9) == 3);
} }
/// Returns the return type `sqrt` will return given an operand of type `T`. /// Returns the return type `sqrt` will return given an operand of type `T`.
pub fn Sqrt(comptime T: type) type { pub fn Sqrt(comptime T: type) type {
return switch (@typeInfo(T)) { return switch (@typeInfo(T)) {
.Int => |int| { .Int => |int| std.meta.Int(.unsigned, (int.bits + 1) / 2),
return switch (int.bits) {
0 => u0,
1 => u1,
else => std.meta.Int(.unsigned, int.bits / 2),
};
},
else => T, else => T,
}; };
} }