tcoratger
diff --git a/‎src/lean_spec/subspecs/poseidon1/permutation.py‎
Lines changed: 74 additions & 38 deletions b/‎src/lean_spec/subspecs/poseidon1/permutation.py‎
Lines changed: 74 additions & 38 deletions
diff --git a/‎tests/lean_spec/subspecs/poseidon1/test_permutation.py‎
Lines changed: 112 additions & 5 deletions b/‎tests/lean_spec/subspecs/poseidon1/test_permutation.py‎
Lines changed: 112 additions & 5 deletions
@@ -4,17 +4,17 @@
 Based on "Poseidon: A New Hash Function for Zero-Knowledge Proof Systems".
 See https://eprint.iacr.org/2019/458.
 
+This is the original Hades-based design.
+
 Uses Numba JIT compilation for native-speed permutation.
 """
 
-from __future__ import annotations
-
 from typing import Self
 
 import numpy as np
 from numba import njit
 from numpy.typing import NDArray
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 
 from ...types import StrictBaseModel
 from ..koalabear.field import Fp, P
@@ -24,18 +24,6 @@
 )
 
 
-def _build_circulant_mds(first_row: list[int], n: int, p: int) -> NDArray[np.int64]:
-    """
-    Expand a circulant matrix from its first row into a dense NxN matrix.
-
-    A circulant matrix C defined by first row [r0, r1, ..., rn-1]:
-    C[i][j] = r[(j - i) mod n]
-
-    Row i is the first row rolled right by i positions.
-    """
-    return np.array([np.roll(first_row, i) for i in range(n)], dtype=np.int64) % p
-
-
 @njit(cache=True)
 def _mds_multiply_jit(
     state: NDArray[np.int64], mds: NDArray[np.int64], p: int
@@ -46,13 +34,26 @@ def _mds_multiply_jit(
     Computes y = MDS * x where MDS is the circulant MDS matrix.
     Each product is reduced mod p before accumulation to prevent overflow.
     """
+    # State length doubles as the matrix dimension since MDS is square.
     n = state.shape[0]
+
+    # Output buffer, written in place by the inner loop.
     result = np.empty(n, dtype=np.int64)
+
+    # One iteration computes a single row of the matrix-vector product.
     for i in range(n):
+        # Accumulator collects n pre-reduced contributions before the final fold.
         s = np.int64(0)
+
         for j in range(n):
+            # Each factor sits below p, so the product fits in 62 bits.
+            #
+            # Without per-product reduction, summing n products would risk int64 overflow.
             s += (mds[i, j] * state[j]) % p
+
+        # Final fold back into the field after n already-reduced terms.
         result[i] = s % p
+
     return result
 
 
