Merge pull request #251 from solidstate-network/sqrt-max-fix

`Math#sqrt` fix for max `uint256` input

Author: ItsNickBarry

contracts/utils/Math.sol

@@ -47,16 +47,27 @@ library Math {
 
     /**
      * @notice estimate square root of number
-     * @dev uses Babylonian method (https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Babylonian_method)
-     * @param x input number
-     * @return y square root
+     * @dev uses Heron's method (https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Heron's_method)
+     * @param n input number
+     * @return root square root of input (rounded down to nearest uint256)
      */
-    function sqrt(uint256 x) internal pure returns (uint256 y) {
-        uint256 z = (x + 1) >> 1;
-        y = x;
-        while (z < y) {
-            y = z;
-            z = (x / z + z) >> 1;
+    function sqrt(uint256 n) internal pure returns (uint256 root) {
+        unchecked {
+            // begin with an upper bound, to be updated each time a better estimate is found
+            // for inputs of 0 and 1, this will be returned as-is
+            root = n;
+            // calculate an overestimate
+            // bitwise-or prevents zero division in the case of input of 1
+            uint256 estimate = (n >> 1) | 1;
+
+            // as long as estimate continues to decrease, it is converging on the square root
+
+            while (estimate < root) {
+                // track the new best estimate as the prospective output value
+                root = estimate;
+                // attempt to find a better estimate
+                estimate = (root + n / root) >> 1;
+            }
         }
     }
 }

test/utils/Math.ts

@@ -70,17 +70,42 @@ describe('Math', () => {
     });
 
     describe('#sqrt(uint256)', () => {
-      it('returns the sqrt of a positive integer from 0 to maxUint256', async () => {
-        expect(await instance.sqrt.staticCall(16)).to.eq(4);
+      it('returns the square root of 0', async () => {
+        expect(await instance.sqrt.staticCall(0n)).to.eq(0n);
+      });
+
+      it('returns the square root of 1', async () => {
+        expect(await instance.sqrt.staticCall(1n)).to.eq(1n);
+      });
+
+      it('returns the square root of 2', async () => {
+        expect(await instance.sqrt.staticCall(2n)).to.eq(1n);
+      });
 
-        for (let i = 10; i < 16; i++) {
-          expect(await instance.sqrt.staticCall(i.toString())).to.eq(3);
+      it('returns the square root of positive integers', async () => {
+        for (let i = 2; i < 16; i++) {
+          expect(await instance.sqrt.staticCall(BigInt(i))).to.eq(
+            Math.floor(Math.sqrt(i)),
+          );
         }
+      });
 
-        expect(await instance.sqrt.staticCall(0)).to.eq(0);
+      it('returns the square root of powers of 2', async () => {
+        for (let i = 0; i < 256; i++) {
+          const input = 2n ** BigInt(i);
+          const output = await instance.sqrt.staticCall(input);
+          expect(output ** 2n).to.be.lte(input);
+          expect((output + 1n) ** 2n).to.be.gt(input);
+        }
+      });
 
+      it('returns the square root of max values', async () => {
         expect(await instance.sqrt.staticCall(ethers.MaxUint256 - 1n)).to.eq(
-          BigInt('340282366920938463463374607431768211455'),
+          340282366920938463463374607431768211455n,
+        );
+
+        expect(await instance.sqrt.staticCall(ethers.MaxUint256)).to.eq(
+          340282366920938463463374607431768211455n,
         );
       });
     });