diff --git a/lib/std/math/sqrt.zig b/lib/std/math/sqrt.zig index 394164ff54..4f32581270 100644 --- a/lib/std/math/sqrt.zig +++ b/lib/std/math/sqrt.zig @@ -39,7 +39,13 @@ pub fn sqrt(x: anytype) Sqrt(@TypeOf(x)) { } } -fn sqrt_int(comptime T: type, value: T) std.meta.Int(.unsigned, @typeInfo(T).Int.bits / 2) { +fn sqrt_int(comptime T: type, value: T) Sqrt(T) { + switch (T) { + u0 => return 0, + u1 => return value, + else => {}, + } + var op = value; var res: T = 0; var one: T = 1 << (@typeInfo(T).Int.bits - 2); @@ -58,11 +64,13 @@ fn sqrt_int(comptime T: type, value: T) std.meta.Int(.unsigned, @typeInfo(T).Int one >>= 2; } - const ResultType = std.meta.Int(.unsigned, @typeInfo(T).Int.bits / 2); + const ResultType = Sqrt(T); return @intCast(ResultType, res); } test "math.sqrt_int" { + expect(sqrt_int(u0, 0) == 0); + expect(sqrt_int(u1, 1) == 1); expect(sqrt_int(u32, 3) == 1); expect(sqrt_int(u32, 4) == 2); expect(sqrt_int(u32, 5) == 2); @@ -74,7 +82,13 @@ test "math.sqrt_int" { /// Returns the return type `sqrt` will return given an operand of type `T`. pub fn Sqrt(comptime T: type) type { return switch (@typeInfo(T)) { - .Int => |int| std.meta.Int(.unsigned, int.bits / 2), + .Int => |int| { + return switch (int.bits) { + 0 => u0, + 1 => u1, + else => std.meta.Int(.unsigned, int.bits / 2), + }; + }, else => T, }; }