autogen

#!/usr/bin/env python3
# Copyright (c) The mlkem-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 = ""

_progress_lock = threading.Lock()

modulus = 3329
root_of_unity = 17
montgomery_factor = pow(2, 16, 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)/mlkem/([^/\s]+?)\.o"
)
_RE_FUNC_SYMBOL = re.compile(r"MLK_ASM_FN_SYMBOL\((.*)\)")
_RE_LABEL = re.compile(r"^(\w+):")
_RE_CONFIG_NAME = re.compile(r"\* Name:\s+(\w+)")
_RE_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()


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 mlkem-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 mlkem-native repository."
    yield " *          Do not modify it directly."
    yield " */"


def gen_header():
    yield "/*"
    yield " * Copyright (c) The mlkem-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 mlkem-native repository."
    yield " *          Do not modify it directly."
    yield " *)"
    yield ""


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


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.parseString(exp, parseAll=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 = []
    for i in range(128):
        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, 7)] for i in range(128))


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 MLK_ALIGN const int16_t mlk_zetas[128] = {"
        yield from map(lambda t: str(t) + ",", gen_c_zetas())
        yield "};"
        yield ""

    update_file("mlkem/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**16) / modulus) // 2
    return root, root_twisted


def gen_aarch64_root_of_unity_for_block(layer, block, inv=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 6
    # in this function.
    log = bitreverse(pow(2, layer) + block, 7)
    if inv is True:
        log = -log
    root, root_twisted = prepare_root_for_barrett(pow(root_of_unity, log, modulus))
    return root, root_twisted


def gen_aarch64_fwd_ntt_zetas_layer12345():
    # 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,7 are merged, but we emit roots for 4,5
    # in separate arrays than those for 6,7
    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 (0, 0)  # Padding


def gen_aarch64_fwd_ntt_zetas_layer67():
    # Layers 4,5,6,7 are merged, but we emit roots for 4,5
    # in separate arrays than those for 6,7
    for block in range(8):

        def double_ith(t, i):
            yield from (t[i], t[i])

        # Ordering of blocks is adjusted to suit the transposed internal
        # presentation of the data
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 0), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 1), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 2), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 3), i
            )
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 0), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 2), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 4), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 6), i
            )
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 1), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 3), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 5), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 7), i
            )


def gen_aarch64_inv_ntt_zetas_layer12345():
    # Layers 4,5,6,7 are merged, but we emit roots for 4,5
    # in separate arrays than those for 6,7
    for block in range(8):  # There are 8 blocks in Layer 4
        yield from gen_aarch64_root_of_unity_for_block(3, block, inv=True)
        yield from gen_aarch64_root_of_unity_for_block(4, 2 * block + 0, inv=True)
        yield from gen_aarch64_root_of_unity_for_block(4, 2 * block + 1, inv=True)
        yield from (0, 0)  # Padding

    # Layers 1,2,3 are merged
    yield from gen_aarch64_root_of_unity_for_block(0, 0, inv=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 (0, 0)  # Padding


def gen_aarch64_inv_ntt_zetas_layer67():
    # Layers 4,5,6,7 are merged, but we emit roots for 4,5
    # in separate arrays than those for 6,7
    for block in range(8):

        def double_ith(t, i):
            yield from (t[i], t[i])

        # Ordering of blocks is adjusted to suit the transposed internal
        # presentation of the data
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 0, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 1, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 2, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(5, 4 * block + 3, inv=True), i
            )
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 0, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 2, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 4, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 6, inv=True), i
            )
        for i in range(2):
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 1, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 3, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 5, inv=True), i
            )
            yield from double_ith(
                gen_aarch64_root_of_unity_for_block(6, 8 * block + 7, inv=True), i
            )


def gen_aarch64_mulcache_twiddles():
    for idx in range(0, 128):
        root = pow(root_of_unity, 2 * bitreverse(idx, 7) + 1, modulus)
        yield prepare_root_for_barrett(root)[0]


def gen_aarch64_mulcache_twiddles_twisted():
    for idx in range(0, 128):
        root = pow(root_of_unity, 2 * bitreverse(idx, 7) + 1, modulus)
        yield prepare_root_for_barrett(root)[1]


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 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_layer12345 = define `ntt_zetas_layer12345:int list = ["
        yield from print_hol_light_array(gen_aarch64_fwd_ntt_zetas_layer12345())
        yield "]`;;"
        yield ""
        yield "let ntt_zetas_layer67 = define `ntt_zetas_layer67:int list = ["
        yield from print_hol_light_array(gen_aarch64_fwd_ntt_zetas_layer67())
        yield "]`;;"
        yield ""
        yield "let intt_zetas_layer12345 = define `intt_zetas_layer12345:int list = ["
        yield from print_hol_light_array(gen_aarch64_inv_ntt_zetas_layer12345())
        yield "]`;;"
        yield ""
        yield "let intt_zetas_layer67 = define `intt_zetas_layer67:int list = ["
        yield from print_hol_light_array(gen_aarch64_inv_ntt_zetas_layer67())
        yield "]`;;"
        yield ""
        yield "let mulcache_zetas = define `mulcache_zetas:int list = ["
        yield from print_hol_light_array(gen_aarch64_mulcache_twiddles())
        yield "]`;;"
        yield ""
        yield ""
        yield "let mulcache_zetas_twisted = define `mulcache_zetas_twisted:int list = ["
        yield from print_hol_light_array(gen_aarch64_mulcache_twiddles_twisted())
        yield "]`;;"
        yield ""

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


def gen_aarch64_zeta_file():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLK_ARITH_BACKEND_AARCH64) && \\"
        yield "    !defined(MLK_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 "MLK_ALIGN const int16_t mlk_aarch64_ntt_zetas_layer12345[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_fwd_ntt_zetas_layer12345())
        yield "};"
        yield ""
        yield "MLK_ALIGN const int16_t mlk_aarch64_ntt_zetas_layer67[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_fwd_ntt_zetas_layer67())
        yield "};"
        yield ""
        yield "MLK_ALIGN const int16_t mlk_aarch64_invntt_zetas_layer12345[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_inv_ntt_zetas_layer12345())
        yield "};"
        yield ""
        yield "MLK_ALIGN const int16_t mlk_aarch64_invntt_zetas_layer67[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_inv_ntt_zetas_layer67())
        yield "};"
        yield ""
        yield "MLK_ALIGN const int16_t mlk_aarch64_zetas_mulcache_native[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_mulcache_twiddles())
        yield "};"
        yield ""
        yield "MLK_ALIGN const int16_t mlk_aarch64_zetas_mulcache_twisted_native[] = {"
        yield from map(lambda t: str(t) + ",", gen_aarch64_mulcache_twiddles_twisted())
        yield "};"
        yield ""
        yield "#else"
        yield ""
        yield "MLK_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_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. 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 -1
        idxs = idxs + [255] * (16 - len(idxs))
        yield idxs


def gen_aarch64_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 " * See autogen for details."
        yield " *)"
        yield ""
        yield "let mlkem_rej_uniform_table = (REWRITE_RULE[MAP] o define)"
        yield "  `mlkem_rej_uniform_table:byte list = MAP word ["
        data = [i for idxs in gen_rej_uniform_table_rows() for i in idxs]
        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/mlkem_rej_uniform_table.ml", "\n".join(gen())
    )


