Support for Modexp functions in huff-math

foundry.toml

@@ -2,5 +2,6 @@
 src = 'src'
 out = 'artifacts'
 libs = ["node_modules", "lib"]
+evm_version = "shanghai"
 ffi=true
 # See more config options https://github.com/foundry-rs/foundry/tree/master/config

src/Math.huff

@@ -42,4 +42,48 @@
     complete:
 }
 
-
+#define macro MODEXP() = takes (3) returns (1) {
+    dup1
+    iszero MODULUS_ZERO jumpi             // Compare it to zero
+    MODULUS_NOT_ZERO jump
+
+    MODULUS_NOT_ZERO:
+        // Setting up correct memory layout
+        0x20 push1 0x0 mstore     // Store length of base (32 bytes) at memory 0x00
+        0x20 0x20 mstore     // Store length of exponent (32 bytes) at memory 0x40
+        0x20 0x40 mstore     // Store length of modulus (32 bytes) at memory 0x80
+
+        0xa0 mstore          // Store 'modulus' at memory 0xa0
+        0x80 mstore          // Store 'exponent' at memory 0xc0
+        0x60 mstore          // Store 'base' at memory 0x60
+
+        // Prepare staticcall to precompiled contract
+
+        0x20                     // Output data size (32 bytes)
+        push1 0x0                // Output data memory
+        0xc0                     // Input data size (192 bytes)
+        push1 0x0                // Input data starting position
+        0x05                     // Address of the modexp precompiled contract (0x05)
+        gas                      // Gas provided
+        staticcall
+
+        // Verify the call success
+        iszero STATIC_CALL_FAILED jumpi // Check if the staticcall was successful
+        CALL_SUCCESSFUL jump           // If staticcall returned 0 (failure), revert
+
+    MODULUS_ZERO:
+        push1 0x0 mload
+        revert         // Revert if modulus is zero
+
+    STATIC_CALL_FAILED:
+        push1 0x0 mload
+        revert         // Revert if staticcall failed
+
+    // Label for call success
+    CALL_SUCCESSFUL:
+        push1 0x0 mload               // Load the result from memory location 0x00
+        END jump                 // Jump to the end of the macro
+
+    END:
+    // The result of base^exponent % modulus is now on top of the stack
+}

src/interfaces/IMath.sol

@@ -12,4 +12,6 @@ interface IMath {
     function divideNumbers(uint256, uint256) external view returns (uint256);
 
     function abs(uint256, uint256) external view returns (uint256);
+
+    function modExp(uint256, uint256, uint256) external view returns (uint256);
 }

src/wrappers/MathWrapper.huff

@@ -6,6 +6,7 @@
 #define function multiplyNumbers(uint256,uint256) nonpayable returns (uint256)
 #define function divideNumbers(uint256,uint256) nonpayable returns (uint256)
 #define function abs(uint256,uint256) nonpayable returns (uint256)
+#define function modExp(uint256,uint256,uint256) nonpayable returns (uint256)
 
 #define macro ADD_WRAPPER() = takes (2) returns (1) {
     0x04 calldataload    // [num1]
@@ -47,7 +48,14 @@
     0x20 0x00 return     // []  
  }
 
-
+  #define macro MODEXP_WRAPPER() = takes (3) returns (1) {
+    0x04 calldataload    // [base]
+    0x24 calldataload    // [exponent, base]
+    0x44 calldataload    // [modulus, exponent, base]
+    MODEXP()             // [base^exponent mod modulus]
+    0x00 mstore          // []
+    0x20 0x00 return     // []  
+ }
 
 
 #define macro MAIN() = takes (0) returns (0) {
@@ -63,6 +71,7 @@
     dup1 0xd3f3cd7b eq  multiplyNumbers jumpi
     dup1 0x8fce12ed eq  divideNumbers   jumpi
     dup1 0xe093a157 eq  abs             jumpi
+    dup1 0x3148f14f eq  modExp          jumpi
 
 
     addNumbers:
@@ -76,6 +85,9 @@
 
      divideNumbers:
         DIVIDE_WRAPPER()
+
+    modExp:
+        MODEXP_WRAPPER()
 
     abs:
         ABS_WRAPPER()

test/foundry/Math.t.sol

@@ -93,4 +93,29 @@ contract MathTest is Test {
         uint256 _result = a > b ? a - b : b - a;
         require(math.abs(a, b) == _result);
     }
+
+    function testModExp() public {
+        // Example test: 2^3 % 5 should equal 3
+        uint256 base = 2;
+        uint256 exponent = 3;
+        uint256 modulus = 5;
+        uint256 expected = 3;
+
+        uint256 result = math.modExp(base, exponent, modulus);
+        assertEq(result, expected, "modExp did not return the expected value");
+    }
+
+    // function testModExp_fuzz(uint256 b, uint256 e, uint256 m) public {
+    //     // To avoid testing with modulus zero, which would revert
+    //     vm.assume(m > 1);
+    //     // To avoid gas issues, cap the exponent
+    //     uint256 exponent = e % 256;
+
+    //     // The actual modExp calculation can be complicated to emulate in Solidity due to gas constraints,
+    //     // so here we just test that the function does not revert and returns a value
+    //     // less than the modulus.
+    //     uint256 result = math.modExp(b, exponent, m);
+
+    //     assertLt(result, m, "modExp result should be less than the modulus");
+    // }
 }

test/foundry/MathForkTest.t.sol

@@ -36,4 +36,15 @@ contract MathForkTest is Test {
     function testAbs() public view {
         require(math.abs(1, 10) == 9);
     }
+
+    function testModExp() public {
+        // Example test: 2^3 % 5 should equal 3
+        uint256 base = 2;
+        uint256 exponent = 3;
+        uint256 modulus = 5;
+        uint256 expected = 3;
+
+        uint256 result = math.modExp(base, exponent, modulus);
+        assertEq(result, expected, "modExp did not return the expected value");
+    }
 }

test/huff/Math.t.huff

@@ -75,3 +75,24 @@
     ASSERT_EQ()                          // [4e18==result]     
 }
 
+#define test TEST_MODEXP() = {
+    // Test case 1: Simple modular exponentiation
+    // Using small numbers for easy manual verification: 2^3 % 5 = 3
+    0x02                                // [base = 2, exponent, modulus]
+    0x03                                // [exponent = 3, modulus]
+    0x05                                // [modulus = 5]
+    MODEXP()                            // [result]
+    0x03                                // [expected = 3, result]
+    ASSERT_EQ()                         // [3 == result]
+
+    // Test case 2: Larger numbers
+    // We need to choose numbers such that we can calculate the expected result manually or with a tool
+    // For example: (0x04)^2 % 0x05 = 0x01
+    0x04                                // [base = 4, exponent, modulus]
+    0x02                                // [exponent = 2, modulus]
+    0x05                                // [modulus = 5]
+    MODEXP()                            // [result]
+    0x01                                // [expected = 4, result]
+    ASSERT_EQ()                         // [4 == result]
+}
+