crypto: poseidon2: Define huff structure and implement simple sbox (#3)

* forge install: foundry-huff

* crypto: poseidon2: roundUtils: Define huff structure and add sbox

.github/workflows/test.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       fail-fast: true
 
-    name: Foundry project
+    name: Foundry checks
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -41,5 +41,5 @@ jobs:
 
       - name: Run Forge tests
         run: |
-          forge test -vvv
+          forge test --ffi -vvv
         id: test

.gitignore

@@ -1,3 +1,6 @@
+# Workspace 
+.vscode
+
 # Compiler files
 cache/
 out/

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "lib/forge-std"]
 	path = lib/forge-std
 	url = https://github.com/foundry-rs/forge-std
+[submodule "lib/foundry-huff"]
+	path = lib/foundry-huff
+	url = https://github.com/huff-language/foundry-huff

foundry.toml

@@ -1,6 +1,7 @@
 [profile.default]
+solc_version = "0.8.20"
+evm_version = "cancun"
 src = "src"
 out = "out"
 libs = ["lib"]
-
-# See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options
+remappings = ["forge-std/=lib/forge-std/src/"]

remappings.txt

@@ -0,0 +1,5 @@
+ds-test/=lib/forge-std/lib/ds-test/src/
+forge-std/=lib/forge-std/src/
+foundry-huff/=lib/foundry-huff/src/
+solidity-stringutils/=lib/foundry-huff/lib/solidity-stringutils/
+stringutils/=lib/foundry-huff/lib/solidity-stringutils/

src/crypto/poseidon2/main.huff

@@ -0,0 +1,9 @@
+/// @title Poseidon2
+/// @author @joeykraut
+/// @notice A Poseidon2 implementation in Huff. See https://eprint.iacr.org/2023/323
+/// for more details. Inspired by https://github.com/zemse/poseidon2-evm/tree/main
+
+#define macro MAIN() = takes(0) returns(0) {
+}
+
+

src/crypto/poseidon2/roundUtils.huff

@@ -0,0 +1,30 @@
+/// Helpers for Poseidon2 round functions
+
+/// @dev The scalar field modulus of BN254
+#define constant PRIME = 0x30644e72e131a029b85045b68181585d2833e84879b9709143e1f593f0000001 
+
+/// @dev Push the prime onto the stack
+#define macro PUSH_PRIME() = {
+    [PRIME]
+}
+
+/// @dev Apply the sbox function to the element on top of the stack,
+/// which is replaced by its fifth power mod p
+/// I.e. transforms the stack [x, ...] -> [x^5 mod p, ...]
+#define macro SBOX() = takes(2) returns(1) {
+    // Takes [x, PRIME]
+    // Prelude: setup primes on the stack in the order they'll be used
+    dup2 dup1       // [PRIME, PRIME, x, PRIME]
+    dup3            // [x, PRIME, PRIME, x, PRIME]
+    dup1            // [x, x, PRIME, PRIME, x, PRIME]
+
+    // Compute x^2 (mod p)
+    mulmod          // [x^2 mod p, PRIME, x, PRIME]
+
+    // Compute x^4 (mod p)
+    dup1            // [x^2 mod p, x^2 mod p, PRIME, x, PRIME]
+    mulmod          // [x^4 mod p, x, PRIME]
+
+    // Compute x^5 (mod p)
+    mulmod          // [x^5 mod p]
+}

test/Poseidon.t.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.0;
+
+import {Test} from "forge-std/Test.sol";
+import {console} from "forge-std/console.sol";
+import {HuffDeployer} from "foundry-huff/HuffDeployer.sol";
+
+contract PoseidonTest is Test {
+    /// @dev The Poseidon main contract
+    PoseidonSuite public poseidonSuite;
+
+    /// @dev The BN254 field modulus from roundUtils.huff
+    uint256 PRIME = 0x30644e72e131a029b85045b68181585d2833e84879b9709143e1f593f0000001;
+
+    /// @dev Deploy the PoseidonSuite contract
+    function setUp() public {
+        poseidonSuite = PoseidonSuite(HuffDeployer.deploy("../test/huff/testPoseidonUtils"));
+    }
+
+    /// @dev Test the sbox function applied to a single input
+    function testSboxSingle() public {
+        uint256 testValue = vm.randomUint();
+        uint256 result = poseidonSuite.testSboxSingle(testValue);
+
+        // Calculate expected x^5 mod p
+        uint256 x2 = mulmod(testValue, testValue, PRIME);
+        uint256 x4 = mulmod(x2, x2, PRIME);
+        uint256 expected = mulmod(testValue, x4, PRIME);
+        assertEq(result, expected, "Expected result to match x^5 mod p");
+    }
+}
+
+interface PoseidonSuite {
+    function testSboxSingle(uint256) external returns (uint256);
+}

test/huff/testPoseidonUtils.huff

@@ -0,0 +1,36 @@
+/// @title Test Poseidon Utils
+/// @author @joeykraut
+/// @notice Test the utils for the Poseidon2 permutation
+
+#include "../../src/crypto/poseidon2/roundUtils.huff"
+
+/// @dev Test the sbox function applied to a single input
+#define function testSboxSingle(uint256) nonpayable returns(uint256) 
+
+/// @dev Entrypoint to the poseidon test suite
+#define macro MAIN() = takes(0) returns(0) {
+    // Get the function selector
+    0x0 calldataload 0xe0 shr // [SELECTOR]
+    __FUNC_SIG(testSboxSingle) eq testSboxSingle jumpi
+
+    // Revert if the function selector is not valid
+    0x1 0x0 mstore
+    0x1 0x0 revert
+
+    testSboxSingle:
+        TEST_SBOX_SINGLE()
+}
+
+/// @notice Test the sbox function applied to a single input
+#define macro TEST_SBOX_SINGLE() = takes(0) returns(0) {
+    // Get the input from calldata 
+    PUSH_PRIME()
+    0x04 calldataload
+
+    // Call the sbox function
+    SBOX()
+
+    // Return the result
+    0x00 mstore
+    0x20 0x00 return
+}