mlkem-sync/test_framework/modules/sample_ntt/gen_vectors.py

"""gen_vectors.py - Test vector generator for sample_ntt module.

Generates random rho seeds, computes expected coefficients with a standard
FIPS 202 SHAKE-128 rejection sampler (hashlib), and writes hex files for
Verilator simulation.

The expected coefficients are FIPS-203-conformant (Algorithm 7 SampleNTT):
SHAKE-128(rho || j || i), consume the full rate, 12-bit rejection sampling.
"""

import os
import random
import sys
import hashlib

# Add test_framework/lib to path for VectorGenerator base class
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "lib"))
from vector_gen import VectorGenerator


def _sample_ntt_ref(seed: bytes) -> list[int]:
    """FIPS 203 Algorithm 7 SampleNTT using standard SHAKE-128.

    Args:
        seed: 34-byte input = rho(32) || j(1) || i(1).

    Returns:
        List of 256 coefficients in [0, 3329).
    """
    # 1344 bytes = 8 full SHAKE-128 rate blocks; ample for 256 accepts.
    stream = hashlib.shake_128(seed).digest(1344)
    coeffs = []
    off = 0
    while len(coeffs) < 256:
        c0, c1, c2 = stream[off], stream[off + 1], stream[off + 2]
        off += 3
        d1 = c0 | ((c1 & 0x0F) << 8)
        d2 = (c1 >> 4) | (c2 << 4)
        if d1 < 3329:
            coeffs.append(d1)
        if d2 < 3329 and len(coeffs) < 256:
            coeffs.append(d2)
    return coeffs


class SampleNTTVectorGenerator(VectorGenerator):
    """Generates test vectors for the sample_ntt_sync module."""

    def generate_one(self, params: dict) -> dict:
        """Generate a single test vector.

        Args:
            params: dict with 'k' key (ML-KEM parameter, 2/3/4).

        Returns:
            dict with 'input' and 'expected' keys.
        """
        k = params.get("k", 2)

        # Random 32-byte rho
        rho = bytes(random.randint(0, 255) for _ in range(32))

        # Choose random (i, j) indices within [0, k-1]
        i = random.randint(0, k - 1)
        j = random.randint(0, k - 1)

        # 34-byte SampleNTT input: rho || j || i (FIPS 203 Alg 13/7)
        seed = rho + bytes([j]) + bytes([i])

        # Expected coefficients from a standard SHAKE-128 sampler
        coeffs = _sample_ntt_ref(seed)

        # rho hex: byte 0 first (matches TB parse_rho byte ordering)
        rho_hex = rho.hex().upper()

        return {
            "input": {
                "rho_hex": rho_hex,
                "k": k,
                "i": i,
                "j": j,
            },
            "expected": {
                "coeffs": list(coeffs),  # 256 coefficients (0 <= c < 3329)
            },
        }

    def write_hex_file(self, vectors: list[dict], filepath: str) -> None:
        """Write input vectors as "RHO_HEX K_HEX I_HEX J_HEX" per line.

        Args:
            vectors: List of vector dicts from generate_one().
            filepath: Path to write the hex file.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, "w") as f:
            for v in vectors:
                inp = v["input"]
                rho_hex = inp["rho_hex"]
                k_hex = format(inp["k"], "X")
                i_hex = format(inp["i"], "X")
                j_hex = format(inp["j"], "X")
                f.write(f"{rho_hex} {k_hex} {i_hex} {j_hex}\n")

    def write_expected_file(self, vectors: list[dict], filepath: str) -> None:
        """Write expected outputs: one coefficient per line (12-bit hex).

        Args:
            vectors: List of vector dicts from generate_one().
            filepath: Path to write the expected hex file.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, "w") as f:
            for v in vectors:
                coeffs = v["expected"]["coeffs"]
                for c in coeffs:
                    f.write(f"{c:03X}\n")

    def compare_results(self, got: list[str], expected_file: str) -> bool:
        """Compare RTL output against expected values.

        Args:
            got: List of hex result strings from simulation ("RESULT: XXX").
            expected_file: Path to expected hex file.

        Returns:
            bool: True if all results match.
        """
        with open(expected_file, "r") as f:
            expected = [
                line.strip()
                for line in f
                if line.strip() and not line.startswith("#")
            ]

        if len(got) != len(expected):
            print(
                f"  MISMATCH: got {len(got)} results, expected {len(expected)}"
            )
            return False

        all_ok = True
        for i, (g, e) in enumerate(zip(got, expected)):
            if g.upper() != e.upper():
                print(f"  MISMATCH[{i}]: got={g}, expected={e}")
                all_ok = False
                if i >= 9:  # Stop after 10 mismatches
                    print(f"  ... too many mismatches, stopping comparison")
                    break

        return all_ok


if __name__ == "__main__":
    # Quick smoke test when run directly
    gen = SampleNTTVectorGenerator()
    for k_val in [2, 4]:
        vec = gen.generate_one({"k": k_val})
        rho = vec["input"]["rho_hex"]
        i = vec["input"]["i"]
        j = vec["input"]["j"]
        coeffs = vec["expected"]["coeffs"]
        print(f"k={k_val}, i={i}, j={j}: {len(coeffs)} coefficients")
        print(f"  rho[0:8]={rho[:8]}...")
        print(f"  coeffs[0:4]={coeffs[:4]}")
        # Verify all coefficients < Q
        assert all(0 <= c < 3329 for c in coeffs), "Coefficient out of range!"
    print("Smoke test PASSED")