Browse Source

Add checked_pow function

est31 9 years ago
parent
commit
3c5ecdc636
1 changed files with 41 additions and 1 deletions
  1. 41 1
      src/lib.rs

+ 41 - 1
src/lib.rs

@@ -1,4 +1,4 @@
-// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// Copyright 2014-2016 The Rust Project Developers. See the COPYRIGHT
 // file at the top-level directory of this distribution and at
 // http://rust-lang.org/COPYRIGHT.
 //
@@ -164,6 +164,46 @@ pub fn pow<T: Clone + One + Mul<T, Output = T>>(mut base: T, mut exp: usize) ->
     acc
 }
 
+/// Raises a value to the power of exp, returning `None` if an overflow occured.
+///
+/// Otherwise same as the `pow` function.
+///
+/// # Example
+///
+/// ```rust
+/// use num;
+///
+/// assert_eq!(num::checked_pow(2i8, 4), Some(16));
+/// assert_eq!(num::checked_pow(7i8, 8), None);
+/// assert_eq!(num::checked_pow(7u32, 8), Some(5_764_801));
+/// ```
+#[inline]
+pub fn checked_pow<T: Clone + One + CheckedMul>(mut base: T, mut exp: usize) -> Option<T> {
+    if exp == 0 { return Some(T::one()) }
+
+    macro_rules! optry {
+        ( $ expr : expr ) => {
+            if let Some(val) = $expr { val } else { return None }
+        }
+    }
+
+    while exp & 1 == 0 {
+        base = optry!(base.checked_mul(&base));
+        exp >>= 1;
+    }
+    if exp == 1 { return Some(base) }
+
+    let mut acc = base.clone();
+    while exp > 1 {
+        exp >>= 1;
+        base = optry!(base.checked_mul(&base));
+        if exp & 1 == 1 {
+            acc = optry!(acc.checked_mul(&base));
+        }
+    }
+    Some(acc)
+}
+
 #[cfg(test)]
 fn hash<T: hash::Hash>(x: &T) -> u64 {
     use std::hash::Hasher;