1// Based on Rust, which is licensed under the MIT license.
2// https://github.com/rust-lang/rust/blob/360432f1e8794de58cd94f34c9c17ad65871e5b5/LICENSE-MIT
3//
4// https://github.com/rust-lang/rust/blob/360432f1e8794de58cd94f34c9c17ad65871e5b5/src/libcore/num/mod.rs#L3423
5
6const std = @import("../std.zig");
7const math = std.math;
8const assert = std.debug.assert;
9const testing = std.testing;
10
11/// Returns the power of x raised by the integer y (x^y).
12///
13/// Errors:
14/// - Overflow: Integer overflow or Infinity
15/// - Underflow: Absolute value of result smaller than 1
16///
17/// Edge case rules ordered by precedence:
18/// - powi(T, x, 0) = 1 unless T is i1, u0
19/// - powi(T, 0, x) = 0 when x > 0
20/// - powi(T, 0, x) = Overflow
21/// - powi(T, 1, y) = 1
22/// - powi(T, -1, y) = -1 for y an odd integer
23/// - powi(T, -1, y) = 1 unless T is i1, u0
24/// - powi(T, -1, y) = Overflow
25/// - powi(T, x, y) = Overflow when y >= @bitSizeOf(x)
26/// - powi(T, x, y) = Underflow when y < 0
27pub fn powi(comptime T: type, x: T, y: T) (error{
28 Overflow,
29 Underflow,
30}!T) {
31 const bit_size = @typeInfo(T).int.bits;
32
33 // `y & 1 == 0` won't compile when `does_one_overflow`.
34 const does_one_overflow = math.maxInt(T) < 1;
35 const is_y_even = !does_one_overflow and y & 1 == 0;
36
37 if (x == 1 or y == 0 or (x == -1 and is_y_even)) {
38 if (does_one_overflow) {
39 return error.Overflow;
40 } else {
41 return 1;
42 }
43 }
44
45 if (x == -1) {
46 return -1;
47 }
48
49 if (x == 0) {
50 if (y > 0) {
51 return 0;
52 } else {
53 // Infinity/NaN, not overflow in strict sense
54 return error.Overflow;
55 }
56 }
57 // x >= 2 or x <= -2 from this point
58 if (y >= bit_size) {
59 return error.Overflow;
60 }
61 if (y < 0) {
62 return error.Underflow;
63 }
64
65 // invariant :
66 // return value = powi(T, base, exp) * acc;
67
68 var base = x;
69 var exp = y;
70 var acc: T = if (does_one_overflow) unreachable else 1;
71
72 while (exp > 1) {
73 if (exp & 1 == 1) {
74 const ov = @mulWithOverflow(acc, base);
75 if (ov[1] != 0) return error.Overflow;
76 acc = ov[0];
77 }
78
79 exp >>= 1;
80
81 const ov = @mulWithOverflow(base, base);
82 if (ov[1] != 0) return error.Overflow;
83 base = ov[0];
84 }
85
86 if (exp == 1) {
87 const ov = @mulWithOverflow(acc, base);
88 if (ov[1] != 0) return error.Overflow;
89 acc = ov[0];
90 }
91
92 return acc;
93}
94
95test powi {
96 try testing.expectError(error.Overflow, powi(i8, -66, 6));
97 try testing.expectError(error.Overflow, powi(i16, -13, 13));
98 try testing.expectError(error.Overflow, powi(i32, -32, 21));
99 try testing.expectError(error.Overflow, powi(i64, -24, 61));
100 try testing.expectError(error.Overflow, powi(i17, -15, 15));
101 try testing.expectError(error.Overflow, powi(i42, -6, 40));
102
103 try testing.expect((try powi(i8, -5, 3)) == -125);
104 try testing.expect((try powi(i16, -16, 3)) == -4096);
105 try testing.expect((try powi(i32, -91, 3)) == -753571);
106 try testing.expect((try powi(i64, -36, 6)) == 2176782336);
107 try testing.expect((try powi(i17, -2, 15)) == -32768);
108 try testing.expect((try powi(i42, -5, 7)) == -78125);
109
110 try testing.expect((try powi(u8, 6, 2)) == 36);
111 try testing.expect((try powi(u16, 5, 4)) == 625);
112 try testing.expect((try powi(u32, 12, 6)) == 2985984);
113 try testing.expect((try powi(u64, 34, 2)) == 1156);
114 try testing.expect((try powi(u17, 16, 3)) == 4096);
115 try testing.expect((try powi(u42, 34, 6)) == 1544804416);
116
117 try testing.expectError(error.Overflow, powi(i8, 120, 7));
118 try testing.expectError(error.Overflow, powi(i16, 73, 15));
119 try testing.expectError(error.Overflow, powi(i32, 23, 31));
120 try testing.expectError(error.Overflow, powi(i64, 68, 61));
121 try testing.expectError(error.Overflow, powi(i17, 15, 15));
122 try testing.expectError(error.Overflow, powi(i42, 121312, 41));
123
124 try testing.expectError(error.Overflow, powi(u8, 123, 7));
125 try testing.expectError(error.Overflow, powi(u16, 2313, 15));
126 try testing.expectError(error.Overflow, powi(u32, 8968, 31));
127 try testing.expectError(error.Overflow, powi(u64, 2342, 63));
128 try testing.expectError(error.Overflow, powi(u17, 2723, 16));
129 try testing.expectError(error.Overflow, powi(u42, 8234, 41));
130
131 const minInt = std.math.minInt;
132 try testing.expect((try powi(i8, -2, 7)) == minInt(i8));
133 try testing.expect((try powi(i16, -2, 15)) == minInt(i16));
134 try testing.expect((try powi(i32, -2, 31)) == minInt(i32));
135 try testing.expect((try powi(i64, -2, 63)) == minInt(i64));
136
137 try testing.expectError(error.Underflow, powi(i8, 6, -2));
138 try testing.expectError(error.Underflow, powi(i16, 5, -4));
139 try testing.expectError(error.Underflow, powi(i32, 12, -6));
140 try testing.expectError(error.Underflow, powi(i64, 34, -2));
141 try testing.expectError(error.Underflow, powi(i17, 16, -3));
142 try testing.expectError(error.Underflow, powi(i42, 34, -6));
143}
144
145test "powi.special" {
146 try testing.expectError(error.Overflow, powi(i8, -2, 8));
147 try testing.expectError(error.Overflow, powi(i16, -2, 16));
148 try testing.expectError(error.Overflow, powi(i32, -2, 32));
149 try testing.expectError(error.Overflow, powi(i64, -2, 64));
150 try testing.expectError(error.Overflow, powi(i17, -2, 17));
151 try testing.expectError(error.Overflow, powi(i17, -2, 16));
152 try testing.expectError(error.Overflow, powi(i42, -2, 42));
153
154 try testing.expect((try powi(i8, -1, 3)) == -1);
155 try testing.expect((try powi(i16, -1, 2)) == 1);
156 try testing.expect((try powi(i32, -1, 16)) == 1);
157 try testing.expect((try powi(i64, -1, 6)) == 1);
158 try testing.expect((try powi(i17, -1, 15)) == -1);
159 try testing.expect((try powi(i42, -1, 7)) == -1);
160
161 try testing.expect((try powi(u8, 1, 2)) == 1);
162 try testing.expect((try powi(u16, 1, 4)) == 1);
163 try testing.expect((try powi(u32, 1, 6)) == 1);
164 try testing.expect((try powi(u64, 1, 2)) == 1);
165 try testing.expect((try powi(u17, 1, 3)) == 1);
166 try testing.expect((try powi(u42, 1, 6)) == 1);
167
168 try testing.expectError(error.Overflow, powi(i8, 2, 7));
169 try testing.expectError(error.Overflow, powi(i16, 2, 15));
170 try testing.expectError(error.Overflow, powi(i32, 2, 31));
171 try testing.expectError(error.Overflow, powi(i64, 2, 63));
172 try testing.expectError(error.Overflow, powi(i17, 2, 16));
173 try testing.expectError(error.Overflow, powi(i42, 2, 41));
174
175 try testing.expectError(error.Overflow, powi(u8, 2, 8));
176 try testing.expectError(error.Overflow, powi(u16, 2, 16));
177 try testing.expectError(error.Overflow, powi(u32, 2, 32));
178 try testing.expectError(error.Overflow, powi(u64, 2, 64));
179 try testing.expectError(error.Overflow, powi(u17, 2, 17));
180 try testing.expectError(error.Overflow, powi(u42, 2, 42));
181
182 try testing.expect((try powi(u8, 6, 0)) == 1);
183 try testing.expect((try powi(u16, 5, 0)) == 1);
184 try testing.expect((try powi(u32, 12, 0)) == 1);
185 try testing.expect((try powi(u64, 34, 0)) == 1);
186 try testing.expect((try powi(u17, 16, 0)) == 1);
187 try testing.expect((try powi(u42, 34, 0)) == 1);
188}
189
190test "powi.narrow" {
191 try testing.expectError(error.Overflow, powi(u0, 0, 0));
192 try testing.expectError(error.Overflow, powi(i1, 0, 0));
193 try testing.expectError(error.Overflow, powi(i1, -1, 0));
194 try testing.expectError(error.Overflow, powi(i1, 0, -1));
195 try testing.expect((try powi(i1, -1, -1)) == -1);
196}