mldsa-native-rs 0.0.1-alpha.6

#!/usr/bin/env python3
# Copyright (c) The mlkem-native project authors
# Copyright (c) The mldsa-native project authors
# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT

import subprocess
import tempfile
import platform
import argparse
import shutil
import pathlib
import re
import sys
import threading
import pyparsing as pp
import os
import yaml
import time

from concurrent.futures import ThreadPoolExecutor
from functools import partial
from rich.console import Console
from rich.progress import (
    Progress,
    BarColumn,
    TextColumn,
    TaskProgressColumn,
    TimeElapsedColumn,
)

console = Console()

# Global progress bar - initialized in _main()
_progress = None
_main_task = None
_current_task = ""

modulus = 8380417
root_of_unity = 1753
montgomery_factor = pow(2, 32, modulus)

# Compiled regex patterns
_RE_DEFINED = re.compile(r"defined\(([^)]+)\)")
_RE_MARKDOWN_CITE = re.compile(r"\[\^(?P<id>\w+)\]")
_RE_C_CITE = re.compile(r"@\[(?P<id>\w+)")
_RE_BYTECODE_START = re.compile(
    r"=== bytecode start: (?:aarch64|x86_64)/mldsa/([^/\s]+?)\.o"
)
_RE_FUNC_SYMBOL = re.compile(r"MLD_ASM_FN_SYMBOL\((.*)\)")
_RE_MACRO_CHECK = re.compile(r"[^_]((?:MLD_|MLDSA_)\w+)(.*)$", re.M)
_RE_MLKEM_MACRO_CHECK = re.compile(r"[^_]((?:MLK_|MLKEM_)\w+)(.*)$", re.M)
_RE_DEFINE = re.compile(r"^\s*#define\s+(\w+)")
_RE_ARGS_COMMENT = re.compile(r"(.*?)(\s*//.*)?$")
_RE_MACRO_DEF = re.compile(r"^\s*\.macro\s+(\w+)")
_RE_MACRO_DEF_ARGS = re.compile(r"^(\s*\.macro\s+\w+)(\s+.*)$")
_RE_LEADING_SPACE = re.compile(r"^(\s*)")

# File cache: {filename: {"content": str, "original": str, "force_format": bool}}
# Caches content of files in preparation/modification to avoid repeated
# read/writes to the file system.
_file_cache = {}
_file_cache_lock = threading.Lock()

_errors = []
_errors_lock = threading.Lock()

_progress_lock = threading.Lock()


def read_file(filename, original=False):
    """Read file content, using cache if available"""
    with _file_cache_lock:
        if filename in _file_cache:
            key = "content" if original is False else "original"
            return _file_cache[filename][key]

        with open(filename, "r") as f:
            content = f.read()
        _file_cache[filename] = {
            "content": content,
            "original": content,
            "force_format": False,
        }
        return content


def update_file(filename, content, force_format=False):
    """Write file content to cache"""
    with _file_cache_lock:
        if filename not in _file_cache:
            try:
                with open(filename, "r") as f:
                    original = f.read()
                _file_cache[filename] = {"original": original}
            except FileNotFoundError:
                _file_cache[filename] = {"original": None}

        e = _file_cache[filename]
        e["content"] = content
        e["force_format"] = e.get("force_format", False) or force_format


def finalize_format_batch(batch):
    """Format a batch of files by passing to clang-format with -i flag"""
    if not batch:
        return

    # Create temp files for each filename in batch
    temp_files = []
    try:
        for filename in batch:
            content = read_file(filename)
            # Skip files scheduled for deletion
            if content is None:
                continue
            with tempfile.NamedTemporaryFile(mode="w", suffix=".c", delete=False) as f:
                f.write(content)
                temp_files.append((f.name, filename))

        # Call clang-format with -i to update files in-place
        clang_format_file = os.path.join(
            os.path.dirname(__file__), "..", ".clang-format"
        )
        p = subprocess.run(
            ["clang-format", "-i", f"-style=file:{clang_format_file}"]
            + [t[0] for t in temp_files],
            capture_output=True,
            text=True,
        )
        if p.returncode != 0:
            print(p.stderr)
            print(
                f"Failed to auto-format autogenerated code (clang-format return code {p.returncode}). Are you running in a nix shell? See CONTRIBUTING.md."
            )
            exit(1)

        # Read formatted files back and update cache
        for temp_path, filename in temp_files:
            with open(temp_path, "r") as f:
                update_file(filename, f.read())
    finally:
        for temp_path, _ in temp_files:
            os.unlink(temp_path)


def finalize_file(item, dry_run):
    """Write a single file or delete it if content is None"""
    filename, data = item

    content_old = data["original"]
    content_new = data["content"]

    if content_old == content_new:
        return

    # Handle deletion (content_new is None)
    if content_new is None:
        if dry_run is False:
            file_updated(filename, removed=True)
            os.remove(filename)
        else:
            error(filename, None)
        return

    if dry_run is False:
        file_updated(filename)
        with open(filename, "w") as f:
            f.write(content_new)
    else:
        filename_new = f"{filename}.new"
        with open(filename_new, "w") as f:
            f.write(content_new)
        error(filename, filename_new)


def format_files(dry_run):
    """Apply formatting to files"""
    to_format = [
        filename
        for filename, data in _file_cache.items()
        if data["force_format"] or filename.endswith((".c", ".h", ".i"))
    ]

    # Group into batches of max 20
    batch_size = 20
    batches = [
        to_format[i : i + batch_size] for i in range(0, len(to_format), batch_size)
    ]

    run_parallel(batches, finalize_format_batch)


def finalize(dry_run):
    """Write dirty files to filesystem"""
    run_parallel(_file_cache.items(), partial(finalize_file, dry_run=dry_run))


# This file re-generated auto-generated source files in mldsa-native.
#
# It currently covers:
# - zeta values for the reference NTT and invNTT
# - lookup tables used for fast rejection sampling
# - source files for monolithic single-CU build
# - simplified assembly sources
# - header guards
# - #undef's for CU-local macros

_step_start_time = time.time()


def high_level_task(msg):
    """Set the current high-level task description"""
    global _current_task
    _current_task = msg
    if _progress:
        _progress.update(_main_task, description=f"[cyan]{msg}[/]")


def high_level_status(msg, skipped=False):
    """Complete a high-level step and print status"""
    global _step_start_time
    elapsed = time.time() - _step_start_time
    if skipped:
        symbol = "[dim]–[/dim]"
    else:
        symbol = "[green]✓[/green]"
    if _progress:
        _progress.print(f"{symbol} {msg} ({elapsed:.1f}s)", highlight=False)
        _progress.advance(_main_task)
    else:
        console.print(f"{symbol} {msg} ({elapsed:.1f}s)", highlight=False)
    _step_start_time = time.time()


def run_parallel(files, func):
    """Run func over files in parallel with progress tracking"""
    if not files:
        return []

    files = list(files)
    total = len(files)
    state = {"completed": 0, "last_file": ""}

    def update_progress():
        if _progress and total > 0:
            suffix = (
                f" {os.path.basename(state['last_file'])}" if state["last_file"] else ""
            )
            _progress.update(
                _main_task,
                description=f"[cyan]{_current_task}[/] [dim][{state['completed']}/{total}]{suffix}[/]",
            )

    def wrapped(f):
        result = func(f)
        with _progress_lock:
            state["completed"] += 1
            state["last_file"] = str(f[0]) if isinstance(f, tuple) else str(f)
            update_progress()
        return result

    with ThreadPoolExecutor() as executor:
        return list(executor.map(wrapped, files))


def error(filename, filename_new):
    with _errors_lock:
        _errors.append((filename, filename_new))


def print_check_errors():
    for filename, filename_new in _errors:
        console.print(f"[red]error[/] {filename}")
        if filename_new is not None:
            console.print(
                f"Autogenerated file {filename} needs updating. Have you called scripts/autogen? Wrote new version to {filename_new}."
            )
            if os.path.exists(filename):
                subprocess.run(["diff", filename, filename_new])
        else:
            console.print(
                f"Autogenerated file {filename} needs removing. Have you called scripts/autogen?"
            )

    return len(_errors) == 0


def file_updated(filename, removed=False):
    if removed is False:
        console.print(f"[bold]updated {filename}[/]")
    else:
        console.print(f"[bold]removed {filename}[/]")


def gen_autogen_warning():
    yield ""
    yield "/*"
    yield " * WARNING: This file is auto-generated from scripts/autogen"
    yield " *          in the mldsa-native repository."
    yield " *          Do not modify it directly."
    yield " */"


def gen_header():
    yield "/*"
    yield " * Copyright (c) The mldsa-native project authors"
    yield " * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT"
    yield " */"
    yield from gen_autogen_warning()
    yield ""


def gen_hol_light_header():
    yield "(*"
    yield " * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved."
    yield " * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0"
    yield " *)"
    yield ""
    yield "(*"
    yield " * WARNING: This file is auto-generated from scripts/autogen"
    yield " *          in the mldsa-native repository."
    yield " *          Do not modify it directly."
    yield " *)"
    yield ""


def gen_yaml_header():
    yield "# Copyright (c) The mldsa-native project authors"
    yield "# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT"
    yield ""


def gen_yaml_autogen_warning():
    yield "# WARNING: This file is auto-generated from scripts/autogen"
    yield "#          in the mldsa-native repository."
    yield "#          Do not modify it directly."


def format_content(content):
    clang_format_file = os.path.join(os.path.dirname(__file__), "..", ".clang-format")
    p = subprocess.run(
        ["clang-format", f"-style=file:{clang_format_file}"],
        capture_output=True,
        input=content,
        text=True,
    )
    if p.returncode != 0:
        print(p.stderr)
        print(
            f"Failed to auto-format autogenerated code (clang-format return code {p.returncode}). Are you running in a nix shell? See CONTRIBUTING.md."
        )
        exit(1)
    return p.stdout


class CondParser:
    """Rudimentary parser for expressions if `#if .. #else ..` directives"""

    def __init__(self):
        c_identifier = pp.common.identifier()
        c_integer_suffix = pp.one_of("U L LU UL LL ULL LLU", caseless=True)
        c_dec_integer = pp.Combine(
            pp.Optional(pp.one_of("+ -"))
            + pp.Word(pp.nums)
            + pp.Optional(c_integer_suffix)
        )
        c_hex_integer = pp.Combine(
            pp.Literal("0x") + pp.Word(pp.hexnums) + pp.Optional(c_integer_suffix)
        )

        self.parser = pp.infix_notation(
            c_identifier | c_hex_integer | c_dec_integer,
            [
                (pp.one_of("!"), 1, pp.opAssoc.RIGHT),
                (pp.one_of("!= == <= >= > <"), 2, pp.opAssoc.LEFT),
                (pp.one_of("&&"), 2, pp.opAssoc.LEFT),
                (pp.one_of("||"), 2, pp.opAssoc.LEFT),
            ],
        )

    @staticmethod
    def connective(res):
        """Extract the top-level connective for the expression"""
        if not isinstance(res, list):
            return None
        elif len(res) == 2:
            # Unary operator (will be "!" in our case)
            return res[0]
        else:
            # Binary operator
            return res[1]

    @staticmethod
    def map_top(f, res):
        """Apply function to arguments of top-level connective"""
        if not isinstance(res, list):
            return res
        else:
            # We expect `f` to do nothing on strings, so it is safe
            # to apply it everywhere, including the connectives.
            return list(map(f, res))

    @staticmethod
    def args(res):
        """Assuming the argument is a binary operation, return all arguments"""
        return res[::2]

    @staticmethod
    def simplify_double_negation(res):
        """Cancel double negations"""
        if CondParser.connective(res) == "!" and CondParser.connective(res[1]) == "!":
            res = res[1][1]
        res = CondParser.map_top(CondParser.simplify_double_negation, res)
        return res

    @staticmethod
    def simplify_not_eq(res):
        """Replace !(x == y) by x != y, and !(x != y) by x == y"""
        if CondParser.connective(res) == "!" and CondParser.connective(res[1]) == "==":
            res = res[1]
            res[1] = "!="
        if CondParser.connective(res) == "!" and CondParser.connective(res[1]) == "!=":
            res = res[1]
            res[1] = "=="
        res = CondParser.map_top(CondParser.simplify_not_eq, res)
        return res

    @staticmethod
    def simplify_neq_chain(res):
        """Check for &&-chains of inequalities followed by an equality
        which implies the inequality. This catches patterns like
        ```
         #if MLKEM_K == 2
         ...
         #elif MLKEM_K == 3
         ...
         #elif MLKEM_K == 4
         ...
         #endif
        ```
        """
        if (
            CondParser.connective(res) == "&&"
            and CondParser.connective(res[-1]) == "=="
        ):
            lhs = res[-1][0]
            rhs = res[-1][2]
            args = []
            for a in CondParser.args(res[:-1]):
                if CondParser.connective(a) == "!=" and a[0] == lhs:
                    args.append(a[2])
                else:
                    args = None
                    break
            if args is None:
                return res
            # Check if all args are numerical and different
            if rhs.isdigit() and all(
                map(lambda a: a.isdigit() and int(a) != int(rhs), args)
            ):
                # Success -- just drop all but the final condition
                return res[-1]
        res = CondParser.map_top(CondParser.simplify_neq_chain, res)
        return res

    @staticmethod
    def print_exp(exp, inner=False):
        conn = CondParser.connective(exp)
        if conn is None:
            return exp
        elif conn == "!":
            res = f"!{CondParser.print_exp(exp[1], inner=True)}"
        else:
            padded_conn = f" {conn} "
            res = padded_conn.join(
                map(lambda e: CondParser.print_exp(e, inner=True), CondParser.args(exp))
            )
        if inner is True and conn in ["&&", "||"]:
            res = f"({res})"
        return res

    def simplify_assoc(exp):
        """Check for unnecesary bracketing and remove it"""
        conn = CondParser.connective(exp)
        if conn in ["&&", "||"]:
            args = CondParser.args(exp)
            new_args = []
            for a in args:
                if CondParser.connective(a) == conn:
                    new_args += CondParser.args(a)
                else:
                    new_args.append(a)
            exp = [x for y in map(lambda x: [x, conn], new_args) for x in y][:-1]
        exp = CondParser.map_top(CondParser.simplify_assoc, exp)
        return exp

    def simplify_all(exp):
        exp = CondParser.simplify_double_negation(exp)
        exp = CondParser.simplify_not_eq(exp)
        exp = CondParser.simplify_neq_chain(exp)
        exp = CondParser.simplify_assoc(exp)
        return exp

    def parse_condition(self, exp, simplify=True):
        try:
            exp = self.parser.parse_string(exp, parse_all=True).as_list()[0]
        except pp.ParseException:
            print(f"WARNING: Ignoring condition '{exp}' I cannot parse")
            return exp
        if simplify is True:
            exp = CondParser.simplify_all(exp)
        return exp

    def normalize_condition(self, exp):
        return CondParser.print_exp(self.parse_condition(exp))


def adjust_preprocessor_comments_for_filename(
    content, source_file, parser, show_status=False
):
    """Automatically add comments to large `#if ... #else ... #endif`
    blocks indicating the guarding conditions.

    For example, a block

    ```c
      #if FOO
      ...
      #else
      ...
      #endif
    ```

    will be transformed into


    ```c
      #if FOO
      ...
      #else /* FOO */
      ...
      #endif /* !FOO */
    ```

    except when the distance between the preprocessor directives is
    very short, and the annotations would be more harmful than useful.

    ```
    """
    content = content.split("\n")
    new_content = []

    # Stack of `#if` statements. Every entry is a tuple
    # `(conds, line_no, if_or_else, has_children)`, where
    # - `conds` is the list of conditions being tested.
    #   In a normal `#if ... #else ...` braach, this is a singleton list
    #   containing the condition being tested. In a chain of
    #   `#if .. #elif ..` it contains all conditions encountered to this point.
    # - `line_no` is the line where it started
    # - `if_or_else` indicates whether we are in the `#if`
    #   or the `#else` branch (if present)
    # - `force_print` indicates if a comment should be omitted
    if_stack = []

    def merge_escaped_lines(line, i):
        while line.endswith("\\"):
            line = line.removesuffix("\\").rstrip() + content[i + 1].lstrip()
            i = i + 1
        return (line, i)

    def merge_commented_lines(line, i):
        # Not very robust, but good enough
        if "/*" not in line or "*/" in line:
            return (line, i)
        i += 1
        while "*/" not in content[i]:
            line += content[i]
            i += 1

        line += content[i]
        return (line, i)

    def should_print(cur_line_no, conds, line_no, force_print):
        line_threshold = 5
        if force_print is True:
            return True

        if cur_line_no - line_no >= line_threshold:
            return True
        return False

    def format_condition(cond):
        cond = _RE_DEFINED.sub(r"\1", cond)
        return parser.normalize_condition(cond)

    def format_conditions(conds, branch):
        prev_conds = list(map(lambda s: f"!({s})", conds[:-1]))
        final_cond = conds[-1]
        if branch is False:
            final_cond = f"!({final_cond})"
        full_cond = "&&".join(prev_conds + [final_cond])
        return format_condition(full_cond)

    def wrap_long_directive(directive, condition, max_len=80):
        """Manually wrap long preprocessor comment lines without subprocess overhead"""
        single_line = directive + " " + condition
        if len(single_line) <= max_len:
            return single_line

        # Wrap condition across multiple lines with backslash continuation
        words = condition.split()
        lines = []
        current = f"{directive} "
        indent = (len(directive) + 4) * " "
        indent_final = (len(directive) + 2) * " "

        for word in words:
            if len(current) + len(word) + 1 <= max_len:
                current += word + " "
            else:
                lines.append(current.rstrip() + " \\")
                if word == "*/":
                    current = indent_final + word
                else:
                    current = indent + word + " "

        lines.append(current.rstrip())

        return "\n".join(lines)

    def adhoc_format(directive, content):
        # .c and .h files are formatted as a whole
        if not source_file.endswith(".S"):
            return directive + " /* " + content + " */"
        # For .S files, manually wrap long lines
        return wrap_long_directive(directive, "/* " + content + " */")

    i = 0
    while i < len(content):
        line = content[i].strip()
        # Replace #ifdef by #if defined(...)
        if line.startswith("#ifdef "):
            line = "#if defined(" + line.removeprefix("#ifdef").strip() + ")"
        if line.startswith("#ifndef "):
            line = "#if !defined(" + line.removeprefix("#ifndef").strip() + ")"
        if line.startswith("#if"):
            line, _ = merge_escaped_lines(line, i)
            cond = line.removeprefix("#if")
            if_stack.append(([cond], i, True, False))
            new_content.append(content[i])
        elif line.startswith("#elif"):
            conds, _, _, force_print = if_stack.pop()
            line, _ = merge_escaped_lines(line, i)
            conds.append(line.removeprefix("#elif"))
            if_stack.append((conds, i, True, force_print))
            new_content.append(content[i])
        elif line.startswith("#else"):
            line, i = merge_escaped_lines(line, i)
            _, i = merge_commented_lines(line, i)
            conds, j, branch, force_print = if_stack.pop()
            assert branch is True
            print_else = should_print(i, cond, j, force_print)
            if_stack.append((conds, i, False, print_else))
            if print_else is True:
                cond = format_conditions(conds, True)
                new_content.append(adhoc_format("#else", cond))
            else:
                new_content.append("#else")
        elif line.startswith("#endif"):
            line, i = merge_escaped_lines(line, i)
            _, i = merge_commented_lines(line, i)
            conds, j, branch, force_print = if_stack.pop()
            print_endif = should_print(i, conds, j, force_print)
            if print_endif is False:
                new_content.append("#endif")
            else:
                cond = format_conditions(conds, branch)
                new_content.append(adhoc_format("#endif", cond))
        else:
            # Skip over multiline comments -- we don't want to
            # handle `#if ...` inside documentation as this would
            # lead to nested `/* ... */`.
            i_old = i
            _, i = merge_commented_lines(line, i_old)
            new_content += content[i_old : i + 1]
        i += 1

    return "\n".join(new_content)


