From 83649f11fa3360d163b111fd9ae9ac2758aa4778 Mon Sep 17 00:00:00 2001
From: Michael Pfaff <michael@pfaff.dev>
Date: Tue, 1 Mar 2022 18:30:30 -0500
Subject: [PATCH] Make const

---
 Cargo.toml |  5 +++--
 src/lib.rs | 57 ++++++++++++++++++++++++++++++++++++++++--------------
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 50c0c44..d19a92f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -2,7 +2,7 @@
 name = "integer-sqrt"
 description = """
 An implementation of integer square root algorithm for primitive rust types"""
-version = "0.1.5"
+version = "0.2.0"
 authors = ["Richard Dodd <richard.o.dodd@gmail.com>", "Joseph Richey", "Sergei Shulepov"]
 include = ["src/**/*.rs", "Cargo.toml"]
 repository = "https://github.com/derekdreery/integer-sqrt-rs"
@@ -10,6 +10,7 @@ readme = "README.md"
 keywords = ["integer", "square", "root", "isqrt", "sqrt"]
 categories = ["algorithms", "no-std"]
 license = "Apache-2.0/MIT"
+edition = "2021"
 
 [dependencies]
-num-traits = { version = "0.2", default-features = false }
+static_assertions = "1.1"
diff --git a/src/lib.rs b/src/lib.rs
index 1feb1e8..1ff3b67 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,6 +16,9 @@
 //!
 //! [`IntegerSquareRoot`]: ./trait.IntegerSquareRoot.html
 #![no_std]
+#![feature(const_trait_impl)]
+#![feature(const_fn_trait_bound)]
+#![feature(const_option)]
 
 /// A trait implementing integer square root.
 pub trait IntegerSquareRoot {
@@ -29,6 +32,7 @@ pub trait IntegerSquareRoot {
     /// For negative numbers (`i` family) this function will panic on negative input
     ///
     /// [wiki_article]: https://en.wikipedia.org/wiki/Integer_square_root
+    #[default_method_body_is_const]
     fn integer_sqrt(&self) -> Self
     where
         Self: Sized,
@@ -44,39 +48,62 @@ pub trait IntegerSquareRoot {
         Self: Sized;
 }
 
-impl<T: num_traits::PrimInt> IntegerSquareRoot for T {
-    fn integer_sqrt_checked(&self) -> Option<Self> {
+macro_rules! integer_sqrt {
+    ($ty:ty as $unsigned_ty:ty, $value:expr) => {{
+        static_assertions::assert_eq_size!($ty, $unsigned_ty);
         use core::cmp::Ordering;
-        match self.cmp(&T::zero()) {
+        match $value {
+            0 => return Some(0),
             // Hopefully this will be stripped for unsigned numbers (impossible condition)
-            Ordering::Less => return None,
-            Ordering::Equal => return Some(T::zero()),
+            v if v < 0 => return None,
             _ => {}
         }
 
         // Compute bit, the largest power of 4 <= n
-        let max_shift: u32 = T::zero().leading_zeros() - 1;
-        let shift: u32 = (max_shift - self.leading_zeros()) & !1;
-        let mut bit = T::one().unsigned_shl(shift);
+        const ZERO: $ty = 0;
+        let max_shift: u32 = ZERO.leading_zeros() - 1;
+        let shift: u32 = (max_shift - $value.leading_zeros()) & !1;
+        let mut bit = (1 as $unsigned_ty << shift) as $ty;
 
         // Algorithm based on the implementation in:
         // https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)
         // Note that result/bit are logically unsigned (even if T is signed).
-        let mut n = *self;
-        let mut result = T::zero();
-        while bit != T::zero() {
+        let mut n = $value;
+        let mut result = 0;
+        while bit != 0 {
             if n >= (result + bit) {
                 n = n - (result + bit);
-                result = result.unsigned_shr(1) + bit;
+                result = (result as $unsigned_ty >> 1) as $ty + bit;
             } else {
-                result = result.unsigned_shr(1);
+                result = (result as $unsigned_ty >> 1) as $ty;
             }
-            bit = bit.unsigned_shr(2);
+            bit = (bit as $unsigned_ty >> 2) as $ty;
         }
         Some(result)
-    }
+    }};
+    (impl const $ty:ty as $unsigned_ty:ty) => {
+        impl const IntegerSquareRoot for $ty {
+            fn integer_sqrt_checked(&self) -> Option<Self> {
+                integer_sqrt!($ty as $unsigned_ty, *self)
+            }
+        }
+    };
 }
 
+integer_sqrt!(impl const u8 as u8);
+integer_sqrt!(impl const u16 as u16);
+integer_sqrt!(impl const u32 as u32);
+integer_sqrt!(impl const u64 as u64);
+integer_sqrt!(impl const u128 as u128);
+integer_sqrt!(impl const usize as usize);
+
+integer_sqrt!(impl const i8 as u8);
+integer_sqrt!(impl const i16 as u16);
+integer_sqrt!(impl const i32 as u32);
+integer_sqrt!(impl const i64 as u64);
+integer_sqrt!(impl const i128 as u128);
+integer_sqrt!(impl const isize as usize);
+
 #[cfg(test)]
 mod tests {
     use super::IntegerSquareRoot;