uint: implement integer_sqrt (#554)

* uint: implement integer_sqrt

* uint: remove wrong assertion

* go brrrrrrrrrrrrrr

* we don't need to move anymore

* impl-num-traits: add integer-sqrt trait support

* add missing test

* fmt

* bump uint

Author: Andronik Ordian

.travis.yml

@@ -61,6 +61,7 @@ script:
   - cd primitive-types/ && cargo test --all-features && cd ..
   - cd primitive-types/ && cargo test --no-default-features --features=serde_no_std && cd ..
   - cd primitive-types/ && cargo test --no-default-features --features=scale-info && cd ..
+  - cd primitive-types/ && cargo test --no-default-features --features=num-traits && cd ..
   - cd rlp/ && cargo test --no-default-features && cargo check --benches && cd ..
   - cd triehash/ && cargo check --benches && cd ..
   - cd kvdb-web/ && wasm-pack test --headless --firefox && cd ..

primitive-types/Cargo.toml

@@ -40,3 +40,7 @@ required-features = ["scale-info"]
 [[test]]
 name = "fp_conversion"
 required-features = ["fp-conversion"]
+
+[[test]]
+name = "num_traits"
+required-features = ["num-traits"]

primitive-types/impls/num-traits/CHANGELOG.md

@@ -5,3 +5,4 @@ The format is based on [Keep a Changelog].
 [Keep a Changelog]: http://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
+- Added `integer-sqrt` trait support. [#554](https://github.com/paritytech/parity-common/pull/554)

primitive-types/impls/num-traits/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "impl-num-traits"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Parity Technologies <[email protected]>"]
 license = "MIT OR Apache-2.0"
 homepage = "https://github.com/paritytech/parity-common"
@@ -9,4 +9,5 @@ edition = "2018"
 
 [dependencies]
 num-traits = { version = "0.2", default-features = false }
+integer-sqrt = "0.1"
 uint = { version = "0.9.0", path = "../../../uint", default-features = false }

primitive-types/impls/num-traits/src/lib.rs

@@ -13,6 +13,9 @@
 #[doc(hidden)]
 pub use num_traits;
 
+#[doc(hidden)]
+pub use integer_sqrt;
+
 #[doc(hidden)]
 pub use uint;
 
@@ -48,5 +51,11 @@ macro_rules! impl_uint_num_traits {
 				Self::from_str_radix(txt, radix)
 			}
 		}
+
+		impl $crate::integer_sqrt::IntegerSquareRoot for $name {
+			fn integer_sqrt_checked(&self) -> Option<Self> {
+				Some(self.integer_sqrt())
+			}
+		}
 	};
 }

primitive-types/tests/num_traits.rs

@@ -0,0 +1,17 @@
+// Copyright 2021 Parity Technologies
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use impl_num_traits::integer_sqrt::IntegerSquareRoot;
+use primitive_types::U256;
+
+#[test]
+fn u256_isqrt() {
+	let x = U256::MAX;
+	let s = x.integer_sqrt_checked().unwrap();
+	assert_eq!(x.integer_sqrt(), s);
+}

uint/CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog].
 [Keep a Changelog]: http://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
