subspec: implement Poseidon2 spec with tests

src/lean_spec/subspecs/poseidon2/__init__.py

@@ -1,4 +1,13 @@
-"""
-Specifications for the Poseidon2 cryptographic permutation for
-zero-knowledge applications.
-"""
+"""Specification for the Poseidon2 permutation."""
+
+from .permutation import (
+    PARAMS_16,
+    PARAMS_24,
+    permute,
+)
+
+__all__ = [
+    "permute",
+    "PARAMS_16",
+    "PARAMS_24",
+]

src/lean_spec/subspecs/poseidon2/permutation.py

@@ -0,0 +1,325 @@
+"""
+A minimal Python specification for the Poseidon2 permutation.
+
+The design is based on the paper "Poseidon2: A Faster Version of the Poseidon
+Hash Function" (https://eprint.iacr.org/2023/323).
+"""
+
+from itertools import chain
+from typing import List, NamedTuple
+
+from ..koalabear.field import Fp
+
+# =================================================================
+# Poseidon2 Parameter Definitions
+# =================================================================
+
+S_BOX_DEGREE = 3
+"""
+The S-box exponent `d`.
+
+For fields where `gcd(d, p-1) = 1`, `x -> x^d` is a permutation.
+
+For KoalaBear, `d=3` is chosen for its low degree.
+"""
+
+
+class Poseidon2Params(NamedTuple):
+    """
+    Encapsulates all necessary parameters for a specific Poseidon2 instance.
+
+    This structure holds the configuration for a given state width, including
+    the number of rounds and the constants for the internal linear layer.
+
+    Attributes:
+        WIDTH (int): The size of the state (t).
+
+        ROUNDS_F (int): The total number of "full" rounds, where the S-box is
+        applied to the entire state.
+
+        ROUNDS_P (int): The number of "partial" rounds, where the S-box is
+        applied to only the first element of the state.
+
+        INTERNAL_DIAG_VECTORS (List[Fp]): The diagonal vectors for the
+        efficient internal linear layer matrix (M_I).
+    """
+
+    WIDTH: int
+    ROUNDS_F: int
+    ROUNDS_P: int
+    INTERNAL_DIAG_VECTORS: List[Fp]
+
+
+def _generate_round_constants(params: Poseidon2Params) -> List[Fp]:
+    """
+    Generates a deterministic list of round constants for the permutation.
+
+    Round constants are added in each round to break symmetries and prevent
+    attacks like slide or interpolation attacks.
+
+    Args:
+        params: The object defining the permutation's configuration.
+
+    Returns:
+        A list of Fp elements to be used as round constants.
+    """
+    # The total number of constants needed for the entire permutation.
+    #
+    # This is the sum of constants for all full rounds and all partial rounds.
+    #   - Full rounds require `WIDTH` constants each
+    #   (one for each state element).
+    #   - Partial rounds require 1 constant each
+    #   (for the first state element).
+    total_constants = (params.ROUNDS_F * params.WIDTH) + params.ROUNDS_P
+
+    # For the specification, we generate the constants as a deterministic d
+    # sequence of integers.
+    #
+    # This is sufficient to define the algorithm's mechanics.
+    #
+    # Real-world implementations would use constants generated from a secure,
+    # pseudo-random source.
+    return [Fp(value=i) for i in range(total_constants)]
+
+
+# Parameters for WIDTH = 16
+PARAMS_16 = Poseidon2Params(
+    WIDTH=16,
+    ROUNDS_F=8,
+    ROUNDS_P=20,
+    INTERNAL_DIAG_VECTORS=[
+        Fp(value=-2),
+        Fp(value=1),
+        Fp(value=2),
+        Fp(value=1) / Fp(value=2),
+        Fp(value=3),
+        Fp(value=4),
+        Fp(value=-1) / Fp(value=2),
+        Fp(value=-3),
+        Fp(value=-4),
+        Fp(value=1) / Fp(value=2**8),
+        Fp(value=1) / Fp(value=8),
+        Fp(value=1) / Fp(value=2**24),
+        Fp(value=-1) / Fp(value=2**8),
+        Fp(value=-1) / Fp(value=8),
+        Fp(value=-1) / Fp(value=16),
+        Fp(value=-1) / Fp(value=2**24),
+    ],
+)
+
+# Parameters for WIDTH = 24
+PARAMS_24 = Poseidon2Params(
+    WIDTH=24,
+    ROUNDS_F=8,
+    ROUNDS_P=23,
+    INTERNAL_DIAG_VECTORS=[
+        Fp(value=-2),
+        Fp(value=1),
+        Fp(value=2),
+        Fp(value=1) / Fp(value=2),
+        Fp(value=3),
+        Fp(value=4),
+        Fp(value=-1) / Fp(value=2),
+        Fp(value=-3),
+        Fp(value=-4),
+        Fp(value=1) / Fp(value=2**8),
+        Fp(value=1) / Fp(value=4),
+        Fp(value=1) / Fp(value=8),
+        Fp(value=1) / Fp(value=16),
+        Fp(value=1) / Fp(value=32),
+        Fp(value=1) / Fp(value=64),
+        Fp(value=1) / Fp(value=2**24),
+        Fp(value=-1) / Fp(value=2**8),
+        Fp(value=-1) / Fp(value=8),
+        Fp(value=-1) / Fp(value=16),
+        Fp(value=-1) / Fp(value=32),
+        Fp(value=-1) / Fp(value=64),
+        Fp(value=-1) / Fp(value=2**7),
+        Fp(value=-1) / Fp(value=2**9),
+        Fp(value=-1) / Fp(value=2**24),
+    ],
+)
+
+# Base 4x4 matrix, used in the external linear layer.
+M4_MATRIX = [
+    [Fp(value=2), Fp(value=3), Fp(value=1), Fp(value=1)],
+    [Fp(value=1), Fp(value=2), Fp(value=3), Fp(value=1)],
+    [Fp(value=1), Fp(value=1), Fp(value=2), Fp(value=3)],
+    [Fp(value=3), Fp(value=1), Fp(value=1), Fp(value=2)],
+]
+
+# =================================================================
+# Linear Layers
+# =================================================================
+
+
+def _apply_m4(chunk: List[Fp]) -> List[Fp]:
+    """
+    Applies the 4x4 M4 MDS matrix to a 4-element chunk of the state.
+    This is a helper function for the external linear layer.
+
+    Args:
+        chunk: A list of 4 Fp elements.
+
+    Returns:
+        The transformed 4-element chunk.
+    """
+    # Initialize the result vector with zeros.
+    result = [Fp(value=0)] * 4
+    # Perform standard matrix-vector multiplication.
+    for i in range(4):
+        for j in range(4):
+            result[i] += M4_MATRIX[i][j] * chunk[j]
+    return result
+
+
+def external_linear_layer(state: List[Fp], width: int) -> List[Fp]:
+    """
+    Applies the external linear layer (M_E).
+
+    This layer provides strong diffusion across the entire state and is used
+    in the full rounds. For a state of size t=4k, it's constructed from the
+    base M4 matrix to form a larger circulant-like matrix, which is efficient
+    while ensuring that a change in any single element affects all other
+    elements after application.
+
+    The process follows Appendix B of the paper.
+
+    Args:
+        state: The current state vector.
+        width: The width `t` of the state.
+
+    Returns:
+        The state vector after applying the external linear layer.
+    """
+    # Apply the M4 matrix to each 4-element chunk of the state.
+    #
+    # This provides strong local diffusion within each block.
+    state_after_m4 = list(
+        chain.from_iterable(
+            _apply_m4(state[i : i + 4]) for i in range(0, width, 4)
+        )
+    )
+
+    # Apply the outer circulant structure for global diffusion.
+    #
+    # We precompute the four sums of elements at the same offset in each chunk.
+    # For each k in 0..4:
+    #     sums[k] = state[k] + state[4 + k] + state[8 + k] + ... up to width
+    sums = [
+        sum((state_after_m4[j + k] for j in range(0, width, 4)), Fp(value=0))
+        for k in range(4)
+    ]
+
+    # Add the corresponding sum to each element of the state.
+    state_after_circulant = [
+        s + sums[i % 4] for i, s in enumerate(state_after_m4)
+    ]
+
+    return state_after_circulant
+
+
+def internal_linear_layer(
+    state: List[Fp], params: Poseidon2Params
+) -> List[Fp]:
+    """
+    Applies the internal linear layer (M_I).
+
+    This layer is used during partial rounds and is optimized for speed. Its
+    matrix is constructed as M_I = J + D, where J is the all-ones matrix and D
+    is a diagonal matrix. This structure allows the matrix-vector product to be
+    computed in O(t) time instead of O(t^2), as M_I * s = J*s + D*s.
+    The term J*s is a vector where each element is the sum of
+    all elements in s.
+
+    Args:
+        state: The current state vector.
+        params: The Poseidon2Params object containing the diagonal vectors.
+
+    Returns:
+        The state vector after applying the internal linear layer.
+    """
+    # Calculate the sum of all elements in the state vector.
+    s_sum = sum(state, Fp(value=0))
+    # For each element s_i, compute s_i' = d_i * s_i + sum(s).
+    # This is the efficient computation of (J + D)s.
+    new_state = [
+        s * d + s_sum
+        for s, d in zip(state, params.INTERNAL_DIAG_VECTORS, strict=False)
+    ]
+    return new_state
+
+
+# =================================================================
+# Core Permutation
+# =================================================================
+
+
+def permute(state: List[Fp], params: Poseidon2Params) -> List[Fp]:
+    """
+    Performs the full Poseidon2 permutation on the given state.
+
+    The permutation follows the structure:
+    Initial Layer -> Full Rounds -> Partial Rounds -> Full Rounds
+
+    Args:
+        state: A list of Fp elements representing the current state.
+        params: The object defining the permutation's configuration.
+
+    Returns:
+        The new state after applying the permutation.
+    """
+    # Ensure the input state has the correct dimensions.
+    if len(state) != params.WIDTH:
+        raise ValueError(f"Input state must have length {params.WIDTH}")
+
+    # Generate the deterministic round constants for this parameter set.
+    round_constants = _generate_round_constants(params)
+    # The number of full rounds is split between the beginning and end.
+    half_rounds_f = params.ROUNDS_F // 2
+    # Initialize index for accessing the flat list of round constants.
+    const_idx = 0
+
+    # 1. Initial Linear Layer
+    #
+    # Another linear layer is applied at the start to prevent certain algebraic
+    # attacks by ensuring the permutation begins with a diffusion layer.
+    state = external_linear_layer(list(state), params.WIDTH)
+
+    # 2. First Half of Full Rounds (R_F / 2)
+    for _r in range(half_rounds_f):
+        # Add round constants to the entire state.
+        state = [
+            s + round_constants[const_idx + i] for i, s in enumerate(state)
+        ]
+        const_idx += params.WIDTH
+        # Apply the S-box (x -> x^d) to the full state.
+        state = [s**S_BOX_DEGREE for s in state]
+        # Apply the external linear layer for diffusion.
+        state = external_linear_layer(state, params.WIDTH)
+
+    # 3. Partial Rounds (R_P)
+    for _r in range(params.ROUNDS_P):
+        # Add a single round constant to the first state element.
+        state[0] += round_constants[const_idx]
+        const_idx += 1
+        # Apply the S-box to the first state element only.
+        #
+        # This is the main optimization of the Hades design.
+        state[0] = state[0] ** S_BOX_DEGREE
+        # Apply the internal linear layer.
+        state = internal_linear_layer(state, params)
+
+    # 4. Second Half of Full Rounds (R_F / 2)
+    for _r in range(half_rounds_f):
+        # Add round constants to the entire state.
+        state = [
+            s + round_constants[const_idx + i] for i, s in enumerate(state)
+        ]
+        const_idx += params.WIDTH
+        # Apply the S-box to the full state.
+        state = [s**S_BOX_DEGREE for s in state]
+        # Apply the external linear layer for diffusion.
+        state = external_linear_layer(state, params.WIDTH)
+
+    return state