def gen_preprocessor_comments_for(parser, source_file):
    content = read_file(source_file)
    new_content = adjust_preprocessor_comments_for_filename(
        content, source_file, parser, show_status=True
    )
    update_file(source_file, new_content)


def gen_preprocessor_comments():
    files = get_c_source_files() + get_asm_source_files() + get_header_files()
    parser = CondParser()
    run_parallel(files, partial(gen_preprocessor_comments_for, parser))


def bitreverse(i, n):
    r = 0
    for _ in range(n):
        r = 2 * r + (i & 1)
        i >>= 1
    return r


def signed_reduce(a):
    """Return signed canonical representative of a mod b"""
    c = a % modulus
    if c >= modulus / 2:
        c -= modulus
    return c


def gen_c_zetas():
    """Generate source and header file for zeta values used in
    the reference NTT and invNTT"""

    # The zeta values are the powers of the chosen root of unity (17),
    # converted to Montgomery form.

    zeta = [0]  # First entry is unused and set to 0
    for i in range(1, 256):
        zeta.append(signed_reduce(pow(root_of_unity, i, modulus) * montgomery_factor))

    # The source code stores the zeta table in bit reversed form
    yield from (zeta[bitreverse(i, 8)] for i in range(256))


def gen_c_zeta_file():
    def gen():
        yield from gen_header()
        yield ""
        yield "/*"
        yield " * Table of zeta values used in the reference NTT and inverse NTT."
        yield " * See autogen for details."
        yield " */"
        yield "static const int32_t mld_zetas[MLDSA_N] = {"
        yield from map(lambda t: str(t) + ",", gen_c_zetas())
        yield "};"
        yield ""

    update_file("mldsa/src/zetas.inc", "\n".join(gen()), force_format=True)


def prepare_root_for_barrett(root):
    """Takes a constant that the code needs to Barrett-multiply with,
    and returns the pair of (a) its signed canonical form, (b) the
    twisted constant used in the high-mul part of the Barrett multiplication."""

    # Signed canonical reduction
    root = signed_reduce(root)

    def round_to_even(t):
        rt = round(t)
        if rt % 2 == 0:
            return rt
        # Make sure to pick a rounding target
        # that's <= 1 away from x in absolute value.
        if rt <= t:
            return rt + 1
        return rt - 1

    root_twisted = round_to_even((root * 2**32) / modulus) // 2
    return root, root_twisted


def gen_aarch64_root_of_unity_for_block(layer, block, inv=False, scale=False):
    # We are computing a negacyclic NTT; the twiddles needed here is
    # the second half of the twiddles for a cyclic NTT of twice the size.
    # For ease of calculating the roots, layers are numbers 0 through 7
    # in this function.
    log = bitreverse(pow(2, layer) + block, 8)
    if inv is True:
        log = -log
    root = pow(root_of_unity, log, modulus)

    if scale is True:
        # Integrate scaling by 2**(-8) and Montgomery factor 2**32 into twiddle
        root = root * pow(2, 32 - 8, modulus)

    root, root_twisted = prepare_root_for_barrett(root)
    return root, root_twisted


def gen_aarch64_fwd_ntt_zetas_layer123456():
    # Layers 1,2,3 are merged
    yield from gen_aarch64_root_of_unity_for_block(0, 0)
    yield from gen_aarch64_root_of_unity_for_block(1, 0)
    yield from gen_aarch64_root_of_unity_for_block(1, 1)
    yield from gen_aarch64_root_of_unity_for_block(2, 0)
    yield from gen_aarch64_root_of_unity_for_block(2, 1)
    yield from gen_aarch64_root_of_unity_for_block(2, 2)
    yield from gen_aarch64_root_of_unity_for_block(2, 3)
    yield from (0, 0)  # Padding

    # Layers 4,5,6 are merged
    for block in range(8):  # There are 8 blocks in Layer 4
        yield from gen_aarch64_root_of_unity_for_block(3, block)
        yield from gen_aarch64_root_of_unity_for_block(4, 2 * block + 0)
        yield from gen_aarch64_root_of_unity_for_block(4, 2 * block + 1)
        yield from gen_aarch64_root_of_unity_for_block(5, 4 * block + 0)
        yield from gen_aarch64_root_of_unity_for_block(5, 4 * block + 1)
        yield from gen_aarch64_root_of_unity_for_block(5, 4 * block + 2)
        yield from gen_aarch64_root_of_unity_for_block(5, 4 * block + 3)
        yield from (0, 0)  # Padding


def gen_aarch64_fwd_ntt_zetas_layer78():
    # Layers 4,5,6,7,8 are merged, but we emit roots for 4,5,6
    # in separate arrays than those for 7,8
    for block in range(8):
        # Ordering of blocks is adjusted to suit the transposed internal
        # presentation of the data

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 0)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 1)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 2)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 3)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 0)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 2)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 4)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 6)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 1)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 3)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 5)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 7)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 0 + 4)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 1 + 4)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 2 + 4)[i]
            yield gen_aarch64_root_of_unity_for_block(6, 8 * block + 3 + 4)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 0 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 2 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 4 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 6 + 8)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 1 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 3 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 5 + 8)[i]
            yield gen_aarch64_root_of_unity_for_block(7, 16 * block + 7 + 8)[i]