+- Added `integer_sqrt` method. [#554](https://github.com/paritytech/parity-common/pull/554)
 
 ## [0.9.0] - 2021-01-05
 - Allow `0x` prefix in `from_str`. [#487](https://github.com/paritytech/parity-common/pull/487)

uint/Cargo.toml

@@ -4,7 +4,7 @@ homepage = "http://parity.io"
 repository = "https://github.com/paritytech/parity-common"
 license = "MIT OR Apache-2.0"
 name = "uint"
-version = "0.9.0"
+version = "0.9.1"
 authors = ["Parity Technologies <[email protected]>"]
 readme = "README.md"
 edition = "2018"

uint/benches/bigint.rs

@@ -44,6 +44,7 @@ criterion_group!(
 	u256_div,
 	u512_div_mod,
 	u256_rem,
+	u256_integer_sqrt,
 	u256_bit_and,
 	u256_bit_or,
 	u256_bit_xor,
@@ -58,6 +59,7 @@ criterion_group!(
 	u512_mul,
 	u512_div,
 	u512_rem,
+	u512_integer_sqrt,
 	u512_mul_u32_vs_u64,
 	mulmod_u512_vs_biguint_vs_gmp,
 	conversions,
@@ -253,6 +255,22 @@ fn u256_rem(c: &mut Criterion) {
 	);
 }
 
+fn u256_integer_sqrt(c: &mut Criterion) {
+	c.bench(
+		"u256_integer_sqrt",
+		ParameterizedBenchmark::new(
+			"",
+			|b, x| b.iter(|| black_box(x.integer_sqrt().0)),
+			vec![
+				U256::from(u64::MAX),
+				U256::from(u128::MAX) + 1,
+				U256::from(u128::MAX - 1) * U256::from(u128::MAX - 1) - 1,
+				U256::MAX,
+			],
+		),
+	);
+}
+
 fn u512_pairs() -> Vec<(U512, U512)> {
 	vec![
 		(U512::from(1u64), U512::from(0u64)),
@@ -286,6 +304,23 @@ fn u512_mul(c: &mut Criterion) {
 	);
 }
 
+fn u512_integer_sqrt(c: &mut Criterion) {
+	c.bench(
+		"u512_integer_sqrt",
+		ParameterizedBenchmark::new(
+			"",
+			|b, x| b.iter(|| black_box(x.integer_sqrt().0)),
+			vec![
+				U512::from(u32::MAX) + 1,
+				U512::from(u64::MAX),
+				(U512::from(u128::MAX) + 1) * (U512::from(u128::MAX) + 1),
+				U256::MAX.full_mul(U256::MAX) - 1,
+				U512::MAX,
+			],
+		),
+	);
+}
+
 fn u512_div(c: &mut Criterion) {
 	let one = U512([
 		8326634216714383706,

uint/fuzz/Cargo.toml

@@ -24,3 +24,7 @@ path = "fuzz_targets/div_mod.rs"
 [[bin]]
 name = "div_mod_word"
 path = "fuzz_targets/div_mod_word.rs"
+
+[[bin]]
+name = "isqrt"
+path = "fuzz_targets/isqrt.rs"

uint/fuzz/fuzz_targets/div_mod.rs

@@ -23,7 +23,7 @@ fn from_gmp(x: Integer) -> U512 {
 }
 
 fuzz_target!(|data: &[u8]| {
-    if data.len() == 128 {
+	if data.len() == 128 {
 		let x = U512::from_little_endian(&data[..64]);
 		let y = U512::from_little_endian(&data[64..]);
 		let x_gmp = Integer::from_digits(&data[..64], Order::LsfLe);
@@ -32,5 +32,5 @@ fuzz_target!(|data: &[u8]| {
 			let (a, b) = x_gmp.div_rem(y_gmp);
 			assert_eq!((from_gmp(a), from_gmp(b)), x.div_mod(y));
 		}
-    }
+	}
 });

uint/fuzz/fuzz_targets/div_mod_word.rs

@@ -57,7 +57,7 @@ fn div_mod_word(hi: u64, lo: u64, y: u64) -> (u64, u64) {
 }
 
 fuzz_target!(|data: &[u8]| {
-    if data.len() == 24 {
+	if data.len() == 24 {
 		let mut buf = [0u8; 8];
 		buf.copy_from_slice(&data[..8]);
 		let x = u64::from_ne_bytes(buf);
@@ -68,5 +68,5 @@ fuzz_target!(|data: &[u8]| {
 		if x < z {
 			assert_eq!(div_mod_word(x, y, z), div_mod_word_u128(x, y, z));
 		}
-    }
+	}
 });

uint/fuzz/fuzz_targets/isqrt.rs

@@ -0,0 +1,50 @@
+// Copyright 2021 Parity Technologies
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+use uint::*;
+
+construct_uint! {
+	pub struct U256(4);
+}
+
+fn isqrt(mut me: U256) -> U256 {
+	let one = U256::one();
+	if me <= one {
+		return me;
+	}
+	// the implementation is based on:
+	// https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)
+
+	// "bit" starts at the highest power of four <= self.
+	let max_shift = 4 * 64 as u32 - 1;
+	let shift: u32 = (max_shift - me.leading_zeros()) & !1;
+	let mut bit = one << shift;
+	let mut result = U256::zero();
+	while !bit.is_zero() {
+		let x = result + bit;
+		result >>= 1;
+		if me >= x {
+			me -= x;
+			result += bit;
+		}
+		bit >>= 2;
+	}
+	result
+}
+
+fuzz_target!(|data: &[u8]| {
+	if data.len() == 32 {
+		let x = U256::from_little_endian(data);
+		let expected = isqrt(x);
+		let got = x.integer_sqrt();
+		assert_eq!(got, expected);
+	}
+});

uint/src/uint.rs

@@ -975,6 +975,28 @@ macro_rules! construct_uint {
 				self.div_mod_knuth(other, n, m)
 			}
 
+			/// Compute the highest `n` such that `n * n <= self`.
+			pub fn integer_sqrt(&self) -> Self {
+				let one = Self::one();
+				if self <= &one {
+					return *self;
+				}
+
+				// the implementation is based on:
+				// https://en.wikipedia.org/wiki/Integer_square_root#Using_only_integer_division
+
+				// Set the initial guess to something higher than √self.
+				let shift: u32 = (self.bits() as u32 + 1) / 2;
+				let mut x_prev = one << shift;
+				loop {
+					let x = (x_prev + self / x_prev) >> 1;
+					if x >= x_prev {
+						return x_prev;
+					}
+					x_prev = x;
+				}
+			}
+
 			/// Fast exponentiation by squaring
 			/// https://en.wikipedia.org/wiki/Exponentiation_by_squaring
 			///

uint/tests/uint_tests.rs

@@ -1140,6 +1140,22 @@ pub mod laws {
 					}
 				}
 
+				quickcheck! {
+					fn isqrt(x: $uint_ty) -> TestResult {
+						let s = x.integer_sqrt();
+						let higher = s + 1;
+						if let Some(y) = higher.checked_mul(higher) {
+							TestResult::from_bool(
+								(s * s <= x) && (y > x)
+							)
+						} else {
+							TestResult::from_bool(
+								s * s <= x
+							)
+						}
+					}
+				}
+
 				quickcheck! {
 					fn pow_mul(x: $uint_ty) -> TestResult {
 						if x.overflowing_pow($uint_ty::from(2)).1 || x.overflowing_pow($uint_ty::from(3)).1 {