Merge pull request #2108 from CosmWasm/chipshort/isqrt-optimization

Optimize `isqrt`

CHANGELOG.md

@@ -48,13 +48,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

@@ -15,17 +15,22 @@ where
     I: Unsigned
         + ops::Add<I, Output = I>
         + ops::Div<I, Output = I>
+        + ops::Shl<u32, Output = I>
         + ops::Shr<u32, Output = I>
         + cmp::PartialOrd
-        + Copy
-        + From<u8>,
+        + Copy,
 {
     /// 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;
+        // sqrt(0) = 0, sqrt(1) = 1
+        if self <= Self::ONE {
+            return self;
+        }
+
+        let mut x0 = Self::ONE << ((self.log_2() / 2) + 1);
 
-        if x0 > 0.into() {
+        if x0 > Self::ZERO {
             let mut x1 = (x0 + self / x0) >> 1;
 
             while x1 < x0 {
@@ -40,17 +45,35 @@ where
 }
 
 /// Marker trait for types that represent unsigned integers.
-pub trait Unsigned {}
-impl Unsigned for u8 {}
-impl Unsigned for u16 {}
-impl Unsigned for u32 {}
-impl Unsigned for u64 {}
-impl Unsigned for u128 {}
-impl Unsigned for Uint64 {}
-impl Unsigned for Uint128 {}
-impl Unsigned for Uint256 {}
-impl Unsigned for Uint512 {}
-impl Unsigned for usize {}
+pub trait Unsigned {
+    const ZERO: Self;
+    const ONE: Self;
+
+    fn log_2(self) -> u32;
+}
+
+macro_rules! impl_unsigned {
+    ($type:ty, $zero:expr, $one:expr) => {
+        impl Unsigned for $type {
+            const ZERO: Self = $zero;
+            const ONE: Self = $one;
+
+            fn log_2(self) -> u32 {
+                self.ilog2()
+            }
+        }
+    };
+}
+impl_unsigned!(u8, 0, 1);
+impl_unsigned!(u16, 0, 1);
+impl_unsigned!(u32, 0, 1);
+impl_unsigned!(u64, 0, 1);
+impl_unsigned!(u128, 0, 1);
+impl_unsigned!(usize, 0, 1);
+impl_unsigned!(Uint64, Self::zero(), Self::one());
+impl_unsigned!(Uint128, Self::zero(), Self::one());
+impl_unsigned!(Uint256, Self::zero(), Self::one());
+impl_unsigned!(Uint512, Self::zero(), Self::one());
 
 #[cfg(test)]
 mod tests {

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.