add sqrt to math contract

.solhint.json

@@ -12,7 +12,7 @@
     ],
     "code-complexity": [
       "error",
-      7
+      10
     ],
     "compiler-version": [
       "error",

contracts/MathHelpers.sol

@@ -115,4 +115,95 @@ contract MathHelpers {
         }
         return a;
     }
+
+/// @notice Calculates the square root of x using the Babylonian method.
+/// Taken from https://github.com/PaulRBerg/prb-math/blob/main/src/Common.sol
+///
+/// @dev See https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Babylonian_method.
+///
+/// Notes:
+/// - If x is not a perfect square, the result is rounded down.
+/// - Credits to OpenZeppelin for the explanations in comments below.
+///
+/// @param x The uint256 number for which to calculate the square root.
+/// @return result The result as a uint256.
+/// @custom:smtchecker abstract-function-nondet
+    function sqrt(uint256 x) external pure returns (uint256 result) {
+    if (x == 0) {
+        return 0;
+    }
+
+    // For our first guess, we calculate the biggest power of 2 which is smaller than the square root of x.
+    //
+    // We know that the "msb" (most significant bit) of x is a power of 2 such that we have:
+    //
+    // $$
+    // msb(x) <= x <= 2*msb(x)$
+    // $$
+    //
+    // We write $msb(x)$ as $2^k$, and we get:
+    //
+    // $$
+    // k = log_2(x)
+    // $$
+    //
+    // Thus, we can write the initial inequality as:
+    //
+    // $$
+    // 2^{log_2(x)} <= x <= 2*2^{log_2(x)+1} \\
+    // sqrt(2^k) <= sqrt(x) < sqrt(2^{k+1}) \\
+    // 2^{k/2} <= sqrt(x) < 2^{(k+1)/2} <= 2^{(k/2)+1}
+    // $$
+    //
+    // Consequently, $2^{log_2(x) /2} is a good first approximation of sqrt(x) with at least one correct bit.
+    uint256 xAux = uint256(x);
+    result = 1;
+    if (xAux >= 2 ** 128) {
+        xAux >>= 128;
+        result <<= 64;
+    }
+    if (xAux >= 2 ** 64) {
+        xAux >>= 64;
+        result <<= 32;
+    }
+    if (xAux >= 2 ** 32) {
+        xAux >>= 32;
+        result <<= 16;
+    }
+    if (xAux >= 2 ** 16) {
+        xAux >>= 16;
+        result <<= 8;
+    }
+    if (xAux >= 2 ** 8) {
+        xAux >>= 8;
+        result <<= 4;
+    }
+    if (xAux >= 2 ** 4) {
+        xAux >>= 4;
+        result <<= 2;
+    }
+    if (xAux >= 2 ** 2) {
+        result <<= 1;
+    }
+
+    // At this point, `result` is an estimation with at least one bit of precision. We know the true value has at
+    // most 128 bits, since it is the square root of a uint256. Newton's method converges quadratically (precision
+    // doubles at every iteration). We thus need at most 7 iteration to turn our partial result with one bit of
+    // precision into the expected uint128 result.
+    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;
+
+        // If x is not a perfect square, round the result toward zero.
+        uint256 roundedResult = x / result;
+        if (result >= roundedResult) {
+            result = roundedResult;
+        }
+    }
+}
 }

test/helpers/MathHelpers.test.ts

@@ -133,5 +133,12 @@ describe('EnsoHelpers', async () => {
 
       expect(await mathHelpers.conditional(condition, method, a, b)).to.eq(expected);
     });
+
+    it('sqrt', async () => {
+      const a = BigNumber.from(4);
+      const expected = BigNumber.from(ethers.constants.Two);
+
+      expect(await mathHelpers.sqrt(a)).to.eq(expected);
+    });
   });
 });