Update Integer Square Root Function

bridge/contracts/libraries/math/Sigmoid.sol

@@ -45,61 +45,64 @@ abstract contract Sigmoid {
         return b_;
     }
 
-    /// @notice Calculates the square root of x, rounding down.
-    /// @dev Uses the Babylonian method https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Babylonian_method.
-    ///
-    /// Caveats:
-    /// - This function does not work with fixed-point numbers.
-    ///
-    /// @param x The uint256 number for which to calculate the square root.
-    /// @return result The result as an uint256.
-    function _sqrt(uint256 x) internal pure returns (uint256 result) {
-        if (x == 0) {
-            return 0;
-        }
-
-        // Set the initial guess to the closest power of two that is higher than x.
-        uint256 xAux = uint256(x);
-        result = 1;
-        if (xAux >= 0x100000000000000000000000000000000) {
-            xAux >>= 128;
-            result <<= 64;
-        }
-        if (xAux >= 0x10000000000000000) {
-            xAux >>= 64;
-            result <<= 32;
-        }
-        if (xAux >= 0x100000000) {
-            xAux >>= 32;
-            result <<= 16;
-        }
-        if (xAux >= 0x10000) {
-            xAux >>= 16;
-            result <<= 8;
-        }
-        if (xAux >= 0x100) {
-            xAux >>= 8;
-            result <<= 4;
-        }
-        if (xAux >= 0x10) {
-            xAux >>= 4;
-            result <<= 2;
-        }
-        if (xAux >= 0x8) {
-            result <<= 1;
-        }
-
-        // The operations can never overflow because the result is max 2^127 when it enters this block.
+    function _sqrt(uint256 x) internal pure returns (uint256) {
         unchecked {
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1;
-            result = (result + x / result) >> 1; // Seven iterations should be enough
-            uint256 roundedDownResult = x / result;
-            return result >= roundedDownResult ? roundedDownResult : result;
+            if (x <= 1) {
+                return x;
+            }
+            if (x >= ((1 << 128) - 1)**2) {
+                return (1 << 128) - 1;
+            }
+            // Here, e represents the bit length;
+            // its value is at most 256, so it could fit in a uint16.
+            uint256 e = 1;
+            // Here, result is a copy of x to compute the bit length
+            uint256 result = x;
+            if (result >= (1 << 128)) {
+                result >>= 128;
+                e += 128;
+            }
+            if (result >= (1 << 64)) {
+                result >>= 64;
+                e += 64;
+            }
+            if (result >= (1 << 32)) {
+                result >>= 32;
+                e += 32;
+            }
+            if (result >= (1 << 16)) {
+                result >>= 16;
+                e += 16;
+            }
+            if (result >= (1 << 8)) {
+                result >>= 8;
+                e += 8;
+            }
+            if (result >= (1 << 4)) {
+                result >>= 4;
+                e += 4;
+            }
+            if (result >= (1 << 2)) {
+                result >>= 2;
+                e += 2;
+            }
+            if (result >= (1 << 1)) {
+                e += 1;
+            }
+            // e is currently bit length; we overwrite it to scale x
+            e = (256 - e) >> 1;
+            // m now satisfies 2**254 <= m < 2**256
+            uint256 m = x << (2 * e);
+            // result now stores the result
+            result = 1 + (m >> 254);
+            result = (result << 1) + (m >> 251) / result;
+            result = (result << 3) + (m >> 245) / result;
+            result = (result << 7) + (m >> 233) / result;
+            result = (result << 15) + (m >> 209) / result;
+            result = (result << 31) + (m >> 161) / result;
+            result = (result << 63) + (m >> 65) / result;
+            result >>= e;
+            return result * result <= x ? result : (result - 1);
         }
     }
 }

bridge/test/contract-mocks/math/MockSigmoid.sol

@@ -3,6 +3,10 @@ pragma solidity ^0.8.0;
 import "contracts/libraries/math/Sigmoid.sol";
 
 contract MockSigmoid is Sigmoid {
+    function sqrtPublic(uint256 x) public returns (uint256) {
+        return Sigmoid._sqrt(x);
+    }
+
     function p(uint256 x) public pure returns (uint256) {
         return Sigmoid._p(x);
     }

bridge/test/math/sigmoid.test.ts

@@ -117,6 +117,147 @@ describe("Sigmoid unit tests", async () => {
       const retSqrt = await sigmoid.sqrt(x);
       expect(retSqrt).to.be.equal(trueSqrt);
     });
+    it("Integer Square Root 17:    (2**128 - 1)**2", async function () {
+      const x = BigNumber.from(
+        "0xfffffffffffffffffffffffffffffffe00000000000000000000000000000001"
+      );
+      const trueSqrt = BigNumber.from("0xffffffffffffffffffffffffffffffff");
+      const retSqrt = await sigmoid.sqrt(x);
+      expect(retSqrt).to.be.equal(trueSqrt);
+    });
+    it("Integer Square Root 18:    (2**128 - 1)**2 - 1", async function () {
+      const x = BigNumber.from(
+        "0xfffffffffffffffffffffffffffffffe00000000000000000000000000000000"
+      );
+      const trueSqrt = BigNumber.from("0xfffffffffffffffffffffffffffffffe");
+      const retSqrt = await sigmoid.sqrt(x);
+      expect(retSqrt).to.be.equal(trueSqrt);
+    });
+  });
+
+  describe("Integer Square Root Tests: Gas", async () => {
+    it("Integer Square Root 0:  0", async function () {
+      const x = 0;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 1:  1", async function () {
+      const x = 1;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 2:  4", async function () {
+      const x = 4;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 3:  5", async function () {
+      const x = 5;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 4:  10", async function () {
+      const x = 10;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 5:  257", async function () {
+      const x = 257;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 6:     2**15  - 19", async function () {
+      const x = 2 ** 15 - 19;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 7:   3*2**16  + 5", async function () {
+      const x = 3 * 2 ** 16 + 5;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 8:   5*2**31  - 27", async function () {
+      const x = 5 * 2 ** 31 - 27;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 9:   7*2**32  + 9", async function () {
+      const x = 7 * 2 ** 32 + 9;
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 10: 11*2**63  - 9", async function () {
+      const x = BigNumber.from("0x57ffffffffffffff7");
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 11: 13*2**64  + 43", async function () {
+      const x = BigNumber.from("0xd000000000000002b");
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 12: 17*2**127 - 23", async function () {
+      const x = BigNumber.from("0x87fffffffffffffffffffffffffffffe9");
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 13: 19*2**128 + 109", async function () {
+      const x = BigNumber.from("0x130000000000000000000000000000006d");
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 14:    2**130 - 5", async function () {
+      const x = BigNumber.from("0x3fffffffffffffffffffffffffffffffb");
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 15:    2**255 - 19", async function () {
+      const x = BigNumber.from(
+        "0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffed"
+      );
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 16:    2**256 - 1", async function () {
+      const x = BigNumber.from(
+        "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
+      );
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 17:    (2**128 - 1**2", async function () {
+      const x = BigNumber.from(
+        "0xfffffffffffffffffffffffffffffffe00000000000000000000000000000001"
+      );
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
+    it("Integer Square Root 18:    (2**128 - 1**2 - 1", async function () {
+      const x = BigNumber.from(
+        "0xfffffffffffffffffffffffffffffffe00000000000000000000000000000000"
+      );
+      const tx = await sigmoid.sqrtPublic(x);
+      const receipt = await tx.wait();
+      expect(receipt.status).to.eq(1);
+    });
   });
 
   describe("safeAbsSub Tests", async () => {