def gen_aarch64_intt_zetas_layer78():
    for block in range(16):
        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(6, block * 4 + 0, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(6, block * 4 + 1, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(6, block * 4 + 2, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(6, block * 4 + 3, inv=True)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 0, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 2, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 4, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 6, inv=True)[i]

        for i in range(2):
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 1, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 3, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 5, inv=True)[i]
            yield gen_aarch64_root_of_unity_for_block(7, block * 8 + 7, inv=True)[i]


def gen_aarch64_intt_zetas_layer123456():
    for i in range(16):
        yield from gen_aarch64_root_of_unity_for_block(4, i, inv=True)
        yield from (0, 0)  # Padding
        yield from gen_aarch64_root_of_unity_for_block(5, i * 2, inv=True)
        yield from gen_aarch64_root_of_unity_for_block(5, i * 2 + 1, inv=True)

    # The last layer has the scaling by 1/256 integrated in the twiddle
    yield from gen_aarch64_root_of_unity_for_block(0, 0, inv=True, scale=True)

    yield from gen_aarch64_root_of_unity_for_block(1, 0, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(1, 1, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(2, 0, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(2, 1, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(2, 2, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(2, 3, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 0, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 1, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 2, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 3, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 4, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 5, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 6, inv=True)
    yield from gen_aarch64_root_of_unity_for_block(3, 7, inv=True)
    yield from (0, 0)  # Padding


def print_hol_light_array(g, as_int=True, entries_per_line=8, pad=0):
    # Format of integer list entries, including `;` separator:
    # - Positive numbers: &42;
    # - Negative numbers: -- &42;
    # If as_int is false, we omit `&` and emit constant as numerals.
    def format_hol_light_int(n):
        prefix = ""
        if n < 0:
            prefix = "-- "
            n = -n
        c = "&" if as_int is True else ""
        return f"{prefix}{c}{n:>{pad}};"

    items = list(map(format_hol_light_int, g))
    # Remove `;` from end of last entry
    items[-1] = items[-1][:-1]

    for i in range(0, len(items), entries_per_line):
        yield "  " + " ".join(items[i : i + entries_per_line])


def emit_c_array(type_str, name, data, formatter=None, suffix=""):
    """Emit a C array definition with explicit size.

    Yields lines for:
        MLD_ALIGN MLD_INTERNAL_DATA_DEFINITION <type>
            <name>[<count>]<suffix> = { <body> };

    data: elements (list or generator; materialized internally).
          If elements are lists/tuples (row-oriented data), count is the
          total number of scalar elements across all rows.
    formatter: callable(data) -> iterable of body lines.
              Default: one "element," per line.
    suffix: extra array dimensions (e.g., "[8]" for 2D arrays).
    """
    data = list(data)
    if data and isinstance(data[0], (list, tuple)):
        count = sum(len(row) for row in data)
    else:
        count = len(data)
    yield f"MLD_ALIGN MLD_INTERNAL_DATA_DEFINITION {type_str}"
    yield f"    {name}[{count}]{suffix} = {{"
    if formatter is not None:
        yield from formatter(data)
    else:
        yield from (f"{elem}," for elem in data)
    yield "};"


def emit_c_array_decl(type_str, name, count, suffix=""):
    """Return a C array forward declaration with explicit size."""
    return f"MLD_INTERNAL_DATA_DECLARATION {type_str} {name}[{count}]{suffix};"


def _fmt_indexed_rows(data):
    """Format rows of data with index comments."""
    for i, row in enumerate(data):
        yield ",".join(map(str, row)) + f" /* {i} */,"


def gen_aarch64_zeta_file():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_AARCH64) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_aarch64.h"'
        yield ""
        yield "/*"
        yield " * Table of zeta values used in the AArch64 forward NTT"
        yield " * See autogen for details."
        yield " */"
        yield from emit_c_array(
            "const int32_t",
            "mld_aarch64_ntt_zetas_layer123456",
            gen_aarch64_fwd_ntt_zetas_layer123456(),
        )
        yield ""
        yield from emit_c_array(
            "const int32_t",
            "mld_aarch64_ntt_zetas_layer78",
            gen_aarch64_fwd_ntt_zetas_layer78(),
        )
        yield ""
        yield from emit_c_array(
            "const int32_t",
            "mld_aarch64_intt_zetas_layer78",
            gen_aarch64_intt_zetas_layer78(),
        )
        yield ""
        yield from emit_c_array(
            "const int32_t",
            "mld_aarch64_intt_zetas_layer123456",
            gen_aarch64_intt_zetas_layer123456(),
        )
        yield ""
        yield "#else"
        yield ""
        yield "MLD_EMPTY_CU(aarch64_zetas)"
        yield ""
        yield "#endif"
        yield ""

    update_file("dev/aarch64_opt/src/aarch64_zetas.c", "\n".join(gen()))

    update_file("dev/aarch64_clean/src/aarch64_zetas.c", "\n".join(gen()))


def gen_aarch64_rej_uniform_eta_table_rows():
    # The index into the lookup table is an 8-bit bitmap, i.e. a number 0..255.
    # Conceptually, the table entry at index i is a vector of 8 16-bit values, of
    # which only the first popcount(i) are set; those are the indices of the set-bits
    # in i. Concretely, we store each 16-bit index as consecutive 8-bit indices.
    def get_set_bits_idxs(i):
        bits = list(map(int, format(i, "08b")))
        bits.reverse()
        return [bit_idx for bit_idx in range(8) if bits[bit_idx] == 1]

    for i in range(256):
        idxs = get_set_bits_idxs(i)
        # Replace each index by two consecutive indices
        idxs = [j for i in idxs for j in [2 * i, 2 * i + 1]]
        # Pad by 255 (invalid index)
        idxs = idxs + [255] * (16 - len(idxs))
        yield idxs


def gen_aarch64_rej_uniform_eta_table():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_AARCH64) &&  \\"
        yield "    !defined(MLD_CONFIG_NO_KEYPAIR_API) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_aarch64.h"'
        yield ""
        yield "/*"
        yield " * Lookup table used by 16-bit rejection sampling (rej_eta)."
        yield " * Adapted from ML-KEM for ML-DSA eta rejection sampling."
        yield " * See autogen for details."
        yield " */"
        yield from emit_c_array(
            "const uint8_t",
            "mld_rej_uniform_eta_table",
            gen_aarch64_rej_uniform_eta_table_rows(),
            formatter=_fmt_indexed_rows,
        )
        yield ""
        yield ("#else /* MLD_ARITH_BACKEND_AARCH64 && !MLD_CONFIG_NO_KEYPAIR_API && \\")
        yield "         !MLD_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLD_EMPTY_CU(aarch64_rej_uniform_eta_table)"
        yield ""
        yield (
            "#endif /* !(MLD_ARITH_BACKEND_AARCH64 && !MLD_CONFIG_NO_KEYPAIR_API && \\"
        )
        yield "          !MLD_CONFIG_MULTILEVEL_NO_SHARED) */"
        yield ""

    update_file("dev/aarch64_opt/src/rej_uniform_eta_table.c", "\n".join(gen()))

    update_file("dev/aarch64_clean/src/rej_uniform_eta_table.c", "\n".join(gen()))

    def gen_aarch64_hol_light_rej_uniform_eta_table():
        yield from gen_hol_light_header()
        yield "(*"
        yield " * Constant table values used in the AArch64 eta rejection sampling."
        yield " * Each entry is 16 bytes. There are 256 entries (one per 8-bit mask),"
        yield " * for a total of 4096 bytes. Entries use 2-byte (16-bit) coefficient"
        yield " * indices since ML-DSA eta samples are stored as 16-bit halfwords"
        yield " * on the stack before being converted to 32-bit in the final copy."
        yield " * See autogen for details."
        yield " *)"
        yield ""
        yield "let mldsa_rej_uniform_eta_table = (REWRITE_RULE[MAP] o define)"
        yield "  `mldsa_rej_uniform_eta_table:byte list = MAP word ["
        data = [v for row in gen_aarch64_rej_uniform_eta_table_rows() for v in row]
        yield from print_hol_light_array(data, as_int=False, entries_per_line=16, pad=3)
        yield "  ]`;;"
        yield ""

    update_file(
        "proofs/hol_light/aarch64/proofs/mldsa_rej_uniform_eta_table.ml",
        "\n".join(gen_aarch64_hol_light_rej_uniform_eta_table()),
    )


def gen_aarch64_rej_uniform_table_rows():
    # The index into the lookup table is an 4-bit bitmap, i.e. a number 0..15.
    # Conceptually, the table entry at index i is a vector of 4-bit values, of
    # which only the first popcount(i) are set; those are the indices of the set-bits
    # in i. Concretely, we store each 32-bit index as consecutive 8-bit indices.
    def get_set_bits_idxs(i):
        bits = list(map(int, format(i, "08b")))
        bits.reverse()
        return [bit_idx for bit_idx in range(8) if bits[bit_idx] == 1]

    for i in range(16):
        idxs = get_set_bits_idxs(i)
        # Replace each index by two consecutive indices
        idxs = [j for i in idxs for j in [4 * i + k for k in range(4)]]
        # Pad by -1
        idxs = idxs + [255] * (16 - len(idxs))
        yield idxs


def gen_aarch64_rej_uniform_table():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_AARCH64) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_aarch64.h"'
        yield ""
        yield "/*"
        yield " * Lookup table used by rejection sampling of the public matrix."
        yield " * See autogen for details."
        yield " */"
        yield from emit_c_array(
            "const uint8_t",
            "mld_rej_uniform_table",
            gen_aarch64_rej_uniform_table_rows(),
            formatter=_fmt_indexed_rows,
        )
        yield ""
        yield "#else"
        yield ""
        yield "MLD_EMPTY_CU(aarch64_rej_uniform_table)"
        yield ""
        yield "#endif"
        yield ""

    update_file("dev/aarch64_opt/src/rej_uniform_table.c", "\n".join(gen()))

    update_file("dev/aarch64_clean/src/rej_uniform_table.c", "\n".join(gen()))


def gen_aarch64_polyz_unpack_indices(bit_width):
    """Generate Neon TBL index bytes for polyz_unpack.

    Each loop iteration loads bit_width * 16 / 8 packed bytes into
    registers v0, v1, v2 and unpacks 16 coefficients. Coefficients
    are unpacked in 4 groups of 4, one group per TBL/TBL2.

    Each coefficient is extracted from 3 consecutive bytes.
    Groups 0,1 index from v0 (and v1); groups 2,3 index from v1 (and v2),
    so their byte offsets are shifted by -16.

    To avoid overreading the input buffer, v0/v1 and v2 are loaded
    separately: v0/v1 via ld1 {v0,v1} at buf, v2 via ld1 {v2} at
    buf + stride - 16 (stride = bit_width * 2). Since v2 starts at
    a different offset than a contiguous 3-register load would give,
    TBL indices referencing v2 (>= 16 within groups 2,3) are shifted
    by 48 - 2 * bit_width.
    """
    v2_shift = 48 - 2 * bit_width
    for group in range(4):
        base_coeff = group * 4
        reg_offset = 0 if group < 2 else 16
        for coeff in range(4):
            i = base_coeff + coeff
            byte_start = (i * bit_width) // 8 - reg_offset
            for j in range(3):
                b = byte_start + j
                if group >= 2 and b >= 16:
                    b += v2_shift
                yield b
            yield 255


def gen_aarch64_polyz_unpack_table():
    # Map gamma1_bits to parameter set guards
    param_set_guards = {
        17: "#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || MLD_CONFIG_PARAMETER_SET == 44",
        19: (
            "#if defined(MLD_CONFIG_MULTILEVEL_WITH_SHARED) || \\\n"
            "    (MLD_CONFIG_PARAMETER_SET == 65 || MLD_CONFIG_PARAMETER_SET == 87)"
        ),
    }

    def format_row(vals):
        return ", ".join(f"{v:>3}" for v in vals) + ","

    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_AARCH64) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_aarch64.h"'
        yield ""
        yield "#if !defined(MLD_CONFIG_NO_SIGN_API) || !defined(MLD_CONFIG_NO_VERIFY_API)"
        yield "/* Table of indices used for tbl instructions in polyz_unpack_{17,19}."
        yield " * See autogen for details. */"
        yield ""
        for gamma1_bits in [17, 19]:
            bit_width = gamma1_bits + 1
            indices = list(gen_aarch64_polyz_unpack_indices(bit_width))
            yield param_set_guards[gamma1_bits]

            def _fmt_polyz(data, _fr=format_row):
                for row_start in range(0, len(data), 16):
                    yield "    " + _fr(data[row_start : row_start + 16])

            yield from emit_c_array(
                "const uint8_t",
                f"mld_polyz_unpack_{gamma1_bits}_indices",
                indices,
                formatter=_fmt_polyz,
            )
            yield "#endif"
            yield ""
        yield "#endif /* !MLD_CONFIG_NO_SIGN_API || !MLD_CONFIG_NO_VERIFY_API */"
        yield ""
        yield "#else /* MLD_ARITH_BACKEND_AARCH64 && !MLD_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLD_EMPTY_CU(aarch64_polyz_unpack_table)"
        yield ""
        yield "#endif /* !(MLD_ARITH_BACKEND_AARCH64 && !MLD_CONFIG_MULTILEVEL_NO_SHARED) */"
        yield ""

    update_file("dev/aarch64_opt/src/polyz_unpack_table.c", "\n".join(gen()))
    update_file("dev/aarch64_clean/src/polyz_unpack_table.c", "\n".join(gen()))


def gen_hol_light_rej_uniform_table():
    def gen():
        yield from gen_hol_light_header()
        yield "(*"
        yield " * Constant table values used in the AArch64 rejection sampling."
        yield " * Each entry is 16 bytes. There are 16 entries (one per 4-bit mask),"
        yield " * for a total of 256 bytes. Entries use 4-byte (32-bit) coefficient"
        yield " * indices since ML-DSA coefficients are 32-bit."
        yield " * See autogen for details."
        yield " *)"
        yield ""
        rows = list(gen_aarch64_rej_uniform_table_rows())
        flat = [v for row in rows for v in row]
        yield "let mldsa_rej_uniform_table = (REWRITE_RULE[MAP] o define)"
        yield "  `mldsa_rej_uniform_table:byte list = MAP word ["
        yield from print_hol_light_array(flat, as_int=False, pad=3, entries_per_line=16)
        yield "  ]`;;"
        yield ""

    update_file(
        "proofs/hol_light/aarch64/proofs/mldsa_rej_uniform_table.ml",
        "\n".join(gen()),
    )


def gen_hol_light_polyz_unpack_table():
    def gen():
        yield from gen_hol_light_header()
        yield "(*"
        yield " * TBL index tables for polyz_unpack_{17,19}"
        yield " * See autogen for details."
        yield " *)"
        yield ""
        for gamma1_bits in [17, 19]:
            bit_width = gamma1_bits + 1
            name = f"mldsa_polyz_unpack_{gamma1_bits}_indices"
            indices = list(gen_aarch64_polyz_unpack_indices(bit_width))
            yield f"let {name} = (REWRITE_RULE[MAP] o define)"
            yield f"  `{name}:byte list = MAP word ["
            yield from print_hol_light_array(indices, as_int=False, pad=3)
            yield "]`;;"
            yield ""

    update_file(
        "proofs/hol_light/aarch64/proofs/mldsa_polyz_unpack_consts.ml",
        "\n".join(gen()),
    )


def gen_avx2_rej_uniform_table_rows():
    # The index into the lookup table is an 8-bit bitmap, i.e. a number 0..255.
    # Conceptually, the table entry at index i is a vector of 8 16-bit values, of
    # which only the first popcount(i) are set; those are the indices of the set-bits
    # in i.
    def get_set_bits_idxs(i):
        bits = list(map(int, format(i, "08b")))
        bits.reverse()
        return [bit_idx for bit_idx in range(8) if bits[bit_idx] == 1]

    for i in range(256):
        idxs = get_set_bits_idxs(i)
        idxs = [i for i in idxs]
        # Pad by 0
        idxs = idxs + [0] * (8 - len(idxs))
        yield "{" + ",".join(map(str, idxs)) + "}"


def gen_avx2_rej_uniform_table():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_X86_64_DEFAULT) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_x86_64.h"'
        yield ""
        yield "/*"
        yield " * Lookup table used by rejection sampling."
        yield " * See autogen for details."
        yield " */"
        yield from emit_c_array(
            "const uint8_t",
            "mld_rej_uniform_table",
            gen_avx2_rej_uniform_table_rows(),
            suffix="[8]",
        )
        yield ""
        yield "#else"
        yield ""
        yield "MLD_EMPTY_CU(avx2_rej_uniform_table)"
        yield ""
        yield "#endif"
        yield ""

    update_file("dev/x86_64/src/rej_uniform_table.c", "\n".join(gen()))


def gen_keccak_round_constants():
    """Generate the 24 Keccak-f[1600] round constants (RC values).

    Yields uint64 values computed via the LFSR-based algorithm
    from the Keccak specification.
    """
    rc = [0] * 24
    lfsr = 1
    for round_idx in range(24):
        rc_val = 0
        for j in range(7):
            if lfsr & 1:
                rc_val ^= 1 << ((1 << j) - 1)
            # x^8 + x^6 + x^5 + x^4 + 1 feedback polynomial
            lfsr = (lfsr << 1) ^ (0x71 if lfsr & 0x80 else 0)
            lfsr &= 0xFF
        rc[round_idx] = rc_val
    yield from rc


def gen_keccak_rho_shuffle(offset):
    """Generate a vpshufb shuffle mask for byte rotation within each qword.

    The Keccak rho step requires 64-bit rotations. Some rotation amounts
    (8 and 56 bits) are multiples of 8 and can be implemented as byte
    shuffles via vpshufb, which is faster than a general-purpose rotation.

    A vpshufb mask for a 256-bit register (YMM) has 32 bytes: 4 groups of 8,
    one per 64-bit lane. Within each lane, byte[i] of the output is sourced
    from byte[src] where src = (i + offset) % 8, shifted to the lane's base.

    offset=-1: each output byte[i] reads from byte[i-1], rotating the qword
               right by 8 bits (equivalently, left by 56 bits).
    offset=+1: each output byte[i] reads from byte[i+1], rotating the qword
               right by 56 bits (equivalently, left by 8 bits).
    """
    for lane in range(4):
        base = lane * 8
        val = 0
        for byte_pos in range(8):
            src = (byte_pos + offset) % 8 + base
            val |= src << (byte_pos * 8)
        yield val


def gen_keccak_rho8():
    """Generate the vpshufb shuffle mask for 8-bit rotation of each qword."""
    yield from gen_keccak_rho_shuffle(-1)


def gen_keccak_rho56():
    """Generate the vpshufb shuffle mask for 56-bit rotation of each qword."""
    yield from gen_keccak_rho_shuffle(+1)


def print_hol_light_word64_list(g, entries_per_line=4):
    """Format a list of uint64 values as a HOL Light int64 word list.

    Analogous to print_hol_light_array but for int64 word lists.
    Yields lines with leading indent, semicolon-separated.
    """
    values = [f"word 0x{v:016X}" for v in g]

    for i in range(0, len(values), entries_per_line):
        row = values[i : i + entries_per_line]
        is_last = i + entries_per_line >= len(values)
        line = "; ".join(row)
        if not is_last:
            line += ";"
        yield "  " + line


def print_c_uint64_array(name, values, entries_per_line=1):
    """Format a C uint64_t array definition with MLD_ALIGN.

    Yields lines for a complete array definition including braces.
    """
    vals = list(values)
    yield "MLD_ALIGN MLD_INTERNAL_DATA_DEFINITION const uint64_t"
    yield f"    {name}[{len(vals)}] = {{"
    for i in range(0, len(vals), entries_per_line):
        row = vals[i : i + entries_per_line]
        yield "        " + ", ".join(f"0x{v:016x}" for v in row) + ","
    yield "};"


def gen_aarch64_keccak_constants_c_file():
    def gen_c():
        yield from gen_header()
        yield '#include "../../../../common.h"'
        yield ""
        yield "#if (defined(MLD_FIPS202_AARCH64_NEED_X1_SCALAR) ||                  \\"
        yield "     defined(MLD_FIPS202_AARCH64_NEED_X1_V84A) ||                    \\"
        yield "     defined(MLD_FIPS202_AARCH64_NEED_X2_V84A) ||                    \\"
        yield "     defined(MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID) ||       \\"
        yield "     defined(MLD_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID)) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "fips202_native_aarch64.h"'
        yield ""
        yield from print_c_uint64_array(
            "mld_keccakf1600_round_constants",
            gen_keccak_round_constants(),
            entries_per_line=3,
        )
        yield ""
        yield "#else /* (MLD_FIPS202_AARCH64_NEED_X1_SCALAR ||                               \\"
        yield "         MLD_FIPS202_AARCH64_NEED_X1_V84A || MLD_FIPS202_AARCH64_NEED_X2_V84A \\"
        yield "         || MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID ||                  \\"
        yield "         MLD_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID) &&               \\"
        yield "         !MLD_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLD_EMPTY_CU(fips202_aarch64_round_constants)"
        yield ""
        yield "#endif /* !((MLD_FIPS202_AARCH64_NEED_X1_SCALAR ||                             \\"
        yield "          MLD_FIPS202_AARCH64_NEED_X1_V84A || MLD_FIPS202_AARCH64_NEED_X2_V84A \\"
        yield "          || MLD_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID ||                  \\"
        yield "          MLD_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID) &&               \\"
        yield "          !MLD_CONFIG_MULTILEVEL_NO_SHARED) */"
        yield ""

    update_file(
        "dev/fips202/aarch64/src/keccakf1600_round_constants.c",
        "\n".join(gen_c()),
        force_format=True,
    )


def gen_avx2_keccak_constants_c_file():
    def gen_c():
        yield from gen_header()
        yield '#include "../../../../common.h"'
        yield "#if defined(MLD_FIPS202_X86_64_NEED_X4_AVX2) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield "#include <stdint.h>"
        yield ""
        yield '#include "fips202_native_x86_64.h"'
        yield ""
        yield from print_c_uint64_array(
            "mld_keccakf1600_round_constants",
            gen_keccak_round_constants(),
            entries_per_line=3,
        )
        yield ""
        yield from print_c_uint64_array("mld_keccak_rho8", gen_keccak_rho8())
        yield ""
        yield from print_c_uint64_array("mld_keccak_rho56", gen_keccak_rho56())
        yield ""
        yield "#else /* MLD_FIPS202_X86_64_NEED_X4_AVX2 && !MLD_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLD_EMPTY_CU(fips202_x86_64_constants)"
        yield ""
        yield "#endif /* !(MLD_FIPS202_X86_64_NEED_X4_AVX2 && !MLD_CONFIG_MULTILEVEL_NO_SHARED) */"
        yield ""

    update_file(
        "dev/fips202/x86_64/src/keccakf1600_constants.c",
        "\n".join(gen_c()),
        force_format=True,
    )


def gen_hol_light_keccak_constants_file():
    def gen():
        yield from gen_hol_light_header()
        yield "(* Keccak-f[1600] round constants RC[i] for i = 0..23. *)"
        yield ""
        yield "let round_constants = define `round_constants:int64 list = ["
        yield from print_hol_light_word64_list(
            gen_keccak_round_constants(), entries_per_line=1
        )
        yield "]`;;"
        yield ""

    update_file(
        "proofs/hol_light/common/keccak_constants.ml",
        "\n".join(gen()),
    )


def gen_avx2_keccak_hol_light_constants_file():
    def gen():
        yield from gen_hol_light_header()
        yield "(* Keccak constants for x86_64 AVX2 implementations. *)"
        yield ""
        yield "let rho8_constant = define `rho8_constant:int64 list = ["
        yield from print_hol_light_word64_list(gen_keccak_rho8(), entries_per_line=1)
        yield "]`;;"
        yield ""
        yield "let rho56_constant = define `rho56_constant:int64 list = ["
        yield from print_hol_light_word64_list(gen_keccak_rho56(), entries_per_line=1)
        yield "]`;;"
        yield ""

    update_file(
        "proofs/hol_light/x86_64/proofs/keccak_f1600_x4_avx2_constants.ml",
        "\n".join(gen()),
    )


def signed_reduce_u32(a):
    """Return signed canonical representative of a mod b"""
    c = a % 2**32
    if c >= 2**31:
        c -= 2**32
    return c


def prepare_root_for_montmul(root, mult):
    """Takes a constant that the code needs to Montgomery-multiply with,
    and returns the pair of (a) the signed canonical representative of its
    Montgomery form, (b) the twisted constant used in the low-mul part of
    the Montgomery multiplication."""

    # Convert to Montgomery form and pick canonical signed representative
    root = signed_reduce(root * montgomery_factor)
    if mult:
        root = signed_reduce_u32(root * pow(modulus, -1, 2**32))
    return root


def gen_avx2_root_of_unity_for_block(layer, block, mult=False):
    # We are computing a negacyclic NTT; the twiddles needed here is
    # the second half of the twiddles for a cyclic NTT of twice the size.
    log = bitreverse(pow(2, layer) + block, 8)
    root = pow(root_of_unity, log, modulus)
    return prepare_root_for_montmul(root, mult)


def gen_avx2_fwd_ntt_zetas(mult=False):
    def gen_twiddles(layer, block, repeat, mult):
        root = gen_avx2_root_of_unity_for_block(layer, block, mult)
        return [root] * repeat

    def gen_twiddles_many(layer, block_base, block_offsets, repeat, mult):
        roots = list(
            map(
                lambda x: gen_twiddles(layer, block_base + x, repeat, mult),
                block_offsets,
            )
        )
        yield from (r for ln in roots for r in ln)

    # embed the scaling of 1/256 and correction of the Montgomery factor
    # from the basemul into last twiddle of the inverse NTT
    # - root^-128 * 2^64/256
    # In the forward NTT this twiddle is unused
    f = signed_reduce(-pow(root_of_unity, -128, modulus) * 2**56)
    if mult:
        f = signed_reduce_u32(f * pow(modulus, -1, 2**32))

    yield f

    # Layers 1 twiddle
    # In the inverse NTT this twiddle is unused
    yield from gen_twiddles_many(0, 0, range(1), 1, mult)

    # Layer 2-8 twiddles
    yield from gen_twiddles_many(1, 0, range(2), 1, mult)
    yield from gen_twiddles_many(2, 0, range(4), 1, mult)
    yield from gen_twiddles_many(3, 0, range(8), 4, mult)
    yield from gen_twiddles_many(4, 0, range(16), 2, mult)
    yield from gen_twiddles_many(5, 0, range(32), 1, mult)
    for i in range(32):
        yield from gen_twiddles_many(6, i * 2, range(1), 1, mult)
    for i in range(32):
        yield from gen_twiddles_many(6, i * 2 + 1, range(1), 1, mult)

    for k in range(4):
        for i in range(32):
            yield from gen_twiddles_many(7, i * 4 + k, range(1), 1, mult)


def gen_avx2_zetas_qdata():
    def cmod(a, mod):
        """Return signed canonical representative of a mod b"""
        c = a % mod
        if c >= mod / 2:
            c -= mod
        return c

    # MLD_AVX2_Q
    q_dup = [modulus] * 8
    # MLD_AVX2_QINV
    qinv_dup = [pow(modulus, -1, 2**32)] * 8
    # MLD_AVX2_DIV_QINV
    div_qinv_dup = [cmod(pow(modulus, -1, 2**32) * pow(2, 64 - 8, modulus), 2**32)] * 8
    # MLD_AVX2_DIV
    div_dup = [pow(2, 64 - 8, modulus)] * 8

    zetas_qinv = list(gen_avx2_fwd_ntt_zetas(mult=True))
    zetas = list(gen_avx2_fwd_ntt_zetas(mult=False))

    q_idx = 0
    qinv_idx = q_idx + len(q_dup)
    div_qinv_idx = qinv_idx + len(qinv_dup)
    div_idx = div_qinv_idx + len(div_qinv_dup)
    zetas_qinv_idx = div_idx + len(div_dup)
    zetas_idx = zetas_qinv_idx + len(zetas_qinv)

    constants = q_dup + qinv_dup + div_qinv_dup + div_dup + zetas_qinv + zetas
    offsets = {
        "q": q_idx,
        "qinv": qinv_idx,
        "div_qinv": div_qinv_idx,
        "div": div_idx,
        "zetas_qinv": zetas_qinv_idx,
        "zetas": zetas_idx,
    }
    return constants, offsets


def gen_aarch64_hol_light_zeta_file():
    def gen():
        yield from gen_hol_light_header()
        yield "(*"
        yield " * Table of zeta values used in the AArch64 NTTs"
        yield " * See autogen for details."
        yield " *)"
        yield ""
        yield "let ntt_zetas_layer012345 = define `ntt_zetas_layer012345:int list = ["
        yield from print_hol_light_array(gen_aarch64_fwd_ntt_zetas_layer123456())
        yield "]`;;"
        yield ""
        yield "let ntt_zetas_layer67 = define `ntt_zetas_layer67:int list = ["
        yield from print_hol_light_array(gen_aarch64_fwd_ntt_zetas_layer78())
        yield "]`;;"
        yield ""
        yield "let intt_zetas_layer78 = define `intt_zetas_layer78:int list = ["
        yield from print_hol_light_array(gen_aarch64_intt_zetas_layer78())
        yield "]`;;"
        yield ""
        yield "let intt_zetas_layer123456 = define `intt_zetas_layer123456:int list = ["
        yield from print_hol_light_array(gen_aarch64_intt_zetas_layer123456())
        yield "]`;;"
        yield ""

    update_file("proofs/hol_light/aarch64/proofs/mldsa_zetas.ml", "\n".join(gen()))


def gen_avx2_hol_light_zeta_file():
    def gen():
        yield from gen_hol_light_header()
        yield "(*"
        yield " * Table of zeta values used in the AVX2 NTTs"
        yield " * See autogen for details."
        yield " *)"
        yield ""
        yield "let mldsa_complete_qdata = define `mldsa_complete_qdata:int list = ["
        constants, _ = gen_avx2_zetas_qdata()
        yield from print_hol_light_array(constants)
        yield "]`;;"
        yield ""

    update_file("proofs/hol_light/x86_64/proofs/mldsa_zetas.ml", "\n".join(gen()))


def gen_avx2_zeta_file():
    constants, offsets = gen_avx2_zetas_qdata()

    def gen_c():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLD_ARITH_BACKEND_X86_64_DEFAULT) && \\"
        yield "    !defined(MLD_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "consts.h"'
        yield ""
        yield "/*"
        yield " * Table of zeta values used in the AVX2 forward and inverse NTT"
        yield " * See autogen for details."
        yield " */"
        yield from emit_c_array("const int32_t", "mld_qdata", constants)
        yield ""
        yield "#else"
        yield ""
        yield "MLD_EMPTY_CU(avx2_consts)"
        yield ""
        yield "#endif"
        yield ""

    def gen_h():
        yield from gen_header()
        yield "#ifndef MLD_NATIVE_X86_64_SRC_CONSTS_H"
        yield "#define MLD_NATIVE_X86_64_SRC_CONSTS_H"
        yield '#include "../../../common.h"'
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_8XQ {offsets['q']}"
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_8XQINV {offsets['qinv']}"
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_8XDIV_QINV {offsets['div_qinv']}"
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_8XDIV {offsets['div']}"
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS_QINV {offsets['zetas_qinv']}"
        yield f"#define MLD_AVX2_BACKEND_DATA_OFFSET_ZETAS {offsets['zetas']}"
        yield ""
        yield "#ifndef __ASSEMBLER__"
        yield "#define mld_qdata MLD_NAMESPACE(qdata)"
        yield emit_c_array_decl("const int32_t", "mld_qdata", len(constants))
        yield "#endif"
        yield ""
        yield "#endif"
        yield ""

    update_file("dev/x86_64/src/consts.c", "\n".join(gen_c()))
    update_file("dev/x86_64/src/consts.h", "\n".join(gen_h()))


def get_c_source_files(main_only=False, core_only=False, strip_mldsa=False):
    if main_only is True:
        return get_files("mldsa/src/**/*.c", strip_mldsa=strip_mldsa)
    elif core_only is True:
        return get_files("mldsa/src/**/*.c", strip_mldsa=strip_mldsa) + get_files(
            "dev/**/*.c"
        )
    else:
        return get_files("**/*.c", strip_mldsa=strip_mldsa)


def get_asm_source_files(main_only=False, core_only=False, strip_mldsa=False):
    if main_only is True:
        return get_files("mldsa/src/**/*.S", strip_mldsa=strip_mldsa)
    elif core_only is True:
        return get_files("mldsa/src/**/*.S", strip_mldsa=strip_mldsa) + get_files(
            "dev/**/*.S", strip_mldsa=strip_mldsa
        )
    else:
        return get_files("**/*.S", strip_mldsa=strip_mldsa)


def get_header_files(main_only=False, core_only=False):
    if main_only is True:
        return get_files("mldsa/*.h") + get_files("mldsa/src/**/*.h")
    elif core_only is True:
        return (
            get_files("mldsa/*.h")
            + get_files("mldsa/src/**/*.h")
            + get_files("dev/**/*.h")
            + get_files("integration/**/*.h")
        )
    else:
        return get_files("**/*.h")


def get_markdown_files(main_only=False):
    return get_files("**/*.md")


def get_files(pattern, strip_mldsa=False):
    def normalize(f):
        return f.removeprefix("mldsa/") if strip_mldsa else f

    fs_files = {f for f in get_all_files() if pathlib.Path(f).is_file()}
    cache_files = set(_file_cache.keys())

    # Convert glob pattern to compiled regex
    pattern = pattern.replace(".", r"\.")
    pattern = pattern.replace("**/", "<<<DOUBLESTAR>>>")
    pattern = pattern.replace("*", "[^/]*")
    pattern = pattern.replace("<<<DOUBLESTAR>>>", "(?:.*/)?")
    regexp = re.compile(f"^{pattern}$")

    # Remove files which are scheduled to be deleted
    res = list(
        map(
            normalize,
            filter(
                lambda f: regexp.match(f) and (read_file(f) is not None),
                fs_files | cache_files,
            ),
        )
    )

    return res


def get_all_files():
    # All git-tracked files, including symlinks
    r = subprocess.run(["git", "ls-files"], capture_output=True, text=True)
    assert r.returncode == 0
    files = r.stdout.split("\n")
    files = filter(lambda s: s != "" and pathlib.Path(s).is_symlink() is False, files)
    return files


