Add Complex::powi and assorted Pow impls

Author: cuviper

src/lib.rs

@@ -38,13 +38,15 @@ use core::str::FromStr;
 #[cfg(feature = "std")]
 use std::error::Error;
 
-use traits::{Inv, MulAdd, Num, One, Signed, Zero};
+use traits::{Inv, MulAdd, Num, One, Pow, Signed, Zero};
 
 #[cfg(feature = "std")]
 use traits::float::Float;
 use traits::float::FloatCore;
 
 mod cast;
+mod pow;
+
 #[cfg(feature = "rand")]
 mod crand;
 #[cfg(feature = "rand")]
@@ -121,6 +123,12 @@ impl<T: Clone + Num> Complex<T> {
     pub fn unscale(&self, t: T) -> Self {
         Self::new(self.re.clone() / t.clone(), self.im.clone() / t)
     }
+
+    /// Raises `self` to an unsigned integer power.
+    #[inline]
+    pub fn powu(&self, exp: u32) -> Self {
+        Pow::pow(self, exp)
+    }
 }
 
 impl<T: Clone + Num + Neg<Output = T>> Complex<T> {
@@ -139,6 +147,12 @@ impl<T: Clone + Num + Neg<Output = T>> Complex<T> {
             -self.im.clone() / norm_sqr,
         )
     }
+
+    /// Raises `self` to a signed integer power.
+    #[inline]
+    pub fn powi(&self, exp: i32) -> Self {
+        Pow::pow(self, exp)
+    }
 }
 
 impl<T: Clone + Signed> Complex<T> {
@@ -1487,6 +1501,25 @@ mod test {
         assert_eq!(_4_2i.l1_norm(), 6.0);
     }
 
+    #[test]
+    fn test_pow() {
+        for c in all_consts.iter() {
+            assert_eq!(c.powi(0), _1_0i);
+            let mut pos = _1_0i;
+            let mut neg = _1_0i;
+            for i in 1i32..20 {
+                pos *= c;
+                assert_eq!(pos, c.powi(i));
+                if c.is_zero() {
+                    assert!(c.powi(-i).is_nan());
+                } else {
+                    neg /= c;
+                    assert_eq!(neg, c.powi(-i));
+                }
+            }
+        }
+    }
+
     #[cfg(feature = "std")]
     mod float {
         use super::*;
@@ -1535,7 +1568,11 @@ mod test {
 
         fn close_to_tol(a: Complex64, b: Complex64, tol: f64) -> bool {
             // returns true if a and b are reasonably close
-            (a == b) || (a - b).norm() < tol
+            let close = (a == b) || (a - b).norm() < tol;
+            if !close {
+                println!("{:?} != {:?}", a, b);
+            }
+            close
         }
 
         #[test]

src/pow.rs

@@ -0,0 +1,101 @@
+use super::Complex;
+
+use std::ops::Neg;
+use traits::{Float, Num, One, Pow};
+
+macro_rules! pow_impl {
+    ($U:ty, $S:ty) => {
+        impl<'a, T: Clone + Num> Pow<$U> for &'a Complex<T> {
+            type Output = Complex<T>;
+
+            #[inline]
+            fn pow(self, mut exp: $U) -> Self::Output {
+                if exp == 0 {
+                    return Complex::one();
+                }
+                let mut base = self.clone();
+
+                while exp & 1 == 0 {
+                    base = base.clone() * base;
+                    exp >>= 1;
+                }
+
+                if exp == 1 {
+                    return base;
+                }
+
+                let mut acc = base.clone();
+                while exp > 1 {
+                    exp >>= 1;
+                    base = base.clone() * base;
+                    if exp & 1 == 1 {
+                        acc = acc * base.clone();
+                    }
+                }
+                acc
+            }
+        }
+
+        impl<'a, 'b, T: Clone + Num> Pow<&'b $U> for &'a Complex<T> {
+            type Output = Complex<T>;
+
+            #[inline]
+            fn pow(self, exp: &$U) -> Self::Output {
+                self.pow(*exp)
+            }
+        }
+
+        impl<'a, T: Clone + Num + Neg<Output = T>> Pow<$S> for &'a Complex<T> {
+            type Output = Complex<T>;
+
+            #[inline]
+            fn pow(self, exp: $S) -> Self::Output {
+                if exp < 0 {
+                    Pow::pow(&self.inv(), exp.wrapping_neg() as $U)
+                } else {
+                    Pow::pow(self, exp as $U)
+                }
+            }
+        }
+
+        impl<'a, 'b, T: Clone + Num + Neg<Output = T>> Pow<&'b $S> for &'a Complex<T> {
+            type Output = Complex<T>;
+
+            #[inline]
+            fn pow(self, exp: &$S) -> Self::Output {
+                self.pow(*exp)
+            }
+        }
+    };
+}
+
+pow_impl!(u8, i8);
+pow_impl!(u16, i16);
+pow_impl!(u32, i32);
+pow_impl!(u64, i64);
+pow_impl!(usize, isize);
+#[cfg(has_i128)]
+pow_impl!(u128, i128);
+
+// Note: the impls above are for `&Complex<T>`, while those below are for `Complex<T>`.  This is
+// fine since `Float: Copy` anyway, but it's also necessary to avoid conflicting implementations.
+// Otherwise rustc would insist that those `Pow<{integer}>` impls could overlap with `Pow<T>` if
+// integers ever implement `Float`, though of course we know they won't...
+
+impl<T: Float> Pow<T> for Complex<T> {
+    type Output = Complex<T>;
+
+    #[inline]
+    fn pow(self, exp: T) -> Self::Output {
+        self.powf(exp)
+    }
+}
+
+impl<T: Float> Pow<Complex<T>> for Complex<T> {
+    type Output = Complex<T>;
+
+    #[inline]
+    fn pow(self, exp: Complex<T>) -> Self::Output {
+        self.powc(exp)
+    }
+}