Optimize isqrt

CHANGELOG.md

@@ -47,13 +47,15 @@ and this project adheres to
 - cosmwasm-schema-derive: Improve emitted error messages ([#2063])
 - cosmwasm-schema: `#[cw_serde]` now doesn't add `#[serde(deny_unknown_fields)]`
   to the expanded code anymore ([#2080])
+- cosmwasm-std: Improve performance of `Uint{64,128,256,512}::isqrt` ([#2108])
 
 [#2044]: https://github.com/CosmWasm/cosmwasm/pull/2044
 [#2051]: https://github.com/CosmWasm/cosmwasm/pull/2051
 [#2059]: https://github.com/CosmWasm/cosmwasm/pull/2059
 [#2063]: https://github.com/CosmWasm/cosmwasm/pull/2063
 [#2070]: https://github.com/CosmWasm/cosmwasm/pull/2070
 [#2080]: https://github.com/CosmWasm/cosmwasm/pull/2080
+[#2108]: https://github.com/CosmWasm/cosmwasm/pull/2108
 
 ## [2.0.1] - 2024-04-03
 

packages/core/src/math/isqrt.rs

@@ -13,8 +13,10 @@ pub trait Isqrt {
 impl<I> Isqrt for I
 where
     I: Unsigned
+        + Log2
         + ops::Add<I, Output = I>
         + ops::Div<I, Output = I>
+        + ops::Shl<u32, Output = I>
         + ops::Shr<u32, Output = I>
         + cmp::PartialOrd
         + Copy
@@ -23,9 +25,13 @@ where
     /// Algorithm adapted from
     /// [Wikipedia](https://en.wikipedia.org/wiki/Integer_square_root#Example_implementation_in_C).
     fn isqrt(self) -> Self {
-        let mut x0 = self >> 1;
+        let zero = Self::from(0);
+        if self == zero {
+            return zero;
+        }
+        let mut x0 = Self::from(1u8) << ((self.log_2() / 2) + 1);
 
-        if x0 > 0.into() {
+        if x0 > zero {
             let mut x1 = (x0 + self / x0) >> 1;
 
             while x1 < x0 {
@@ -52,6 +58,29 @@ impl Unsigned for Uint256 {}
 impl Unsigned for Uint512 {}
 impl Unsigned for usize {}
 
+trait Log2 {
+    fn log_2(self) -> u32;
+}
+macro_rules! impl_log2 {
+    ($type:ty) => {
+        impl Log2 for $type {
+            fn log_2(self) -> u32 {
+                self.ilog2()
+            }
+        }
+    };
+}
+impl_log2!(u8);
+impl_log2!(u16);
+impl_log2!(u32);
+impl_log2!(u64);
+impl_log2!(u128);
+impl_log2!(usize);
+impl_log2!(Uint64);
+impl_log2!(Uint128);
+impl_log2!(Uint256);
+impl_log2!(Uint512);
+
 #[cfg(test)]
 mod tests {
     use super::*;

packages/core/src/math/uint128.rs

@@ -96,6 +96,16 @@ impl Uint128 {
         self.0.pow(exp).into()
     }
 
+    /// Returns the base 2 logarithm of the number, rounded down.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if `self` is zero.
+    #[must_use = "this returns the result of the operation, without modifying the original"]
+    pub fn ilog2(self) -> u32 {
+        self.0.checked_ilog2().unwrap()
+    }
+
     /// Returns `self * numerator / denominator`.
     ///
     /// Due to the nature of the integer division involved, the result is always floored.

packages/core/src/math/uint256.rs

@@ -168,6 +168,16 @@ impl Uint256 {
         Self(self.0.pow(exp))
     }
 
+    /// Returns the base 2 logarithm of the number, rounded down.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if `self` is zero.
+    #[must_use = "this returns the result of the operation, without modifying the original"]
+    pub fn ilog2(self) -> u32 {
+        self.0.checked_ilog2().unwrap()
+    }
+
     /// Returns `self * numerator / denominator`.
     ///
     /// Due to the nature of the integer division involved, the result is always floored.

packages/core/src/math/uint512.rs

@@ -193,6 +193,16 @@ impl Uint512 {
         Self(self.0.pow(exp))
     }
 
+    /// Returns the base 2 logarithm of the number, rounded down.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if `self` is zero.
+    #[must_use = "this returns the result of the operation, without modifying the original"]
+    pub fn ilog2(self) -> u32 {
+        self.0.checked_ilog2().unwrap()
+    }
+
     pub fn checked_add(self, other: Self) -> Result<Self, OverflowError> {
         self.0
             .checked_add(other.0)

packages/core/src/math/uint64.rs

@@ -90,6 +90,16 @@ impl Uint64 {
         self.0.pow(exp).into()
     }
 
+    /// Returns the base 2 logarithm of the number, rounded down.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if `self` is zero.
+    #[must_use = "this returns the result of the operation, without modifying the original"]
+    pub fn ilog2(self) -> u32 {
+        self.0.checked_ilog2().unwrap()
+    }
+
     /// Returns `self * numerator / denominator`.
     ///
     /// Due to the nature of the integer division involved, the result is always floored.