def get_defines_from_file(c):
    for line in read_file(c).split("\n"):
        m = _RE_DEFINE.match(line)
        if m:
            yield (c, m.group(1))


def get_defines(all=False):
    if all is False:
        files = get_header_files(main_only=True)
    else:
        files = get_header_files() + get_c_source_files() + get_asm_source_files()

    for results in run_parallel(files, get_defines_from_file):
        yield from results


def get_checked_defines():
    allow_list = [("__contract__", "cbmc.h"), ("__loop__", "cbmc.h")]

    def is_allowed(d, c):
        for d0, c0 in allow_list:
            if c.endswith(c0) is True and d0 == d:
                return True
        return False

    for c, d in get_defines():
        if d.startswith("_") and is_allowed(d, c) is False:
            raise Exception(
                f"{d} from {c}: starts with an underscore, which is not allowed for mldsa-native macros. "
                f"If this is an mldsa-native specific macro, please pick a different name. "
                f"If this is an external macro, it likely needs removing from `gen_monolithic_undef_all_core()` in `scripts/autogen` -- check this!"
            )
        yield (c, d)


def gen_monolithic_undef_all_core(filt=None, desc=""):
    if filt is None:

        def filt(c):
            return True

    if desc != "":
        yield "/*"
        yield f" * Undefine macros from {desc}"
        yield " */"

    defines = list(set(get_checked_defines()))
    defines.sort()

    last_filename = None
    for filename, d in defines:
        if filt(filename) is False:
            continue
        if last_filename != filename:
            yield f"/* {filename} */"
            last_filename = filename
        yield f"#undef {d}"


def native(c):
    return "/native/" in c


def fips202(c):
    return "/fips202/" in c


def aarch64(c):
    return "/aarch64/" in c


def x86_64(c):
    return "/x86_64/" in c


def riscv64(c):
    return "/riscv64/" in c


def armv81m(c):
    return "/armv81m/" in c


def native_fips202(c):
    return native(c) and fips202(c)


def native_arith(c):
    return native(c) and not fips202(c)


def native_fips202_aarch64(c):
    return native_fips202(c) and aarch64(c)


def native_fips202_x86_64(c):
    return native_fips202(c) and x86_64(c)


def native_fips202_armv81m(c):
    return native_fips202(c) and armv81m(c)


def native_fips202_core(c):
    return (
        native_fips202(c)
        and not native_fips202_x86_64(c)
        and not native_fips202_aarch64(c)
        and not native_fips202_armv81m(c)
    )


def native_arith_aarch64(c):
    return native_arith(c) and aarch64(c)


def native_arith_x86_64(c):
    return native_arith(c) and x86_64(c)


def native_arith_riscv64(c):
    return native_arith(c) and riscv64(c)


def native_arith_core(c):
    return (
        native_arith(c)
        and not native_arith_x86_64(c)
        and not native_arith_aarch64(c)
        and not native_arith_riscv64(c)
    )


# List of level-specific source files
# All other files only need including and building once
# in multi-level build.
def k_specific(c):
    k_specific_sources = [
        "mldsa_native.h",
        "params.h",
        "common.h",
        "packing.c",
        "packing.h",
        "poly_kl.c",
        "poly_kl.h",
        "polyvec.c",
        "polyvec.h",
        "polyvec_lazy.h",
        "rounding.h",
        "sign.c",
        "sign.h",
    ]
    for f in k_specific_sources:
        if c.endswith(f):
            return True
    return False


def k_generic(c):
    return not k_specific(c) and c != "mldsa/mldsa_native_config.h"


def gen_macro_undefs(extra_notes=None):
    if extra_notes is None:
        extra_notes = []

    yield "/* Macro #undef's"
    yield " *"
    yield " * The following undefines macros from headers"
    yield " * included by the source files imported above."
    yield " *"
    yield " * This is to allow building and linking multiple builds"
    yield " * of mldsa-native for varying parameter sets through concatenation"
    yield " * of this file, as if the files had been compiled separately."
    yield " * If this is not relevant to you, you may remove the following."
    for e in extra_notes:
        yield f" * {e}"
    yield " */"
    yield ""
    yield from gen_monolithic_undef_all_core(
        filt=k_specific, desc="MLD_CONFIG_PARAMETER_SET-specific files"
    )
    yield ""
    yield "#if !defined(MLD_CONFIG_MONOBUILD_KEEP_SHARED_HEADERS)"
    yield from gen_monolithic_undef_all_core(
        filt=lambda c: (
            not native(c) and k_generic(c) and not fips202(c) and "cbmc.h" not in c
        ),
        desc="MLD_CONFIG_PARAMETER_SET-generic files",
    )
    # Handle cbmc.h manually -- most #define's therein are only defined when CBMC is set
    # and need not be #undef'ed. In fact, #undef'ing them is risky since their names may
    # well already be occupied.
    yield "/* mldsa/src/cbmc.h */"
    yield "#undef MLD_CBMC_H"
    yield "#undef __contract__"
    yield "#undef __loop__"
    yield ""
    yield "#if !defined(MLD_CONFIG_FIPS202_CUSTOM_HEADER)"
    yield from gen_monolithic_undef_all_core(
        filt=lambda c: not native(c) and k_generic(c) and fips202(c),
        desc="FIPS-202 files",
    )
    yield "#endif"
    yield ""
    yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
    yield from gen_monolithic_undef_all_core(filt=native_fips202_core, desc="")
    yield "#if defined(MLD_SYS_AARCH64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_fips202_aarch64, desc="native code (FIPS202, AArch64)"
    )
    yield "#endif"
    yield "#if defined(MLD_SYS_X86_64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_fips202_x86_64, desc="native code (FIPS202, x86_64)"
    )
    yield "#endif"
    yield "#if defined(MLD_SYS_ARMV81M_MVE)"
    yield from gen_monolithic_undef_all_core(
        filt=native_fips202_armv81m, desc="native code (FIPS202, Armv8.1-M)"
    )
    yield "#endif"
    yield "#endif"
    yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_ARITH)"
    yield from gen_monolithic_undef_all_core(filt=native_arith_core, desc="")
    yield "#if defined(MLD_SYS_AARCH64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_arith_aarch64, desc="native code (Arith, AArch64)"
    )
    yield "#endif"
    yield "#if defined(MLD_SYS_X86_64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_arith_x86_64, desc="native code (Arith, X86_64)"
    )
    yield "#endif"
    yield "#endif"
    yield "#endif"
    yield ""


