1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

//!
//! This module contains the single trait [`IntegerSquareRoot`] and implements it for primitive
//! integer types.
//!
//! # Example
//!
//! ```
//! extern crate integer_sqrt;
//! // `use` trait to get functionality
//! use integer_sqrt::IntegerSquareRoot;
//!
//! # fn main() {
//! assert_eq!(4u8.integer_sqrt(), 2);
//! ```
//!
//! [`IntegerSquareRoot`]: ./trait.IntegerSquareRoot.html

/// A trait implementing integer square root.
pub trait IntegerSquareRoot {
    /// Find the integer square root.
    ///
    /// See [Integer_square_root on wikipedia][wiki_article] for more information (and also the
    /// source of this algorithm)
    ///
    /// # Panics
    ///
    /// For negative numbers (`i` family) this function will panic on negative input
    ///
    /// [wiki_article]: https://en.wikipedia.org/wiki/Integer_square_root
    fn integer_sqrt(&self) -> Self where Self: Sized {
        self.integer_sqrt_checked().expect("cannot calculate square root of negative number")
    }

    /// Find the integer square root, returning `None` if the number is negative (this can never
    /// happen for unsigned types).
    fn integer_sqrt_checked(&self) -> Option<Self> where Self: Sized;
}

// This could be more optimized
macro_rules! impl_isqrt {
    () => ();
    ($t:ty) => {impl_isqrt!($t,);};
    ($t:ty, $($e:tt)*) => {
        impl IntegerSquareRoot for $t {
            #[allow(unused_comparisons)]
            fn integer_sqrt_checked(&self) -> Option<Self> {
                // Hopefully this will be stripped for unsigned numbers (impossible condition)
                if *self < 0 {
                    return None
                }
                // Find greatest shift
                let mut shift = 2;
                let mut n_shifted = *self >> shift;
                // We check for n_shifted being self, since some implementations of logical
                // right shifting shift modulo the word size.
                while n_shifted != 0 && n_shifted != *self {
                    shift = shift + 2;
                    n_shifted = self.wrapping_shr(shift);
                }
                shift = shift - 2;

                // Find digits of result.
                let mut result = 0;
                loop {
                    result = result << 1;
                    let candidate_result = result + 1;
                    if candidate_result * candidate_result <= *self >> shift {
                        result = candidate_result;
                    }
                    if shift == 0 {
                        break;
                    }
                    shift = shift.saturating_sub(2);
                }

                Some(result)
            }
        }

        impl_isqrt!($($e)*);
    };
}


impl_isqrt!(usize, u64, u32, u16, u8, isize, i64, i32, i16, i8);


#[cfg(test)]
mod tests {
    use super::IntegerSquareRoot;
    use std::{u8, u16, u64, i8};

    #[test]
    fn u8_sqrt() {
        let tests = [
            (0u8, 0u8),
            (4, 2),
            (7, 2),
            (81, 9),
            (80, 8),
            (u8::MAX, (u8::MAX as f64).sqrt() as u8),
        ];
        for &(in_, out) in tests.iter() {
            assert_eq!(in_.integer_sqrt(), out, "in {}", in_);
        }
    }

    #[test]
    fn i8_sqrt() {
        let tests = [
            (0i8, 0i8),
            (4, 2),
            (7, 2),
            (81, 9),
            (80, 8),
            (i8::MAX, (i8::MAX as f64).sqrt() as i8),
        ];
        for &(in_, out) in tests.iter() {
            assert_eq!(in_.integer_sqrt(), out, "in {}", in_);
        }
    }

    #[test]
    #[should_panic]
    fn i8_sqrt_negative() {
        (-12i8).integer_sqrt();
    }

    #[test]
    fn u16_sqrt() {
        let tests = [
            (0u16, 0u16),
            (4, 2),
            (7, 2),
            (81, 9),
            (80, 8),
            (u16::MAX, (u16::MAX as f64).sqrt() as u16),
        ];
        for &(in_, out) in tests.iter() {
            assert_eq!(in_.integer_sqrt(), out, "in {}", in_);
        }
    }

    #[test]
    fn u64_sqrt() {
        let sqrt_max = 4_294_967_295;
        let tests = [
            (0u64, 0u64),
            (4, 2),
            (7, 2),
            (81, 9),
            (80, 8),
            (u64::MAX, sqrt_max),
        ];
        for &(in_, out) in tests.iter() {
            assert_eq!(in_.integer_sqrt(), out, "in {}", in_);
        }
        // checks to make sure we have the right number for u64::MAX.integer_sqrt()
        // we can't use the same strategy as in previous tests as f64 is now not returning the
        // correct floored integer
        assert!(sqrt_max * sqrt_max <= u64::MAX);
        // check that the next number's square overflows
        assert!((sqrt_max + 1).checked_mul(sqrt_max + 1).is_none());
    }
}