Add Math.modExp and a Panic library

.changeset/shiny-poets-whisper.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`Math`: Add `modExp` function that exposes the `EIP-198` precompile.

.changeset/silver-swans-promise.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`Panic`: Add a library for reverting with panic codes.

contracts/utils/Panic.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.20;
+
+/**
+ * @dev Helper library for emitting standardized panic codes.
+ */
+// slither-disable-next-line unused-state
+library Panic {
+    uint256 internal constant ASSERTION_ERROR = 0x1;
+    uint256 internal constant ARITHMETIC_UNDER_OR_OVERFLOW = 0x11;
+    uint256 internal constant DIVISION_BY_ZERO = 0x12;
+    uint256 internal constant ENUM_CONVERSION_OUT_OF_BOUNDS = 0x21;
+    uint256 internal constant INCORRECTLY_ENCODED_STORAGE_BYTE_ARRAY = 0x22;
+    uint256 internal constant POP_ON_EMPTY_ARRAY = 0x31;
+    uint256 internal constant ARRAY_ACCESS_OUT_OF_BOUNDS = 0x32;
+    uint256 internal constant TOO_MUCH_MEMORY_ALLOCATED = 0x41;
+    uint256 internal constant ZERO_INITIALIZED_VARIABLE = 0x51;
+
+    function panic(uint256 code) internal pure {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, shl(0xe0, 0x4e487b71))
+            mstore(0x04, code)
+            revert(0x00, 0x24)
+        }
+    }
+}

contracts/utils/math/Math.sol

@@ -3,15 +3,12 @@
 
 pragma solidity ^0.8.20;
 
