stdlib: Add isqrt (integer square root) to core::num

Returns floor(sqrt(x)) for u256, matching Vyper's isqrt builtin and
OpenZeppelin's Math.sqrt: seed with 1 << (log2(x) / 2) via branchy
bit-length estimate, run exactly seven Newton iterations, then round
down with min(z, x / z). Deterministic gas, no input-dependent loops.

Author: cburgdorf

crates/fe/tests/fixtures/fe_test/isqrt.fe

@@ -0,0 +1,60 @@
+// Tests for `core::num::isqrt` — floor of the integer square root of a u256.
+
+use core::num::isqrt
+
+#[test]
+fn isqrt_small_values() {
+    assert(isqrt(0) == 0)
+    assert(isqrt(1) == 1)
+    assert(isqrt(2) == 1)
+    assert(isqrt(3) == 1)
+    assert(isqrt(4) == 2)
+    assert(isqrt(5) == 2)
+    assert(isqrt(8) == 2)
+    assert(isqrt(9) == 3)
+    assert(isqrt(15) == 3)
+    assert(isqrt(16) == 4)
+}
+
+// For a perfect square k*k, the floor contract pins the results at the
+// square and its immediate neighbors: k*k - 1 -> k - 1, k*k -> k, k*k + 1 -> k.
+#[test]
+fn isqrt_perfect_squares_small() {
+    let k: u256 = 10
+    assert(isqrt(k * k - 1) == k - 1)
+    assert(isqrt(k * k) == k)
+    assert(isqrt(k * k + 1) == k)
+}
+
+#[test]
+fn isqrt_perfect_squares_mid() {
+    let k: u256 = 1000003
+    assert(isqrt(k * k - 1) == k - 1)
+    assert(isqrt(k * k) == k)
+    assert(isqrt(k * k + 1) == k)
+}
+
+#[test]
+fn isqrt_perfect_squares_huge() {
+    // k = 2^127; k*k = 2^254 still fits in a u256.
+    let k: u256 = 170141183460469231731687303715884105728
+    assert(isqrt(k * k - 1) == k - 1)
+    assert(isqrt(k * k) == k)
+    assert(isqrt(k * k + 1) == k)
+}
+
+#[test]
+fn isqrt_u256_max() {
+    // floor(sqrt(2^256 - 1)) == 2^128 - 1
+    assert(isqrt(u256::max()) == 340282366920938463463374607431768211455)
+}
+
+// Arbitrary values cross-checked against precomputed floor square roots.
+#[test]
+fn isqrt_arbitrary_values() {
+    assert(isqrt(1000000000000000000) == 1000000000)
+    assert(isqrt(999999999999999999) == 999999999)
+    assert(isqrt(12345678901234567890) == 3513641828)
+    // 2^255
+    assert(isqrt(57896044618658097711785492504343953926634992332820282019728792003956564819968) == 240615969168004511545033772477625056927)
+}

ingots/core/src/num.fe

@@ -919,6 +919,61 @@ pub const fn sign_extend(word: u256, _ mask: u256, _ sign_bit: u256) -> u256 {
     }
 }
 
+/// Returns floor(sqrt(x)), i.e. the largest `z` such that `z * z <= x`.
+#[arithmetic(unchecked)]
+pub fn isqrt(_ x: u256) -> u256 {
+    if x == 0 {
+        return 0
+    }
+
+    // Seed with `1 << (log2(x) / 2)`, a power of two within a factor of two
+    // of the true root, so a fixed number of Newton iterations suffices.
+    let mut y = x
+    let mut z: u256 = 1
+    if y >= 1 << 128 {
+        y >>= 128
+        z <<= 64
+    }
+    if y >= 1 << 64 {
+        y >>= 64
+        z <<= 32
+    }
+    if y >= 1 << 32 {
+        y >>= 32
+        z <<= 16
+    }
+    if y >= 1 << 16 {
+        y >>= 16
+        z <<= 8
+    }
+    if y >= 1 << 8 {
+        y >>= 8
+        z <<= 4
+    }
+    if y >= 1 << 4 {
+        y >>= 4
+        z <<= 2
+    }
+    if y >= 1 << 2 {
+        z <<= 1
+    }
+
+    // Newton's method doubles the number of correct bits per step, so seven
+    // iterations refine the seed past the 128 bits a root can occupy.
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+    z = (z + x / z) >> 1
+
+    // Newton converges from above and can settle on `floor + 1`; subtract
+    // one exactly when `z` overshoots (`z > x / z`) to round down to the floor.
+    let q = x / z
+    z - ((z > q) as u256)
+}
+
 const fn pow_checked<T>(_ base: own T, _ exp: own T) -> T
     where T: IntWord, T: Mul<Output = T>
 {