Improved integer square root.

contracts/utils/math/Math.sol

@@ -339,38 +339,69 @@ library Math {
      * @dev Returns the square root of a number. If the number is not a perfect square, the value is rounded
      * towards zero.
      *
-     * Inspired by Henry S. Warren, Jr.'s "Hacker's Delight" (Chapter 11).
+     * This method is based on Newton's method for computing square roots; the algorithm is restricted to only
+     * using integer operations.
      */
     function sqrt(uint256 a) internal pure returns (uint256) {
-        if (a == 0) {
-            return 0;
-        }
-
-        // For our first guess, we get the biggest power of 2 which is smaller than the square root of the target.
-        //
-        // We know that the "msb" (most significant bit) of our target number `a` is a power of 2 such that we have
-        // `msb(a) <= a < 2*msb(a)`. This value can be written `msb(a)=2**k` with `k=log2(a)`.
-        //
-        // This can be rewritten `2**log2(a) <= a < 2**(log2(a) + 1)`
-        // → `sqrt(2**k) <= sqrt(a) < sqrt(2**(k+1))`
-        // → `2**(k/2) <= sqrt(a) < 2**((k+1)/2) <= 2**(k/2 + 1)`
-        //
-        // Consequently, `2**(log2(a) / 2)` is a good first approximation of `sqrt(a)` with at least 1 correct bit.
-        uint256 result = 1 << (log2(a) >> 1);
-
-        // At this point `result` is an estimation with one bit of precision. We know the true value is a uint128,
-        // 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 + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            result = (result + a / result) >> 1;
-            return min(result, a / result);
+            // Take care of easy edge cases when a == 0 or a == 1
+            if (a <= 1) {
+                return a;
+            }
+
+            uint256 aAux = a;
+            uint256 result = 1;
+
+            // For the first guess of `result` (e), we get the biggest power of 2 which is smaller than sqrt(a)
+            // (i.e. 2^e <= sqrt(a)). We know that e is at most 127 given (2^128)^2 overflows an uint256.
+            // Thus, we approximate e by iterating 2^{i/2} where i starts at 128, and applying the exponent
+            // to e if the result is still smaller than a (up to e == 127).
+            if (aAux >= (1 << 128)) {
+                aAux >>= 128;
+                result <<= 64;
+            }
+            if (aAux >= (1 << 64)) {
+                aAux >>= 64;
+                result <<= 32;
+            }
+            if (aAux >= (1 << 32)) {
+                aAux >>= 32;
+                result <<= 16;
+            }
+            if (aAux >= (1 << 16)) {
+                aAux >>= 16;
+                result <<= 8;
+            }
+            if (aAux >= (1 << 8)) {
+                aAux >>= 8;
+                result <<= 4;
+            }
+            if (aAux >= (1 << 4)) {
+                aAux >>= 4;
+                result <<= 2;
+            }
+            if (aAux >= (1 << 2)) {
+                result <<= 1;
+            }
+
+            // We can use the fact that 2^e <= sqrt(a) to improve the estimation
+            // by computing the arithmetic mean between the current estimation and
+            // the next one (result * 2), ensuring that result - sqrt(a) <= 2^{e-2}.
+            result = (3 * result) >> 1;
+
+            // Each Newton iteration will have f(x) = (x + a / x) / 2.
+            // Given the error (ε) is defined by x - sqrt(a), then we know that
+            // ε + 1 == ε^2 / 2x <= ε^2 / 2 * sqrt(a).
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-4.5}
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-9}
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-18}
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-36}
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-72}
+            result = (result + a / result) >> 1; // err := result - sqrt(a) <= 2^{e-144}
+
+            // After 6 iterations, no more precision can be obtained since the max result is 127.
+            // result is either sqrt(a) or sqrt(a) + 1.
+            return result - SafeCast.toUint(result > a / result);
         }
     }