1const std = @import("../std.zig");
 2const math = std.math;
 3const expect = std.testing.expect;
 4const TypeId = std.builtin.TypeId;
 5const maxInt = std.math.maxInt;
 6
 7/// Returns the square root of x.
 8///
 9/// Special Cases:
10///  - sqrt(+inf)  = +inf
11///  - sqrt(+-0)   = +-0
12///  - sqrt(x)     = nan if x < 0
13///  - sqrt(nan)   = nan
14/// TODO Decide if all this logic should be implemented directly in the @sqrt builtin function.
15pub fn sqrt(x: anytype) Sqrt(@TypeOf(x)) {
16    const T = @TypeOf(x);
17    switch (@typeInfo(T)) {
18        .float, .comptime_float => return @sqrt(x),
19        .comptime_int => comptime {
20            if (x > maxInt(u128)) {
21                @compileError("sqrt not implemented for comptime_int greater than 128 bits");
22            }
23            if (x < 0) {
24                @compileError("sqrt on negative number");
25            }
26            return @as(T, sqrt_int(u128, x));
27        },
28        .int => |IntType| switch (IntType.signedness) {
29            .signed => @compileError("sqrt not implemented for signed integers"),
30            .unsigned => return sqrt_int(T, x),
31        },
32        else => @compileError("sqrt not implemented for " ++ @typeName(T)),
33    }
34}
35
36fn sqrt_int(comptime T: type, value: T) Sqrt(T) {
37    if (@typeInfo(T).int.bits <= 2) {
38        return if (value == 0) 0 else 1; // shortcut for small number of bits to simplify general case
39    } else {
40        const bits = @typeInfo(T).int.bits;
41        const max = math.maxInt(T);
42        const minustwo = (@as(T, 2) ^ max) + 1; // unsigned int cannot represent -2
43        var op = value;
44        var res: T = 0;
45        var one: T = 1 << ((bits - 1) & minustwo); // highest power of four that fits into T
46
47        // "one" starts at the highest power of four <= than the argument.
48        while (one > op) {
49            one >>= 2;
50        }
51
52        while (one != 0) {
53            const c = op >= res + one;
54            if (c) op -= res + one;
55            res >>= 1;
56            if (c) res += one;
57            one >>= 2;
58        }
59
60        return @as(Sqrt(T), @intCast(res));
61    }
62}
63
64test sqrt_int {
65    try expect(sqrt_int(u32, 3) == 1);
66    try expect(sqrt_int(u32, 4) == 2);
67    try expect(sqrt_int(u32, 5) == 2);
68    try expect(sqrt_int(u32, 8) == 2);
69    try expect(sqrt_int(u32, 9) == 3);
70    try expect(sqrt_int(u32, 10) == 3);
71
72    try expect(sqrt_int(u0, 0) == 0);
73    try expect(sqrt_int(u1, 1) == 1);
74    try expect(sqrt_int(u2, 3) == 1);
75    try expect(sqrt_int(u3, 4) == 2);
76    try expect(sqrt_int(u4, 8) == 2);
77    try expect(sqrt_int(u4, 9) == 3);
78}
79
80/// Returns the return type `sqrt` will return given an operand of type `T`.
81pub fn Sqrt(comptime T: type) type {
82    return switch (@typeInfo(T)) {
83        .int => |int| @Int(.unsigned, (int.bits + 1) / 2),
84        else => T,
85    };
86}