vendor/num-integer/src/roots.rs

   1 use core;
   2 use core::mem;
   3 use traits::checked_pow;
   4 use traits::PrimInt;
   5 use Integer;
   6
   7 /// Provides methods to compute an integer's square root, cube root,
   8 /// and arbitrary `n`th root.
   9 pub trait Roots: Integer {
  10     /// Returns the truncated principal `n`th root of an integer
  11     /// -- `if x >= 0 { ⌊ⁿ√x⌋ } else { ⌈ⁿ√x⌉ }`
  12     ///
  13     /// This is solving for `r` in `rⁿ = x`, rounding toward zero.
  14     /// If `x` is positive, the result will satisfy `rⁿ ≤ x < (r+1)ⁿ`.
  15     /// If `x` is negative and `n` is odd, then `(r-1)ⁿ < x ≤ rⁿ`.
  16     ///
  17     /// # Panics
  18     ///
  19     /// Panics if `n` is zero:
  20     ///
  21     /// ```should_panic
  22     /// # use num_integer::Roots;
  23     /// println!("can't compute ⁰√x : {}", 123.nth_root(0));
  24     /// ```
  25     ///
  26     /// or if `n` is even and `self` is negative:
  27     ///
  28     /// ```should_panic
  29     /// # use num_integer::Roots;
  30     /// println!("no imaginary numbers... {}", (-1).nth_root(10));
  31     /// ```
  32     ///
  33     /// # Examples
  34     ///
  35     /// ```
  36     /// use num_integer::Roots;
  37     ///
  38     /// let x: i32 = 12345;
  39     /// assert_eq!(x.nth_root(1), x);
  40     /// assert_eq!(x.nth_root(2), x.sqrt());
  41     /// assert_eq!(x.nth_root(3), x.cbrt());
  42     /// assert_eq!(x.nth_root(4), 10);
  43     /// assert_eq!(x.nth_root(13), 2);
  44     /// assert_eq!(x.nth_root(14), 1);
  45     /// assert_eq!(x.nth_root(std::u32::MAX), 1);
  46     ///
  47     /// assert_eq!(std::i32::MAX.nth_root(30), 2);
  48     /// assert_eq!(std::i32::MAX.nth_root(31), 1);
  49     /// assert_eq!(std::i32::MIN.nth_root(31), -2);
  50     /// assert_eq!((std::i32::MIN + 1).nth_root(31), -1);
  51     ///
  52     /// assert_eq!(std::u32::MAX.nth_root(31), 2);
  53     /// assert_eq!(std::u32::MAX.nth_root(32), 1);
  54     /// ```
  55     fn nth_root(&self, n: u32) -> Self;
  56
  57     /// Returns the truncated principal square root of an integer -- `⌊√x⌋`
  58     ///
  59     /// This is solving for `r` in `r² = x`, rounding toward zero.
  60     /// The result will satisfy `r² ≤ x < (r+1)²`.
  61     ///
  62     /// # Panics
  63     ///
  64     /// Panics if `self` is less than zero:
  65     ///
  66     /// ```should_panic
  67     /// # use num_integer::Roots;
  68     /// println!("no imaginary numbers... {}", (-1).sqrt());
  69     /// ```
  70     ///
  71     /// # Examples
  72     ///
  73     /// ```
  74     /// use num_integer::Roots;
  75     ///
  76     /// let x: i32 = 12345;
  77     /// assert_eq!((x * x).sqrt(), x);
  78     /// assert_eq!((x * x + 1).sqrt(), x);
  79     /// assert_eq!((x * x - 1).sqrt(), x - 1);
  80     /// ```
  81     #[inline]
  82     fn sqrt(&self) -> Self {
  83         self.nth_root(2)
  84     }
  85
  86     /// Returns the truncated principal cube root of an integer --
  87     /// `if x >= 0 { ⌊∛x⌋ } else { ⌈∛x⌉ }`
  88     ///
  89     /// This is solving for `r` in `r³ = x`, rounding toward zero.
  90     /// If `x` is positive, the result will satisfy `r³ ≤ x < (r+1)³`.
  91     /// If `x` is negative, then `(r-1)³ < x ≤ r³`.
  92     ///
  93     /// # Examples
  94     ///
  95     /// ```
  96     /// use num_integer::Roots;
  97     ///
  98     /// let x: i32 = 1234;
  99     /// assert_eq!((x * x * x).cbrt(), x);
 100     /// assert_eq!((x * x * x + 1).cbrt(), x);
 101     /// assert_eq!((x * x * x - 1).cbrt(), x - 1);
 102     ///
 103     /// assert_eq!((-(x * x * x)).cbrt(), -x);
 104     /// assert_eq!((-(x * x * x + 1)).cbrt(), -x);
 105     /// assert_eq!((-(x * x * x - 1)).cbrt(), -(x - 1));
 106     /// ```
 107     #[inline]
 108     fn cbrt(&self) -> Self {
 109         self.nth_root(3)
 110     }
 111 }
 112
 113 /// Returns the truncated principal square root of an integer --
 114 /// see [Roots::sqrt](trait.Roots.html#method.sqrt).
 115 #[inline]
 116 pub fn sqrt<T: Roots>(x: T) -> T {
 117     x.sqrt()
 118 }
 119
 120 /// Returns the truncated principal cube root of an integer --
 121 /// see [Roots::cbrt](trait.Roots.html#method.cbrt).
 122 #[inline]
 123 pub fn cbrt<T: Roots>(x: T) -> T {
 124     x.cbrt()
 125 }
 126
 127 /// Returns the truncated principal `n`th root of an integer --
 128 /// see [Roots::nth_root](trait.Roots.html#tymethod.nth_root).
 129 #[inline]
 130 pub fn nth_root<T: Roots>(x: T, n: u32) -> T {
 131     x.nth_root(n)
 132 }
 133
 134 macro_rules! signed_roots {
 135     ($T:ty, $U:ty) => {
 136         impl Roots for $T {
 137             #[inline]
 138             fn nth_root(&self, n: u32) -> Self {
 139                 if *self >= 0 {
 140                     (*self as $U).nth_root(n) as Self
 141                 } else {
 142                     assert!(n.is_odd(), "even roots of a negative are imaginary");
 143                     -((self.wrapping_neg() as $U).nth_root(n) as Self)
 144                 }
 145             }
 146
 147             #[inline]
 148             fn sqrt(&self) -> Self {
 149                 assert!(*self >= 0, "the square root of a negative is imaginary");
 150                 (*self as $U).sqrt() as Self
 151             }
 152
 153             #[inline]
 154             fn cbrt(&self) -> Self {
 155                 if *self >= 0 {
 156                     (*self as $U).cbrt() as Self
 157                 } else {
 158                     -((self.wrapping_neg() as $U).cbrt() as Self)
 159                 }
 160             }
 161         }
 162     };
 163 }
 164
 165 signed_roots!(i8, u8);
 166 signed_roots!(i16, u16);
 167 signed_roots!(i32, u32);
 168 signed_roots!(i64, u64);
 169 #[cfg(has_i128)]
 170 signed_roots!(i128, u128);
 171 signed_roots!(isize, usize);
 172
 173 #[inline]
 174 fn fixpoint<T, F>(mut x: T, f: F) -> T
 175 where
 176     T: Integer + Copy,
 177     F: Fn(T) -> T,
 178 {
 179     let mut xn = f(x);
 180     while x < xn {
 181         x = xn;
 182         xn = f(x);
 183     }
 184     while x > xn {
 185         x = xn;
 186         xn = f(x);
 187     }
 188     x
 189 }
 190
 191 #[inline]
 192 fn bits<T>() -> u32 {
 193     8 * mem::size_of::<T>() as u32
 194 }
 195
 196 #[inline]
 197 fn log2<T: PrimInt>(x: T) -> u32 {
 198     debug_assert!(x > T::zero());
 199     bits::<T>() - 1 - x.leading_zeros()
 200 }
 201
 202 macro_rules! unsigned_roots {
 203     ($T:ident) => {
 204         impl Roots for $T {
 205             fn nth_root(&self, n: u32) -> Self {
 206                 // Specialize small roots
 207                 match n {
 208                     0 => panic!("can't find a root of degree 0!"),
 209                     1 => return *self,
 210                     2 => return self.sqrt(),
 211                     3 => return self.cbrt(),
 212                     _ => (),
 213                 }
 214
 215                 // The root of values less than 2ⁿ can only be 0 or 1.
 216                 if bits::<$T>() <= n || *self < (1 << n) {
 217                     return (*self > 0) as $T;
 218                 }
 219
 220                 if bits::<$T>() > 64 {
 221                     // 128-bit division is slow, so do a bitwise `nth_root` until it's small enough.
 222                     return if *self <= core::u64::MAX as $T {
 223                         (*self as u64).nth_root(n) as $T
 224                     } else {
 225                         let lo = (self >> n).nth_root(n) << 1;
 226                         let hi = lo + 1;
 227                         // 128-bit `checked_mul` also involves division, but we can't always
 228                         // compute `hiⁿ` without risking overflow.  Try to avoid it though...
 229                         if hi.next_power_of_two().trailing_zeros() * n >= bits::<$T>() {
 230                             match checked_pow(hi, n as usize) {
 231                                 Some(x) if x <= *self => hi,
 232                                 _ => lo,
 233                             }
 234                         } else {
 235                             if hi.pow(n) <= *self {
 236                                 hi
 237                             } else {
 238                                 lo
 239                             }
 240                         }
 241                     };
 242                 }
 243
 244                 #[cfg(feature = "std")]
 245                 #[inline]
 246                 fn guess(x: $T, n: u32) -> $T {
 247                     // for smaller inputs, `f64` doesn't justify its cost.
 248                     if bits::<$T>() <= 32 || x <= core::u32::MAX as $T {
 249                         1 << ((log2(x) + n - 1) / n)
 250                     } else {
 251                         ((x as f64).ln() / f64::from(n)).exp() as $T
 252                     }
 253                 }
 254
 255                 #[cfg(not(feature = "std"))]
 256                 #[inline]
 257                 fn guess(x: $T, n: u32) -> $T {
 258                     1 << ((log2(x) + n - 1) / n)
 259                 }
 260
 261                 // https://en.wikipedia.org/wiki/Nth_root_algorithm
 262                 let n1 = n - 1;
 263                 let next = |x: $T| {
 264                     let y = match checked_pow(x, n1 as usize) {
 265                         Some(ax) => self / ax,
 266                         None => 0,
 267                     };
 268                     (y + x * n1 as $T) / n as $T
 269                 };
 270                 fixpoint(guess(*self, n), next)
 271             }
 272
 273             fn sqrt(&self) -> Self {
 274                 if bits::<$T>() > 64 {
 275                     // 128-bit division is slow, so do a bitwise `sqrt` until it's small enough.
 276                     // https://en.wikipedia.org/wiki/Integer_square_root#Using_bitwise_operations
 277                     return if *self <= core::u64::MAX as $T {
 278                         (*self as u64).sqrt() as $T
 279                     } else {
 280                         let lo = (self >> 2u32).sqrt() << 1;
 281                         let hi = lo + 1;
 282                         if hi * hi <= *self {
 283                             hi
 284                         } else {
 285                             lo
 286                         }
 287                     };
 288                 }
 289
 290                 if *self < 4 {
 291                     return (*self > 0) as Self;
 292                 }
 293
 294                 #[cfg(feature = "std")]
 295                 #[inline]
 296                 fn guess(x: $T) -> $T {
 297                     (x as f64).sqrt() as $T
 298                 }
 299
 300                 #[cfg(not(feature = "std"))]
 301                 #[inline]
 302                 fn guess(x: $T) -> $T {
 303                     1 << ((log2(x) + 1) / 2)
 304                 }
 305
 306                 // https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Babylonian_method
 307                 let next = |x: $T| (self / x + x) >> 1;
 308                 fixpoint(guess(*self), next)
 309             }
 310
 311             fn cbrt(&self) -> Self {
 312                 if bits::<$T>() > 64 {
 313                     // 128-bit division is slow, so do a bitwise `cbrt` until it's small enough.
 314                     return if *self <= core::u64::MAX as $T {
 315                         (*self as u64).cbrt() as $T
 316                     } else {
 317                         let lo = (self >> 3u32).cbrt() << 1;
 318                         let hi = lo + 1;
 319                         if hi * hi * hi <= *self {
 320                             hi
 321                         } else {
 322                             lo
 323                         }
 324                     };
 325                 }
 326
 327                 if bits::<$T>() <= 32 {
 328                     // Implementation based on Hacker's Delight `icbrt2`
 329                     let mut x = *self;
 330                     let mut y2 = 0;
 331                     let mut y = 0;
 332                     let smax = bits::<$T>() / 3;
 333                     for s in (0..smax + 1).rev() {
 334                         let s = s * 3;
 335                         y2 *= 4;
 336                         y *= 2;
 337                         let b = 3 * (y2 + y) + 1;
 338                         if x >> s >= b {
 339                             x -= b << s;
 340                             y2 += 2 * y + 1;
 341                             y += 1;
 342                         }
 343                     }
 344                     return y;
 345                 }
 346
 347                 if *self < 8 {
 348                     return (*self > 0) as Self;
 349                 }
 350                 if *self <= core::u32::MAX as $T {
 351                     return (*self as u32).cbrt() as $T;
 352                 }
 353
 354                 #[cfg(feature = "std")]
 355                 #[inline]
 356                 fn guess(x: $T) -> $T {
 357                     (x as f64).cbrt() as $T
 358                 }
 359
 360                 #[cfg(not(feature = "std"))]
 361                 #[inline]
 362                 fn guess(x: $T) -> $T {
 363                     1 << ((log2(x) + 2) / 3)
 364                 }
 365
 366                 // https://en.wikipedia.org/wiki/Cube_root#Numerical_methods
 367                 let next = |x: $T| (self / (x * x) + x * 2) / 3;
 368                 fixpoint(guess(*self), next)
 369             }
 370         }
 371     };
 372 }
 373
 374 unsigned_roots!(u8);
 375 unsigned_roots!(u16);
 376 unsigned_roots!(u32);
 377 unsigned_roots!(u64);
 378 #[cfg(has_i128)]
 379 unsigned_roots!(u128);
 380 unsigned_roots!(usize);