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

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

Generates random rho seeds, calls the Python reference sampleNTT to compute
expected coefficients, and writes hex files for Verilator simulation.

Matches the Python reference (sample.py / SHA_3.py) bit-exactly.
"""

import os
import random
import sys
import hashlib

# Add the Python reference implementation to path
_REF_DIR = os.path.expanduser(
    "~/Dev/server_code/python_project/PQC_2025/A_ML_KEM_v0"
)
if _REF_DIR not in sys.path:
    sys.path.insert(0, _REF_DIR)

import utils
import sample as sample_ref

# 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


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)

        # Generate random 32-byte rho (as binary string, matching the reference)
        rho_bin = utils.random_Generator(8 * 32)  # 256-bit binary string

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

        # Build the 34-byte input for sampleNTT: rho || j || i
        # Each component is a binary string (LSB at index 0)
        s_j_bin = utils.dec_to_binary_little_endian(j)  # 8-bit binary string
        s_i_bin = utils.dec_to_binary_little_endian(i)  # 8-bit binary string
        B = rho_bin + s_j_bin + s_i_bin  # 272-bit binary string

        # Compute expected coefficients using Python reference
        coeffs = sample_ref.sampleNTT(B)  # numpy array of 256 ints

        # Convert rho to hex for the input file
        # binary_to_hex_little_endian produces MSB-first hex per byte
        rho_hex = utils.binary_to_hex_little_endian(rho_bin)

        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")