def gen_monolithic_source_file():
    def gen():
        c_sources = get_c_source_files(main_only=True, strip_mldsa=True)
        yield from gen_header()

        yield "/******************************************************************************"
        yield " *"
        yield " * Single compilation unit (SCU) for fixed-level build of mldsa-native"
        yield " *"
        yield " * This compilation unit bundles together all source files for a build"
        yield " * of mldsa-native for a fixed security level (MLDSA-44/65/87)."
        yield " *"
        yield " * # API"
        yield " *"
        yield " * The API exposed by this file is described in mldsa_native.h."
        yield " *"
        yield " * # Multi-level build"
        yield " *"
        yield " * If you want an SCU build of mldsa-native with support for multiple security"
        yield " * levels, you need to include this file multiple times, and set"
        yield " * MLD_CONFIG_MULTILEVEL_WITH_SHARED and MLD_CONFIG_MULTILEVEL_NO_SHARED"
        yield " * appropriately. This is exemplified in examples/monolithic_build_multilevel"
        yield " * and examples/monolithic_build_multilevel_native."
        yield " *"
        yield " * # Configuration"
        yield " *"
        yield " * The following options from the mldsa-native configuration are relevant:"
        yield " *"
        yield " * - MLD_CONFIG_FIPS202_CUSTOM_HEADER"
        yield " *   Set this option if you use a custom FIPS202 implementation."
        yield " *"
        yield " * - MLD_CONFIG_USE_NATIVE_BACKEND_ARITH"
        yield " *   Set this option if you want to include the native arithmetic backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202"
        yield " *   Set this option if you want to include the native FIPS202 backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLD_CONFIG_MONOBUILD_KEEP_SHARED_HEADERS"
        yield " *   Set this option if you want to keep the directives defined in"
        yield " *   level-independent headers. This is needed for a multi-level build."
        yield " */"
        yield " "
        yield " /* If parts of the mldsa-native source tree are not used,"
        yield " * consider reducing this header via `unifdef`."
        yield " *"
        yield " * Example:"
        yield " * ```bash"
        yield " * unifdef -UMLD_CONFIG_USE_NATIVE_BACKEND_ARITH mldsa_native.c"
        yield " * ```"
        yield " */"
        yield ""
        yield '#include "src/common.h"'
        yield ""
        for c in filter(lambda c: not native(c) and not fips202(c), c_sources):
            yield f'#include "{c}"'
        yield ""
        yield "#if !defined(MLD_CONFIG_FIPS202_CUSTOM_HEADER)"
        for c in filter(lambda c: not native(c) and fips202(c), c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield ""
        yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_ARITH)"
        yield "#if defined(MLD_SYS_AARCH64)"
        for c in filter(native_arith_aarch64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_X86_64)"
        for c in filter(native_arith_x86_64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
        yield "#if defined(MLD_SYS_AARCH64)"
        for c in filter(native_fips202_aarch64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_X86_64)"
        for c in filter(native_fips202_x86_64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_ARMV81M_MVE)"
        for c in filter(native_fips202_armv81m, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        yield from gen_macro_undefs()

    update_file("mldsa/mldsa_native.c", "\n".join(gen()))


def gen_monolithic_asm_file():
    def gen():
        asm_sources = get_asm_source_files(main_only=True, strip_mldsa=True)
        yield from gen_header()
        yield "/******************************************************************************"
        yield " *"
        yield " * Single assembly unit for fixed-level build of mldsa-native"
        yield " *"
        yield " * This assembly unit bundles together all assembly files for a build"
        yield " * of mldsa-native for a fixed security level (MLDSA-44/65/87)."
        yield " *"
        yield " * # Multi-level build"
        yield " *"
        yield " * If you want an SCU build of mldsa-native with support for multiple security"
        yield " * levels, you should include this file once with MLD_CONFIG_MULTILEVEL_WITH_SHARED set."
        yield " *"
        yield " * (You could also follow the same pattern as for mldsa_native_monobuild.c"
        yield " *  and include it for every level, setting MLD_CONFIG_MULTILEVEL_NO_SHARED"
        yield " *  for all but one. For builds with MLD_CONFIG_MULTILEVEL_NO_SHARED, this"
        yield " *  file will then be ignored.)"
        yield " *"
        yield " * # Configuration"
        yield " *"
        yield " * The following options from the mldsa-native configuration are relevant:"
        yield " *"
        yield " * - MLD_CONFIG_FIPS202_CUSTOM_HEADER"
        yield " *   Set this option if you use a custom FIPS202 implementation."
        yield " *"
        yield " * - MLD_CONFIG_USE_NATIVE_BACKEND_ARITH"
        yield " *   Set this option if you want to include the native arithmetic backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202"
        yield " *   Set this option if you want to include the native FIPS202 backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLD_CONFIG_MONOBUILD_KEEP_SHARED_HEADERS"
        yield " *   Set this option if you want to keep the directives defined in"
        yield " *   level-independent headers. This is needed for a multi-level build."
        yield " */"
        yield ""
        yield "/* If parts of the mldsa-native source tree are not used,"
        yield " * consider reducing this header via `unifdef`."
        yield " *"
        yield " * Example:"
        yield " * ```bash"
        yield " * unifdef -UMLD_CONFIG_USE_NATIVE_BACKEND_ARITH mldsa_native_asm.S"
        yield " * ```"
        yield " */"
        yield ""
        yield '#include "src/common.h"'
        yield ""
        yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_ARITH)"
        yield "#if defined(MLD_SYS_AARCH64)"
        for c in filter(native_arith_aarch64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_X86_64)"
        for c in filter(native_arith_x86_64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        yield "#if defined(MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
        yield "#if defined(MLD_SYS_AARCH64)"
        for c in filter(native_fips202_aarch64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_X86_64)"
        for c in filter(native_fips202_x86_64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLD_SYS_ARMV81M_MVE)"
        for c in filter(native_fips202_armv81m, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        # We generate #undef's for all headers, even though most are not
        # included by the assembly files. This does not harm, and avoids
        # having to trace which headers are being pulled in from common.h
        # included from the assembly files.
        yield ""
        extra = [
            "",
            "NOTE: This is not needed for the assembly SCU since, at present,",
            "there is no need to include it multiple times.",
            "We keep it for uniformity with mldsa_native.c only.",
            "",
            "NOTE: To avoid having to distinguish between which headers are included",
            "from the assembly files, we #undef the same set of directives",
            "as in mldsa_native.c",
        ]
        yield from gen_macro_undefs(extra_notes=extra)

    update_file("mldsa/mldsa_native_asm.S", "\n".join(gen()), force_format=True)


def get_config_options():
    content = read_file("mldsa/mldsa_native_config.h")
    # Match `#define X`, `/* #define X */`, `#ifndef X`, `#ifdef X` lines.
    config_pattern = r"(?m)^\s*(?:/\*\s*)?#\s*(?:define|ifn?def)\s+(MLD_CONFIG_\w+)\b"

    configs = list(set(re.findall(config_pattern, content)))

    configs += [
        "MLD_FORCE_AARCH64",
        "MLD_FORCE_AARCH64_EB",
        "MLD_FORCE_X86_64",
        "MLD_FORCE_PPC64LE",
        "MLD_FORCE_RISCV64",
        "MLD_FORCE_RISCV32",
        "MLD_SYS_AARCH64_SLOW_BARREL_SHIFTER",
        "MLD_SYS_AARCH64_FAST_SHA3",
        "MLDSA_DEBUG",  # TODO: Rename?
        "MLD_BREAK_PCT",  # Use in PCT breakage test
        "MLD_ALLOW_NONCOMPLIANT_SIGNING_BOUND",  # Internal testing escape hatch
        "MLD_CHECK_APIS",
        "MLD_ERR_XXX",
        "MLD_USE_NATIVE_XXX",
        "MLD_CONFIG_XXX",
        "MLD_PREHASH_",
        "MLD_UNIT_TEST",
    ]

    return configs


def check_macro_typos_in_file(filename, macro_check):
    """Checks for typos in MLD_XXX and MLDSA_XXX identifiers."""
    content = read_file(filename)

    # Separate check for wrongly ported MLK/MLKEM macros
    for m in _RE_MLKEM_MACRO_CHECK.finditer(content):
        txt = m.group(1)
        rest = m.group(2)
        line_no = content[: m.start()].count("\n") + 1
        if filename != "scripts/autogen":
            raise Exception(
                f"Likely typo {txt} in {filename}:{line_no}? wrongly ported MLK_XXX / MLKEM_XXX macros from mlkem-native."
            )

    # Check MLD/MLDSA macros
    for m in _RE_MACRO_CHECK.finditer(content):
        txt = m.group(1)
        rest = m.group(2)
        if macro_check(txt, rest, filename) is False:
            line_no = content[: m.start()].count("\n") + 1
            raise Exception(
                f"Likely typo {txt} in {filename}:{line_no}? Not a defined macro."
            )


def get_syscaps():
    return [
        "MLD_SYS_CAP_X86_64_AVX2",
        "MLD_SYS_CAP_AARCH64_SHA3",
        "MLD_SYS_CAP_ARMV81M_MVE",
    ]


def check_macro_typos():
    files = get_all_files()
    syscaps = get_syscaps()

    macros = set(map(lambda t: t[1], get_defines(all=True)))

    # Add configuration options to the list of allows macro names
    macros.update(get_config_options())

    def macro_check(m, rest, filename):
        if m in macros:
            return True

        # Ignore alloc macros only defined in mldsa_native.h
        if m.startswith("MLD_TOTAL_ALLOC") or m.startswith("MLD_MAX_TOTAL_ALLOC"):
            return True

        is_autogen = filename == "scripts/autogen"

        # Exclude system capabilities, which are enum values
        if m in syscaps:
            return True

        #
        # Register some file-specific exceptions
        #

        # 1. Makefiles use MLD_SOURCE_XXX to list source files
        if is_autogen or filename.endswith("/Makefile"):
            if m.startswith("MLD_SOURCE") or m.startswith("MLD_OBJ"):
                return True

        # 2. libOQS specific identifier
        if is_autogen or filename.startswith("integration/liboqs"):
            if m.startswith("MLDSA_NATIVE_MLDSA") or m in ["MLDSA_DIR"]:
                return True

        # 3. Exclude HOL-Light proof scripts
        if is_autogen or filename.startswith("proofs/hol_light"):
            if filename.endswith(".ml"):
                return True

        # 3b. Isabelle .thy files: <name>_def is the auto-generated
        # unfolding-rule name for a `definition <name> = ...`. Allow it.
        if is_autogen or filename.startswith("proofs/isabelle"):
            if filename.endswith(".thy") and m.endswith("_def"):
                return True

        # 4. Exclude regexp patterns in `autogen`
        if is_autogen:
            if rest.startswith("\\") or m in ["MLD_XXX", "MLD_SOURCE_XXX"]:
                return True

        # 5. AWS-LC post-import patch
        if is_autogen or filename == "integration/aws-lc/post_import.patch":
            return True

        if (
            is_autogen
            or filename == "mldsa/src/common.h"
            or filename == "mldsa/src/context.h"
        ):
            if m == "MLD_CONTEXT_PARAMETERS_n":
                return True

        return False

    run_parallel(
        list(files), partial(check_macro_typos_in_file, macro_check=macro_check)
    )


def check_asm_register_aliases_for_file(filename):
    """Checks that `filename` has no mismatching or dangling register aliases"""

    def get_alias_def(line):
        s = list(filter(lambda s: s != "", line.strip().split(" ")))
        if len(s) < 3 or s[1] != ".req":
            return None
        return s[0]

    def get_alias_undef(line):
        if line.strip().startswith(".unreq") is False:
            return None
        return list(filter(lambda s: s != "", line.strip().split(" ")))[1]

    content = read_file(filename)
    aliases = {}
    for i, line in enumerate(content.split("\n")):
        alias_def = get_alias_def(line)
        alias_undef = get_alias_undef(line)
        if alias_def is not None:
            if alias_def in aliases.keys():
                raise Exception(
                    f"Invalid assembly file {filename}: Duplicate .req directive for {alias_def} at line {i}"
                )
            aliases[alias_def] = i
        elif alias_undef is not None:
            if alias_undef not in aliases.keys():
                raise Exception(
                    f"Invalid assembly file {filename}: .unreq without prior .req for {alias_undef} at line {i}"
                )
            del aliases[alias_undef]

    if len(aliases) > 0:
        fixup_suggestion = [
            "/****************** REGISTER DEALLOCATIONS *******************/"
        ]
        dangling = list(aliases.items())
        # Sort by line number of .req
        dangling.sort(key=lambda s: s[1])

        for a, _ in dangling:
            fixup_suggestion.append(f"    .unreq {a}")
        fixup_suggestion.append("")
        fixup_suggestion = "\n".join(fixup_suggestion)

        raise Exception(
            f"Invalid assembly file {filename}: Dangling .req directives {aliases}.\n\nTry adding this?\n\n{fixup_suggestion}"
        )


def check_asm_register_aliases():
    run_parallel(get_asm_source_files(), check_asm_register_aliases_for_file)


def check_asm_loop_labels_for_file(filename):
    """Checks that all labels in `filename` are prefixed according to the function defined in the file"""

    content = read_file(filename)

    # Find function symbol name
    res = _RE_FUNC_SYMBOL.search(content)
    if res is None:
        raise Exception(f"Could not find function symbol in assembly file {filename}")
    funcname = res.group(1)
    lbl_prefix = re.sub(r"(_(aarch64|avx2|mve))?_asm$", "", funcname) + "_"
    lbl_pattern = r"^(\w+):"
    for m in re.finditer(lbl_pattern, content, flags=re.M):
        lbl = m.group(1)
        if not lbl.startswith(lbl_prefix):
            raise Exception(
                f"Please change label {lbl} in {filename} to be prefixed by {lbl_prefix}"
            )


def check_asm_loop_labels():
    # Operate on assembly files in dev/ only. The ones in mlkem/ are autogenerated
    # from that and don't have the original MLK_ASM_FN_SYMBOL marker anymore.
    files = list(filter(lambda s: s.startswith("dev/"), get_asm_source_files()))
    run_parallel(files, check_asm_loop_labels_for_file)


def normalize_comma_separated_args(args_str):
    """Convert whitespace-separated args to comma-separated, add spaces after commas"""
    # Extract and preserve comment
    match = _RE_ARGS_COMMENT.match(args_str)
    args_only = match.group(1).rstrip()
    comment = match.group(2) or ""

    # If already has commas, just normalize spacing
    if "," in args_only:
        result = re.sub(r",(?! )", ", ", args_only)
    else:
        # Split on whitespace and join with commas
        args = args_only.split()
        if not args:
            return args_str
        result = ", ".join(args)

    return result + comment


def normalize_asm_macro_syntax_for_file(filename):
    """Normalize macro definitions and invocations to use commas with spaces"""

    content = read_file(filename)
    lines = content.split("\n")

    # First pass: collect macro names
    macro_names = set()
    for line in lines:
        macro_match = _RE_MACRO_DEF.match(line)
        if macro_match:
            macro_names.add(macro_match.group(1))

    # Second pass: normalize syntax
    new_lines = []
    for line in lines:
        # Normalize .macro definitions
        macro_def_match = _RE_MACRO_DEF_ARGS.match(line)
        if macro_def_match:
            prefix = macro_def_match.group(1)
            args_with_space = macro_def_match.group(2)
            # Preserve leading whitespace, normalize the rest
            leading_space = _RE_LEADING_SPACE.match(args_with_space).group(1)
            args = args_with_space.lstrip()
            normalized_args = normalize_comma_separated_args(args)
            line = prefix + leading_space + normalized_args
        else:
            # Normalize macro invocations
            for macro_name in macro_names:
                # Match: whitespace + macro_name + whitespace + args
                pattern = r"^(\s*" + re.escape(macro_name) + r")(\s+.*)$"
                invocation_match = re.match(pattern, line)
                if not invocation_match:
                    continue
                prefix = invocation_match.group(1)
                args_with_space = invocation_match.group(2)
                # Preserve leading whitespace, normalize the rest
                leading_space = _RE_LEADING_SPACE.match(args_with_space).group(1)
                args = args_with_space.lstrip()
                normalized_args = normalize_comma_separated_args(args)
                line = prefix + leading_space + normalized_args
                break

        new_lines.append(line)

    update_file(filename, "\n".join(new_lines))


def normalize_asm_macro_syntax():
    """Normalize macro syntax in all assembly files"""
    # Operate on assembly files in dev/ only. The ones in mldsa/ are autogenerated.
    files = list(filter(lambda s: s.startswith("dev/"), get_asm_source_files()))
    files += list(filter(lambda s: s.endswith(".inc"), get_files("dev/**/*.inc")))
    run_parallel(files, normalize_asm_macro_syntax_for_file)


# Architectures autogen knows how to (cross-)compile assembly for. Used to
# expand a bare `--force-cross` (no value) into "force all architectures".
FORCE_CROSS_ALL_ARCHES = {"aarch64", "x86_64", "armv81m"}


def resolve_force_cross(value):
    """Normalize the --force-cross CLI value into a set of architectures.

    Cross-compilation always happens for an architecture when its cross
    toolchain is found; `force_cross` only controls what happens when the
    toolchain is *missing*: architectures in the set raise an error, the
    rest are silently skipped.

    - None (flag absent)          -> skip all missing toolchains
    - []   (flag without a value) -> require all architectures
    - [a, b] (explicit list)      -> require exactly those architectures
    """
    if value is None:
        return set()
    if value == []:
        return set(FORCE_CROSS_ALL_ARCHES)
    return set(value)


def update_via_simpasm(
    infile_full,
    outdir,
    outfile=None,
    cflags=None,
    preserve_header=True,
    force_cross=(),
    x86_64_syntax="att",
):
    _, infile = os.path.split(infile_full)
    if outfile is None:
        outfile = infile
    outfile_full = os.path.join(outdir, outfile)

    if cflags is None:
        cflags = ""
    cflags += " -Imldsa"

    # Check if we need to use a cross-compiler
    if "aarch64" in infile_full:
        source_arch = "aarch64"
    elif "x86_64" in infile_full:
        source_arch = "x86_64"
    elif "armv81m" in infile_full:
        source_arch = "armv81m"
    else:
        raise Exception(f"Could not detect architecture of source file {infile_full}.")
    # Check native architecture
    if platform.machine().lower() in ["arm64", "aarch64"]:
        native_arch = "aarch64"
    else:
        native_arch = "x86_64"

    # Armv8.1-M is always cross-compiled
    if source_arch == "armv81m":
        cross_prefix = "arm-none-eabi-"
        cross_gcc = cross_prefix + "gcc"
        if shutil.which(cross_gcc) is None:
            if source_arch not in force_cross:
                return
            raise Exception(f"Could not find cross toolchain {cross_prefix}")
    elif native_arch != source_arch:
        cross_prefix = f"{source_arch}-unknown-linux-gnu-"
        cross_gcc = cross_prefix + "gcc"
        # Check if cross-compiler is present
        if shutil.which(cross_gcc) is None:
            if source_arch not in force_cross:
                return
            raise Exception(f"Could not find cross toolchain {cross_prefix}")
    else:
        cross_prefix = None

    with tempfile.NamedTemporaryFile(suffix=".S") as tmp:
        try:
            # Determine architecture from filename
            if "aarch64" in infile_full:
                arch = "aarch64"
            elif "armv81m" in infile_full:
                arch = "armv81m"
            else:
                arch = "x86_64"

            cmd = [
                "./scripts/simpasm",
                "--objdump=llvm-objdump",
                "--arch=" + arch,
                "-i",
                infile_full,
                "-o",
                tmp.name,
            ]
            cmd += ["--cfify"]
            if cross_prefix is not None:
                # Stick with llvm-objdump for disassembly
                cmd += ["--cc", cross_prefix + "gcc"]
                cmd += ["--nm", cross_prefix + "nm"]
            if cflags is not None:
                cmd += [f'--cflags="{cflags}"']
            if preserve_header is True:
                cmd += ["-p"]
            # Add syntax option for x86_64
            if arch == "x86_64" and x86_64_syntax != "att":
                cmd += ["--syntax", x86_64_syntax]
            subprocess.run(
                cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
                check=True,
                text=True,
            )
        except subprocess.CalledProcessError as e:
            print(f"Command failed: {' '.join(cmd)}")
            print(f"Exit code: {e.returncode}")
            print(f"stderr: {e.stderr}")
            raise Exception("Failed to run simpasm") from e
        tmp.seek(0)
        new_contents = tmp.read().decode()

    update_file(outfile_full, new_contents)


def gen_hol_light_asm_file(job):
    infile, indir, cflags, arch = job
    # Override the default namespace (PQCP_MLDSA_NATIVE_MLDSA44) with a short
    # `mld_` prefix in the proofs/hol_light/ assembly drops. The exported symbol
    # is determined by MLD_ASM_NAMESPACE() at preprocessing time, which expands
    # via MLD_CONFIG_NAMESPACE_PREFIX with a trailing underscore appended.
    update_via_simpasm(
        f"{indir}/{infile}",
        "proofs/hol_light/" + arch + "/mldsa",
        cflags=cflags + " -DMLD_CONFIG_NAMESPACE_PREFIX=mld",
        preserve_header=False,
    )


def extract_yaml_from_assembly(assembly_file):
    """Extract YAML metadata from assembly file."""
    with open(assembly_file, "r") as f:
        content = f.read()

    yaml_match = re.search(r"/\*yaml\s*\n(.*?)\n\*/", content, re.DOTALL)
    if not yaml_match:
        raise ValueError(f"No YAML metadata found in {assembly_file}")

    return yaml.safe_load(yaml_match.group(1))


# Map dev/ source directory to the corresponding upstream under mldsa/.
def upstream_src_dir_for_dev_dir(dev_dir):
    if dev_dir.startswith("dev/fips202/"):
        # dev/fips202/<arch>/src -> mldsa/src/fips202/native/<arch>/src
        arch = dev_dir.split("/")[2]
        return f"mldsa/src/fips202/native/{arch}/src"
    if dev_dir.startswith("dev/aarch64_"):
        # dev/aarch64_opt/src and dev/aarch64_clean/src map to the same tree
        return "mldsa/src/native/aarch64/src"
    if dev_dir.startswith("dev/x86_64/"):
        return "mldsa/src/native/x86_64/src"
    raise ValueError(
        f"Unrecognised dev/ source directory {dev_dir!r}; "
        "extend upstream_src_dir_for_dev_dir() to map it to mldsa/src/..."
    )


def hol_light_asm_joblist():
    aarch64_flags = "-march=armv8.4-a+sha3"
    joblist_aarch64 = [
        (
            "ntt_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "intt_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_caddq_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_chknorm_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_decompose_32_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_decompose_88_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_use_hint_32_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "poly_use_hint_88_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "mld_polyvecl_pointwise_acc_montgomery_l4_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "mld_polyvecl_pointwise_acc_montgomery_l5_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "mld_polyvecl_pointwise_acc_montgomery_l7_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "keccak_f1600_x1_scalar_aarch64_asm.S",
            "dev/fips202/aarch64/src",
            f"-Imldsa/src/fips202/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "keccak_f1600_x1_v84a_aarch64_asm.S",
            "dev/fips202/aarch64/src",
            f"-Imldsa/src/fips202/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "keccak_f1600_x2_v84a_aarch64_asm.S",
            "dev/fips202/aarch64/src",
            f"-Imldsa/src/fips202/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "keccak_f1600_x4_v8a_v84a_scalar_hybrid_aarch64_asm.S",
            "dev/fips202/aarch64/src",
            f"-Imldsa/src/fips202/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "keccak_f1600_x4_v8a_scalar_hybrid_aarch64_asm.S",
            "dev/fips202/aarch64/src",
            f"-Imldsa/src/fips202/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "pointwise_montgomery_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "polyz_unpack_17_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "polyz_unpack_19_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "rej_uniform_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "rej_uniform_eta2_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
        (
            "rej_uniform_eta4_aarch64_asm.S",
            "dev/aarch64_opt/src",
            f"-Imldsa/src/native/aarch64/src {aarch64_flags}",
            "aarch64",
        ),
    ]

    x86_64_flags = "-mavx2 -mbmi2 -msse4 -fcf-protection=full"
    joblist_x86_64 = [
        (
            "poly_caddq_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_chknorm_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_use_hint_32_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_use_hint_88_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "ntt_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "intt_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "nttunpack_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "pointwise_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompose_32_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompose_88_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "polyz_unpack_17_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "polyz_unpack_19_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "pointwise_acc_l4_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "pointwise_acc_l5_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "pointwise_acc_l7_avx2_asm.S",
            "dev/x86_64/src",
            f"-Imldsa/src/native/x86_64/src -Imldsa/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "keccak_f1600_x4_avx2_asm.S",
            "dev/fips202/x86_64/src",
            f"-Idev/fips202/x86_64/src -Imldsa/src/fips202/native/x86_64/src {x86_64_flags}",
            "x86_64",
        ),
    ]

    return joblist_aarch64 + joblist_x86_64


def gen_hol_light_asm():
    run_parallel(hol_light_asm_joblist(), gen_hol_light_asm_file)


def check_hol_light_asm_autogenerated():
    """Check that every .S file in proofs/hol_light is autogenerated from dev/."""
    expected = {
        os.path.join("proofs/hol_light", arch, "mldsa", os.path.basename(infile))
        for infile, _indir, _cflags, arch in hol_light_asm_joblist()
    }
    stray = sorted(
        os.path.join(d, fn)
        for arch in ("aarch64", "x86_64")
        for d in [os.path.join("proofs/hol_light", arch, "mldsa")]
        if os.path.isdir(d)
        for fn in os.listdir(d)
        if fn.endswith(".S") and os.path.join(d, fn) not in expected
    )
    if stray:
        raise Exception(
            "HOL-Light assembly files not autogenerated from dev/:\n  "
            + "\n  ".join(stray)
            + "\n\nAdd the corresponding source to the gen_hol_light_asm joblist."
        )


def update_via_copy(infile_full, outfile_full, transform=None):
    content = read_file(infile_full)

    if transform is not None:
        content = transform(content)

    update_file(outfile_full, content)


def update_via_remove(filename):
    update_file(filename, None)


# Only synchronize sources, but not README.md, Makefile and so on
SYNCHRONIZED_EXTENSIONS = (".c", ".h", ".i", ".inc", ".S")


def synchronize_file(f, in_dir, out_dir, delete=False, no_simplify=False, **kwargs):
    if not f.endswith(SYNCHRONIZED_EXTENSIONS):
        return None

    basename = os.path.basename(f)

    if delete is True:
        return basename

    if no_simplify is False and f.endswith(".S"):
        update_via_simpasm(f, out_dir, **kwargs)
    else:
        # Update via copy
        _, infile = os.path.split(f)
        outfile_full = os.path.join(out_dir, infile)
        # The header guards will also be checked later, but if we
        # don't do it here, the dry-run would fail because of a
        # mismatching intermediate file
        if f.endswith(".h"):

            def transform(c):
                return adjust_header_guard_for_filename(c, outfile_full)
        else:
            transform = None
        update_via_copy(f, outfile_full, transform=transform)

    return basename


def synchronize_backend(in_dir, out_dir, delete=False, no_simplify=False, **kwargs):
    copied = []

    files = get_files(os.path.join(in_dir, "*"))
    pool_results = run_parallel(
        files,
        partial(
            synchronize_file,
            in_dir=in_dir,
            out_dir=out_dir,
            delete=delete,
            no_simplify=no_simplify,
            **kwargs,
        ),
    )

    copied = [r for r in pool_results if r is not None]

    if delete is False:
        return

    # Check for files in the target directory that have not been copied
    for f in get_files(os.path.join(out_dir, "*")):
        if os.path.basename(f) in copied:
            continue
        if not f.endswith(SYNCHRONIZED_EXTENSIONS):
            continue
        # Otherwise, remove it
        update_via_remove(f)


def synchronize_backends(
    *,
    force_cross=(),
    clean=False,
    delete=False,
    no_simplify=False,
    x86_64_syntax="att",
):
    if clean is False:
        ty = "opt"
    else:
        ty = "clean"

    if delete is False:
        # We may switch the AArch64 arithmetic backend, so adjust the metadata file
        update_via_copy(
            f"dev/aarch64_{ty}/meta.h",
            "mldsa/src/native/aarch64/meta.h",
            transform=lambda c: adjust_header_guard_for_filename(
                c, "mldsa/src/native/aarch64/meta.h"
            ),
        )

        update_via_copy(
            "dev/x86_64/meta.h",
            "mldsa/src/native/x86_64/meta.h",
            transform=lambda c: adjust_header_guard_for_filename(
                c, "mldsa/src/native/x86_64/meta.h"
            ),
        )

    synchronize_backend(
        f"dev/aarch64_{ty}/src",
        "mldsa/src/native/aarch64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Imldsa/src/native/aarch64/src",
    )
    synchronize_backend(
        "dev/fips202/aarch64/src",
        "mldsa/src/fips202/native/aarch64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Imldsa/src/fips202/native/aarch64/src -march=armv8.4-a+sha3",
    )
    synchronize_backend(
        "dev/fips202/aarch64",
        "mldsa/src/fips202/native/aarch64",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Imldsa/src/fips202/native/aarch64 -march=armv8.4-a+sha3",
    )
    synchronize_backend(
        "dev/x86_64/src",
        "mldsa/src/native/x86_64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        x86_64_syntax=x86_64_syntax,
        # Turn off control-flow protection (CET) explicitly. Newer versions of
        # clang turn it on by default and insert endbr64 instructions at every
        # global symbol.
        # We insert endbr64 instruction manually via the MLD_ASM_FN_SYMBOL
        # macro.
        # This leads to duplicate endbr64 instructions causing a failure when
        # comparing the object code before and after simplification.
        cflags="-Imldsa/src/native/x86_64/src/ -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )
    synchronize_backend(
        "dev/fips202/x86_64/src",
        "mldsa/src/fips202/native/x86_64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        x86_64_syntax=x86_64_syntax,
        # Turn off control-flow protection (CET) explicitly. Newer versions of
        # clang turn it on by default and insert endbr64 instructions at every
        # global symbol.
        # We insert endbr64 instruction manually via the MLD_ASM_FN_SYMBOL
        # macro.
        # This leads to duplicate endbr64 instructions causing a failure when
        # comparing the object code before and after simplification.
        cflags="-Idev/fips202/x86_64/src -Imldsa/src/fips202/native/x86_64/src -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )
    synchronize_backend(
        "dev/fips202/x86_64",
        "mldsa/src/fips202/native/x86_64",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        x86_64_syntax=x86_64_syntax,
        cflags="-Idev/fips202/x86_64 -Imldsa/src/fips202/native/x86_64 -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )
    synchronize_backend(
        "dev/fips202/armv81m/src",
        "mldsa/src/fips202/native/armv81m/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/armv81m/src -Imldsa/src/fips202/native/armv81m/src -march=armv8.1-m.main+mve -mthumb",
    )
    synchronize_backend(
        "dev/fips202/armv81m",
        "mldsa/src/fips202/native/armv81m",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/armv81m -Imldsa/src/fips202/native/armv81m -march=armv8.1-m.main+mve -mthumb",
    )


def adjust_header_guard_for_filename(content, header_file):
    content = content.split("\n")
    exceptions = {
        "mldsa/mldsa_native.h": "MLD_H",
        "mldsa/mldsa_native_config.h": "MLD_CONFIG_H",
    }

    # Use full filename as the header guard, with '/' and '.' replaced by '_'
    guard_name = (
        header_file.removeprefix("mldsa/src/")
        .replace("/", "_")
        .replace(".", "_")
        .upper()
    )
    guard_name = "MLD_" + guard_name

    if header_file in exceptions.keys():
        guard_name = exceptions[header_file]

    def gen_guard():
        yield f"#ifndef {guard_name}"
        yield f"#define {guard_name}"

    def gen_footer():
        yield "#endif"
        yield ""

    guard = list(gen_guard())
    footer = list(gen_footer())

    # Skip over initial commentary
    insert_at = None
    for i, line in enumerate(content):
        if line.strip() == "" or line.startswith(("/*", " *")):
            continue
        insert_at = i
        break

    i = insert_at
    while content[i].strip() == "":
        i += 1
    # Check if header file has some guard -- if so, drop it
    if content[i].strip().startswith("#if !defined") or content[i].strip().startswith(
        "#ifndef"
    ):
        del content[i]
        if content[i].strip().startswith("#define"):
            del content[i]
        has_guard = True
    else:
        has_guard = False
    # Add standardized guard
    content = content[:i] + guard + content[i:]
    # Check if header has some footer
    if (
        has_guard is True
        and content[-1] == ""
        and content[-2].strip().startswith("#endif")
    ):
        del content[-2:]
    # Add standardized footer
    content = content + footer

    return "\n".join(content)


def gen_header_guard(header_file):
    content = read_file(header_file)
    new_content = adjust_header_guard_for_filename(content, header_file)
    update_file(header_file, new_content)


def gen_header_guards():
    run_parallel(get_header_files(main_only=True), gen_header_guard)


def gen_source_undefs(source_file):
    # Get list of #define's clauses in this source file (ignore filename)
    undef_list = list(map(lambda c: c[1], get_defines_from_file(source_file)))
    if not undef_list:
        return

    # Get define clauses from header files, as dict
    header_defs = {d: c for (c, d) in get_defines()}

    undefs = []
    ignored = []
    for d in undef_list:
        if d not in header_defs.keys():
            undefs.append(f"#undef {d}")
        else:
            ignored.append((d, header_defs[d]))

    if len(ignored) != 0:
        undefs.append(
            "/* Some macros are kept because they are also defined in a header. */"
        )
        for d, c in ignored:
            undefs.append(f"/* Keep: {d} ({c.split('/')[-1]}) */")

    content = read_file(source_file).split("\n")

    # Find simpasm footer if present (search from end)
    footer_start = None
    if source_file.endswith(".S"):
        footer_start_marker = "simpasm: footer-start"
        for i in range(len(content) - 1, -1, -1):
            if footer_start_marker in content[i]:
                footer_start = i
                break

    if footer_start is not None:
        simpasm_footer = content[footer_start:]
        content = content[:footer_start]
    else:
        simpasm_footer = []

    # Strip trailing undefs and empty lines
    while content and (
        content[-1].startswith("#undef")
        or content[-1].startswith("/* Keep:")
        or content[-1].startswith("/* Some macros")
        or content[-1] == ""
    ):
        content.pop()

    footer = [
        "",
        "/* To facilitate single-compilation-unit (SCU) builds, undefine all macros.",
        " * Don't modify by hand -- this is auto-generated by scripts/autogen. */",
    ]

    # Remove existing footer if present
    if len(content) >= len(footer) and content[-len(footer) :] == footer:
        content = content[: -len(footer)]

    content.extend(footer)
    content.extend(undefs)
    content.append("")

    new_content = "\n".join(content + simpasm_footer)
    update_file(source_file, new_content)


def gen_undefs():
    files = get_c_source_files(core_only=True) + get_asm_source_files(core_only=True)
    run_parallel(files, gen_source_undefs)


def gen_slothy(funcs):
    if not isinstance(funcs, list):
        return

    targets = list(map(lambda s: s + ".S", funcs))

    for t in targets:
        if t.startswith("keccak"):
            base = "dev/fips202/aarch64/src"
        else:
            base = "dev/aarch64_opt/src"

        # Remove file(s) to be re-generated
        if t.endswith(".S"):
            subprocess.run(["rm", "-f", f"{base}/{t}"])

        p = subprocess.run(["make", t] + ["-C", base])
        if p.returncode != 0:
            print(f"Failed to run SLOTHY on {t}!")
            exit(1)


class BibliographyEntry:
    def __init__(self, raw_dict):
        self._raw = raw_dict
        self._usages = []

    def register_usages(self, lst):
        self._usages += lst

    @property
    def usages(self):
        return self._usages

    @property
    def name(self):
        return self._raw["name"]

    @property
    def short(self):
        if "short" in self._raw.keys():
            return self._raw["short"]
        return self.name

    @property
    def id(self):
        return self._raw["id"]

    @property
    def url(self):
        return self._raw["url"]

    @staticmethod
    def full_name(name):
        if "," not in name:
            return name
        surname, forename = name.split(",")
        return forename.strip() + " " + surname.strip()

    @property
    def authors(self):
        authors = self._raw["author"]
        if not isinstance(authors, list):
            authors = [authors]
        authors = list(map(BibliographyEntry.full_name, authors))
        return authors

    @property
    def authors_text(self):
        authors = self._raw["author"]
        if not isinstance(authors, list):
            authors = [authors]

        def surname(name):
            return name.split(",")[0].strip()

        if len(authors) > 1:
            authors = ", ".join(map(surname, authors))
        else:
            authors = BibliographyEntry.full_name(authors[0])
        return authors


def gen_markdown_citations_for(filename, bibliography):
    # Skip BIBLIOGRAPHY.md
    if filename == "BIBLIOGRAPHY.md":
        return

    content = read_file(filename)
    if not _RE_MARKDOWN_CITE.search(content):
        return

    content = content.split("\n")

    # Lookup all citations in style `[^ID]`
    citations = {}
    for i, line in enumerate(content):
        for m in _RE_MARKDOWN_CITE.finditer(line):
            cite_id = m.group("id")
            uses = citations.get(cite_id, [])
            uses.append((filename, i))
            citations[cite_id] = uses

    # Find and remove any existing citation footnotes
    footnote_footer_start = "<!--- bibliography --->"
    try:
        i = content.index(footnote_footer_start)
        content = content[:i]
    except ValueError:
        pass

    # Find explicit footnote definitions outside the bibliography block
    # These are lines matching `[^ID]: ...` and should be excluded from autogeneration
    explicit_footnotes = set()
    for line in content:
        m = re.match(r"^\[\^(\w+)\]:", line)
        if m:
            explicit_footnotes.add(m.group(1))

    # Add footnotes for all citations found (excluding explicit ones)
    cite_ids = [c for c in citations.keys() if c not in explicit_footnotes]
    cite_ids.sort()
    if len(cite_ids) > 0:
        content.append(footnote_footer_start)
    for cite_id in cite_ids:
        uses = citations[cite_id]
        entry = bibliography.get(cite_id, None)
        if entry is None:
            raise Exception(
                f"Could not find bibliography entry {cite_id} referenced in {filename}. Known entries: {list(bibliography.keys())}"
            )
        content.append(
            f"[^{cite_id}]: {entry.authors_text}: {entry.name}, [{entry.url}]({entry.url})"
        )

        # Remember this usage of the bibliography entry
        entry.register_usages(uses)

    if len(cite_ids) > 0:
        content.append("")

    update_file(filename, "\n".join(content))


def gen_c_citations_for(filename, bibliography):
    content = read_file(filename)

    if not _RE_C_CITE.search(content):
        return

    references_start = [
        "/* References",
        " * ==========",
    ]
    references_end = [" */"]

    # Find and remove any existing reference section
    ref_pattern = r"/\* (# )?References.*?\*/\n+"
    content = re.sub(ref_pattern, "", content, flags=re.DOTALL)

    content = content.split("\n")

    # Lookup all citations in style `@[ID]`
    citations = {}
    for i, line in enumerate(content):
        for m in _RE_C_CITE.finditer(line):
            cite_id = m.group("id")
            uses = citations.get(cite_id, [])
            # Remember usage. +1 because line counting starts at 1
            uses.append((filename, i + 1))
            citations[cite_id] = uses

    # Add references section
    references = []
    references += references_start

    cite_ids = list(citations.keys())
    cite_ids.sort()
    for cite_id in cite_ids:
        uses = citations[cite_id]
        entry = bibliography.get(cite_id, None)
        if entry is None:
            raise Exception(
                f"Could not find bibliography entry {cite_id} referenced in {filename}"
            )
        references.append(" *")
        references.append(f" * - [{cite_id}]")
        prefix = " *   "
        # Wrap long lines at 80 chars
        for line in [entry.name, entry.authors_text, entry.url]:
            if len(line) + len(prefix) <= 80:
                references.append(f"{prefix}{line}")
            else:
                words = line.split()
                current = prefix
                for word in words:
                    if len(current) + len(word) <= 80 or current == prefix:
                        current += word + " "
                    else:
                        references.append(current.rstrip())
                        current = prefix + word + " "
                if current.rstrip() != prefix.rstrip():
                    references.append(current.rstrip())

    references += references_end

    # Fix indentation for comment lines
    references = "\n".join(references)
    references = references.split("\n")
    references = [""] + references

    if len(cite_ids) > 0:
        # Add references to file after initial header section
        # Skip over copyright
        insert_at = None
        for i, line in enumerate(content):
            if line.startswith(("/*", " *")):
                continue
            insert_at = i
            break
        content = content[:insert_at] + references + content[insert_at:]

    # Remember uses -- needs to happen after insertion of references
    # since we need to adjust the line count
    for cite_id in cite_ids:
        uses = citations[cite_id]
        entry = bibliography.get(cite_id, None)

        # Adjust line count after insertion of references
        def bump_line_count(x):
            return (x[0], x[1] + len(references))

        uses = list(map(bump_line_count, uses))

        # Remember this usage of the bibliography entry
        entry.register_usages(uses)

    update_file(filename, "\n".join(content))


def gen_citations_for(filename, bibliography):
    if filename.endswith(".md"):
        gen_markdown_citations_for(filename, bibliography)
    elif filename.endswith((".c", ".h", ".S")):
        gen_c_citations_for(filename, bibliography)
    else:
        raise Exception(f"Unexpected file extension in {filename}")


def gen_bib_file(bibliography):
    content = [
        "[//]: # (SPDX-License-Identifier: CC-BY-4.0)",
        "[//]: # (This file is auto-generated from BIBLIOGRAPHY.yml)",
        "[//]: # (Do not modify it directly)",
        "",
        "# Bibliography",
        "",
        "This file lists the citations made throughout the mldsa-native ",
        "source code and documentation.",
        "",
    ]

    cite_ids = list(bibliography.keys())
    cite_ids.sort()

    for cite_id in cite_ids:
        entry = bibliography[cite_id]
        content.append(f"### `{cite_id}`")
        content.append("")
        content.append(f"* {entry.name}")
        content.append("* Author(s):")
        for author in entry.authors:
            content.append(f"  - {author}")
        content.append(f"* URL: {entry.url}")
        content.append("* Referenced from:")
        # Usages are pairs of (filename, line_count)
        # Ignore line_count for now, as it would require `autogen` after
        # a change to source files.
        files = list(set(map(lambda x: x[0], entry.usages)))
        files.sort()
        for filename in files:
            content.append(f"  - [{filename}]({filename})")
        content.append("")

    update_file("BIBLIOGRAPHY.md", "\n".join(content))


def get_oqs_shared_sources(backend):
    """Get shared source files for OQS integration"""
    mldsa_dir = "mldsa/src/"

    # add files mldsa/*
    sources = [
        f"mldsa/src/{f}"
        for f in os.listdir(mldsa_dir)
        if os.path.isfile(f"{mldsa_dir}/{f}")
        and not f.endswith(".o")
        and not f == "mldsa_native.h"
    ]

    if backend != "ref":
        # add files mldsa/native/* (API definitions)
        sources += [
            f"mldsa/src/native/{f}"
            for f in os.listdir(f"{mldsa_dir}/native")
            if os.path.isfile(f"{mldsa_dir}/native/{f}")
        ]
    # Add FIPS202 glue code
    sources += [
        "integration/liboqs/fips202_glue.h",
        "integration/liboqs/fips202x4_glue.h",
    ]
    # Add custom config
    if backend == "ref":
        backend = "c"
    sources.append(f"integration/liboqs/config_{backend.lower()}.h")

    return sources


def get_oqs_native_sources(backend):
    """Get native source files for OQS integration"""
    return [f"mldsa/src/native/{backend}"]


def gen_oqs_meta_file(filename):
    """Generate OQS META.yml file with updated source lists"""

    content = read_file(filename)

    # Parse YAML while preserving structure
    yml_data = yaml.safe_load(content)

    for impl in yml_data["implementations"]:
        name = impl["name"]

        sources = get_oqs_shared_sources(name)

        # NOTE: Sorting at the end causes the libOQS importer to fail.
        # Somehow, the native directory cannot be imported too early.
        sources.sort()

        if name != "ref":
            sources += get_oqs_native_sources(name)
        impl["sources"] = " ".join(sources)

    # Convert back to YAML string with standard copyright header
    yaml_header = "\n".join(gen_yaml_header())

    new_content = yaml.dump(
        yml_data,
        default_flow_style=False,
        sort_keys=False,
        allow_unicode=True,
        encoding=None,
    )

    # Combine copyright header with new YAML content
    new_content = yaml_header + new_content

    update_file(filename, new_content)


def gen_oqs_meta_files():
    """Generate all OQS META.yml files"""
    meta_files = [
        "integration/liboqs/ML-DSA-44_META.yml",
        "integration/liboqs/ML-DSA-65_META.yml",
        "integration/liboqs/ML-DSA-87_META.yml",
    ]

    for meta_file in meta_files:
        gen_oqs_meta_file(meta_file)


def gen_citations():
    # Load bibliography
    with open("BIBLIOGRAPHY.yml", "r") as f:
        bibliography_raw = yaml.safe_load(f.read())

    bibliography = {}
    for r in bibliography_raw:
        cite_id = r["id"]
        bibliography[cite_id] = BibliographyEntry(r)

    files = (
        get_markdown_files()
        + get_asm_source_files()
        + get_c_source_files()
        + get_header_files()
    )
    run_parallel(files, partial(gen_citations_for, bibliography=bibliography))

    # Check that every bibliography entry has been used as least once
    for e in bibliography.values():
        if len(e.usages) == 0:
            raise Exception(
                f"Bibliography entry {e.id} is unused! "
                "Add a citation or remove from BIBLIOGRAPHY.yml."
            )

    gen_bib_file(bibliography)


def extract_bytecode_from_output(output_text):
    """Convert output of proofs/hol_light/{aarch64,x86_64}/proofs/dump_bytecode.native
    into a dictionary mapping function names to byte code strings."""
    bytecode_dict = {}

    lines = output_text.split("\n")
    i = 0

    while i < len(lines):
        line = lines[i]
        match = _RE_BYTECODE_START.search(line)
        if match:
            filename = match.group(1)

            # Collect bytecode until end marker
            bytecode_lines = []
            i += 1
            while i < len(lines) and "==== bytecode end" not in lines[i]:
                bytecode_lines.append(lines[i])
                i += 1

            bytecode = "\n".join(bytecode_lines).strip()
            bytecode_dict[filename] = bytecode
        i += 1
    return bytecode_dict


def update_bytecode_in_proof_script(filepath, bytecode):
    content = read_file(filepath)

    # Check if markers exist
    start_marker = "(*** BYTECODE START ***)"
    end_marker = "(*** BYTECODE END ***)"

    if start_marker not in content or end_marker not in content:
        raise Exception(f"Could not find BYTECODE START/END markers in {filepath}")

    # Replace content between markers
    pattern = rf"{re.escape(start_marker)}.*?{re.escape(end_marker)}"
    replacement = f"{start_marker}\n{bytecode}\n{end_marker}"

    updated_content = re.sub(pattern, replacement, content, flags=re.DOTALL)

    update_file(filepath, updated_content)


def update_hol_light_bytecode_for_arch(arch, force_cross=()):
    source_arch = arch
    if platform.machine().lower() in ["arm64", "aarch64"]:
        native_arch = "aarch64"
    else:
        native_arch = "x86_64"

    if native_arch != source_arch:
        cross_prefix = f"{source_arch}-unknown-linux-gnu-"
        cross_gcc = cross_prefix + "gcc"
        # Check if cross-compiler is present
        if shutil.which(cross_gcc) is None:
            if source_arch not in force_cross:
                return
            raise Exception(f"Could not find cross toolchain {cross_prefix}")

    # Run make to get bytecode output
    result = subprocess.run(
        ["make", "-C", "proofs/hol_light/" + arch, "dump_bytecode"],
        capture_output=True,
        text=True,
        check=True,
    )
    output_text = result.stdout

    # Extract bytecode
    bytecode_dict = extract_bytecode_from_output(output_text)

    # Update each .ml file
    for obj_name, bytecode in bytecode_dict.items():
        ml_file = "proofs/hol_light/" + arch + "/proofs/" + obj_name + ".ml"
        update_bytecode_in_proof_script(ml_file, bytecode)


def update_hol_light_bytecode(force_cross=()):
    """Update HOL Light proof files with bytecode from make dump_bytecode."""
    update_hol_light_bytecode_for_arch("aarch64", force_cross=force_cross)
    update_hol_light_bytecode_for_arch("x86_64", force_cross=force_cross)


def gen_test_config(config_path, config_spec, default_config_content):
    """Generate a config file by modifying the default config."""

    # Start with the default config
    lines = default_config_content.split("\n")

    # Find copyright and reference header
    references_start = None
    references_end = None

    # NOTE: This needs further work if any custom config contains citations
    # not included in the default configuration. In this case, the reference
    # section needs to be updated.

    for i, line in enumerate(lines):
        if "/* References" in line:
            references_start = i
        elif (
            references_start is not None
            and references_end is None
            and line.strip() == "*/"
        ):
            references_end = i + 1
            break

    header = lines[:references_end]
    header += list(gen_autogen_warning())

    header.append("")
    header.append("/*")
    header.append(f" * Test configuration: {config_spec['description']}")
    header.append(" *")
    header.append(
        " * This configuration differs from the default mldsa/mldsa_native_config.h in the following places:"
    )

    def spec_has_value(opt_value):
        if not isinstance(opt_value, dict):
            return True
        else:
            return "value" in opt_value.keys() or "content" in opt_value.keys()

    for opt_name, opt_value in config_spec["defines"].items():
        if not spec_has_value(opt_value):
            continue
        header.append(f" *   - {opt_name}")
    header.append(" */")
    header.append("")

    # Combine: new header + config body
    lines = header + lines[references_end:]

    def locate_config_option(lines, opt_name):
        """Locate configuration option in lines. Returns (start, end) line indices"""
        # Find the line that defines (or comments out) opt_name, e.g.
        #   #define <opt>
        #   #ifndef <opt> / #ifdef <opt>
        #   /* #define <opt> ... */
        define_re = re.compile(
            rf"^\s*(/\*\s*)?#\s*(define|ifn?def)\s+{re.escape(opt_name)}\b"
        )
        impl_line = None
        for i, line in enumerate(lines):
            if define_re.search(line):
                impl_line = i
                break
        if impl_line is None:
            raise Exception(
                f"Could not find config option {opt_name} in default config"
            )

        # Walk back to the closing */ of the preceding /** ... */ doxygen block.
        j = impl_line - 1
        while j >= 0 and not lines[j].strip().endswith("*/"):
            j -= 1
        if j < 0:
            raise Exception(
                f"Could not find doxygen block before config option {opt_name}"
            )

        start = j + 1  # Replacement starts right after */

        # Replacement ends at the next /** block (option or section divider).
        i = start
        while i < len(lines):
            if lines[i].strip().startswith("/**"):
                while i > start and lines[i - 1].strip() == "":
                    i -= 1
                return (start, i)
            i += 1
        return (start, len(lines))

    def replace_config_option(lines, opt_name, opt_value):
        """Replace config option with new value"""
        block_start, block_end = locate_config_option(lines, opt_name)

        def content_from_value(value):
            if value is True:
                return [f"#define {opt_name}"]
            elif value is False:
                return [f"/* #define {opt_name} */"]
            else:
                return [f"#define {opt_name} {str(value)}"]

        if isinstance(opt_value, dict):
            if "content" in opt_value:
                content = opt_value.get("content").split("\n")
            elif "value" in opt_value:
                content = content_from_value(opt_value["value"])
            else:
                # Use original content
                content = lines[block_start:block_end]
            if "comment" in opt_value:
                comment = opt_value.get("comment").split("\n")
            else:
                comment = []
        else:
            content = content_from_value(opt_value)
            comment = []

        lines[block_start:block_end] = comment + content

    # Apply modifications for each defined option
    for opt_name, opt_value in config_spec["defines"].items():
        replace_config_option(lines, opt_name, opt_value)

    content = "\n".join(lines)
    update_file(config_path, content)


def gen_test_configs():
    """Generate all test configuration files from metadata."""
    # Load metadata
    metadata = yaml.safe_load(read_file("test/configs/configs.yml"))

    # Load default config
    default_config = read_file("mldsa/mldsa_native_config.h")

    # Generate each test config
    for config_spec in metadata["configs"]:
        gen_test_config(config_spec["path"], config_spec, default_config)


def gen_test_vectors(msg, ctx):
    """Generate expected test vectors using ACVP binaries and notrandombytes.

    A single 32-byte value is drawn from notrandombytes (using the default
    pi-based seed) and reused as both the keygen seed and the signing rnd
    for all parameter sets (44, 65, 87).
    """
    nrb_script = os.path.join(os.path.dirname(__file__), "notrandombytes")

    def notrandombytes(num_bytes):
        result = subprocess.run(
            [sys.executable, nrb_script, str(num_bytes)],
            capture_output=True,
            encoding="utf-8",
        )
        if result.returncode != 0:
            raise RuntimeError(f"notrandombytes failed: {result.stderr}")
        return result.stdout.strip().upper()

    def run_acvp(lvl, *args):
        acvp_bin = f"./test/build/mldsa{lvl}/bin/acvp_mldsa{lvl}"
        if not os.path.isfile(acvp_bin):
            console.print(
                f"[red]Error:[/] ACVP binary not found: {acvp_bin}\n"
                f"Build it first: [bold]make acvp OPT=0[/]"
            )
            sys.exit(1)
        result = subprocess.run(
            [acvp_bin, *args],
            capture_output=True,
            encoding="utf-8",
        )
        if result.returncode != 0:
            raise RuntimeError(f"acvp_mldsa{lvl} {args[0]} failed: {result.stderr}")
        out = {}
        for line in result.stdout.splitlines():
            k, v = line.split("=", 1)
            out[k] = v
        return out

    levels = [44, 65, 87]
    msg_hex = msg.encode().hex().upper()
    ctx_hex = ctx.encode().hex().upper()

    # Fixed 64-byte mu for extmu test vectors (all zeros)
    mu_hex = "00" * 64
    # All randomness drawn from notrandombytes with the default pi seed
    rnd = notrandombytes(32)

    results = {}
    for lvl in levels:
        keygen_out = run_acvp(lvl, "keyGen", f"seed={rnd}")

        sig_out = run_acvp(
            lvl,
            "sigGen",
            f"message={msg_hex}",
            f"sk={keygen_out['sk']}",
            f"context={ctx_hex}",
            f"rnd={rnd}",
        )

        extmu_out = run_acvp(
            lvl,
            "sigGenInternal",
            f"message={mu_hex}",
            f"sk={keygen_out['sk']}",
            "externalMu=1",
            f"rnd={rnd}",
        )

        prehash_out = run_acvp(
            lvl,
            "sigGenPreHashShake256",
            f"message={msg_hex}",
            f"context={ctx_hex}",
            f"sk={keygen_out['sk']}",
            f"rnd={rnd}",
        )

        results[lvl] = {
            "pk": keygen_out["pk"],
            "sk": keygen_out["sk"],
            "sig": sig_out["signature"],
            "sig_extmu": extmu_out["signature"],
            "sig_pre_hash_shake256": prehash_out["signature"],
        }

    def hex_to_c_array(hex_str, indent="    "):
        """Convert hex string to C array initializer with 12 bytes per line."""
        raw = bytes.fromhex(hex_str)
        lines = []
        for i in range(0, len(raw), 12):
            chunk = raw[i : i + 12]
            elems = ", ".join(f"0x{b:02x}" for b in chunk)
            lines.append(f"{indent}{elems},")
        # Remove trailing comma from last line
        if lines:
            lines[-1] = lines[-1].rstrip(",")
        return "\n".join(lines)

    def c_string_literal(s):
        """Escape a string for use in a C string literal."""
        return s.replace("\\", "\\\\").replace('"', '\\"')

    def hex_len(hex_str):
        return len(hex_str) // 2

    def gen_msg_ctx_defines():
        yield f'#define TEST_VECTOR_MSG "{c_string_literal(msg)}"'
        yield "#define TEST_VECTOR_MSG_LEN (sizeof(TEST_VECTOR_MSG) - 1)"
        yield f'#define TEST_VECTOR_CTX "{c_string_literal(ctx)}"'
        yield "#define TEST_VECTOR_CTX_LEN (sizeof(TEST_VECTOR_CTX) - 1)"
        yield ""
        yield "/* notrandombytes(32) */"
        yield f"static const uint8_t test_vector_rnd[{hex_len(rnd)}] = {{"
        yield hex_to_c_array(rnd)
        yield "};"
        yield "/* 64 zero bytes */"
        yield "static const uint8_t test_vector_mu[64] = {0};"

    def gen_description():
        yield "/*"
        yield " * Generated with:"
        yield f' *   message: "{c_string_literal(msg)}"'
        yield f' *   context: "{c_string_literal(ctx)}"'
        yield " *   rnd:     notrandombytes(32)"
        yield " *   mu:      64 zero bytes"
        yield " */"

    # Generate expected_test_vectors.h (single-level, with #if guards)
    def gen_single():
        yield from gen_header()
        yield from gen_description()
        yield ""
        yield "#ifndef EXPECTED_TEST_VECTORS_H"
        yield "#define EXPECTED_TEST_VECTORS_H"
        yield ""
        yield "#include <stdint.h>"
        yield ""
        yield from gen_msg_ctx_defines()
        yield ""

        first = True
        for lvl in levels:
            r = results[lvl]
            directive = "#if" if first else "#elif"
            yield f"{directive} MLD_CONFIG_PARAMETER_SET == {lvl}"
            yield f"static const uint8_t test_vector_pk[{hex_len(r['pk'])}] = {{"
            yield hex_to_c_array(r["pk"])
            yield "};"
            yield f"static const uint8_t test_vector_sk[{hex_len(r['sk'])}] = {{"
            yield hex_to_c_array(r["sk"])
            yield "};"
            yield f"static const uint8_t test_vector_sig[{hex_len(r['sig'])}] = {{"
            yield hex_to_c_array(r["sig"])
            yield "};"
            yield f"static const uint8_t test_vector_sig_extmu[{hex_len(r['sig_extmu'])}] = {{"
            yield hex_to_c_array(r["sig_extmu"])
            yield "};"
            yield f"static const uint8_t test_vector_sig_pre_hash_shake256[{hex_len(r['sig_pre_hash_shake256'])}] = {{"
            yield hex_to_c_array(r["sig_pre_hash_shake256"])
            yield "};"
            first = False

        yield f"#endif /* MLD_CONFIG_PARAMETER_SET == {levels[-1]} */"
        yield ""
        yield "#endif /* !EXPECTED_TEST_VECTORS_H */"
        yield ""

    update_file(
        "test/test_vectors/expected_test_vectors.h",
        "\n".join(gen_single()),
    )

    # Generate expected_test_vectors_multilevel.h (all levels, suffixed names)
    def gen_multi():
        yield from gen_header()
        yield "/*"
        yield " * Expected test vectors for multilevel builds."
        yield " * Generated with per-level array names."
        yield " */"
        yield ""
        yield from gen_description()
        yield ""
        yield "#ifndef EXPECTED_TEST_VECTORS_MULTILEVEL_H"
        yield "#define EXPECTED_TEST_VECTORS_MULTILEVEL_H"
        yield ""
        yield "#include <stdint.h>"
        yield ""
        yield from gen_msg_ctx_defines()
        yield ""

        for lvl in levels:
            r = results[lvl]
            yield f"/* ML-DSA-{lvl} */"
            yield f"static const uint8_t test_vector_pk_{lvl}[{hex_len(r['pk'])}] = {{"
            yield hex_to_c_array(r["pk"])
            yield "};"
            yield f"static const uint8_t test_vector_sk_{lvl}[{hex_len(r['sk'])}] = {{"
            yield hex_to_c_array(r["sk"])
            yield "};"
            yield f"static const uint8_t test_vector_sig_{lvl}[{hex_len(r['sig'])}] = {{"
            yield hex_to_c_array(r["sig"])
            yield "};"
            yield f"static const uint8_t test_vector_sig_extmu_{lvl}[{hex_len(r['sig_extmu'])}] = {{"
            yield hex_to_c_array(r["sig_extmu"])
            yield "};"
            yield f"static const uint8_t test_vector_sig_pre_hash_shake256_{lvl}[{hex_len(r['sig_pre_hash_shake256'])}] = {{"
            yield hex_to_c_array(r["sig_pre_hash_shake256"])
            yield "};"
            yield ""

        yield "#endif /* !EXPECTED_TEST_VECTORS_MULTILEVEL_H */"
        yield ""

    update_file(
        "test/test_vectors/expected_test_vectors_multilevel.h",
        "\n".join(gen_multi()),
    )


# ABI checker capability table.
#
# Maps each `ABI.Features:` tag from the kernel YAMLs to:
#   - 'asm_subdir':  the arch this cap belongs to; selects the per-arch
#                    test/abicheck/<asm_subdir>/abicheck_<asm_subdir>.mk it is
#                    emitted into.
#   - 'guard':       if non-None, wrap the generated check in #if defined(...).
#   - 'syscap':      runtime capability; check skips if the host lacks it.
#   - 'cflags':      for the autogeneration of the per-arch abicheck_<arch>.mk.
#   - 'description': shown when a check is skipped.
#   - 'aux_files':   extra .S files needing these CFLAGS though not generated
#                    kernel exports (hand-written abicheck stubs/selftests, and
#                    sources whose YAML doesn't list them as ABI-tested kernels
#                    but still need the per-cap toolchain flags).
ABI_CAPS = {
    "AVX2": {
        "asm_subdir": "x86_64",
        "guard": "__AVX2__",
        "syscap": "MLD_SYS_CAP_X86_64_AVX2",
        "description": "AVX2",
        "cflags": ["-mavx2", "-mbmi2"],
        "aux_files": [],
    },
    "SHA3": {
        "asm_subdir": "aarch64",
        "guard": "__ARM_FEATURE_SHA3",
        "syscap": "MLD_SYS_CAP_AARCH64_SHA3",
        "description": "Armv8.4-A SHA3 (eor3, rax1, xar, bcax)",
        "cflags": ["-march=armv8.4-a+sha3"],
        "aux_files": [],
    },
    "MVE": {
        "asm_subdir": "armv81m",
        "guard": "__ARM_FEATURE_MVE",
        "syscap": "MLD_SYS_CAP_ARMV81M_MVE",
        "description": "Armv8.1-M MVE",
        "cflags": ["-march=armv8.1-m.main+mve", "-mthumb"],
        "aux_files": [
            "test/abicheck/armv81m/callstub_armv81m.S",
            "test/abicheck/armv81m/selftest_armv81m.S",
            # Shipped MVE helpers pulled in by the abicheck wildcard but
            # not registered as ABI-tested kernels - they still need MVE
            # CFLAGS to assemble.
            "mldsa/src/fips202/native/armv81m/src/state_extract_bytes_x4_mve.S",
            "mldsa/src/fips202/native/armv81m/src/state_xor_bytes_x4_mve.S",
        ],
    },
}


# ABI checker config, keyed by the (ABI.Architecture, ABI.CallingConvention)
# pair each kernel YAML declares.
ARCH_CALLINGS = {
    ("aarch64", "AAPCS64"): {
        "arch_guard": "MLD_SYS_AARCH64",
        "extra_guards": [],
        "asm_subdir": "aarch64",
        "state_type": "struct aarch64_register_state",
        # GPR-slot integer type. Used to cast pointer-arg buffers into the
        # input state without introducing an intermediate type (avoiding
        # uintptr_t / size_t which we don't use elsewhere in mldsa-native).
        "gpr_int_type": "uint64_t",
        "fn_qualifier": "",
        "init_func": "init_aarch64_register_state",
        "stub_func": "asm_call_stub_aarch64",
        "check_func": "check_aarch64_aapcs_compliance",
        # YAML register xN maps to gpr[N].
        "reg_to_field": lambda reg: f"gpr[{int(reg[1:])}]",
    },
    ("x86_64", "SysV"): {
        "arch_guard": "MLD_SYS_X86_64",
        "extra_guards": ["MLD_SYSV_ABI_SUPPORTED"],
        "asm_subdir": "x86_64",
        "state_type": "struct x86_64_register_state",
        "gpr_int_type": "uint64_t",
        # MLD_SYSV_ABI is `__attribute__((sysv_abi))` on Windows-MinGW
        # (where Microsoft x64 is the default convention) and empty on
        # ELF/Mach-O x86_64. Matches the library's own kernel decls in
        # mldsa/src/native/x86_64/src/arith_native_x86_64.h.
        "fn_qualifier": "MLD_SYSV_ABI",
        "init_func": "init_x86_64_register_state",
        "stub_func": "asm_call_stub_x86_64_sysv",
        "check_func": "check_x86_64_sysv_compliance",
        "reg_to_field": lambda reg: reg,
    },
    ("armv81m", "AAPCS32"): {
        "arch_guard": "MLD_SYS_ARMV81M_MVE",
        "extra_guards": [],
        "asm_subdir": "armv81m",
        "state_type": "struct armv81m_register_state",
        "gpr_int_type": "uint32_t",
        "fn_qualifier": "",
        "init_func": "init_armv81m_register_state",
        "stub_func": "asm_call_stub_armv81m",
        "check_func": "check_armv81m_aapcs32_compliance",
        # AAPCS32 GPRs r0-r12 map directly to gpr[0..12].
        "reg_to_field": lambda reg: f"gpr[{int(reg[1:])}]",
    },
}


def _abi_block(yaml_data, dev_source):
    """Return the YAML 'ABI:' block, raising on missing/wrong type."""
    abi = yaml_data.get("ABI")
    if abi is None:
        raise ValueError(f"{dev_source}: YAML missing 'ABI:' block")
    if not isinstance(abi, dict):
        raise ValueError(
            f"{dev_source}: YAML 'ABI:' must be a mapping, got {type(abi).__name__}"
        )
    return abi


def kernel_arch_calling(yaml_data, dev_source):
    """Return the (arch, calling) pair from a kernel's YAML; both are required
    and must form a key in ARCH_CALLINGS."""
    abi = _abi_block(yaml_data, dev_source)
    arch = abi.get("Architecture")
    calling = abi.get("CallingConvention")
    if arch is None or calling is None:
        raise ValueError(
            f"{dev_source}: YAML 'ABI:' must set both 'Architecture:' and "
            f"'CallingConvention:' (got Architecture={arch!r}, "
            f"CallingConvention={calling!r})"
        )
    if (arch, calling) not in ARCH_CALLINGS:
        known = ", ".join(f"({a},{c})" for (a, c) in sorted(ARCH_CALLINGS))
        raise ValueError(
            f"{dev_source}: unknown (Architecture, CallingConvention) pair "
            f"({arch!r}, {calling!r}); expected one of: {known}. "
            f"Add it to ARCH_CALLINGS in scripts/autogen."
        )
    return (arch, calling)


def kernel_features(yaml_data, dev_source):
    """Return the 'ABI.Features:' list (capability tags, keys of ABI_CAPS);
    missing/empty means no requirement, unknown tags raise."""
    abi = _abi_block(yaml_data, dev_source)
    feats = abi.get("Features") or []
    if not isinstance(feats, list):
        raise ValueError(
            f"{dev_source}: YAML 'ABI.Features:' must be a list, "
            f"got {type(feats).__name__}"
        )
    for r in feats:
        if r not in ABI_CAPS:
            known = ", ".join(sorted(ABI_CAPS.keys()))
            raise ValueError(
                f"{dev_source}: unknown feature {r!r} in 'ABI.Features:'; "
                f"expected one of: {known}. Add it to ABI_CAPS in scripts/autogen."
            )
    return feats


def kernel_registers(yaml_data, dev_source):
    """Return the per-register entries in the kernel's 'ABI:' block: the direct
    children carrying a `type:` field (e.g. `x0:`, `rdi:`), skipping the scalar
    metadata keys."""
    abi = _abi_block(yaml_data, dev_source)
    abi_metadata_keys = {"Architecture", "CallingConvention", "Features"}
    regs = {}
    for key, value in abi.items():
        if key in abi_metadata_keys:
            continue
        if not isinstance(value, dict) or "type" not in value:
            raise ValueError(
                f"{dev_source}: unexpected key {key!r} in 'ABI:' block; "
                f"expected one of {sorted(abi_metadata_keys)} or a register "
                f"entry with a 'type:' field"
            )
        regs[key] = value
    return regs


def resolve_buffer_size(reg_info, abi_data):
    """Resolve buffer size from register info, handling cross-references."""
    size = reg_info.get("size_bytes")
    if isinstance(size, str):
        if size in abi_data:
            ref_reg = abi_data[size]
            if ref_reg.get("type") == "scalar" and "test_with" in ref_reg:
                return ref_reg["test_with"]
            raise ValueError(f"Buffer {reg_info} references non-scalar register {size}")
        return int(size)
    elif isinstance(size, int):
        return size
    raise ValueError(f"Cannot resolve buffer size from {reg_info}")


def gen_abicheck():
    """Generate ABI checker tests"""

    # Sources the ABI checker covers in addition to those exported for
    # HOL-Light proofs (armv81m). Shape matches hol_light_asm_joblist
    # (basename, dev_dir, _cflags, arch); the cflags slot is unused here.
    joblist_abicheck_extra = [
        ("keccak_f1600_x4_mve.S", "dev/fips202/armv81m/src", "", "armv81m"),
    ]
    joblist = hol_light_asm_joblist() + joblist_abicheck_extra

    # Bucket joblist entries by their YAML-declared (arch, calling) pair so a
    # single dev source file is fully self-describing.
    jobs_by_arch_calling = {key: [] for key in ARCH_CALLINGS}
    # Map cap -> sorted list of .S files that require it. Drives the
    # generated per-arch abicheck_<arch>.mk at the end of this function.
    cap_to_files = {cap: [] for cap in ABI_CAPS}

    for entry in sorted(joblist, key=lambda j: j[0]):
        assembly_basename, dev_src_dir, _cflags, joblist_arch = entry
        dev_source = f"{dev_src_dir}/{assembly_basename}"
        yaml_data = extract_yaml_from_assembly(dev_source)
        arch_calling = kernel_arch_calling(yaml_data, dev_source)
        # Sanity: the joblist directory architecture should match the YAML.
        if arch_calling[0] != joblist_arch:
            raise ValueError(
                f"{dev_source}: YAML 'ABI.Architecture: {arch_calling[0]}' "
                f"disagrees with joblist arch {joblist_arch!r}"
            )
        upstream_src_dir = upstream_src_dir_for_dev_dir(dev_src_dir)
        for cap in kernel_features(yaml_data, dev_source):
            cap_to_files[cap].append(f"{upstream_src_dir}/{assembly_basename}")
        jobs_by_arch_calling[arch_calling].append((entry, yaml_data))

    # Per-(arch, calling) generated metadata used to assemble checks_all.h:
    # list of (check_name, arch_flags, extra_guards) tuples plus the
    # arch_guard / extra_guards keying.
    all_groups = []  # list of (arch_calling, arch_guard, extra_guards, [(name, arch_flags)])

    # Per-arch-subdir set of check_*.c basenames just generated. Used after the
    # main loop to delete stale check_*.c left over from a renamed/removed
    # kernel (mirrors synchronize_backend's delete= path).
    generated_basenames_by_subdir = {}

    for arch_calling, jobs in jobs_by_arch_calling.items():
        config = ARCH_CALLINGS[arch_calling]
        arch_guard = config["arch_guard"]
        extra_guards = config["extra_guards"]
        generated_functions = []
        generated_basenames_by_subdir.setdefault(config["asm_subdir"], set())

        for (assembly_basename, dev_src_dir, _cflags, _joblist_arch), yaml_data in jobs:
            dev_source = f"{dev_src_dir}/{assembly_basename}"
            features = kernel_features(yaml_data, dev_source)
            registers = kernel_registers(yaml_data, dev_source)
            # Compile-time guards derived from Features. Each capability with a
            # 'guard' (e.g. __AVX2__ for AVX2, __ARM_FEATURE_SHA3 for SHA3)
            # contributes an #if term mirroring the backend's own compile-time
            # gate; runtime gating via mld_sys_check_capability is
            # additional, not a replacement.
            feature_guards = [
                ABI_CAPS[f]["guard"] for f in features if ABI_CAPS[f]["guard"]
            ]

            function_name = yaml_data.get("Name")
            c_function_name = "mld_" + function_name
            # The YAML Name already encodes the architecture (e.g.
            # ntt_aarch64_asm, ntt_avx2_asm), so use it directly rather than
            # appending the architecture suffix again.
            check_name = function_name

            def gen_c_test(
                function_name=function_name,
                check_name=check_name,
                c_function_name=c_function_name,
                yaml_data=yaml_data,
                features=features,
                registers=registers,
                feature_guards=feature_guards,
                config=config,
                arch_guard=arch_guard,
                extra_guards=extra_guards,
            ):
                yield from gen_header()
                yield "#include <stdio.h>"
                yield ""
                yield f'#include "../abicheck_{config["asm_subdir"]}.h"'
                yield f'#include "../checks_{config["asm_subdir"]}_all.h"'
                yield ""

                # Full guard set: architecture macro, then any calling-
                # convention-specific guards (e.g. MLD_SYSV_ABI_SUPPORTED for
                # x86_64 SysV), then any feature-derived guards (e.g.
                # __ARM_FEATURE_SHA3 for SHA3 kernels).
                conds = (
                    [f"defined({arch_guard})"]
                    + [f"defined({g})" for g in extra_guards]
                    + [f"defined({g})" for g in feature_guards]
                )
                cond = " && ".join(conds)

                yield f"#if {cond}"
                yield ""
                yield '#include "../../../notrandombytes/notrandombytes.h"'
                yield ""
                yield f"typedef {config['state_type']} reg_state;"
                yield ""
                qualifier = config["fn_qualifier"]
                yield f"{qualifier} {yaml_data.get('Signature')};"
                yield ""
                yield f"int check_{check_name}(void)"
                yield "{"

                yield "  int test_iter;"
                yield "  reg_state input_state, output_state;"
                yield "  int violations;"

                sorted_registers = sorted(registers.items())

                buffer_info = []
                for reg_name, reg_info in sorted_registers:
                    if reg_info.get("type") == "buffer":
                        size_bytes = resolve_buffer_size(reg_info, registers)
                        buffer_name = f"buf_{reg_name}"
                        description = reg_info.get("description", "")
                        yield f"  MLD_ALIGN uint8_t {buffer_name}[{size_bytes}]; /* {description} */"
                        buffer_info.append((reg_name, buffer_name, size_bytes))

                # Emit a runtime mld_sys_check_capability() skip (returning
                # MLD_ABICHECK_SKIPPED) per Features cap that has a syscap; caps
                # without a syscap (build-time-only) are no-ops here.
                # Emitted after the buffer decls to keep C90 mixed-decl-and-code happy.
                for cap in features:
                    cap_enum = ABI_CAPS[cap]["syscap"]
                    if cap_enum is None:
                        continue
                    cap_desc = ABI_CAPS[cap]["description"]
                    yield ""
                    yield f"  if (!mld_sys_check_capability({cap_enum}))"
                    yield "  {"
                    yield (
                        f'    fprintf(stderr, "ABI check {check_name}: '
                        f'host lacks {cap_desc}, skipping\\n");'
                    )
                    yield "    return MLD_ABICHECK_SKIPPED;"
                    yield "  }"

                yield ""
                yield "  for (test_iter = 0; test_iter < MLD_ABICHECK_NUM_TESTS; test_iter++)"
                yield "  {"
                yield "    /* Initialize random register state */"
                yield f"    {config['init_func']}(&input_state);"
                yield ""

                for reg_name, buffer_name, size_bytes in buffer_info:
                    yield f"    randombytes({buffer_name}, {size_bytes});"

                yield ""
                yield "    /* Set up register state for function arguments */"

                reg_to_field = config["reg_to_field"]
                for reg_name, reg_info in sorted_registers:
                    field = reg_to_field(reg_name)
                    if reg_info.get("type") == "buffer":
                        # Cast pointer to gpr_int_type (uint64_t / uint32_t on
                        # armv81m) so the conversion lands at the slot width.
                        yield f"    input_state.{field} = ({config['gpr_int_type']}){f'buf_{reg_name}'};"
                    elif reg_info.get("type") == "scalar":
                        test_with = reg_info.get("test_with")
                        if test_with:
                            yield f"    input_state.{field} = {test_with};"

                yield ""
                yield "    /* Call function through ABI test stub */"
                yield (
                    f"    {config['stub_func']}(&input_state, &output_state, "
                    f"({qualifier} void (*)(void)){c_function_name});"
                )
                yield ""
                yield "    /* Check ABI compliance */"
                yield (
                    f"    violations = {config['check_func']}("
                    f"&input_state, &output_state, MLD_ABICHECK_VERBOSE);"
                )
                yield "    if (violations > 0) {"
                yield f'      fprintf(stderr, "ABI test FAILED for {function_name} (iteration %d): %d violations\\n",'
                yield "             test_iter + 1, violations);"
                yield "      return MLD_ABICHECK_FAILED;"
                yield "    }"
                yield "  }"
                yield ""
                yield "  return MLD_ABICHECK_PASSED;"
                yield "}"
                yield ""
                # No #else: components.mk only globs check_*.c from the
                # active arch's checks/ directory, and the forward decl in
                # checks_<arch>_all.h is gated on the same `cond`, so an
                # inactive-arch fallback would be unreachable.
                yield f"#endif /* {cond} */"
                yield ""

            output_file = (
                f"test/abicheck/{config['asm_subdir']}/checks/check_{check_name}.c"
            )
            update_file(output_file, "\n".join(gen_c_test()), force_format=True)
            generated_functions.append((check_name, feature_guards))
            generated_basenames_by_subdir[config["asm_subdir"]].add(
                f"check_{check_name}.c"
            )

        all_groups.append((arch_calling, arch_guard, extra_guards, generated_functions))

    # Prune stale check_*.c left from a renamed/removed kernel. components.mk
    # globs check_*.c into the abicheck build, so a stale file would silently
    # compile and link against an out-of-date YAML signature. Mirrors the
    # delete= path in synchronize_backend.
    for asm_subdir, generated_basenames in generated_basenames_by_subdir.items():
        subdir = f"test/abicheck/{asm_subdir}/checks"
        for f in get_files(f"{subdir}/check_*.c"):
            if os.path.basename(f) not in generated_basenames:
                update_via_remove(f)

    # Generate per-arch checks_<arch>_all.h headers, one per arch_guard. Each
    # header lives at test/abicheck/<asm_subdir>/checks_<asm_subdir>_all.h (one
    # level above that arch's check_*.c sources in checks/) and is selected
    # from abicheck.c via a single #if-chain on the arch guards.
    #
    # Group the per-(arch, calling) rows by arch_guard, so kernels for the
    # same architecture but different calling conventions share one
    # all_checks[] array on that arch (gated per-entry on the calling
    # convention's extra_guards). Today each arch has only one calling
    # convention, but this keeps the door open to e.g. x86_64 + SysV plus
    # x86_64 + MSx64 coexisting in the same build.
    groups_by_arch = {}
    asm_subdir_by_arch = {}
    for arch_calling, arch_guard, extra_guards, functions in all_groups:
        groups_by_arch.setdefault(arch_guard, []).append(
            (arch_calling, extra_guards, functions)
        )
        asm_subdir_by_arch[arch_guard] = ARCH_CALLINGS[arch_calling]["asm_subdir"]

    def gen_checks_arch_header(arch_guard, groups, asm_subdir):
        guard_macro = f"MLD_TEST_ABICHECK_CHECKS_{asm_subdir.upper()}_ALL_H"
        yield from gen_header()
        yield ""
        yield f"#ifndef {guard_macro}"
        yield f"#define {guard_macro}"
        yield ""
        yield "#include <stddef.h>"
        yield '#include "../abicheck_common.h"'
        yield ""

        # Wrap the whole body in the arch guard so accidental inclusion on a
        # foreign architecture is harmless (this header is selected by an
        # #if-chain in abicheck.c, but defense in depth).
        yield f"#if defined({arch_guard})"
        yield ""

        # Emit one row per function, consolidating consecutive functions that
        # share the same feature_guards into a single #if/#endif block (rather
        # than one block per function). This collapses a uniform-feature group
        # (e.g. x86_64, where every kernel is gated on __AVX2__) into a single
        # block, while still splitting where guards differ (e.g. aarch64, where
        # SHA3 and non-SHA3 kernels are interleaved).
        def emit_feature_grouped(functions, render):
            i = 0
            while i < len(functions):
                guards = functions[i][1]
                j = i
                while j < len(functions) and functions[j][1] == guards:
                    j += 1
                for g in guards:
                    yield f"#if defined({g})"
                for func_name, _ in functions[i:j]:
                    yield render(func_name)
                for _ in guards:
                    yield "#endif"
                i = j

        # Forward-declare each check_<kernel>(void) for kernels in THIS arch.
        # Gate on the same extra+feature guards as the registry entry below;
        # this matches the #if guarding the function definition in
        # check_<kernel>.c, so the inactive-arch fallback can be omitted there.
        for _arch_calling, extra_guards, functions in groups:
            if extra_guards:
                inner = " && ".join(f"defined({g})" for g in extra_guards)
                yield f"#if {inner}"
            yield from emit_feature_grouped(
                functions, lambda fn: f"int check_{fn}(void);"
            )
            if extra_guards:
                yield f"#endif /* {inner} */"

        if any(functions for _, _, functions in groups):
            yield ""
            yield "static const abicheck_entry_t all_checks[] = {"

            for _arch_calling, extra_guards, functions in groups:
                if not functions:
                    continue
                if extra_guards:
                    inner = " && ".join(f"defined({g})" for g in extra_guards)
                    yield f"#if {inner}"
                yield from emit_feature_grouped(
                    functions, lambda fn: f'    {{"{fn}", check_{fn}}},'
                )
                if extra_guards:
                    yield f"#endif /* {inner} */"

            yield "    {NULL, NULL} /* Sentinel */"
            yield "};"

        yield ""
        yield f"#endif /* defined({arch_guard}) */"
        yield ""
        yield f"#endif /* !{guard_macro} */"
        yield ""

    for arch_guard, groups in groups_by_arch.items():
        asm_subdir = asm_subdir_by_arch[arch_guard]
        update_file(
            f"test/abicheck/{asm_subdir}/checks_{asm_subdir}_all.h",
            "\n".join(gen_checks_arch_header(arch_guard, groups, asm_subdir)),
            force_format=True,
        )

    # Emit per-capability CFLAGS injection for the build, split per arch:
    # one test/abicheck/<arch>/abicheck_<arch>.mk per arch with caps, plus a
    # top-level test/abicheck/abicheck.mk that includes them all. Single source
    # of truth alongside the C-time guard (arch_flags) and runtime cap check.
    caps_by_subdir = {}
    for cap in sorted(ABI_CAPS):
        caps_by_subdir.setdefault(ABI_CAPS[cap]["asm_subdir"], []).append(cap)

    def gen_features_mk(asm_subdir, caps):
        yield from gen_yaml_header()
        yield from gen_yaml_autogen_warning()
        yield "#"
        yield "# Edit the YAML 'ABI.Features:' list in dev/<arch>/src/<kernel>.S"
        yield "# and re-run scripts/autogen instead."
        yield "#"
        yield "# For each capability declared by a kernel's ABI.Features list, this"
        yield "# file appends the capability's CFLAGS to that kernel's .S object"
        yield "# under mldsa/src/."
        yield ""
        yield "# Default each cap's file list to empty so the unconditional appends"
        yield "# below are safe even when a cap has no kernels on this arch."
        for cap in caps:
            yield f"ABICHECK_REQ_{cap}_FILES :="
        for cap in caps:
            # Files needing this cap's CFLAGS are: the kernel exports declared
            # by YAML, plus any always-on auxiliary files (callstub/selftest)
            # listed in ABI_CAPS[cap]['aux_files'].
            files = sorted(
                set(cap_to_files[cap]) | set(ABI_CAPS[cap].get("aux_files", []))
            )
            if not files:
                continue
            cflags = " ".join(ABI_CAPS[cap]["cflags"])
            yield ""
            yield f"# {cap}: {ABI_CAPS[cap]['description']}"
            yield f"ABICHECK_REQ_{cap}_FILES := \\"
            for i, f in enumerate(files):
                sep = " \\" if i + 1 < len(files) else ""
                yield f"  {f}{sep}"
            yield (
                f"ABICHECK_REQ_{cap}_OBJS := "
                f"$(call MAKE_OBJS,$(ABICHECK_DIR),$(ABICHECK_REQ_{cap}_FILES))"
            )
            yield f"$(ABICHECK_REQ_{cap}_OBJS): CFLAGS += {cflags}"
        yield ""

    for asm_subdir, caps in caps_by_subdir.items():
        update_file(
            f"test/abicheck/{asm_subdir}/abicheck_{asm_subdir}.mk",
            "\n".join(gen_features_mk(asm_subdir, caps)),
        )

    def gen_abicheck_mk():
        yield from gen_yaml_header()
        yield from gen_yaml_autogen_warning()
        yield "#"
        yield "# Includes the per-arch abicheck_<arch>.mk, each of which appends"
        yield "# its capabilities' CFLAGS to the matching .S objects."
        yield ""
        for asm_subdir in sorted(caps_by_subdir):
            yield f"include test/abicheck/{asm_subdir}/abicheck_{asm_subdir}.mk"
        yield ""

    update_file("test/abicheck/abicheck.mk", "\n".join(gen_abicheck_mk()))


def _main():
    slothy_choices = [
        "ntt_aarch64_asm",
        "intt_aarch64_asm",
        "mld_polyvecl_pointwise_acc_montgomery_l4_aarch64_asm",
        "mld_polyvecl_pointwise_acc_montgomery_l5_aarch64_asm",
        "mld_polyvecl_pointwise_acc_montgomery_l7_aarch64_asm",
        "pointwise_montgomery_aarch64_asm",
        "poly_caddq_aarch64_asm",
        "poly_chknorm_aarch64_asm",
        "poly_decompose_32_aarch64_asm",
        "poly_decompose_88_aarch64_asm",
        "poly_use_hint_32_aarch64_asm",
        "poly_use_hint_88_aarch64_asm",
        "polyz_unpack_17_aarch64_asm",
        "polyz_unpack_19_aarch64_asm",
        "rej_uniform_aarch64_asm",
        "rej_uniform_eta2_aarch64_asm",
        "rej_uniform_eta4_aarch64_asm",
    ]

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--dry-run", default=False, action="store_true")
    parser.add_argument(
        "--update-hol-light-bytecode", default=False, action="store_true"
    )
    parser.add_argument("--slothy", nargs="*", default=None, choices=slothy_choices)
    parser.add_argument("--aarch64-clean", default=False, action="store_true")
    parser.add_argument("--no-simplify", default=False, action="store_true")
    parser.add_argument(
        "--force-cross",
        nargs="*",
        default=None,
        choices=sorted(FORCE_CROSS_ALL_ARCHES),
        metavar="ARCH",
        help="Treat a missing cross toolchain as an error instead of silently "
        "skipping. Without a value, requires all architectures "
        f"({', '.join(sorted(FORCE_CROSS_ALL_ARCHES))}); or pass a "
        "space-separated list to require only those, "
        "e.g. --force-cross aarch64 x86_64.",
    )
    parser.add_argument(
        "--x86-64-syntax",
        type=str,
        choices=["att", "intel"],
        default="att",
        help="Assembly syntax for x86_64 disassembly output (att or intel)",
    )
    parser.add_argument(
        "--test-vectors",
        default=False,
        action="store_true",
        help="Generate expected test vectors using ACVP binaries",
    )
    parser.add_argument(
        "--test-vectors-msg",
        type=str,
        default="This is a test message for ML-DSA digital signature algorithm!",
        help="Message to use for test vector generation",
    )
    parser.add_argument(
        "--test-vectors-ctx",
        type=str,
        default="test_context_123",
        help="Context to use for test vector generation",
    )

    args = parser.parse_args()

    os.chdir(os.path.join(os.path.dirname(__file__), ".."))

    if args.slothy == []:
        args.slothy = slothy_choices

    force_cross = resolve_force_cross(args.force_cross)

    def sync_backends():
        synchronize_backends(
            clean=args.aarch64_clean,
            no_simplify=args.no_simplify,
            force_cross=force_cross,
            x86_64_syntax=args.x86_64_syntax,
        )

    def sync_backends_final():
        synchronize_backends(
            clean=args.aarch64_clean,
            delete=True,
            force_cross=force_cross,
            no_simplify=args.no_simplify,
            x86_64_syntax=args.x86_64_syntax,
        )

    def gen_zeta_tables():
        gen_c_zeta_file()
        gen_aarch64_zeta_file()
        gen_aarch64_hol_light_zeta_file()
        gen_aarch64_rej_uniform_table()
        gen_hol_light_rej_uniform_table()
        gen_aarch64_rej_uniform_eta_table()
        gen_aarch64_polyz_unpack_table()
        gen_hol_light_polyz_unpack_table()
        gen_avx2_hol_light_zeta_file()
        gen_avx2_zeta_file()
        gen_avx2_rej_uniform_table()
        gen_hol_light_keccak_constants_file()
        gen_aarch64_keccak_constants_c_file()
        gen_avx2_keccak_constants_c_file()
        gen_avx2_keccak_hol_light_constants_file()

    def gen_monolithic():
        gen_monolithic_source_file()
        gen_monolithic_asm_file()

    hol_light_asm_supported = platform.machine().lower() in [
        "arm64",
        "aarch64",
        "x86_64",
    ]

    # Build step list: (description, function, enabled)
    # If enabled is False, step is skipped
    steps = [
        ("Generate citations", gen_citations),
        ("Generate OQS META.yml files", gen_oqs_meta_files),
        ("Normalize assembly macro syntax", normalize_asm_macro_syntax),
        (
            "Generate SLOTHY optimized assembly",
            lambda: gen_slothy(args.slothy),
            args.slothy is not None and not args.dry_run,
        ),
        ("Check assembly register aliases", check_asm_register_aliases),
        ("Check assembly loop labels", check_asm_loop_labels),
        ("Generate zeta and lookup tables", gen_zeta_tables),
        ("Generate HOL Light assembly", gen_hol_light_asm, hol_light_asm_supported),
        (
            "Check HOL Light assembly is autogenerated",
            check_hol_light_asm_autogenerated,
        ),
        ("Synchronize backends", sync_backends),
        ("Generate header guards", gen_header_guards),
        ("Complete final backend synchronization", sync_backends_final),
        (
            "Update HOL Light bytecode",
            partial(update_hol_light_bytecode, force_cross=force_cross),
            args.update_hol_light_bytecode,
        ),
        ("Generate monolithic source files", gen_monolithic),
        ("Generate undefs", gen_undefs),
        ("Generate test configs", gen_test_configs),
        (
            "Generate test vectors",
            lambda: gen_test_vectors(args.test_vectors_msg, args.test_vectors_ctx),
            args.test_vectors,
        ),
        ("Generate ABI checker tests", gen_abicheck),
        ("Check macro typos", check_macro_typos),
        ("Generate preprocessor comments", gen_preprocessor_comments),
        # Formatting should be the last step
        ("Format files", lambda: format_files(args.dry_run)),
    ]

    # When --test-vectors is set, only run test vector generation + formatting
    if args.test_vectors:
        steps = [s for s in steps if s[0] in ("Generate test vectors", "Format files")]

    global _progress, _main_task
    with Progress(
        BarColumn(),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        TextColumn("{task.description}"),
        console=console,
    ) as progress:
        _progress = progress
        _main_task = progress.add_task("autogen", total=len(steps) + 1)

        for step in steps:
            desc, func = step[0], step[1]
            enabled = step[2] if len(step) > 2 else True
            high_level_task(desc)
            if enabled:
                func()
            high_level_status(desc, skipped=not enabled)

        high_level_task("Write files")
        finalize(args.dry_run)
        txt = (
            "Finalize and check files are up to date"
            if args.dry_run
            else "Finalize and write files"
        )
        high_level_status(txt)

        _progress.update(_main_task, description="[green] Done ✓[/]")
        _progress = None

    return print_check_errors()


if __name__ == "__main__":
    sys.exit(0 if _main() else 1)