Add 512bits add and mult operations (#5035)

.changeset/blue-nails-give.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`Math`: Add `add512`, `mul512` and `mulShr`.

contracts/utils/math/Math.sol

@@ -17,6 +17,34 @@ library Math {
         Expand // Away from zero
     }
 
+    /**
+     * @dev Return the 512-bit addition of two uint256.
+     *
+     * The result is stored in two 256 variables such that sum = high * 2²⁵⁶ + low.
+     */
+    function add512(uint256 a, uint256 b) internal pure returns (uint256 high, uint256 low) {
+        assembly ("memory-safe") {
+            low := add(a, b)
+            high := lt(low, a)
+        }
+    }
+
+    /**
+     * @dev Return the 512-bit multiplication of two uint256.
+     *
+     * The result is stored in two 256 variables such that product = high * 2²⁵⁶ + low.
+     */
+    function mul512(uint256 a, uint256 b) internal pure returns (uint256 high, uint256 low) {
+        // 512-bit multiply [high low] = x * y. Compute the product mod 2²⁵⁶ and mod 2²⁵⁶ - 1, then use
+        // the Chinese Remainder Theorem to reconstruct the 512 bit result. The result is stored in two 256
+        // variables such that product = high * 2²⁵⁶ + low.
+        assembly ("memory-safe") {
+            let mm := mulmod(a, b, not(0))
+            low := mul(a, b)
+            high := sub(sub(mm, low), lt(mm, low))
+        }
+    }
+
     /**
      * @dev Returns the addition of two unsigned integers, with an success flag (no overflow).
      */
@@ -143,26 +171,18 @@ library Math {
      */
     function mulDiv(uint256 x, uint256 y, uint256 denominator) internal pure returns (uint256 result) {
         unchecked {
-            // 512-bit multiply [prod1 prod0] = x * y. Compute the product mod 2²⁵⁶ and mod 2²⁵⁶ - 1, then use
-            // the Chinese Remainder Theorem to reconstruct the 512 bit result. The result is stored in two 256
-            // variables such that product = prod1 * 2²⁵⁶ + prod0.
-            uint256 prod0 = x * y; // Least significant 256 bits of the product
-            uint256 prod1; // Most significant 256 bits of the product
-            assembly {
-                let mm := mulmod(x, y, not(0))
-                prod1 := sub(sub(mm, prod0), lt(mm, prod0))
-            }
+            (uint256 high, uint256 low) = mul512(x, y);
 
             // Handle non-overflow cases, 256 by 256 division.
-            if (prod1 == 0) {
+            if (high == 0) {
                 // Solidity will revert if denominator == 0, unlike the div opcode on its own.
                 // The surrounding unchecked block does not change this fact.
                 // See https://docs.soliditylang.org/en/latest/control-structures.html#checked-or-unchecked-arithmetic.
-                return prod0 / denominator;
+                return low / denominator;
             }
 
             // Make sure the result is less than 2²⁵⁶. Also prevents denominator == 0.
-            if (denominator <= prod1) {
+            if (denominator <= high) {
                 Panic.panic(ternary(denominator == 0, Panic.DIVISION_BY_ZERO, Panic.UNDER_OVERFLOW));
             }
 
@@ -170,15 +190,15 @@ library Math {
             // 512 by 256 division.
             ///////////////////////////////////////////////
 
-            // Make division exact by subtracting the remainder from [prod1 prod0].
+            // Make division exact by subtracting the remainder from [high low].
             uint256 remainder;
             assembly {
                 // Compute remainder using mulmod.
                 remainder := mulmod(x, y, denominator)
 
                 // Subtract 256 bit number from 512 bit number.
-                prod1 := sub(prod1, gt(remainder, prod0))
-                prod0 := sub(prod0, remainder)
+                high := sub(high, gt(remainder, low))
+                low := sub(low, remainder)
             }
 
             // Factor powers of two out of denominator and compute largest power of two divisor of denominator.
@@ -189,15 +209,15 @@ library Math {
                 // Divide denominator by twos.
                 denominator := div(denominator, twos)
 
-                // Divide [prod1 prod0] by twos.
-                prod0 := div(prod0, twos)
+                // Divide [high low] by twos.
+                low := div(low, twos)
 
                 // Flip twos such that it is 2²⁵⁶ / twos. If twos is zero, then it becomes one.
                 twos := add(div(sub(0, twos), twos), 1)
             }
 
-            // Shift in bits from prod1 into prod0.
-            prod0 |= prod1 * twos;
+            // Shift in bits from high into low.
+            low |= high * twos;
 
             // Invert denominator mod 2²⁵⁶. Now that denominator is an odd number, it has an inverse modulo 2²⁵⁶ such
             // that denominator * inv ≡ 1 mod 2²⁵⁶. Compute the inverse by starting with a seed that is correct for
@@ -215,9 +235,9 @@ library Math {
 
             // Because the division is now exact we can divide by multiplying with the modular inverse of denominator.
             // This will give us the correct result modulo 2²⁵⁶. Since the preconditions guarantee that the outcome is
-            // less than 2²⁵⁶, this is the final result. We don't need to compute the high bits of the result and prod1
+            // less than 2²⁵⁶, this is the final result. We don't need to compute the high bits of the result and high
             // is no longer required.
-            result = prod0 * inverse;
+            result = low * inverse;
             return result;
         }
     }
@@ -229,6 +249,26 @@ library Math {
         return mulDiv(x, y, denominator) + SafeCast.toUint(unsignedRoundsUp(rounding) && mulmod(x, y, denominator) > 0);
     }
 
+    /**
+     * @dev Calculates floor(x * y >> n) with full precision. Throws if result overflows a uint256.
+     */
+    function mulShr(uint256 x, uint256 y, uint8 n) internal pure returns (uint256 result) {
+        unchecked {
+            (uint256 high, uint256 low) = mul512(x, y);
+            if (high >= 1 << n) {
+                Panic.panic(Panic.UNDER_OVERFLOW);
+            }
+            return (high << (256 - n)) | (low >> n);
+        }
+    }
+
+    /**
+     * @dev Calculates x * y >> n with full precision, following the selected rounding direction.
+     */
+    function mulShr(uint256 x, uint256 y, uint8 n, Rounding rounding) internal pure returns (uint256) {
+        return mulShr(x, y, n) + SafeCast.toUint(unsignedRoundsUp(rounding) && mulmod(x, y, 1 << n) > 0);
+    }
+
     /**
      * @dev Calculate the modular multiplicative inverse of a number in Z/nZ.
      *

test/helpers/enums.js

@@ -1,12 +1,14 @@
-function Enum(...options) {
-  return Object.fromEntries(options.map((key, i) => [key, BigInt(i)]));
-}
+const { ethers } = require('ethers');
+
+const Enum = (...options) => Object.fromEntries(options.map((key, i) => [key, BigInt(i)]));
+const EnumTyped = (...options) => Object.fromEntries(options.map((key, i) => [key, ethers.Typed.uint8(i)]));
 
 module.exports = {
   Enum,
+  EnumTyped,
   ProposalState: Enum('Pending', 'Active', 'Canceled', 'Defeated', 'Succeeded', 'Queued', 'Expired', 'Executed'),
   VoteType: Object.assign(Enum('Against', 'For', 'Abstain'), { Parameters: 255n }),
-  Rounding: Enum('Floor', 'Ceil', 'Trunc', 'Expand'),
+  Rounding: EnumTyped('Floor', 'Ceil', 'Trunc', 'Expand'),
   OperationState: Enum('Unset', 'Waiting', 'Ready', 'Done'),
-  RevertType: Enum('None', 'RevertWithoutMessage', 'RevertWithMessage', 'RevertWithCustomError', 'Panic'),
+  RevertType: EnumTyped('None', 'RevertWithoutMessage', 'RevertWithMessage', 'RevertWithCustomError', 'Panic'),
 };

test/utils/math/Math.t.sol

@@ -11,6 +11,48 @@ contract MathTest is Test {
         assertEq(Math.ternary(f, a, b), f ? a : b);
     }
 
+    // ADD512 & MUL512
+    function testAdd512(uint256 a, uint256 b) public pure {
+        (uint256 high, uint256 low) = Math.add512(a, b);
+
+        // test against tryAdd
+        (bool success, uint256 result) = Math.tryAdd(a, b);
+        if (success) {
+            assertEq(high, 0);
+            assertEq(low, result);
+        } else {
+            assertEq(high, 1);
+        }
+
+        // test against unchecked
+        unchecked {
+            assertEq(low, a + b); // unchecked allow overflow
+        }
+    }
+
+    function testMul512(uint256 a, uint256 b) public pure {
+        (uint256 high, uint256 low) = Math.mul512(a, b);
+
+        // test against tryMul
+        (bool success, uint256 result) = Math.tryMul(a, b);
+        if (success) {
+            assertEq(high, 0);
+            assertEq(low, result);
+        } else {
+            assertGt(high, 0);
+        }
+
+        // test against unchecked
+        unchecked {
+            assertEq(low, a * b); // unchecked allow overflow
+        }
+
+        // test against alternative method
+        (uint256 _high, uint256 _low) = _mulKaratsuba(a, b);
+        assertEq(high, _high);
+        assertEq(low, _low);
+    }
+
     // MIN & MAX
     function testSymbolicMinMax(uint256 a, uint256 b) public pure {
         assertEq(Math.min(a, b), a < b ? a : b);
@@ -184,7 +226,7 @@ contract MathTest is Test {
     // MULDIV
     function testMulDiv(uint256 x, uint256 y, uint256 d) public pure {
         // Full precision for x * y
-        (uint256 xyHi, uint256 xyLo) = _mulHighLow(x, y);
+        (uint256 xyHi, uint256 xyLo) = Math.mul512(x, y);
 
         // Assume result won't overflow (see {testMulDivDomain})
         // This also checks that `d` is positive
@@ -194,9 +236,9 @@ contract MathTest is Test {
         uint256 q = Math.mulDiv(x, y, d);
 
         // Full precision for q * d
-        (uint256 qdHi, uint256 qdLo) = _mulHighLow(q, d);
+        (uint256 qdHi, uint256 qdLo) = Math.mul512(q, d);
         // Add remainder of x * y / d (computed as rem = (x * y % d))
-        (uint256 qdRemLo, uint256 c) = _addCarry(qdLo, mulmod(x, y, d));
+        (uint256 c, uint256 qdRemLo) = Math.add512(qdLo, mulmod(x, y, d));
         uint256 qdRemHi = qdHi + c;
 
         // Full precision check that x * y = q * d + rem
@@ -206,7 +248,7 @@ contract MathTest is Test {
 
     /// forge-config: default.allow_internal_expect_revert = true
     function testMulDivDomain(uint256 x, uint256 y, uint256 d) public {
-        (uint256 xyHi, ) = _mulHighLow(x, y);
+        (uint256 xyHi, ) = Math.mul512(x, y);
 
         // Violate {testMulDiv} assumption (covers d is 0 and result overflow)
         vm.assume(xyHi >= d);
@@ -266,26 +308,13 @@ contract MathTest is Test {
         }
     }
 
-    function _nativeModExp(uint256 b, uint256 e, uint256 m) private pure returns (uint256) {
-        if (m == 1) return 0;
-        uint256 r = 1;
-        while (e > 0) {
-            if (e % 2 > 0) {
-                r = mulmod(r, b, m);
-            }
-            b = mulmod(b, b, m);
-            e >>= 1;
-        }
-        return r;
-    }
-
     // Helpers
     function _asRounding(uint8 r) private pure returns (Math.Rounding) {
         vm.assume(r < uint8(type(Math.Rounding).max));
         return Math.Rounding(r);
     }
 
-    function _mulHighLow(uint256 x, uint256 y) private pure returns (uint256 high, uint256 low) {
+    function _mulKaratsuba(uint256 x, uint256 y) private pure returns (uint256 high, uint256 low) {
         (uint256 x0, uint256 x1) = (x & type(uint128).max, x >> 128);
         (uint256 y0, uint256 y1) = (y & type(uint128).max, y >> 128);
 
@@ -305,10 +334,16 @@ contract MathTest is Test {
         }
     }
 
-    function _addCarry(uint256 x, uint256 y) private pure returns (uint256 res, uint256 carry) {
-        unchecked {
-            res = x + y;
+    function _nativeModExp(uint256 b, uint256 e, uint256 m) private pure returns (uint256) {
+        if (m == 1) return 0;
+        uint256 r = 1;
+        while (e > 0) {
+            if (e % 2 > 0) {
+                r = mulmod(r, b, m);
+            }
+            b = mulmod(b, b, m);
+            e >>= 1;
         }
-        carry = res < x ? 1 : 0;
+        return r;
     }
 }

test/utils/math/Math.test.js

@@ -16,10 +16,13 @@ const uint256 = value => ethers.Typed.uint256(value);
 bytes.zero = '0x';
 uint256.zero = 0n;
 
-async function testCommutative(fn, lhs, rhs, expected, ...extra) {
-  expect(await fn(lhs, rhs, ...extra)).to.deep.equal(expected);
-  expect(await fn(rhs, lhs, ...extra)).to.deep.equal(expected);
-}
+const testCommutative = (fn, lhs, rhs, expected, ...extra) =>
+  Promise.all([
+    expect(fn(lhs, rhs, ...extra)).to.eventually.deep.equal(expected),
+    expect(fn(rhs, lhs, ...extra)).to.eventually.deep.equal(expected),
+  ]);
+
+const splitHighLow = n => [n / (1n << 256n), n % (1n << 256n)];
 
 async function fixture() {
   const mock = await ethers.deployContract('$Math');
@@ -39,6 +42,24 @@ describe('Math', function () {
     Object.assign(this, await loadFixture(fixture));
   });
 
+  describe('add512', function () {
+    it('adds correctly without reverting', async function () {
+      const values = [0n, 1n, 17n, 42n, ethers.MaxUint256 - 1n, ethers.MaxUint256];
+      for (const [a, b] of product(values, values)) {
+        await expect(this.mock.$add512(a, b)).to.eventually.deep.equal(splitHighLow(a + b));
+      }
+    });
+  });
+
+  describe('mul512', function () {
+    it('multiplies correctly without reverting', async function () {
+      const values = [0n, 1n, 17n, 42n, ethers.MaxUint256 - 1n, ethers.MaxUint256];
+      for (const [a, b] of product(values, values)) {
+        await expect(this.mock.$mul512(a, b)).to.eventually.deep.equal(splitHighLow(a * b));
+      }
+    });
+  });
+
   describe('tryAdd', function () {
     it('adds correctly', async function () {
       const a = 5678n;
@@ -57,13 +78,13 @@ describe('Math', function () {
     it('subtracts correctly', async function () {
       const a = 5678n;
       const b = 1234n;
-      expect(await this.mock.$trySub(a, b)).to.deep.equal([true, a - b]);
+      await expect(this.mock.$trySub(a, b)).to.eventually.deep.equal([true, a - b]);
     });
 
     it('reverts if subtraction result would be negative', async function () {
       const a = 1234n;
       const b = 5678n;
-      expect(await this.mock.$trySub(a, b)).to.deep.equal([false, 0n]);
+      await expect(this.mock.$trySub(a, b)).to.eventually.deep.equal([false, 0n]);
     });
   });
 
@@ -91,25 +112,25 @@ describe('Math', function () {
     it('divides correctly', async function () {
       const a = 5678n;
       const b = 5678n;
-      expect(await this.mock.$tryDiv(a, b)).to.deep.equal([true, a / b]);
+      await expect(this.mock.$tryDiv(a, b)).to.eventually.deep.equal([true, a / b]);
     });
 
     it('divides zero correctly', async function () {
       const a = 0n;
       const b = 5678n;
-      expect(await this.mock.$tryDiv(a, b)).to.deep.equal([true, a / b]);
+      await expect(this.mock.$tryDiv(a, b)).to.eventually.deep.equal([true, a / b]);
     });
 
     it('returns complete number result on non-even division', async function () {
       const a = 7000n;
       const b = 5678n;
-      expect(await this.mock.$tryDiv(a, b)).to.deep.equal([true, a / b]);
+      await expect(this.mock.$tryDiv(a, b)).to.eventually.deep.equal([true, a / b]);
     });
 
     it('reverts on division by zero', async function () {
       const a = 5678n;
       const b = 0n;
-      expect(await this.mock.$tryDiv(a, b)).to.deep.equal([false, 0n]);
+      await expect(this.mock.$tryDiv(a, b)).to.eventually.deep.equal([false, 0n]);
     });
   });
 
@@ -118,32 +139,32 @@ describe('Math', function () {
       it('when the dividend is smaller than the divisor', async function () {
         const a = 284n;
         const b = 5678n;
-        expect(await this.mock.$tryMod(a, b)).to.deep.equal([true, a % b]);
+        await expect(this.mock.$tryMod(a, b)).to.eventually.deep.equal([true, a % b]);
       });
 
       it('when the dividend is equal to the divisor', async function () {
         const a = 5678n;
         const b = 5678n;
-        expect(await this.mock.$tryMod(a, b)).to.deep.equal([true, a % b]);
+        await expect(this.mock.$tryMod(a, b)).to.eventually.deep.equal([true, a % b]);
       });
 
       it('when the dividend is larger than the divisor', async function () {
         const a = 7000n;
         const b = 5678n;
-        expect(await this.mock.$tryMod(a, b)).to.deep.equal([true, a % b]);
+        await expect(this.mock.$tryMod(a, b)).to.eventually.deep.equal([true, a % b]);
       });
 
       it('when the dividend is a multiple of the divisor', async function () {
         const a = 17034n; // 17034 == 5678 * 3
         const b = 5678n;
-        expect(await this.mock.$tryMod(a, b)).to.deep.equal([true, a % b]);
+        await expect(this.mock.$tryMod(a, b)).to.eventually.deep.equal([true, a % b]);
       });
     });
 
     it('reverts with a 0 divisor', async function () {
       const a = 5678n;
       const b = 0n;
-      expect(await this.mock.$tryMod(a, b)).to.deep.equal([false, 0n]);
+      await expect(this.mock.$tryMod(a, b)).to.eventually.deep.equal([false, 0n]);
     });
   });
 
@@ -163,24 +184,24 @@ describe('Math', function () {
     it('is correctly calculated with two odd numbers', async function () {
       const a = 57417n;
       const b = 95431n;
-      expect(await this.mock.$average(a, b)).to.equal((a + b) / 2n);
+      await expect(this.mock.$average(a, b)).to.eventually.equal((a + b) / 2n);
     });
 
     it('is correctly calculated with two even numbers', async function () {
       const a = 42304n;
       const b = 84346n;
-      expect(await this.mock.$average(a, b)).to.equal((a + b) / 2n);
+      await expect(this.mock.$average(a, b)).to.eventually.equal((a + b) / 2n);
     });
 
     it('is correctly calculated with one even and one odd number', async function () {
       const a = 57417n;
       const b = 84346n;
-      expect(await this.mock.$average(a, b)).to.equal((a + b) / 2n);
+      await expect(this.mock.$average(a, b)).to.eventually.equal((a + b) / 2n);
     });
 
     it('is correctly calculated with two max uint256 numbers', async function () {
       const a = ethers.MaxUint256;
-      expect(await this.mock.$average(a, a)).to.equal(a);
+      await expect(this.mock.$average(a, a)).to.eventually.equal(a);
     });
   });
 
@@ -196,35 +217,35 @@ describe('Math', function () {
       const a = 0n;
       const b = 2n;
       const r = 0n;
-      expect(await this.mock.$ceilDiv(a, b)).to.equal(r);
+      await expect(this.mock.$ceilDiv(a, b)).to.eventually.equal(r);
     });
 
     it('does not round up on exact division', async function () {
       const a = 10n;
       const b = 5n;
       const r = 2n;
-      expect(await this.mock.$ceilDiv(a, b)).to.equal(r);
+      await expect(this.mock.$ceilDiv(a, b)).to.eventually.equal(r);
     });
 
     it('rounds up on division with remainders', async function () {
       const a = 42n;
       const b = 13n;
       const r = 4n;
-      expect(await this.mock.$ceilDiv(a, b)).to.equal(r);
+      await expect(this.mock.$ceilDiv(a, b)).to.eventually.equal(r);
     });
 
     it('does not overflow', async function () {
       const a = ethers.MaxUint256;
       const b = 2n;
       const r = 1n << 255n;
-      expect(await this.mock.$ceilDiv(a, b)).to.equal(r);
+      await expect(this.mock.$ceilDiv(a, b)).to.eventually.equal(r);
     });
 
     it('correctly computes max uint256 divided by 1', async function () {
       const a = ethers.MaxUint256;
       const b = 1n;
       const r = ethers.MaxUint256;
-      expect(await this.mock.$ceilDiv(a, b)).to.equal(r);
+      await expect(this.mock.$ceilDiv(a, b)).to.eventually.equal(r);
     });
   });
 
@@ -248,28 +269,30 @@ describe('Math', function () {
     describe('does round down', function () {
       it('small values', async function () {
         for (const rounding of RoundingDown) {
-          expect(await this.mock.$mulDiv(3n, 4n, 5n, rounding)).to.equal(2n);
-          expect(await this.mock.$mulDiv(3n, 5n, 5n, rounding)).to.equal(3n);
+          await expect(this.mock.$mulDiv(3n, 4n, 5n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$mulDiv(3n, 5n, 5n, rounding)).to.eventually.equal(3n);
         }
       });
 
       it('large values', async function () {
         for (const rounding of RoundingDown) {
-          expect(await this.mock.$mulDiv(42n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding)).to.equal(41n);
+          await expect(this.mock.$mulDiv(42n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding)).to.eventually.equal(
+            41n,
+          );
 
-          expect(await this.mock.$mulDiv(17n, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.equal(17n);
+          await expect(this.mock.$mulDiv(17n, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.eventually.equal(17n);
 
-          expect(
-            await this.mock.$mulDiv(ethers.MaxUint256 - 1n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
-          ).to.equal(ethers.MaxUint256 - 2n);
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256 - 1n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256 - 2n);
 
-          expect(
-            await this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
-          ).to.equal(ethers.MaxUint256 - 1n);
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256 - 1n);
 
-          expect(await this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.equal(
-            ethers.MaxUint256,
-          );
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256);
         }
       });
     });
@@ -277,28 +300,91 @@ describe('Math', function () {
     describe('does round up', function () {
       it('small values', async function () {
         for (const rounding of RoundingUp) {
-          expect(await this.mock.$mulDiv(3n, 4n, 5n, rounding)).to.equal(3n);
-          expect(await this.mock.$mulDiv(3n, 5n, 5n, rounding)).to.equal(3n);
+          await expect(this.mock.$mulDiv(3n, 4n, 5n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$mulDiv(3n, 5n, 5n, rounding)).to.eventually.equal(3n);
         }
       });
 
       it('large values', async function () {
         for (const rounding of RoundingUp) {
-          expect(await this.mock.$mulDiv(42n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding)).to.equal(42n);
+          await expect(this.mock.$mulDiv(42n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding)).to.eventually.equal(
+            42n,
+          );
+
+          await expect(this.mock.$mulDiv(17n, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.eventually.equal(17n);
+
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256 - 1n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256 - 1n);
+
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256 - 1n);
+
+          await expect(
+            this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256, ethers.MaxUint256, rounding),
+          ).to.eventually.equal(ethers.MaxUint256);
+        }
+      });
+    });
+  });
+
+  describe('mulShr', function () {
+    it('reverts with result higher than 2 ^ 256', async function () {
+      const a = 5n;
+      const b = ethers.MaxUint256;
+      const c = 1n;
+      await expect(this.mock.$mulShr(a, b, c, Rounding.Floor)).to.be.revertedWithPanic(
+        PANIC_CODES.ARITHMETIC_UNDER_OR_OVERFLOW,
+      );
+    });
+
+    describe('does round down', function () {
+      it('small values', async function () {
+        for (const rounding of RoundingDown) {
+          await expect(this.mock.$mulShr(3n, 5n, 1n, rounding)).to.eventually.equal(7n);
+          await expect(this.mock.$mulShr(3n, 5n, 2n, rounding)).to.eventually.equal(3n);
+        }
+      });
+
+      it('large values', async function () {
+        for (const rounding of RoundingDown) {
+          await expect(this.mock.$mulShr(42n, ethers.MaxUint256, 255n, rounding)).to.eventually.equal(83n);
+
+          await expect(this.mock.$mulShr(17n, ethers.MaxUint256, 255n, rounding)).to.eventually.equal(33n);
+
+          await expect(this.mock.$mulShr(ethers.MaxUint256, ethers.MaxInt256 + 1n, 255n, rounding)).to.eventually.equal(
+            ethers.MaxUint256,
+          );
+
+          await expect(this.mock.$mulShr(ethers.MaxUint256, ethers.MaxInt256, 255n, rounding)).to.eventually.equal(
+            ethers.MaxUint256 - 2n,
+          );
+        }
+      });
+    });
 
-          expect(await this.mock.$mulDiv(17n, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.equal(17n);
+    describe('does round up', function () {
+      it('small values', async function () {
+        for (const rounding of RoundingUp) {
+          await expect(this.mock.$mulShr(3n, 5n, 1n, rounding)).to.eventually.equal(8n);
+          await expect(this.mock.$mulShr(3n, 5n, 2n, rounding)).to.eventually.equal(4n);
+        }
+      });
 
-          expect(
-            await this.mock.$mulDiv(ethers.MaxUint256 - 1n, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
-          ).to.equal(ethers.MaxUint256 - 1n);
+      it('large values', async function () {
+        for (const rounding of RoundingUp) {
+          await expect(this.mock.$mulShr(42n, ethers.MaxUint256, 255n, rounding)).to.eventually.equal(84n);
 
-          expect(
-            await this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256 - 1n, ethers.MaxUint256, rounding),
-          ).to.equal(ethers.MaxUint256 - 1n);
+          await expect(this.mock.$mulShr(17n, ethers.MaxUint256, 255n, rounding)).to.eventually.equal(34n);
 
-          expect(await this.mock.$mulDiv(ethers.MaxUint256, ethers.MaxUint256, ethers.MaxUint256, rounding)).to.equal(
+          await expect(this.mock.$mulShr(ethers.MaxUint256, ethers.MaxInt256 + 1n, 255n, rounding)).to.eventually.equal(
             ethers.MaxUint256,
           );
+
+          await expect(this.mock.$mulShr(ethers.MaxUint256, ethers.MaxInt256, 255n, rounding)).to.eventually.equal(
+            ethers.MaxUint256 - 1n,
+          );
         }
       });
     });
@@ -320,8 +406,8 @@ describe('Math', function () {
 
       describe(`using p=${p} which is ${p > 1 && factors.length > 1 ? 'not ' : ''}a prime`, function () {
         it('trying to inverse 0 returns 0', async function () {
-          expect(await this.mock.$invMod(0, p)).to.equal(0n);
-          expect(await this.mock.$invMod(p, p)).to.equal(0n); // p is 0 mod p
+          await expect(this.mock.$invMod(0, p)).to.eventually.equal(0n);
+          await expect(this.mock.$invMod(p, p)).to.eventually.equal(0n); // p is 0 mod p
         });
 
         if (p != 0) {
@@ -349,7 +435,7 @@ describe('Math', function () {
           const e = 200n;
           const m = 50n;
 
-          expect(await this.mock.$modExp(type(b), type(e), type(m))).to.equal(type(b ** e % m).value);
+          await expect(this.mock.$modExp(type(b), type(e), type(m))).to.eventually.equal(type(b ** e % m).value);
         });
 
         it('is correctly reverting when modulus is zero', async function () {
@@ -373,7 +459,9 @@ describe('Math', function () {
         it(`calculates b ** e % m (b=2**${log2b}+1) (e=2**${log2e}+1) (m=2**${log2m}+1)`, async function () {
           const mLength = ethers.dataLength(ethers.toBeHex(m));
 
-          expect(await this.mock.$modExp(bytes(b), bytes(e), bytes(m))).to.equal(bytes(modExp(b, e, m), mLength).value);
+          await expect(this.mock.$modExp(bytes(b), bytes(e), bytes(m))).to.eventually.equal(
+            bytes(modExp(b, e, m), mLength).value,
+          );
         });
       }
     });
@@ -387,7 +475,10 @@ describe('Math', function () {
           const e = 200n;
           const m = 50n;
 
-          expect(await this.mock.$tryModExp(type(b), type(e), type(m))).to.deep.equal([true, type(b ** e % m).value]);
+          await expect(this.mock.$tryModExp(type(b), type(e), type(m))).to.eventually.deep.equal([
+            true,
+            type(b ** e % m).value,
+          ]);
         });
 
         it('is correctly reverting when modulus is zero', async function () {
@@ -395,7 +486,7 @@ describe('Math', function () {
           const e = 200n;
           const m = 0n;
 
-          expect(await this.mock.$tryModExp(type(b), type(e), type(m))).to.deep.equal([false, type.zero]);
+          await expect(this.mock.$tryModExp(type(b), type(e), type(m))).to.eventually.deep.equal([false, type.zero]);
         });
       });
     }
@@ -409,7 +500,7 @@ describe('Math', function () {
         it(`calculates b ** e % m (b=2**${log2b}+1) (e=2**${log2e}+1) (m=2**${log2m}+1)`, async function () {
           const mLength = ethers.dataLength(ethers.toBeHex(m));
 
-          expect(await this.mock.$tryModExp(bytes(b), bytes(e), bytes(m))).to.deep.equal([
+          await expect(this.mock.$tryModExp(bytes(b), bytes(e), bytes(m))).to.eventually.deep.equal([
             true,
             bytes(modExp(b, e, m), mLength).value,
           ]);
@@ -421,35 +512,39 @@ describe('Math', function () {
   describe('sqrt', function () {
     it('rounds down', async function () {
       for (const rounding of RoundingDown) {
-        expect(await this.mock.$sqrt(0n, rounding)).to.equal(0n);
-        expect(await this.mock.$sqrt(1n, rounding)).to.equal(1n);
-        expect(await this.mock.$sqrt(2n, rounding)).to.equal(1n);
-        expect(await this.mock.$sqrt(3n, rounding)).to.equal(1n);
-        expect(await this.mock.$sqrt(4n, rounding)).to.equal(2n);
-        expect(await this.mock.$sqrt(144n, rounding)).to.equal(12n);
-        expect(await this.mock.$sqrt(999999n, rounding)).to.equal(999n);
-        expect(await this.mock.$sqrt(1000000n, rounding)).to.equal(1000n);
-        expect(await this.mock.$sqrt(1000001n, rounding)).to.equal(1000n);
-        expect(await this.mock.$sqrt(1002000n, rounding)).to.equal(1000n);
-        expect(await this.mock.$sqrt(1002001n, rounding)).to.equal(1001n);
-        expect(await this.mock.$sqrt(ethers.MaxUint256, rounding)).to.equal(340282366920938463463374607431768211455n);
+        await expect(this.mock.$sqrt(0n, rounding)).to.eventually.equal(0n);
+        await expect(this.mock.$sqrt(1n, rounding)).to.eventually.equal(1n);
+        await expect(this.mock.$sqrt(2n, rounding)).to.eventually.equal(1n);
+        await expect(this.mock.$sqrt(3n, rounding)).to.eventually.equal(1n);
+        await expect(this.mock.$sqrt(4n, rounding)).to.eventually.equal(2n);
+        await expect(this.mock.$sqrt(144n, rounding)).to.eventually.equal(12n);
+        await expect(this.mock.$sqrt(999999n, rounding)).to.eventually.equal(999n);
+        await expect(this.mock.$sqrt(1000000n, rounding)).to.eventually.equal(1000n);
+        await expect(this.mock.$sqrt(1000001n, rounding)).to.eventually.equal(1000n);
+        await expect(this.mock.$sqrt(1002000n, rounding)).to.eventually.equal(1000n);
+        await expect(this.mock.$sqrt(1002001n, rounding)).to.eventually.equal(1001n);
+        await expect(this.mock.$sqrt(ethers.MaxUint256, rounding)).to.eventually.equal(
+          340282366920938463463374607431768211455n,
+        );
       }
     });
 
     it('rounds up', async function () {
       for (const rounding of RoundingUp) {
-        expect(await this.mock.$sqrt(0n, rounding)).to.equal(0n);
-        expect(await this.mock.$sqrt(1n, rounding)).to.equal(1n);
-        expect(await this.mock.$sqrt(2n, rounding)).to.equal(2n);
-        expect(await this.mock.$sqrt(3n, rounding)).to.equal(2n);
-        expect(await this.mock.$sqrt(4n, rounding)).to.equal(2n);
-        expect(await this.mock.$sqrt(144n, rounding)).to.equal(12n);
-        expect(await this.mock.$sqrt(999999n, rounding)).to.equal(1000n);
-        expect(await this.mock.$sqrt(1000000n, rounding)).to.equal(1000n);
-        expect(await this.mock.$sqrt(1000001n, rounding)).to.equal(1001n);
-        expect(await this.mock.$sqrt(1002000n, rounding)).to.equal(1001n);
-        expect(await this.mock.$sqrt(1002001n, rounding)).to.equal(1001n);
-        expect(await this.mock.$sqrt(ethers.MaxUint256, rounding)).to.equal(340282366920938463463374607431768211456n);
+        await expect(this.mock.$sqrt(0n, rounding)).to.eventually.equal(0n);
+        await expect(this.mock.$sqrt(1n, rounding)).to.eventually.equal(1n);
+        await expect(this.mock.$sqrt(2n, rounding)).to.eventually.equal(2n);
+        await expect(this.mock.$sqrt(3n, rounding)).to.eventually.equal(2n);
+        await expect(this.mock.$sqrt(4n, rounding)).to.eventually.equal(2n);
+        await expect(this.mock.$sqrt(144n, rounding)).to.eventually.equal(12n);
+        await expect(this.mock.$sqrt(999999n, rounding)).to.eventually.equal(1000n);
+        await expect(this.mock.$sqrt(1000000n, rounding)).to.eventually.equal(1000n);
+        await expect(this.mock.$sqrt(1000001n, rounding)).to.eventually.equal(1001n);
+        await expect(this.mock.$sqrt(1002000n, rounding)).to.eventually.equal(1001n);
+        await expect(this.mock.$sqrt(1002001n, rounding)).to.eventually.equal(1001n);
+        await expect(this.mock.$sqrt(ethers.MaxUint256, rounding)).to.eventually.equal(
+          340282366920938463463374607431768211456n,
+        );
       }
     });
   });
@@ -458,33 +553,33 @@ describe('Math', function () {
     describe('log2', function () {
       it('rounds down', async function () {
         for (const rounding of RoundingDown) {
-          expect(await this.mock.$log2(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log2(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log2(2n, rounding)).to.equal(1n);
-          expect(await this.mock.$log2(3n, rounding)).to.equal(1n);
-          expect(await this.mock.$log2(4n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(5n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(6n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(7n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(8n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(9n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(ethers.MaxUint256, rounding)).to.equal(255n);
+          await expect(this.mock.$log2(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log2(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log2(2n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log2(3n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log2(4n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(5n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(6n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(7n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(8n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(9n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(ethers.MaxUint256, rounding)).to.eventually.equal(255n);
         }
       });
 
       it('rounds up', async function () {
         for (const rounding of RoundingUp) {
-          expect(await this.mock.$log2(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log2(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log2(2n, rounding)).to.equal(1n);
-          expect(await this.mock.$log2(3n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(4n, rounding)).to.equal(2n);
-          expect(await this.mock.$log2(5n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(6n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(7n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(8n, rounding)).to.equal(3n);
-          expect(await this.mock.$log2(9n, rounding)).to.equal(4n);
-          expect(await this.mock.$log2(ethers.MaxUint256, rounding)).to.equal(256n);
+          await expect(this.mock.$log2(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log2(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log2(2n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log2(3n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(4n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log2(5n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(6n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(7n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(8n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log2(9n, rounding)).to.eventually.equal(4n);
+          await expect(this.mock.$log2(ethers.MaxUint256, rounding)).to.eventually.equal(256n);
         }
       });
     });
@@ -492,37 +587,37 @@ describe('Math', function () {
     describe('log10', function () {
       it('rounds down', async function () {
         for (const rounding of RoundingDown) {
-          expect(await this.mock.$log10(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(2n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(9n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(10n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(11n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(99n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(100n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(101n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(999n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(1000n, rounding)).to.equal(3n);
-          expect(await this.mock.$log10(1001n, rounding)).to.equal(3n);
-          expect(await this.mock.$log10(ethers.MaxUint256, rounding)).to.equal(77n);
+          await expect(this.mock.$log10(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(2n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(9n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(10n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(11n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(99n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(100n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(101n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(999n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(1000n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log10(1001n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log10(ethers.MaxUint256, rounding)).to.eventually.equal(77n);
         }
       });
 
       it('rounds up', async function () {
         for (const rounding of RoundingUp) {
-          expect(await this.mock.$log10(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log10(2n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(9n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(10n, rounding)).to.equal(1n);
-          expect(await this.mock.$log10(11n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(99n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(100n, rounding)).to.equal(2n);
-          expect(await this.mock.$log10(101n, rounding)).to.equal(3n);
-          expect(await this.mock.$log10(999n, rounding)).to.equal(3n);
-          expect(await this.mock.$log10(1000n, rounding)).to.equal(3n);
-          expect(await this.mock.$log10(1001n, rounding)).to.equal(4n);
-          expect(await this.mock.$log10(ethers.MaxUint256, rounding)).to.equal(78n);
+          await expect(this.mock.$log10(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log10(2n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(9n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(10n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log10(11n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(99n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(100n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log10(101n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log10(999n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log10(1000n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log10(1001n, rounding)).to.eventually.equal(4n);
+          await expect(this.mock.$log10(ethers.MaxUint256, rounding)).to.eventually.equal(78n);
         }
       });
     });
@@ -530,31 +625,31 @@ describe('Math', function () {
     describe('log256', function () {
       it('rounds down', async function () {
         for (const rounding of RoundingDown) {
-          expect(await this.mock.$log256(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(2n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(255n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(256n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(257n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(65535n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(65536n, rounding)).to.equal(2n);
-          expect(await this.mock.$log256(65537n, rounding)).to.equal(2n);
-          expect(await this.mock.$log256(ethers.MaxUint256, rounding)).to.equal(31n);
+          await expect(this.mock.$log256(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(2n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(255n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(256n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(257n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(65535n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(65536n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log256(65537n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log256(ethers.MaxUint256, rounding)).to.eventually.equal(31n);
         }
       });
 
       it('rounds up', async function () {
         for (const rounding of RoundingUp) {
-          expect(await this.mock.$log256(0n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(1n, rounding)).to.equal(0n);
-          expect(await this.mock.$log256(2n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(255n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(256n, rounding)).to.equal(1n);
-          expect(await this.mock.$log256(257n, rounding)).to.equal(2n);
-          expect(await this.mock.$log256(65535n, rounding)).to.equal(2n);
-          expect(await this.mock.$log256(65536n, rounding)).to.equal(2n);
-          expect(await this.mock.$log256(65537n, rounding)).to.equal(3n);
-          expect(await this.mock.$log256(ethers.MaxUint256, rounding)).to.equal(32n);
+          await expect(this.mock.$log256(0n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(1n, rounding)).to.eventually.equal(0n);
+          await expect(this.mock.$log256(2n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(255n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(256n, rounding)).to.eventually.equal(1n);
+          await expect(this.mock.$log256(257n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log256(65535n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log256(65536n, rounding)).to.eventually.equal(2n);
+          await expect(this.mock.$log256(65537n, rounding)).to.eventually.equal(3n);
+          await expect(this.mock.$log256(ethers.MaxUint256, rounding)).to.eventually.equal(32n);
         }
       });
     });