num_traits/
pow.rs

1use core::num::Wrapping;
2use core::ops::Mul;
3use {CheckedMul, One};
4
5/// Binary operator for raising a value to a power.
6pub trait Pow<RHS> {
7    /// The result after applying the operator.
8    type Output;
9
10    /// Returns `self` to the power `rhs`.
11    ///
12    /// # Examples
13    ///
14    /// ```
15    /// use num_traits::Pow;
16    /// assert_eq!(Pow::pow(10u32, 2u32), 100);
17    /// ```
18    fn pow(self, rhs: RHS) -> Self::Output;
19}
20
21macro_rules! pow_impl {
22    ($t:ty) => {
23        pow_impl!($t, u8);
24        pow_impl!($t, usize);
25
26        // FIXME: these should be possible
27        // pow_impl!($t, u16);
28        // pow_impl!($t, u32);
29        // pow_impl!($t, u64);
30    };
31    ($t:ty, $rhs:ty) => {
32        pow_impl!($t, $rhs, usize, pow);
33    };
34    ($t:ty, $rhs:ty, $desired_rhs:ty, $method:expr) => {
35        impl Pow<$rhs> for $t {
36            type Output = $t;
37            #[inline]
38            fn pow(self, rhs: $rhs) -> $t {
39                ($method)(self, <$desired_rhs>::from(rhs))
40            }
41        }
42
43        impl<'a> Pow<&'a $rhs> for $t {
44            type Output = $t;
45            #[inline]
46            fn pow(self, rhs: &'a $rhs) -> $t {
47                ($method)(self, <$desired_rhs>::from(*rhs))
48            }
49        }
50
51        impl<'a> Pow<$rhs> for &'a $t {
52            type Output = $t;
53            #[inline]
54            fn pow(self, rhs: $rhs) -> $t {
55                ($method)(*self, <$desired_rhs>::from(rhs))
56            }
57        }
58
59        impl<'a, 'b> Pow<&'a $rhs> for &'b $t {
60            type Output = $t;
61            #[inline]
62            fn pow(self, rhs: &'a $rhs) -> $t {
63                ($method)(*self, <$desired_rhs>::from(*rhs))
64            }
65        }
66    };
67}
68
69pow_impl!(u8, u8, u32, u8::pow);
70pow_impl!(u8, u16, u32, u8::pow);
71pow_impl!(u8, u32, u32, u8::pow);
72pow_impl!(u8, usize);
73pow_impl!(i8, u8, u32, i8::pow);
74pow_impl!(i8, u16, u32, i8::pow);
75pow_impl!(i8, u32, u32, i8::pow);
76pow_impl!(i8, usize);
77pow_impl!(u16, u8, u32, u16::pow);
78pow_impl!(u16, u16, u32, u16::pow);
79pow_impl!(u16, u32, u32, u16::pow);
80pow_impl!(u16, usize);
81pow_impl!(i16, u8, u32, i16::pow);
82pow_impl!(i16, u16, u32, i16::pow);
83pow_impl!(i16, u32, u32, i16::pow);
84pow_impl!(i16, usize);
85pow_impl!(u32, u8, u32, u32::pow);
86pow_impl!(u32, u16, u32, u32::pow);
87pow_impl!(u32, u32, u32, u32::pow);
88pow_impl!(u32, usize);
89pow_impl!(i32, u8, u32, i32::pow);
90pow_impl!(i32, u16, u32, i32::pow);
91pow_impl!(i32, u32, u32, i32::pow);
92pow_impl!(i32, usize);
93pow_impl!(u64, u8, u32, u64::pow);
94pow_impl!(u64, u16, u32, u64::pow);
95pow_impl!(u64, u32, u32, u64::pow);
96pow_impl!(u64, usize);
97pow_impl!(i64, u8, u32, i64::pow);
98pow_impl!(i64, u16, u32, i64::pow);
99pow_impl!(i64, u32, u32, i64::pow);
100pow_impl!(i64, usize);
101
102#[cfg(has_i128)]
103pow_impl!(u128, u8, u32, u128::pow);
104#[cfg(has_i128)]
105pow_impl!(u128, u16, u32, u128::pow);
106#[cfg(has_i128)]
107pow_impl!(u128, u32, u32, u128::pow);
108#[cfg(has_i128)]
109pow_impl!(u128, usize);
110
111#[cfg(has_i128)]
112pow_impl!(i128, u8, u32, i128::pow);
113#[cfg(has_i128)]
114pow_impl!(i128, u16, u32, i128::pow);
115#[cfg(has_i128)]
116pow_impl!(i128, u32, u32, i128::pow);
117#[cfg(has_i128)]
118pow_impl!(i128, usize);
119
120pow_impl!(usize, u8, u32, usize::pow);
121pow_impl!(usize, u16, u32, usize::pow);
122pow_impl!(usize, u32, u32, usize::pow);
123pow_impl!(usize, usize);
124pow_impl!(isize, u8, u32, isize::pow);
125pow_impl!(isize, u16, u32, isize::pow);
126pow_impl!(isize, u32, u32, isize::pow);
127pow_impl!(isize, usize);
128pow_impl!(Wrapping<u8>);
129pow_impl!(Wrapping<i8>);
130pow_impl!(Wrapping<u16>);
131pow_impl!(Wrapping<i16>);
132pow_impl!(Wrapping<u32>);
133pow_impl!(Wrapping<i32>);
134pow_impl!(Wrapping<u64>);
135pow_impl!(Wrapping<i64>);
136#[cfg(has_i128)]
137pow_impl!(Wrapping<u128>);
138#[cfg(has_i128)]
139pow_impl!(Wrapping<i128>);
140pow_impl!(Wrapping<usize>);
141pow_impl!(Wrapping<isize>);
142
143// FIXME: these should be possible
144// pow_impl!(u8, u64);
145// pow_impl!(i16, u64);
146// pow_impl!(i8, u64);
147// pow_impl!(u16, u64);
148// pow_impl!(u32, u64);
149// pow_impl!(i32, u64);
150// pow_impl!(u64, u64);
151// pow_impl!(i64, u64);
152// pow_impl!(usize, u64);
153// pow_impl!(isize, u64);
154
155#[cfg(any(feature = "std", feature = "libm"))]
156mod float_impls {
157    use super::Pow;
158    use Float;
159
160    pow_impl!(f32, i8, i32, <f32 as Float>::powi);
161    pow_impl!(f32, u8, i32, <f32 as Float>::powi);
162    pow_impl!(f32, i16, i32, <f32 as Float>::powi);
163    pow_impl!(f32, u16, i32, <f32 as Float>::powi);
164    pow_impl!(f32, i32, i32, <f32 as Float>::powi);
165    pow_impl!(f64, i8, i32, <f64 as Float>::powi);
166    pow_impl!(f64, u8, i32, <f64 as Float>::powi);
167    pow_impl!(f64, i16, i32, <f64 as Float>::powi);
168    pow_impl!(f64, u16, i32, <f64 as Float>::powi);
169    pow_impl!(f64, i32, i32, <f64 as Float>::powi);
170    pow_impl!(f32, f32, f32, <f32 as Float>::powf);
171    pow_impl!(f64, f32, f64, <f64 as Float>::powf);
172    pow_impl!(f64, f64, f64, <f64 as Float>::powf);
173}
174
175/// Raises a value to the power of exp, using exponentiation by squaring.
176///
177/// Note that `0⁰` (`pow(0, 0)`) returns `1`. Mathematically this is undefined.
178///
179/// # Example
180///
181/// ```rust
182/// use num_traits::pow;
183///
184/// assert_eq!(pow(2i8, 4), 16);
185/// assert_eq!(pow(6u8, 3), 216);
186/// assert_eq!(pow(0u8, 0), 1); // Be aware if this case affects you
187/// ```
188#[inline]
189pub fn pow<T: Clone + One + Mul<T, Output = T>>(mut base: T, mut exp: usize) -> T {
190    if exp == 0 {
191        return T::one();
192    }
193
194    while exp & 1 == 0 {
195        base = base.clone() * base;
196        exp >>= 1;
197    }
198    if exp == 1 {
199        return base;
200    }
201
202    let mut acc = base.clone();
203    while exp > 1 {
204        exp >>= 1;
205        base = base.clone() * base;
206        if exp & 1 == 1 {
207            acc = acc * base.clone();
208        }
209    }
210    acc
211}
212
213/// Raises a value to the power of exp, returning `None` if an overflow occurred.
214///
215/// Note that `0⁰` (`checked_pow(0, 0)`) returns `Some(1)`. Mathematically this is undefined.
216///
217/// Otherwise same as the `pow` function.
218///
219/// # Example
220///
221/// ```rust
222/// use num_traits::checked_pow;
223///
224/// assert_eq!(checked_pow(2i8, 4), Some(16));
225/// assert_eq!(checked_pow(7i8, 8), None);
226/// assert_eq!(checked_pow(7u32, 8), Some(5_764_801));
227/// assert_eq!(checked_pow(0u32, 0), Some(1)); // Be aware if this case affect you
228/// ```
229#[inline]
230pub fn checked_pow<T: Clone + One + CheckedMul>(mut base: T, mut exp: usize) -> Option<T> {
231    if exp == 0 {
232        return Some(T::one());
233    }
234
235    macro_rules! optry {
236        ($expr:expr) => {
237            if let Some(val) = $expr {
238                val
239            } else {
240                return None;
241            }
242        };
243    }
244
245    while exp & 1 == 0 {
246        base = optry!(base.checked_mul(&base));
247        exp >>= 1;
248    }
249    if exp == 1 {
250        return Some(base);
251    }
252
253    let mut acc = base.clone();
254    while exp > 1 {
255        exp >>= 1;
256        base = optry!(base.checked_mul(&base));
257        if exp & 1 == 1 {
258            acc = optry!(acc.checked_mul(&base));
259        }
260    }
261    Some(acc)
262}