+import {Panic} from "../Panic.sol";
+
 /**
  * @dev Standard math utilities missing in the Solidity language.
  */
 library Math {
-    /**
-     * @dev Muldiv operation overflow.
-     */
-    error MathOverflowedMulDiv();
-
     enum Rounding {
         Floor, // Toward negative infinity
         Ceil, // Toward positive infinity
@@ -107,7 +104,7 @@ library Math {
     function ceilDiv(uint256 a, uint256 b) internal pure returns (uint256) {
         if (b == 0) {
             // Guarantee the same behavior as in a regular Solidity division.
-            return a / b;
+            Panic.panic(Panic.DIVISION_BY_ZERO);
         }
 
         // The following calculation ensures accurate ceiling division without overflow.
@@ -149,7 +146,7 @@ library Math {
 
             // Make sure the result is less than 2^256. Also prevents denominator == 0.
             if (denominator <= prod1) {
-                revert MathOverflowedMulDiv();
+                Panic.panic(denominator == 0 ? Panic.DIVISION_BY_ZERO : Panic.ARITHMETIC_UNDER_OR_OVERFLOW);
             }
 
             ///////////////////////////////////////////////
@@ -226,6 +223,9 @@ library Math {
      * If n is not a prime, then Z/nZ is not a field, and some elements might not be inversible.
      *
      * If the input value is not inversible, 0 is returned.
+     *
+     * Note: If you know for sure that n is (big) a prime, it may be cheaper to use Ferma's little theorem and get the
+     * inverse using `Math.modExp(a, n - 2, n)`.
      */
     function invMod(uint256 a, uint256 n) internal pure returns (uint256) {
         unchecked {
@@ -275,6 +275,33 @@ library Math {
         }
     }
 
+    /**
+     * @dev Returns the modular exponentiation of the specified base, exponent and modulus (b ** e % m)
+     *
+     * Requirements:
+     * - modulus can't be zero
+     * - result should be obtained successfully
+     */
+    function modExp(uint256 b, uint256 e, uint256 m) internal view returns (uint256 result) {
+        if (m == 0) {
+            Panic.panic(Panic.DIVISION_BY_ZERO);
+        }
+        /// @solidity memory-safe-assembly
+        assembly {
+            let ptr := mload(0x40)
+            mstore(ptr, 0x20)
+            mstore(add(ptr, 0x20), 0x20)
+            mstore(add(ptr, 0x40), 0x20)
+            mstore(add(ptr, 0x60), b)
+            mstore(add(ptr, 0x80), e)
+            mstore(add(ptr, 0xa0), m)
+            if iszero(staticcall(gas(), 0x05, ptr, 0xc0, ptr, 0x20)) {
+                revert(0, 0)
+            }
+            result := mload(ptr)
+        }
+    }
+
     /**
      * @dev Returns the square root of a number. If the number is not a perfect square, the value is rounded
      * towards zero.

scripts/generate/templates/Checkpoints.t.js

@@ -26,14 +26,14 @@ function _bound${capitalize(opts.keyTypeName)}(
     ${opts.keyTypeName} x,
     ${opts.keyTypeName} min,
     ${opts.keyTypeName} max
-) internal view returns (${opts.keyTypeName}) {
+) internal pure returns (${opts.keyTypeName}) {
     return SafeCast.to${capitalize(opts.keyTypeName)}(bound(uint256(x), uint256(min), uint256(max)));
 }
 
 function _prepareKeys(
     ${opts.keyTypeName}[] memory keys,
     ${opts.keyTypeName} maxSpread
-) internal view {
+) internal pure {
     ${opts.keyTypeName} lastKey = 0;
     for (uint256 i = 0; i < keys.length; ++i) {
         ${opts.keyTypeName} key = _bound${capitalize(opts.keyTypeName)}(keys[i], lastKey, lastKey + maxSpread);

test/utils/Panic.test.js

@@ -0,0 +1,26 @@
+const { ethers } = require('hardhat');
+const { expect } = require('chai');
+const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers');
+const { PANIC_CODES } = require('@nomicfoundation/hardhat-chai-matchers/panic');
+
+async function fixture() {
+  return { mock: await ethers.deployContract('$Panic') };
+}
+
+describe('Panic', function () {
+  beforeEach(async function () {
+    Object.assign(this, await loadFixture(fixture));
+  });
+
+  for (const [name, code] of Object.entries(PANIC_CODES)) {
+    describe(name, function () {
+      it('exposes panic code as constant', async function () {
+        expect(await this.mock.getFunction(`$${name}`)()).to.equal(code);
+      });
+
+      it(`throw panic code ${ethers.toBeHex(code)}`, async function () {
+        await expect(this.mock.$panic(code)).to.be.revertedWithPanic(code);
+      });
+    });
+  }
+});

test/utils/math/Math.t.sol

@@ -186,7 +186,7 @@ contract MathTest is Test {
         // Full precision for q * d
         (uint256 qdHi, uint256 qdLo) = _mulHighLow(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 qdRemLo, uint256 c) = _addCarry(qdLo, mulmod(x, y, d));
         uint256 qdRemHi = qdHi + c;
 
         // Full precision check that x * y = q * d + rem
@@ -201,14 +201,31 @@ contract MathTest is Test {
         vm.assume(xyHi >= d);
 
         // we are outside the scope of {testMulDiv}, we expect muldiv to revert
-        try this.muldiv(x, y, d) returns (uint256) {
-            fail();
-        } catch {}
+        vm.expectRevert();
+        Math.mulDiv(x, y, d);
     }
 
-    // External call
-    function muldiv(uint256 x, uint256 y, uint256 d) external pure returns (uint256) {
-        return Math.mulDiv(x, y, d);
+    // MOD EXP
+    function testModExp(uint256 b, uint256 e, uint256 m) public {
+        if (m == 0) {
+            vm.expectRevert(); // panic 0x12
+        }
+        uint256 result = Math.modExp(b, e, m);
+        assertTrue(result < m);
+        assertEq(result, _nativeModExp(b, e, m));
+    }
+
+    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
@@ -217,12 +234,6 @@ contract MathTest is Test {
         return Math.Rounding(r);
     }
 
-    function _mulmod(uint256 x, uint256 y, uint256 z) private pure returns (uint256 r) {
-        assembly {
-            r := mulmod(x, y, z)
-        }
-    }
-
     function _mulHighLow(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);

test/utils/math/Math.test.js

@@ -222,7 +222,7 @@ describe('Math', function () {
     });
   });
 
-  describe('muldiv', function () {
+  describe('mulDiv', function () {
     it('divide by 0', async function () {
       const a = 1n;
       const b = 1n;
@@ -234,9 +234,8 @@ describe('Math', function () {
       const a = 5n;
       const b = ethers.MaxUint256;
       const c = 2n;
-      await expect(this.mock.$mulDiv(a, b, c, Rounding.Floor)).to.be.revertedWithCustomError(
-        this.mock,
-        'MathOverflowedMulDiv',
+      await expect(this.mock.$mulDiv(a, b, c, Rounding.Floor)).to.be.revertedWithPanic(
+        PANIC_CODES.ARITHMETIC_UNDER_OR_OVERFLOW,
       );
     });
 
@@ -336,6 +335,24 @@ describe('Math', function () {
     }
   });
 
+  describe('modExp', function () {
+    it('is correctly calculating modulus', async function () {
+      const base = 3n;
+      const exponent = 200n;
+      const modulus = 50n;
+
+      expect(await this.mock.$modExp(base, exponent, modulus)).to.equal(base ** exponent % modulus);
+    });
+
+    it('is correctly reverting when modulus is zero', async function () {
+      const base = 3n;
+      const exponent = 200n;
+      const modulus = 0n;
+
+      await expect(this.mock.$modExp(base, exponent, modulus)).to.be.revertedWithPanic(PANIC_CODES.DIVISION_BY_ZERO);
+    });
+  });
+
   describe('sqrt', function () {
     it('rounds down', async function () {
       for (const rounding of RoundingDown) {

test/utils/structs/Checkpoints.t.sol

@@ -17,11 +17,11 @@ contract CheckpointsTrace224Test is Test {
     Checkpoints.Trace224 internal _ckpts;
 
     // helpers
-    function _boundUint32(uint32 x, uint32 min, uint32 max) internal view returns (uint32) {
+    function _boundUint32(uint32 x, uint32 min, uint32 max) internal pure returns (uint32) {
         return SafeCast.toUint32(bound(uint256(x), uint256(min), uint256(max)));
     }
 
-    function _prepareKeys(uint32[] memory keys, uint32 maxSpread) internal view {
+    function _prepareKeys(uint32[] memory keys, uint32 maxSpread) internal pure {
         uint32 lastKey = 0;
         for (uint256 i = 0; i < keys.length; ++i) {
             uint32 key = _boundUint32(keys[i], lastKey, lastKey + maxSpread);
@@ -125,11 +125,11 @@ contract CheckpointsTrace208Test is Test {
     Checkpoints.Trace208 internal _ckpts;
 
     // helpers
-    function _boundUint48(uint48 x, uint48 min, uint48 max) internal view returns (uint48) {
+    function _boundUint48(uint48 x, uint48 min, uint48 max) internal pure returns (uint48) {
         return SafeCast.toUint48(bound(uint256(x), uint256(min), uint256(max)));
     }
 
-    function _prepareKeys(uint48[] memory keys, uint48 maxSpread) internal view {
+    function _prepareKeys(uint48[] memory keys, uint48 maxSpread) internal pure {
         uint48 lastKey = 0;
         for (uint256 i = 0; i < keys.length; ++i) {
             uint48 key = _boundUint48(keys[i], lastKey, lastKey + maxSpread);
@@ -233,11 +233,11 @@ contract CheckpointsTrace160Test is Test {
     Checkpoints.Trace160 internal _ckpts;
 
     // helpers
-    function _boundUint96(uint96 x, uint96 min, uint96 max) internal view returns (uint96) {
+    function _boundUint96(uint96 x, uint96 min, uint96 max) internal pure returns (uint96) {
         return SafeCast.toUint96(bound(uint256(x), uint256(min), uint256(max)));
     }
 
-    function _prepareKeys(uint96[] memory keys, uint96 maxSpread) internal view {
+    function _prepareKeys(uint96[] memory keys, uint96 maxSpread) internal pure {
         uint96 lastKey = 0;
         for (uint256 i = 0; i < keys.length; ++i) {
             uint96 key = _boundUint96(keys[i], lastKey, lastKey + maxSpread);