zig

commit 2ebd6bd7060dab3347a736ffd8a0ca5914624dc4 (tree)
parent ebf97627fd4e1ee46c1b446108e4b4e3bf4b5769
Author: LemonBoy <thatlemon@gmail.com>
Date:   Wed, 14 Apr 2021 17:52:24 +0200

std: Fix sqrt for u0/u1 input types

Diffstat:
M
lib/std/math/sqrt.zig
 | 
20
+++++++++++++++++---

1 file changed, 17 insertions(+), 3 deletions(-)
diff --git 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,
     };
 }