Add sqrt for math

CHANGELOG.md

@@ -17,6 +17,7 @@
  * `ERC721`, `ERC1155`: simplified revert reasons. ([#3254](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3254), ([#3438](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3438)))
  * `ERC721`: removed redundant require statement. ([#3434](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3434))
  * `PaymentSplitter`: add `releasable` getters. ([#3350](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3350))
+ * `Math`: Add a `sqrt` function. ([#3242](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3242))
 
 ## 4.6.0 (2022-04-26)
 

contracts/mocks/MathMock.sol

@@ -29,4 +29,8 @@ contract MathMock {
     ) public pure returns (uint256) {
         return Math.mulDiv(a, b, denominator, direction);
     }
+
+    function sqrt(uint256 a) public pure returns (uint256) {
+        return Math.sqrt(a);
+    }
 }

contracts/utils/math/Math.sol

@@ -149,4 +149,62 @@ library Math {
         }
         return result;
     }
+
+    /**
+     * @dev Returns the square root of a number.
+     *
+     * Inspired by Henry S. Warren, Jr.'s "Hacker's Delight" (Chapter 11).
+     */
+    function sqrt(uint256 a) internal pure returns (uint256) {
+        if (a == 0) {
+            return 0;
+        }
+
+        // For our first guess, we use the log2 of the square root. We do that by shifting `a` and only counting half
+        // of the bits. the partial result produced is a power of two that verifies `result <= sqrt(a) < 2 * result`
+        uint256 result = 1;
+        uint256 x = a;
+        if (x > 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF) {
+            x >>= 128;
+            result <<= 64;
+        }
+        if (x > 0xFFFFFFFFFFFFFFFF) {
+            x >>= 64;
+            result <<= 32;
+        }
+        if (x > 0xFFFFFFFF) {
+            x >>= 32;
+            result <<= 16;
+        }
+        if (x > 0xFFFF) {
+            x >>= 16;
+            result <<= 8;
+        }
+        if (x > 0xFF) {
+            x >>= 8;
+            result <<= 4;
+        }
+        if (x > 0xF) {
+            x >>= 4;
+            result <<= 2;
+        }
+        if (x > 0x3) {
+            result <<= 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);
+        }
+    }
 }

test/utils/math/Math.test.js

@@ -182,4 +182,20 @@ contract('Math', function (accounts) {
       });
     });
   });
+
+  describe('sqrt', function () {
+    it('rounds down', async function () {
+      expect(await this.math.sqrt(new BN('0'))).to.be.bignumber.equal('0');
+      expect(await this.math.sqrt(new BN('1'))).to.be.bignumber.equal('1');
+      expect(await this.math.sqrt(new BN('2'))).to.be.bignumber.equal('1');
+      expect(await this.math.sqrt(new BN('3'))).to.be.bignumber.equal('1');
+      expect(await this.math.sqrt(new BN('4'))).to.be.bignumber.equal('2');
+      expect(await this.math.sqrt(new BN('144'))).to.be.bignumber.equal('12');
+      expect(await this.math.sqrt(new BN('1000000'))).to.be.bignumber.equal('1000');
+      expect(await this.math.sqrt(new BN('1000001'))).to.be.bignumber.equal('1000');
+      expect(await this.math.sqrt(new BN('1002000'))).to.be.bignumber.equal('1000');
+      expect(await this.math.sqrt(new BN('1002001'))).to.be.bignumber.equal('1001');
+      expect(await this.math.sqrt(MAX_UINT256)).to.be.bignumber.equal('340282366920938463463374607431768211455');
+    });
+  });
 });