@@ -72,44 +73,51 @@ def _permute_jit(
     Modifies state array in-place.
     S-box: x^3 computed as (x*x % p) * x % p to avoid int64 overflow.
 
-    Round structure: AddRoundConstants -> S-box -> MDS multiply.
-    No initial linear layer is applied before the round structure begins.
+    - Round structure: AddRoundConstants -> S-box -> MDS multiply.
+    - No initial linear layer is applied before the round structure begins.
+    - This matches the original Poseidon1 design.
     """
     const_idx = 0
 
-    # 1. First half of full rounds.
+    # Phase 1: opening full rounds.
     #
-    # Full rounds apply the S-box to every state element.
-    # Note: for S_BOX_DEGREE=3, state**3 would overflow int64 before modulo.
-    # Expand S-box to `(state*state % P) * state % P` to stay in range.
+    # Full S-boxes at the boundary maximize non-linearity where
+    # attacker control over inputs is highest.
     for _ in range(half_rounds_f):
         # Add round constants to entire state.
         state[:] = (state + round_constants[const_idx : const_idx + width]) % p
         const_idx += width
 
         # Apply S-box (x -> x^d) to full state.
+        #
+        # Cubing in one shot would overflow int64 inside Numba.
+        # Splitting into two modular multiplies keeps each intermediate below p squared.
         state[:] = (state * state % p) * state % p
 
         # Apply dense MDS multiply for diffusion.
         state[:] = _mds_multiply_jit(state, mds, p)
 
-    # 2. Partial rounds.
+    # Phase 2: partial rounds.
     #
-    # Partial rounds add constants to ALL state elements but apply
-    # the S-box only to state[0]. The same dense MDS matrix is used.
+    # AddRoundConstants and the MDS multiply still run on the entire state.
+    # Only the S-box layer is partial.
+    # Applying the S-box to only one element is the central Hades optimization.
+    # It still saturates algebraic degree while cutting SNARK constraint cost.
     for _ in range(rounds_p):
         # Add round constants to entire state.
         state[:] = (state + round_constants[const_idx : const_idx + width]) % p
         const_idx += width
 
         # Apply S-box to first element only.
-        # This is the main optimization of the Hades design.
         state[0] = (state[0] * state[0] % p) * state[0] % p
 
         # Apply dense MDS multiply.
         state[:] = _mds_multiply_jit(state, mds, p)
 
-    # 3. Second half of full rounds.
+    # Phase 3: closing full rounds.
+    #
+    # A second wall of full S-boxes blocks algebraic attacks that could
+    # unwind the partial-round middle.
     for _ in range(half_rounds_f):
         # Add round constants to entire state.
         state[:] = (state + round_constants[const_idx : const_idx + width]) % p
@@ -123,7 +131,12 @@ def _permute_jit(
 
 
 class Poseidon1Params(StrictBaseModel):
-    """Parameters for a specific Poseidon1 instance."""
+    """Parameters for a specific Poseidon1 instance.
+
+    - The paper requires at least 6 full rounds for statistical-attack security.
+    - Some regimes raise this bound to 10 per Eq. 2 of the paper.
+    - This minimum is not enforced here. Callers must choose secure parameters.
+    """
 
     width: int = Field(gt=0, description="The size of the state (t).")
     rounds_f: int = Field(gt=0, description="Total number of 'full' rounds.")
@@ -137,6 +150,19 @@ class Poseidon1Params(StrictBaseModel):
         description="The list of pre-computed constants for all rounds.",
     )
 
+    @field_validator("rounds_f")
+    @classmethod
+    def _rounds_f_must_be_even(cls, value: int) -> int:
+        """Require an even full-round count.
+
+        - The permutation runs equal halves of full rounds before and after the partial middle.
+        - An odd count silently drops one full round and orphans a width-sized block of constants.
+        - The original Poseidon design assumes an even split.
+        """
+        if value % 2 != 0:
+            raise ValueError("Full-round count must be even.")
+        return value
+
     @model_validator(mode="after")
     def check_lengths(self) -> Self:
         """Ensures vector lengths match the configuration."""
@@ -151,12 +177,7 @@ def check_lengths(self) -> Self:
 
 
 class Poseidon1:
-    """
-    Optimized execution engine for Poseidon1.
-
-    Pre-processes parameters into numpy arrays during initialization.
-    Minimizes overhead during permute calls.
-    """
+    """Execution engine for Poseidon1."""
 
     __slots__ = ("_width", "_half_rounds_f", "_rounds_p", "_mds", "_round_constants")
 
@@ -186,9 +207,14 @@ def __init__(self, params: Poseidon1Params) -> None:
         self._half_rounds_f = params.rounds_f // 2
         self._rounds_p = params.rounds_p
 
-        # Build the dense circulant MDS matrix from first row.
-        first_row_ints = [int(fp) for fp in params.mds_first_row]
-        self._mds = _build_circulant_mds(first_row_ints, params.width, P)
+        # Expand the n-by-n circulant MDS matrix from its first row r.
+        #
+        # Row i is r rolled right by i positions.
+        # Equivalently, C[i][j] = r[(j - i) mod n].
+        first_row = [int(fp) for fp in params.mds_first_row]
+        self._mds = (
+            np.array([np.roll(first_row, i) for i in range(self._width)], dtype=np.int64) % P
+        )
 
         # Pre-convert round constants to numpy array.
         self._round_constants = np.array([int(fp) for fp in params.round_constants], dtype=np.int64)
@@ -228,7 +254,12 @@ def permute(self, current_state: list[Fp]) -> list[Fp]:
 
 
 _MDS_FIRST_ROW_16: list[int] = [1, 1, 51, 1, 11, 17, 2, 1, 101, 63, 15, 2, 67, 22, 13, 3]
-"""MDS first row for width-16 circulant matrix. From Plonky3: koala-bear/src/mds.rs."""
+"""MDS first row for width-16 circulant matrix.
+
+- From Plonky3: https://github.com/Plonky3/Plonky3/blob/main/koala-bear/src/mds.rs
+- The paper recommends Cauchy matrices over the circulant family used here.
+- The matrix must avoid invariant subspace trails per GRS21.
+"""
 
 _MDS_FIRST_ROW_24: list[int] = [
     0x2D0AAAAB,
@@ -256,7 +287,12 @@ def permute(self, current_state: list[Fp]) -> list[Fp]:
     0x17E118F6,
     0x0878A07F,
 ]
-"""MDS first row for width-24 circulant matrix. From Plonky3: koala-bear/src/mds.rs."""
+"""MDS first row for width-24 circulant matrix.
+
+- From Plonky3: https://github.com/Plonky3/Plonky3/blob/main/koala-bear/src/mds.rs
+- The paper recommends Cauchy matrices over the circulant family used here.
+- The matrix must avoid invariant subspace trails per GRS21.
+"""
 
 PARAMS_16 = Poseidon1Params(
     width=16,
 
@@ -5,8 +5,9 @@
 """
 
 import pytest
+from pydantic import ValidationError
 
-from lean_spec.subspecs.koalabear.field import Fp
+from lean_spec.subspecs.koalabear.field import Fp, P
 from lean_spec.subspecs.poseidon1.permutation import (
     PARAMS_16,
     PARAMS_24,
@@ -106,7 +107,7 @@ class TestPoseidon1ParamsValidation:
 
     def test_invalid_mds_first_row_length(self) -> None:
         """Raises error when mds_first_row length doesn't match width."""
-        with pytest.raises(ValueError, match="Length of mds_first_row must equal width"):
+        with pytest.raises(ValueError, match=r"Length of mds_first_row must equal width\."):
             Poseidon1Params(
                 width=3,
                 rounds_f=8,
@@ -117,7 +118,7 @@ def test_invalid_mds_first_row_length(self) -> None:
 
     def test_invalid_round_constants_count(self) -> None:
         """Raises error when round_constants count is incorrect."""
-        with pytest.raises(ValueError, match="Incorrect number of round constants"):
+        with pytest.raises(ValueError, match=r"Incorrect number of round constants provided\."):
             Poseidon1Params(
                 width=3,
                 rounds_f=8,
@@ -126,20 +127,92 @@ def test_invalid_round_constants_count(self) -> None:
                 round_constants=[Fp(1)] * 20,
             )
 
+    def test_width_must_be_positive(self) -> None:
+        """Rejects a non-positive width."""
+        with pytest.raises(ValidationError, match=r"Input should be greater than 0"):
+            Poseidon1Params(
+                width=0,
+                rounds_f=8,
+                rounds_p=20,
+                mds_first_row=[Fp(1), Fp(2), Fp(3)],
+                round_constants=[Fp(1)] * 84,
+            )
+
+    def test_rounds_f_must_be_positive(self) -> None:
+        """Rejects a non-positive full-round count before the even-check runs."""
+        with pytest.raises(ValidationError, match=r"Input should be greater than 0"):
+            Poseidon1Params(
+                width=3,
+                rounds_f=0,
+                rounds_p=20,
+                mds_first_row=[Fp(1), Fp(2), Fp(3)],
+                round_constants=[Fp(1)] * 60,
+            )
+
+    def test_rounds_p_must_be_non_negative(self) -> None:
+        """Rejects a negative partial-round count."""
+        with pytest.raises(ValidationError, match=r"Input should be greater than or equal to 0"):
+            Poseidon1Params(
+                width=3,
+                rounds_f=8,
+                rounds_p=-1,
+                mds_first_row=[Fp(1), Fp(2), Fp(3)],
+                round_constants=[Fp(1)] * 21,
+            )
+
+    def test_mds_first_row_must_be_non_empty(self) -> None:
+        """Rejects an empty MDS first row."""
+        with pytest.raises(
+            ValidationError,
+            match=r"List should have at least 1 item after validation, not 0",
+        ):
+            Poseidon1Params(
+                width=3,
+                rounds_f=8,
+                rounds_p=20,
+                mds_first_row=[],
+                round_constants=[Fp(1)] * 84,
+            )
+
+    def test_round_constants_must_be_non_empty(self) -> None:
+        """Rejects an empty round-constants list."""
+        with pytest.raises(
+            ValidationError,
+            match=r"List should have at least 1 item after validation, not 0",
+        ):
+            Poseidon1Params(
+                width=3,
+                rounds_f=8,
+                rounds_p=20,
+                mds_first_row=[Fp(1), Fp(2), Fp(3)],
+                round_constants=[],
+            )
+
+    def test_rounds_f_must_be_even(self) -> None:
+        """Rejects odd full-round counts that would leave constants unused."""
+        with pytest.raises(ValidationError, match=r"Full-round count must be even\."):
+            Poseidon1Params(
+                width=3,
+                rounds_f=7,
+                rounds_p=20,
+                mds_first_row=[Fp(1)] * 3,
+                round_constants=[Fp(1)] * 81,
+            )
+
 
 class TestPoseidon1Engine:
     """Tests for Poseidon1 engine."""
 
     def test_permute_wrong_state_length_too_short(self) -> None:
         """Raises error when input state is too short."""
         engine = Poseidon1(PARAMS_16)
-        with pytest.raises(ValueError, match="Input state must have length 16"):
+        with pytest.raises(ValueError, match=r"^Input state must have length 16$"):
             engine.permute([Fp(1)] * 10)
 
     def test_permute_wrong_state_length_too_long(self) -> None:
         """Raises error when input state is too long."""
         engine = Poseidon1(PARAMS_16)
-        with pytest.raises(ValueError, match="Input state must have length 16"):
+        with pytest.raises(ValueError, match=r"^Input state must have length 16$"):
             engine.permute([Fp(1)] * 20)
 
     def test_permute_determinism(self) -> None:
@@ -160,3 +233,37 @@ def test_permute_output_differs_from_input(self) -> None:
         output = engine.permute(state)
 
         assert output != state
+
+    @pytest.mark.parametrize(
+        "params, input_state",
+        [
+            (PARAMS_16, INPUT_16),
+            (PARAMS_24, INPUT_24),
+        ],
+        ids=["width_16", "width_24"],
+    )
+    def test_permute_output_in_field(self, params: Poseidon1Params, input_state: list[Fp]) -> None:
+        """Every output element lies strictly below the field modulus."""
+        engine = Poseidon1(params)
+
+        output = engine.permute(input_state)
+
+        assert all(int(x) < P for x in output)
+
+    def test_permute_all_zero_input(self) -> None:
+        """All-zero input produces an in-field output of the expected width."""
+        engine = Poseidon1(PARAMS_16)
+
+        output = engine.permute([Fp(0)] * 16)
+
+        assert len(output) == 16
+        assert all(int(x) < P for x in output)
+
+    def test_permute_field_boundary_input(self) -> None:
+        """Maximum-value input stays in-field and exposes int64 regressions."""
+        engine = Poseidon1(PARAMS_16)
+
+        output = engine.permute([Fp(value=P - 1)] * 16)
+
+        assert len(output) == 16
+        assert all(int(x) < P for x in output)