| ... | ... | @@ -39,55 +39,50 @@ pub fn sqrt(x: anytype) Sqrt(@TypeOf(x)) { |
| 39 | 39 | } |
| 40 | 40 | |
| 41 | 41 | fn sqrt_int(comptime T: type, value: T) Sqrt(T) { |
| 42 | | switch (T) { |
| 43 | | u0 => return 0, |
| 44 | | u1 => return value, |
| 45 | | else => {}, |
| 46 | | } |
| 47 | | |
| 48 | | var op = value; |
| 49 | | var res: T = 0; |
| 50 | | var one: T = 1 << (@typeInfo(T).Int.bits - 2); |
| 42 | if (@typeInfo(T).Int.bits <= 2) { |
| 43 | return if (value == 0) 0 else 1; // shortcut for small number of bits to simplify general case |
| 44 | } else { |
| 45 | var op = value; |
| 46 | var res: T = 0; |
| 47 | var one: T = 1 << ((@typeInfo(T).Int.bits - 1) & -2); // highest power of four that fits into T |
| 51 | 48 | |
| 52 | | // "one" starts at the highest power of four <= than the argument. |
| 53 | | while (one > op) { |
| 54 | | one >>= 2; |
| 55 | | } |
| 49 | // "one" starts at the highest power of four <= than the argument. |
| 50 | while (one > op) { |
| 51 | one >>= 2; |
| 52 | } |
| 56 | 53 | |
| 57 | | while (one != 0) { |
| 58 | | if (op >= res + one) { |
| 59 | | op -= res + one; |
| 60 | | res += 2 * one; |
| 54 | while (one != 0) { |
| 55 | var c = op >= res + one; |
| 56 | if (c) op -= res + one; |
| 57 | res >>= 1; |
| 58 | if (c) res += one; |
| 59 | one >>= 2; |
| 61 | 60 | } |
| 62 | | res >>= 1; |
| 63 | | one >>= 2; |
| 64 | | } |
| 65 | 61 | |
| 66 | | const ResultType = Sqrt(T); |
| 67 | | return @intCast(ResultType, res); |
| 62 | return @intCast(Sqrt(T), res); |
| 63 | } |
| 68 | 64 | } |
| 69 | 65 | |
| 70 | 66 | test "math.sqrt_int" { |
| 71 | | try expect(sqrt_int(u0, 0) == 0); |
| 72 | | try expect(sqrt_int(u1, 1) == 1); |
| 73 | 67 | try expect(sqrt_int(u32, 3) == 1); |
| 74 | 68 | try expect(sqrt_int(u32, 4) == 2); |
| 75 | 69 | try expect(sqrt_int(u32, 5) == 2); |
| 76 | 70 | try expect(sqrt_int(u32, 8) == 2); |
| 77 | 71 | try expect(sqrt_int(u32, 9) == 3); |
| 78 | 72 | try expect(sqrt_int(u32, 10) == 3); |
| 73 | |
| 74 | try expect(sqrt_int(u0, 0) == 0); |
| 75 | try expect(sqrt_int(u1, 1) == 1); |
| 76 | try expect(sqrt_int(u2, 3) == 1); |
| 77 | try expect(sqrt_int(u3, 4) == 2); |
| 78 | try expect(sqrt_int(u4, 8) == 2); |
| 79 | try expect(sqrt_int(u4, 9) == 3); |
| 79 | 80 | } |
| 80 | 81 | |
| 81 | 82 | /// Returns the return type `sqrt` will return given an operand of type `T`. |
| 82 | 83 | pub fn Sqrt(comptime T: type) type { |
| 83 | 84 | return switch (@typeInfo(T)) { |
| 84 | | .Int => |int| { |
| 85 | | return switch (int.bits) { |
| 86 | | 0 => u0, |
| 87 | | 1 => u1, |
| 88 | | else => std.meta.Int(.unsigned, int.bits / 2), |
| 89 | | }; |
| 90 | | }, |
| 85 | .Int => |int| std.meta.Int(.unsigned, (int.bits + 1) / 2), |
| 91 | 86 | else => T, |
| 92 | 87 | }; |
| 93 | 88 | } |