def gen_aarch64_rej_uniform_table():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLK_ARITH_BACKEND_AARCH64) && \\"
        yield "    !defined(MLK_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 "MLK_ALIGN const uint8_t mlk_rej_uniform_table[] = {"
        for i, idxs in enumerate(gen_rej_uniform_table_rows()):
            yield ",".join(map(str, idxs)) + f" /* {i} */,"
        yield "};"
        yield ""
        yield "#else"
        yield ""
        yield "MLK_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_avx2_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 " * See autogen for details."
        yield " *)"
        yield ""
        yield "let mlkem_rej_uniform_table = (REWRITE_RULE[MAP] o define)"
        yield "  `mlkem_rej_uniform_table:byte list = MAP word ["
        data = [i for idxs in gen_rej_uniform_table_rows() for i in idxs]
        yield from print_hol_light_array(data, as_int=False, entries_per_line=16, pad=3)
        yield "]`;;"
        yield ""

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


def gen_avx2_rej_uniform_table():
    def gen():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \\"
        yield "    !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "arith_native_x86_64.h"'
        yield ""
        yield "/*"
        yield " * Lookup table used by rejection sampling of the public matrix."
        yield " * See autogen for details."
        yield " */"
        yield "MLK_ALIGN const uint8_t mlk_rej_uniform_table[] = {"
        for i, idxs in enumerate(gen_rej_uniform_table_rows()):
            yield ",".join(map(str, idxs)) + f" /* {i} */,"
        yield "};"
        yield ""
        yield "#else"
        yield ""
        yield "MLK_EMPTY_CU(avx2_rej_uniform_table)"
        yield ""
        yield "#endif"
        yield ""

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


def signed_reduce_u16(x):
    x = x % 2**16
    if x >= 2**15:
        x -= 2**16
    return x


def prepare_root_for_montmul(root):
    """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)
    root_twisted = signed_reduce_u16(root * pow(modulus, -1, 2**16))
    return root, root_twisted


def gen_avx2_root_of_unity_for_block(layer, block, inv=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, 7)
    if inv is True:
        log = -log
    root, root_twisted = prepare_root_for_montmul(pow(root_of_unity, log, modulus))
    return root, root_twisted


def gen_avx2_fwd_ntt_zetas():
    def gen_twiddles(layer, block, repeat):
        """Generates twisted twiddle, then twiddle, for given layer and block.
        Repeat both the given number of times."""
        root, root_twisted = gen_avx2_root_of_unity_for_block(layer, block)
        return [root] * repeat, [root_twisted] * repeat

    def gen_twiddles_many(layer, block_base, block_offsets, repeat):
        """Generates twisted twiddles, then twiddles, of each (layer, block_base + i)
        pair for i in block_offsets. Each twiddle is repeated `repeat` times."""
        root_pairs = list(
            map(lambda x: gen_twiddles(layer, block_base + x, repeat), block_offsets)
        )
        yield from (r for pair in root_pairs for r in pair[1])
        yield from (r for pair in root_pairs for r in pair[0])

    # Layers 1 twiddle
    yield from gen_twiddles_many(0, 0, range(1), 4)
    # Padding so that the subsequent twiddles are 16-byte aligned
    yield from [0] * 8

    # Layer 2-7 twiddles, separated by whether they belong to the upper or lower half
    for i in range(2):
        yield from gen_twiddles_many(1, i * (2**0), range(1), 16)
        yield from gen_twiddles_many(2, i * (2**1), range(2), 8)
        yield from gen_twiddles_many(3, i * (2**2), range(4), 4)
        yield from gen_twiddles_many(4, i * (2**3), range(8), 2)
        yield from gen_twiddles_many(5, i * (2**4), range(16), 1)
        yield from gen_twiddles_many(6, i * (2**5), range(0, 32, 2), 1)
        yield from gen_twiddles_many(6, i * (2**5), range(1, 32, 2), 1)


def gen_avx2_mulcache_twiddles():
    for i in range(2):
        for idx in range(16):
            root, root_twisted = prepare_root_for_montmul(
                pow(root_of_unity, bitreverse(64 + 32 * i + 2 * idx, 7), modulus)
            )
            yield root

        for idx in range(16):
            root, root_twisted = prepare_root_for_montmul(
                pow(root_of_unity, bitreverse(64 + 32 * i + 2 * idx + 1, 7), modulus)
            )
            yield root

    for i in range(2):
        for idx in range(16):
            root, root_twisted = prepare_root_for_montmul(
                pow(root_of_unity, bitreverse(64 + 32 * i + 2 * idx, 7), modulus)
            )
            yield root_twisted

        for idx in range(16):
            root, root_twisted = prepare_root_for_montmul(
                pow(root_of_unity, bitreverse(64 + 32 * i + 2 * idx + 1, 7), modulus)
            )
            yield root_twisted


def gen_avx2_zetas_qdata():
    # Twiddle permutation indices for INTT (used with vpshufb and vpermd)
    def offset_revidxb():
        # vpshufb control: reverses 16-bit elements within each 128-bit lane
        # Each int16 = (hi_byte_idx << 8) | lo_byte_idx
        for _ in range(2):  # Same for both lanes
            for i in range(8):
                lo = 2 * (7 - i)
                yield ((lo + 1) << 8) | lo

    def offset_revidxd():
        # vpermd control: reverses 32-bit dwords across YMM register
        # Stored as int16 pairs (idx, 0) forming 32-bit indices
        for i in range(8):
            yield 7 - i
            yield 0

    revidxb = list(offset_revidxb())
    revidxd = list(offset_revidxd())
    ntt_zetas = list(gen_avx2_fwd_ntt_zetas())
    mulcache_twiddles = list(gen_avx2_mulcache_twiddles())

    revidxb_idx = 0
    revidxd_idx = revidxb_idx + len(revidxb)
    ntt_zetas_idx = revidxd_idx + len(revidxd)
    mulcache_twiddles_idx = ntt_zetas_idx + len(ntt_zetas)

    constants = revidxb + revidxd + ntt_zetas + mulcache_twiddles
    offsets = {
        "revidxb": revidxb_idx,
        "revidxd": revidxd_idx,
        "zetas": ntt_zetas_idx,
        "mulcache_twiddles": mulcache_twiddles_idx,
    }
    return constants, offsets


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 qdata_full = define `qdata_full:int list = ["
        constants, _ = gen_avx2_zetas_qdata()
        yield from print_hol_light_array(constants)
        yield "]`;;"
        yield ""

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


# =============================================================================
# AVX2 Compression/Decompression Constants
# =============================================================================
#
# Constants used by the AVX2 polynomial compression and decompression functions.
# These include shuffle indices, masks, and shift values for the various bit
# widths (d4, d5, d10, d11).
#
# Each function has a single merged constant array (e.g., mlk_compress_d5_data,
# mlk_decompress_d11_data) containing all required constants at fixed offsets.


def gen_decompress_d4_shufbidx():
    # Shuffle index for poly_decompress_d4_avx2.
    # Each input byte contains two 4-bit coefficients. To unpack, we replicate
    # each byte 4 times (will extract high/low nibbles with mask afterwards).
    # Pattern: byte[i] -> positions [4i, 4i+1, 4i+2, 4i+3]
    # Only 32 bytes needed; assembly duplicates to both lanes via vinserti128.
    for byte_idx in range(8):
        for _ in range(4):
            yield byte_idx


def gen_compress_d4_permdidx():
    # Permutation index for poly_compress_d4_avx2 (vpermd).
    # Corresponds to _mm256_set_epi32(7, 3, 6, 2, 5, 1, 4, 0),
    # which in memory order (low to high dword) is: 0, 4, 1, 5, 2, 6, 3, 7.
    # Stored as bytes in little-endian order.
    for dword in [0, 4, 1, 5, 2, 6, 3, 7]:
        yield from dword.to_bytes(4, "little")


def gen_compress_d5_shufbidx():
    # Shuffle index for poly_compress_d5_avx2.
    # 5-bit compression: 8 coefficients -> 5 bytes (40 bits).
    # Also used as vpblendvb mask (high bit selects source).
    # 255 (0xFF) means "zero this byte" in vpshufb / "select src2" in vpblendvb.
    yield from [0, 1, 2, 3, 4, 255, 255, 255, 255, 255, 8, 9, 10, 11, 12, 255]
    yield from [9, 10, 11, 12, 255, 0, 1, 2, 3, 4, 255, 255, 255, 255, 255, 8]


def gen_compress_d10_shufbidx():
    # Shuffle index for poly_compress_d10_avx2.
    # 10-bit compression: 4 coefficients -> 5 bytes (40 bits).
    # Lane 0: extract bytes 0-4 and 8-12, Lane 1: wrap-around for blending
    # 255 (0xFF) means "zero this byte" in vpshufb
    yield from [0, 1, 2, 3, 4, 8, 9, 10, 11, 12, 255, 255, 255, 255, 255, 255]
    yield from [9, 10, 11, 12, 255, 255, 255, 255, 255, 255, 0, 1, 2, 3, 4, 8]


def gen_compress_d11_srlvqidx():
    # Variable qword right shift counts for poly_compress_d11_avx2.
    # From _mm256_set_epi64x(30, 10, 30, 10)
    for v in [10, 30, 10, 30]:
        yield v & 0xFF
        yield (v >> 8) & 0xFF
        yield (v >> 16) & 0xFF
        yield (v >> 24) & 0xFF
        yield 0
        yield 0
        yield 0
        yield 0


def gen_compress_d11_shufbidx():
    # Shuffle index for poly_compress_d11_avx2.
    # 11-bit compression: 8 coefficients -> 11 bytes (88 bits).
    # Also used as vpblendvb mask (high bit selects source).
    # 255 (0xFF) means "zero this byte" in vpshufb / "select src2" in vpblendvb.
    # From _mm256_set_epi8(4,3,2,1,0,0,-1,-1,-1,-1,10,9,8,7,6,5,
    #                      -1,-1,-1,-1,-1,10,9,8,7,6,5,4,3,2,1,0)
    yield from [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 255, 255, 255, 255, 255]
    yield from [5, 6, 7, 8, 9, 10, 255, 255, 255, 255, 0, 0, 1, 2, 3, 4]


def gen_decompress_d10_shufbidx():
    # Shuffle index for poly_decompress_d10_avx2.
    # 10-bit decompression: coefficient i uses byte floor(10*i/8).
    # Pattern per lane selects source bytes for each 16-bit output slot.
    yield from [0, 1, 1, 2, 2, 3, 3, 4, 5, 6, 6, 7, 7, 8, 8, 9]
    yield from [2, 3, 3, 4, 4, 5, 5, 6, 7, 8, 8, 9, 9, 10, 10, 11]


def gen_decompress_d5_shufbidx():
    # Shuffle index for poly_decompress_d5_avx2.
    # 5-bit decompression: coefficient i at bit position 5*i uses byte floor(5*i/8).
    # Pattern per lane selects source bytes for each 16-bit output slot.
    yield from [0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4, 4]
    yield from [5, 5, 5, 6, 6, 6, 6, 7, 7, 8, 8, 8, 8, 9, 9, 9]


def gen_decompress_d5_mask():
    # Bit masks for poly_decompress_d5_avx2.
    # Each 16-bit mask isolates 5 bits at different alignments within 16-bit words.
    # Masks: 31, 992, 124, 3968, 496, 62, 1984, 248 (repeated for both lanes)
    masks_16bit = [31, 992, 124, 3968, 496, 62, 1984, 248]
    for m in masks_16bit:
        yield m & 0xFF
        yield (m >> 8) & 0xFF
    for m in masks_16bit:
        yield m & 0xFF
        yield (m >> 8) & 0xFF


def gen_decompress_d5_shift():
    # Shift multipliers for poly_decompress_d5_avx2.
    # Powers of 2 to align extracted 5-bit values for mulhrs scaling.
    # Shifts: 1024, 32, 256, 8, 64, 512, 16, 128 (repeated for both lanes)
    shifts_16bit = [1024, 32, 256, 8, 64, 512, 16, 128]
    for s in shifts_16bit:
        yield s & 0xFF
        yield (s >> 8) & 0xFF
    for s in shifts_16bit:
        yield s & 0xFF
        yield (s >> 8) & 0xFF


def gen_decompress_d11_shufbidx():
    # Shuffle index for poly_decompress_d11_avx2.
    # 11-bit decompression: coefficient i uses bytes at positions determined
    # by floor(11*i/8). Pattern per lane selects source bytes for each
    # 16-bit output slot.
    # From _mm256_set_epi8(13,12,12,11,10,9,9,8,8,7,6,5,5,4,4,3,
    #                      10,9,9,8,7,6,6,5,5,4,3,2,2,1,1,0)
    yield from [0, 1, 1, 2, 2, 3, 4, 5, 5, 6, 6, 7, 8, 9, 9, 10]
    yield from [3, 4, 4, 5, 5, 6, 7, 8, 8, 9, 9, 10, 11, 12, 12, 13]


def gen_decompress_d11_srlvdidx():
    # Variable dword right shift counts for poly_decompress_d11_avx2.
    # From _mm256_set_epi32(0, 0, 1, 0, 0, 0, 1, 0)
    for v in [0, 1, 0, 0, 0, 1, 0, 0]:
        yield v & 0xFF
        yield (v >> 8) & 0xFF
        yield (v >> 16) & 0xFF
        yield (v >> 24) & 0xFF


def gen_decompress_d11_srlvqidx():
    # Variable qword right shift counts for poly_decompress_d11_avx2.
    # From _mm256_set_epi64x(2, 0, 2, 0)
    for v in [0, 2, 0, 2]:
        yield v & 0xFF
        yield (v >> 8) & 0xFF
        yield (v >> 16) & 0xFF
        yield (v >> 24) & 0xFF
        yield 0
        yield 0
        yield 0
        yield 0


def gen_decompress_d11_shift():
    # Shift multipliers for poly_decompress_d11_avx2.
    # From _mm256_set_epi16(4,32,1,8,32,1,4,32, 4,32,1,8,32,1,4,32)
    shifts_16bit = [32, 4, 1, 32, 8, 1, 32, 4]
    for s in shifts_16bit:
        yield s & 0xFF
        yield (s >> 8) & 0xFF
    for s in shifts_16bit:
        yield s & 0xFF
        yield (s >> 8) & 0xFF


def gen_avx2_compress_consts():
    # Generate merged constant arrays, one per function.
    # Returns dict of {name: (sections, dtype, layout_comment, level)}
    # where sections is a list of (section_name, values) tuples.
    return {
        "compress_d4_data": (
            [("permdidx", list(gen_compress_d4_permdidx()))],
            "uint8_t",
            "permdidx",
            "d4_d10",
        ),
        "decompress_d4_data": (
            [("shufbidx", list(gen_decompress_d4_shufbidx()))],
            "uint8_t",
            "shufbidx",
            "d4_d10",
        ),
        "compress_d10_data": (
            [("shufbidx", list(gen_compress_d10_shufbidx()))],
            "uint8_t",
            "shufbidx",
            "d4_d10",
        ),
        "decompress_d10_data": (
            [("shufbidx", list(gen_decompress_d10_shufbidx()))],
            "uint8_t",
            "shufbidx",
            "d4_d10",
        ),
        "compress_d5_data": (
            [("shufbidx", list(gen_compress_d5_shufbidx()))],
            "uint8_t",
            "shufbidx",
            "d5_d11",
        ),
        "decompress_d5_data": (
            [
                ("shufbidx", list(gen_decompress_d5_shufbidx())),
                ("mask", list(gen_decompress_d5_mask())),
                ("shift", list(gen_decompress_d5_shift())),
            ],
            "uint8_t",
            "shufbidx, mask, shift",
            "d5_d11",
        ),
        "compress_d11_data": (
            [
                ("srlvqidx", list(gen_compress_d11_srlvqidx())),
                ("shufbidx", list(gen_compress_d11_shufbidx())),
            ],
            "uint8_t",
            "srlvqidx, shufbidx",
            "d5_d11",
        ),
        "decompress_d11_data": (
            [
                ("shufbidx", list(gen_decompress_d11_shufbidx())),
                ("srlvdidx", list(gen_decompress_d11_srlvdidx())),
                ("srlvqidx", list(gen_decompress_d11_srlvqidx())),
                ("shift", list(gen_decompress_d11_shift())),
            ],
            "uint8_t",
            "shufbidx, srlvdidx, srlvqidx, shift",
            "d5_d11",
        ),
    }


def gen_avx2_compress_consts_file():
    consts = gen_avx2_compress_consts()

    def gen_const_def(name, sections, dtype, layout):
        """Generate a merged constant array definition."""
        # Calculate total size
        total_size = sum(len(vals) for _, vals in sections)
        # Generate layout comment for multi-section arrays
        if len(sections) > 1:
            offset = 0
            parts = []
            for sec_name, vals in sections:
                parts.append(f"{sec_name}[{offset}:{offset + len(vals)}]")
                offset += len(vals)
            yield f"/* {', '.join(parts)} */"
        yield f"MLK_ALIGN const {dtype} mlk_{name}[{total_size}] = {{"
        for sec_name, vals in sections:
            for i in range(0, len(vals), 8):
                row = vals[i : i + 8]
                is_last_row = i + 8 >= len(vals)
                row_str = ", ".join(f"{v:>3}" for v in row) + ","
                # Add section comment at end of last row to anchor formatting
                if is_last_row:
                    yield f"    {row_str} /* {sec_name} */"
                else:
                    yield f"    {row_str}"
        yield "};"

    def gen_c():
        yield from gen_header()
        yield '#include "../../../common.h"'
        yield ""
        yield "#if defined(MLK_ARITH_BACKEND_X86_64_DEFAULT)"
        yield ""
        yield '#include "compress_consts.h"'
        yield ""

        # d4/d10 constants (K == 2 or 3)
        yield "#if !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED) && \\"
        yield "    (defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_K == 2 || MLKEM_K == 3)"
        yield ""
        for name, (sections, dtype, layout, level) in consts.items():
            if level == "d4_d10":
                yield from gen_const_def(name, sections, dtype, layout)
                yield ""
        yield "#endif /* d4/d10 constants */"
        yield ""

        # d5/d11 constants (K == 4)
        yield "#if !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED) && \\"
        yield "    (defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_K == 4)"
        yield ""
        for name, (sections, dtype, layout, level) in consts.items():
            if level == "d5_d11":
                yield from gen_const_def(name, sections, dtype, layout)
                yield ""
        yield "#endif /* d5/d11 constants */"
        yield ""

        yield "#else"
        yield ""
        yield "MLK_EMPTY_CU(avx2_compress_consts)"
        yield ""
        yield "#endif /* MLK_ARITH_BACKEND_X86_64_DEFAULT */"
        yield ""

    def gen_h():
        yield from gen_header()
        yield "#ifndef MLK_DEV_X86_64_SRC_COMPRESS_CONSTS_H"
        yield "#define MLK_DEV_X86_64_SRC_COMPRESS_CONSTS_H"
        yield ""
        yield '#include "../../../common.h"'
        yield ""
        yield "#ifndef __ASSEMBLER__"
        yield ""

        # Declare all constants
        for name, (sections, dtype, layout, level) in consts.items():
            total_size = sum(len(vals) for _, vals in sections)
            yield f"#define mlk_{name} MLK_NAMESPACE({name})"
            yield f"extern const {dtype} mlk_{name}[{total_size}];"
            yield ""

        yield "#endif /* !__ASSEMBLER__ */"
        yield ""
        yield "#endif /* !MLK_DEV_X86_64_SRC_COMPRESS_CONSTS_H */"
        yield ""

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


def gen_avx2_hol_light_compress_consts_file():
    consts = gen_avx2_compress_consts()

    def gen():
        yield from gen_hol_light_header()
        yield "(* Compression/decompression constants for AVX2 implementations. *)"
        yield ""

        for name, (sections, dtype, layout, level) in consts.items():
            # Flatten all sections into one list
            all_values = []
            for _, vals in sections:
                all_values.extend(vals)
            yield f"(* {layout} *)"
            yield f"let {name} = define `{name}:int list = ["
            yield from print_hol_light_array(all_values)
            yield "]`;;"
            yield ""

    update_file(
        "proofs/hol_light/x86_64/proofs/mlkem_compress_consts.ml", "\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.
    """
    items = [f"word 0x{v:016X}" for v in g]

    for i in range(0, len(items), entries_per_line):
        row = items[i : i + entries_per_line]
        is_last = i + entries_per_line >= len(items)
        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 MLK_ALIGN.

    Yields lines for a complete array definition including braces.
    """
    yield f"MLK_ALIGN const uint64_t {name}[] = {{"
    vals = list(values)
    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(MLK_FIPS202_AARCH64_NEED_X1_SCALAR) ||                  \\"
        yield "     defined(MLK_FIPS202_AARCH64_NEED_X1_V84A) ||                    \\"
        yield "     defined(MLK_FIPS202_AARCH64_NEED_X2_V84A) ||                    \\"
        yield "     defined(MLK_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID) ||       \\"
        yield "     defined(MLK_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID)) && \\"
        yield "    !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "fips202_native_aarch64.h"'
        yield ""
        yield from print_c_uint64_array(
            "mlk_keccakf1600_round_constants",
            gen_keccak_round_constants(),
            entries_per_line=3,
        )
        yield ""
        yield "#else /* (MLK_FIPS202_AARCH64_NEED_X1_SCALAR ||                               \\"
        yield "         MLK_FIPS202_AARCH64_NEED_X1_V84A || MLK_FIPS202_AARCH64_NEED_X2_V84A \\"
        yield "         || MLK_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID ||                  \\"
        yield "         MLK_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID) &&               \\"
        yield "         !MLK_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLK_EMPTY_CU(fips202_aarch64_round_constants)"
        yield ""
        yield "#endif /* !((MLK_FIPS202_AARCH64_NEED_X1_SCALAR ||                             \\"
        yield "          MLK_FIPS202_AARCH64_NEED_X1_V84A || MLK_FIPS202_AARCH64_NEED_X2_V84A \\"
        yield "          || MLK_FIPS202_AARCH64_NEED_X4_V8A_SCALAR_HYBRID ||                  \\"
        yield "          MLK_FIPS202_AARCH64_NEED_X4_V8A_V84A_SCALAR_HYBRID) &&               \\"
        yield "          !MLK_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(MLK_FIPS202_X86_64_NEED_X4_AVX2) && \\"
        yield "    !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield "#include <stdint.h>"
        yield ""
        yield '#include "fips202_native_x86_64.h"'
        yield ""
        yield from print_c_uint64_array(
            "mlk_keccakf1600_round_constants",
            gen_keccak_round_constants(),
            entries_per_line=3,
        )
        yield ""
        yield from print_c_uint64_array("mlk_keccak_rho8", gen_keccak_rho8())
        yield ""
        yield from print_c_uint64_array("mlk_keccak_rho56", gen_keccak_rho56())
        yield ""
        yield "#else /* MLK_FIPS202_X86_64_NEED_X4_AVX2 && !MLK_CONFIG_MULTILEVEL_NO_SHARED */"
        yield ""
        yield "MLK_EMPTY_CU(fips202_x86_64_constants)"
        yield ""
        yield "#endif /* !(MLK_FIPS202_X86_64_NEED_X4_AVX2 && !MLK_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 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(MLK_ARITH_BACKEND_X86_64_DEFAULT) && \\"
        yield "    !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)"
        yield ""
        yield '#include "consts.h"'
        yield ""
        yield "/*"
        yield " * Table of zeta values used in the AVX2 NTTs"
        yield " * See autogen for details."
        yield " */"
        yield f"MLK_ALIGN const int16_t mlk_qdata[{len(constants)}] = {{"
        yield from map(lambda t: str(t) + ",", constants)
        yield "};"
        yield ""
        yield "#else"
        yield ""
        yield "MLK_EMPTY_CU(avx2_consts)"
        yield ""
        yield "#endif"
        yield ""

    def gen_h():
        yield from gen_header()
        yield "#ifndef MLK_DEV_X86_64_SRC_CONSTS_H"
        yield "#define MLK_DEV_X86_64_SRC_CONSTS_H"
        yield '#include "../../../common.h"'
        yield f"#define MLK_AVX2_BACKEND_DATA_OFFSET_REVIDXB {offsets['revidxb']}"
        yield f"#define MLK_AVX2_BACKEND_DATA_OFFSET_REVIDXD {offsets['revidxd']}"
        yield f"#define MLK_AVX2_BACKEND_DATA_OFFSET_ZETAS_EXP {offsets['zetas']}"
        yield f"#define MLK_AVX2_BACKEND_DATA_OFFSET_MULCACHE_TWIDDLES {offsets['mulcache_twiddles']}"
        yield ""
        yield "#ifndef __ASSEMBLER__"
        yield "#define mlk_qdata MLK_NAMESPACE(qdata)"
        yield f"extern const int16_t mlk_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 gen_riscv64_root_of_unity_for_block(layer, block, inv=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 6
    # in this function.
    log = bitreverse(pow(2, layer) + block, 7)
    if inv is True:
        log = -log
    root = signed_reduce(pow(root_of_unity, log, modulus) * montgomery_factor)
    return root


def gen_riscv64_zetas():
    zetas_0123 = [
        signed_reduce(montgomery_factor),
        gen_riscv64_root_of_unity_for_block(0, 0),
    ]

    zetas_0123 += [gen_riscv64_root_of_unity_for_block(1, i) for i in range(2)]
    zetas_0123 += [gen_riscv64_root_of_unity_for_block(2, i) for i in range(4)]
    zetas_0123 += [gen_riscv64_root_of_unity_for_block(3, i) for i in range(8)]

    for j in range(8):
        yield from [zetas_0123[j * 2]]
        yield from [zetas_0123[j * 2 + 1]]
        for i in range(2):
            yield from [gen_riscv64_root_of_unity_for_block(4, j * 2 + i)]
        for i in range(4):
            yield from [gen_riscv64_root_of_unity_for_block(5, j * 4 + i)]
        for i in range(8):
            yield from [gen_riscv64_root_of_unity_for_block(6, j * 8 + i)]


def gen_riscv64_izetas():
    zetas_0123 = [
        signed_reduce(montgomery_factor),
        gen_riscv64_root_of_unity_for_block(0, 0, inv=True),
    ]

    zetas_0123 += [
        gen_riscv64_root_of_unity_for_block(1, i, inv=True) for i in range(2)
    ]
    zetas_0123 += [
        gen_riscv64_root_of_unity_for_block(2, i, inv=True) for i in range(4)
    ]
    zetas_0123 += [
        gen_riscv64_root_of_unity_for_block(3, i, inv=True) for i in range(8)
    ]

    for j in range(8):
        yield from [zetas_0123[j * 2]]
        yield from [zetas_0123[j * 2 + 1]]
        for i in range(2):
            yield from [gen_riscv64_root_of_unity_for_block(4, j * 2 + i, inv=True)]
        for i in range(4):
            yield from [gen_riscv64_root_of_unity_for_block(5, j * 4 + i, inv=True)]
        for i in range(8):
            yield from [gen_riscv64_root_of_unity_for_block(6, j * 8 + i, inv=True)]


def gen_riscv64_basemul_roots():
    R = signed_reduce(montgomery_factor)
    for i in range(64):
        yield from [
            R,
            signed_reduce(
                pow(root_of_unity, bitreverse(64 + i, 7), modulus) * montgomery_factor
            ),
        ]
        yield from [
            R,
            signed_reduce(
                -pow(root_of_unity, bitreverse(64 + i, 7), modulus) * montgomery_factor
            ),
        ]


def gen_riscv64_zeta_files():
    """Generate all RISC-V 64 zeta files"""

    # Generate rv64v_zetas.inc
    def gen_zetas():
        yield from gen_header()
        yield '#include "arith_native_riscv64.h"'
        yield ""
        yield "const int16_t zeta[] = {"
        yield from map(lambda t: str(t) + ",", gen_riscv64_zetas())
        yield "};"
        yield ""

    update_file(
        "mlkem/src/native/riscv64/src/rv64v_zetas.inc",
        "\n".join(gen_zetas()),
        force_format=True,
    )

    # Generate rv64v_izetas.inc
    def gen_izetas():
        yield from gen_header()
        yield '#include "arith_native_riscv64.h"'
        yield ""
        yield "const int16_t izeta[] = {"
        yield from map(lambda t: str(t) + ",", gen_riscv64_izetas())
        yield "};"
        yield ""

    update_file(
        "mlkem/src/native/riscv64/src/rv64v_izetas.inc",
        "\n".join(gen_izetas()),
        force_format=True,
    )

    # Generate rv64v_zetas_basemul.inc
    def gen_basemul():
        yield from gen_header()
        yield '#include "arith_native_riscv64.h"'
        yield "  "
        yield "const int16_t roots[] = {"
        yield from map(lambda t: str(t) + ",", gen_riscv64_basemul_roots())
        yield "};"
        yield ""

    update_file(
        "mlkem/src/native/riscv64/src/rv64v_zetas_basemul.inc",
        "\n".join(gen_basemul()),
        force_format=True,
    )


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


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


def get_header_files(main_only=False, core_only=False):
    if main_only is True:
        return get_files("mlkem/*.h") + get_files("mlkem/src/**/*.h")
    elif core_only is True:
        return (
            get_files("mlkem/*.h")
            + get_files("mlkem/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_mlkem=False):
    def normalize(f):
        return f.removeprefix("mlkem/") if strip_mlkem 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()

    with ThreadPoolExecutor() as executor:
        for results in executor.map(get_defines_from_file, files):
            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 mlkem-native macros. "
                f"If this is an mlkem-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 = [
        "mlkem_native.h",
        "params.h",
        # Deliberately omit config.h, which is not #undef'ed
        "common.h",
        "indcpa.c",
        "indcpa.h",
        "kem.c",
        "kem.h",
        "poly_k.c",
        "poly_k.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 != "mlkem/mlkem_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 mlkem-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="MLK_CONFIG_PARAMETER_SET-specific files"
    )
    yield ""
    yield "#if !defined(MLK_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="MLK_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 "/* mlkem/src/cbmc.h */"
    yield "#undef MLK_CBMC_H"
    yield "#undef __contract__"
    yield "#undef __loop__"
    yield ""
    yield "#if !defined(MLK_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(MLK_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
    yield from gen_monolithic_undef_all_core(filt=native_fips202_core, desc="")
    yield "#if defined(MLK_SYS_AARCH64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_fips202_aarch64, desc="native code (FIPS202, AArch64)"
    )
    yield "#endif"
    yield "#if defined(MLK_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(MLK_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(MLK_CONFIG_USE_NATIVE_BACKEND_ARITH)"
    yield from gen_monolithic_undef_all_core(filt=native_arith_core, desc="")
    yield "#if defined(MLK_SYS_AARCH64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_arith_aarch64, desc="native code (Arith, AArch64)"
    )
    yield "#endif"
    yield "#if defined(MLK_SYS_X86_64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_arith_x86_64, desc="native code (Arith, X86_64)"
    )
    yield "#endif"
    yield "#if defined(MLK_SYS_RISCV64)"
    yield from gen_monolithic_undef_all_core(
        filt=native_arith_riscv64, desc="native code (Arith, RISC-V 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_mlkem=True)
        yield from gen_header()

        yield "/******************************************************************************"
        yield " *"
        yield " * Single compilation unit (SCU) for fixed-level build of mlkem-native"
        yield " *"
        yield " * This compilation unit bundles together all source files for a build"
        yield " * of mlkem-native for a fixed security level (MLKEM-512/768/1024)."
        yield " *"
        yield " * # API"
        yield " *"
        yield " * The API exposed by this file is described in mlkem_native.h."
        yield " *"
        yield " * # Multi-level build"
        yield " *"
        yield " * If you want an SCU build of mlkem-native with support for multiple security"
        yield " * levels, you need to include this file multiple times, and set"
        yield " * MLK_CONFIG_MULTILEVEL_WITH_SHARED and MLK_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 mlkem-native configuration are relevant:"
        yield " *"
        yield " * - MLK_CONFIG_FIPS202_CUSTOM_HEADER"
        yield " *   Set this option if you use a custom FIPS202 implementation."
        yield " *"
        yield " * - MLK_CONFIG_USE_NATIVE_BACKEND_ARITH"
        yield " *   Set this option if you want to include the native arithmetic backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLK_CONFIG_USE_NATIVE_BACKEND_FIPS202"
        yield " *   Set this option if you want to include the native FIPS202 backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLK_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 mlkem-native source tree are not used,"
        yield " * consider reducing this header via `unifdef`."
        yield " *"
        yield " * Example:"
        yield " * ```bash"
        yield " * unifdef -UMLK_CONFIG_USE_NATIVE_BACKEND_ARITH mlkem_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(MLK_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(MLK_CONFIG_USE_NATIVE_BACKEND_ARITH)"
        yield "#if defined(MLK_SYS_AARCH64)"
        for c in filter(native_arith_aarch64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_X86_64)"
        for c in filter(native_arith_x86_64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_RISCV64)"
        for c in filter(native_arith_riscv64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        yield "#if defined(MLK_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
        yield "#if defined(MLK_SYS_AARCH64)"
        for c in filter(native_fips202_aarch64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_X86_64)"
        for c in filter(native_fips202_x86_64, c_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_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("mlkem/mlkem_native.c", "\n".join(gen()))


def gen_monolithic_asm_file():
    def gen():
        asm_sources = get_asm_source_files(main_only=True, strip_mlkem=True)
        yield from gen_header()
        yield "/******************************************************************************"
        yield " *"
        yield " * Single assembly unit for fixed-level build of mlkem-native"
        yield " *"
        yield " * This assembly unit bundles together all assembly files for a build"
        yield " * of mlkem-native for a fixed security level (MLKEM-512/768/1024)."
        yield " *"
        yield " * # Multi-level build"
        yield " *"
        yield " * If you want an SCU build of mlkem-native with support for multiple security"
        yield " * levels, you should include this file once with MLK_CONFIG_MULTILEVEL_WITH_SHARED set."
        yield " *"
        yield " * (You could also follow the same pattern as for mlkem_native_monobuild.c"
        yield " *  and include it for every level, setting MLK_CONFIG_MULTILEVEL_NO_SHARED"
        yield " *  for all but one. For builds with MLK_CONFIG_MULTILEVEL_NO_SHARED, this"
        yield " *  file will then be ignored.)"
        yield " *"
        yield " * # Configuration"
        yield " *"
        yield " * The following options from the mlkem-native configuration are relevant:"
        yield " *"
        yield " * - MLK_CONFIG_FIPS202_CUSTOM_HEADER"
        yield " *   Set this option if you use a custom FIPS202 implementation."
        yield " *"
        yield " * - MLK_CONFIG_USE_NATIVE_BACKEND_ARITH"
        yield " *   Set this option if you want to include the native arithmetic backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLK_CONFIG_USE_NATIVE_BACKEND_FIPS202"
        yield " *   Set this option if you want to include the native FIPS202 backends"
        yield " *   in your build."
        yield " *"
        yield " * - MLK_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 mlkem-native source tree are not used,"
        yield " * consider reducing this header via `unifdef`."
        yield " *"
        yield " * Example:"
        yield " * ```bash"
        yield " * unifdef -UMLK_CONFIG_USE_NATIVE_BACKEND_ARITH mlkem_native_asm.S"
        yield " * ```"
        yield " */"
        yield ""
        yield '#include "src/common.h"'
        yield ""
        yield "#if defined(MLK_CONFIG_USE_NATIVE_BACKEND_ARITH)"
        yield "#if defined(MLK_SYS_AARCH64)"
        for c in filter(native_arith_aarch64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_X86_64)"
        for c in filter(native_arith_x86_64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_RISCV64)"
        for c in filter(native_arith_riscv64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#endif"
        yield ""
        yield "#if defined(MLK_CONFIG_USE_NATIVE_BACKEND_FIPS202)"
        yield "#if defined(MLK_SYS_AARCH64)"
        for c in filter(native_fips202_aarch64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_SYS_X86_64)"
        for c in filter(native_fips202_x86_64, asm_sources):
            yield f'#include "{c}"'
        yield "#endif"
        yield "#if defined(MLK_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 mlkem_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 mlkem_native.c",
        ]
        yield from gen_macro_undefs(extra_notes=extra)

    update_file("mlkem/mlkem_native_asm.S", "\n".join(gen()), force_format=True)


def get_config_options():
    content = read_file("mlkem/mlkem_native_config.h")
    config_pattern = r"Name:\s*(MLK_CONFIG_\w+)"

    configs = re.findall(config_pattern, content)

    configs += [
        "MLK_FORCE_AARCH64",
        "MLK_FORCE_AARCH64_EB",
        "MLK_FORCE_X86_64",
        "MLK_FORCE_PPC64LE",
        "MLK_FORCE_RISCV64",
        "MLK_FORCE_RISCV32",
        "MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER",
        "MLKEM_DEBUG",  # TODO: Rename?
        "MLK_BREAK_PCT",  # Use in PCT breakage test]
        "MLK_CHECK_APIS",
        "MLK_CONFIG_XXX",
        "MLK_CONFIG_API_XXX",
        "MLK_USE_NATIVE_XXX",
    ]

    return configs


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

    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 ["MLK_SYS_CAP_AVX2", "MLK_SYS_CAP_SHA3", "MLK_SYS_CAP_DUMMY"]


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 mlkem_native.h
        if m.startswith("MLK_TOTAL_ALLOC") or m.startswith("MLK_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
        #

        # Makefiles use MLK_SOURCE_XXX to list source files
        if is_autogen or filename.endswith("/Makefile"):
            if m.startswith("MLK_SOURCE") or m.startswith("MLK_OBJ"):
                return True

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

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

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

        # AWS-LC importer patch
        if is_autogen or filename == "integration/awslc/awslc.patch":
            return True

        if is_autogen or filename == "mlkem/src/common.h":
            if m == "MLK_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 = funcname.replace("_asm", "") + "_"
    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 mlkem/ 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)


def update_via_simpasm(
    infile_full,
    outdir,
    outfile=None,
    cflags=None,
    preserve_header=True,
    force_cross=False,
    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 += " -Imlkem"

    # 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 force_cross is False:
                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 force_cross is False:
                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, outfile, indir, cflags, arch = job
    update_via_simpasm(
        f"{indir}/{infile}",
        "proofs/hol_light/" + arch + "/mlkem",
        outfile=outfile,
        cflags=cflags + " -Imlkem",
        preserve_header=False,
    )


def gen_hol_light_asm():
    joblist_aarch64 = [
        (
            "ntt.S",
            "mlkem_ntt.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "intt.S",
            "mlkem_intt.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "poly_tomont_asm.S",
            "mlkem_poly_tomont.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "poly_tobytes_asm.S",
            "mlkem_poly_tobytes.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "poly_reduce_asm.S",
            "mlkem_poly_reduce.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "poly_mulcache_compute_asm.S",
            "mlkem_poly_mulcache_compute.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k2.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k2.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k3.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k3.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k4.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k4.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "rej_uniform_asm.S",
            "mlkem_rej_uniform.S",
            "dev/aarch64_opt/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/native/aarch64/src",
            "aarch64",
        ),
        (
            "keccak_f1600_x1_scalar_asm.S",
            "keccak_f1600_x1_scalar.S",
            "dev/fips202/aarch64/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/fips202/native/aarch64/src/",
            "aarch64",
        ),
        (
            "keccak_f1600_x1_v84a_asm.S",
            "keccak_f1600_x1_v84a.S",
            "dev/fips202/aarch64/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/fips202/native/aarch64/src/",
            "aarch64",
        ),
        (
            "keccak_f1600_x2_v84a_asm.S",
            "keccak_f1600_x2_v84a.S",
            "dev/fips202/aarch64/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/fips202/native/aarch64/src/",
            "aarch64",
        ),
        (
            "keccak_f1600_x4_v8a_v84a_scalar_hybrid_asm.S",
            "keccak_f1600_x4_v8a_v84a_scalar.S",
            "dev/fips202/aarch64/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/fips202/native/aarch64/src/",
            "aarch64",
        ),
        (
            "keccak_f1600_x4_v8a_scalar_hybrid_asm.S",
            "keccak_f1600_x4_v8a_scalar.S",
            "dev/fips202/aarch64/src",
            "-march=armv8.4-a+sha3 -Imlkem/src/fips202/native/aarch64/src/",
            "aarch64",
        ),
    ]

    x86_64_flags = "-mavx2 -mbmi2 -msse4 -fcf-protection=full"
    joblist_x86_64 = [
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k2.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k2.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k3.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k3.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "polyvec_basemul_acc_montgomery_cached_asm_k4.S",
            "mlkem_poly_basemul_acc_montgomery_cached_k4.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "ntt.S",
            "mlkem_ntt.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "intt.S",
            "mlkem_intt.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "reduce.S",
            "mlkem_reduce.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "ntttobytes.S",
            "mlkem_tobytes.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "rej_uniform_asm.S",
            "mlkem_rej_uniform.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "nttfrombytes.S",
            "mlkem_frombytes.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "tomont.S",
            "mlkem_tomont.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "nttunpack.S",
            "mlkem_unpack.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "mulcache_compute.S",
            "mlkem_mulcache_compute.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_compress_d4.S",
            "mlkem_poly_compress_d4.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompress_d4.S",
            "mlkem_poly_decompress_d4.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_compress_d5.S",
            "mlkem_poly_compress_d5.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompress_d5.S",
            "mlkem_poly_decompress_d5.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_compress_d10.S",
            "mlkem_poly_compress_d10.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompress_d10.S",
            "mlkem_poly_decompress_d10.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_compress_d11.S",
            "mlkem_poly_compress_d11.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "poly_decompress_d11.S",
            "mlkem_poly_decompress_d11.S",
            "dev/x86_64/src",
            f"-Idev/x86_64/src -Imlkem/src/native/x86_64/src -Imlkem/src/common.h {x86_64_flags}",
            "x86_64",
        ),
        (
            "keccak_f1600_x4_avx2.S",
            "keccak_f1600_x4_avx2.S",
            "dev/fips202/x86_64/src",
            f"-Idev/fips202/x86_64/src -Imlkem/src/fips202/native/x86_64/src {x86_64_flags}",
            "x86_64",
        ),
    ]

    joblist = joblist_aarch64 + joblist_x86_64
    run_parallel(joblist, gen_hol_light_asm_file)


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

    # Skip .inc files when simplification is enabled (they get inlined by simpasm)
    if no_simplify is False and f.endswith(".inc"):
        return None

    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=False,
    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",
            "mlkem/src/native/aarch64/meta.h",
            transform=lambda c: adjust_header_guard_for_filename(
                c, "mlkem/src/native/aarch64/meta.h"
            ),
        )

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

    synchronize_backend(
        f"dev/aarch64_{ty}/src",
        "mlkem/src/native/aarch64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags=f"-Idev/aarch64_{ty}/src -Imlkem/src/native/aarch64/src",
    )
    synchronize_backend(
        "dev/fips202/aarch64/src",
        "mlkem/src/fips202/native/aarch64/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/aarch64/src -Imlkem/src/fips202/native/aarch64/src -march=armv8.4-a+sha3",
    )
    synchronize_backend(
        "dev/fips202/aarch64",
        "mlkem/src/fips202/native/aarch64",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/aarch64 -Imlkem/src/fips202/native/aarch64 -march=armv8.4-a+sha3",
    )
    synchronize_backend(
        "dev/x86_64/src",
        "mlkem/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 MLK_ASM_FN_SYMBOL
        # macro.
        # This leads to duplicate endbr64 instructions causing a failure when
        # comparing the object code before and after simplification.
        cflags="-Idev/x86_64/src/ -Imlkem/src/native/x86_64/src/ -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )
    synchronize_backend(
        "dev/fips202/armv81m/src",
        "mlkem/src/fips202/native/armv81m/src",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/armv81m/src -Imlkem/src/fips202/native/armv81m/src -march=armv8.1-m.main+mve -mthumb",
    )
    synchronize_backend(
        "dev/fips202/armv81m",
        "mlkem/src/fips202/native/armv81m",
        delete=delete,
        force_cross=force_cross,
        no_simplify=no_simplify,
        cflags="-Idev/fips202/armv81m -Imlkem/src/fips202/native/armv81m -march=armv8.1-m.main+mve -mthumb",
    )
    synchronize_backend(
        "dev/fips202/x86_64/src",
        "mlkem/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 MLK_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 -Imlkem/src/fips202/native/x86_64/src -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )
    synchronize_backend(
        "dev/fips202/x86_64",
        "mlkem/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 -Imlkem/src/fips202/native/x86_64 -mavx2 -mbmi2 -msse4 -fcf-protection=none",
    )


def adjust_header_guard_for_filename(content, header_file):
    content = content.split("\n")
    exceptions = {
        "mlkem/mlkem_native.h": "MLK_H",
        "mlkem/mlkem_native_config.h": "MLK_CONFIG_H",
    }

    # Use full filename as the header guard, with '/' and '.' replaced by '_'
    guard_name = (
        header_file.removeprefix("mlkem/src/")
        .replace("/", "_")
        .replace(".", "_")
        .upper()
    )
    guard_name = "MLK_" + 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(core_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 mlkem-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"""
    mlkem_dir = "mlkem/src"

    # add files mlkem/src/*
    sources = [
        f"mlkem/src/{f}"
        for f in os.listdir(mlkem_dir)
        if os.path.isfile(f"{mlkem_dir}/{f}") and not f.endswith(".o")
    ]
    # add files mlkem/src/native/* (API definitions)
    sources += [
        f"mlkem/src/native/{f}"
        for f in os.listdir(f"{mlkem_dir}/native")
        if os.path.isfile(f"{mlkem_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"mlkem/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-KEM-512_META.yml",
        "integration/liboqs/ML-KEM-768_META.yml",
        "integration/liboqs/ML-KEM-1024_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=False):
    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 force_cross is False:
                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=False):
    """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 mlkem/mlkem_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"""
        i = 0
        while i < len(lines):
            if re.search(rf"\* Name:\s+{re.escape(opt_name)}\b", lines[i]):
                # Skip to the end of this comment block (find the closing */)
                while i < len(lines) and not lines[i].strip().endswith("*/"):
                    i += 1
                i += 1  # Skip the closing */

                start = i
                # Find the next config option (starts with /*****)
                while i < len(lines):
                    if lines[i].strip().startswith("/****"):
                        # Back up to exclude empty line before next option
                        while i > start and lines[i - 1].strip() == "":
                            i -= 1
                        return (start, i)
                    i += 1
                # If no next option found, go to end
                return (start, len(lines))
            i += 1
        raise Exception(f"Could not find config option {opt_name} in default config")

    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("mlkem/mlkem_native_config.h")

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


def _main():
    slothy_choices = [
        "ntt",
        "intt",
        "poly_tobytes_asm",
        "poly_tomont_asm",
        "poly_reduce_asm",
        "poly_mulcache_compute_asm",
        "polyvec_basemul_acc_montgomery_cached_asm_k2",
        "polyvec_basemul_acc_montgomery_cached_asm_k3",
        "polyvec_basemul_acc_montgomery_cached_asm_k4",
        "rej_uniform_asm",
        "keccak_f1600_x1_scalar_asm",
        "keccak_f1600_x4_v8a_scalar_hybrid_asm",
        "keccak_f1600_x4_v8a_v84a_scalar_hybrid_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", default=False, action="store_true")
    parser.add_argument(
        "--x86-64-syntax",
        type=str,
        choices=["att", "intel"],
        default="att",
        help="Assembly syntax for x86_64 disassembly output (att or intel)",
    )

    args = parser.parse_args()

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

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

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

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

    # 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),
        (
            "Generate SLOTHY optimized assembly",
            lambda: gen_slothy(args.slothy),
            args.slothy is not None and not args.dry_run,
        ),
    ]

    def gen_zeta_tables():
        gen_c_zeta_file()
        gen_aarch64_hol_light_zeta_file()
        gen_aarch64_hol_light_rej_uniform_table()
        gen_aarch64_zeta_file()
        gen_aarch64_rej_uniform_table()
        gen_avx2_hol_light_zeta_file()
        gen_avx2_hol_light_rej_uniform_table()
        gen_avx2_zeta_file()
        gen_avx2_compress_consts_file()
        gen_avx2_hol_light_compress_consts_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()
        gen_riscv64_zeta_files()

    def gen_monolithic():
        gen_monolithic_source_file()
        gen_monolithic_asm_file()

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

    steps += [
        ("Check assembly register aliases", check_asm_register_aliases),
        ("Check assembly loop labels", check_asm_loop_labels),
        ("Normalize assembly macro syntax", normalize_asm_macro_syntax),
        ("Generate zeta and lookup tables", gen_zeta_tables),
        ("Generate HOL Light assembly", gen_hol_light_asm, hol_light_asm_supported),
        ("Synchronize backends", sync_backends),
        ("Generate header guards", gen_header_guards),
        ("Complete final backend synchronization", sync_backends_final),
        (
            "Update HOL Light bytecode",
            lambda: update_hol_light_bytecode(args.force_cross),
            args.update_hol_light_bytecode,
        ),
        ("Generate monolithic source files", gen_monolithic),
        ("Generate undefs", gen_undefs),
        ("Generate test configs", gen_test_configs),
        ("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)),
    ]

    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)