mldsa-native-rs 0.0.1-alpha.6

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

"""Rewrite simplified assembly to include CFI (Call Frame Information)
directives for stack unwinding.

Reads assembly from stdin (or --input) and emits the same assembly with
architecture-appropriate `.cfi_*` directives inserted at prologue and
epilogue boundaries. Supports aarch64, x86_64, and armv81m targets.

Invoked by scripts/simpasm when --cfify is passed; see that script for
the round-trip validation that guards correctness."""

import sys
import re
import argparse


# -----------------------------------------------------------------------------
# aarch64 module-scope constants
# -----------------------------------------------------------------------------
AARCH64_SIMD_PAIRS = ((8, 9), (10, 11), (12, 13), (14, 15))
AARCH64_GPR_PAIRS = (
    (19, 20),
    (21, 22),
    (23, 24),
    (25, 26),
    (27, 28),
    (29, 30),
)

AARCH64_SIMD_SAVE_PATTERN = re.compile(
    r"(\s*)stp\s+d8,\s*d9,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+d10,\s*d11,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+d12,\s*d13,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+d14,\s*d15,\s*\[sp(?:,\s*#([^]]+))?\]",
    re.IGNORECASE,
)
AARCH64_SIMD_RESTORE_PATTERN = re.compile(
    r"(\s*)ldp\s+d8,\s*d9,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+d10,\s*d11,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+d12,\s*d13,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+d14,\s*d15,\s*\[sp(?:,\s*#[^]]+)?\]",
    re.IGNORECASE,
)
AARCH64_GPR_SAVE_PATTERN = re.compile(
    r"(\s*)stp\s+x19,\s*x20,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+x21,\s*x22,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+x23,\s*x24,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+x25,\s*x26,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+x27,\s*x28,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
    r"\s*stp\s+x29,\s*x30,\s*\[sp(?:,\s*#([^]]+))?\]",
    re.IGNORECASE,
)
AARCH64_GPR_RESTORE_PATTERN = re.compile(
    r"(\s*)ldp\s+x19,\s*x20,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+x21,\s*x22,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+x23,\s*x24,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+x25,\s*x26,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+x27,\s*x28,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
    r"\s*ldp\s+x29,\s*x30,\s*\[sp(?:,\s*#[^]]+)?\]",
    re.IGNORECASE,
)
AARCH64_ADD_SP_PATTERN = re.compile(
    r"(\s*)add\s+sp,\s*sp,\s*#(0x[0-9a-fA-F]+|\d+)", re.IGNORECASE
)
AARCH64_SUB_SP_PATTERN = re.compile(
    r"(\s*)sub\s+sp,\s*sp,\s*#(0x[0-9a-fA-F]+|\d+)", re.IGNORECASE
)
AARCH64_RET_PATTERN = re.compile(r"(\s*)ret\s*$", re.IGNORECASE)


# -----------------------------------------------------------------------------
# x86_64 module-scope constants
# -----------------------------------------------------------------------------
X86_64_LABEL_PATTERN = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_]*):$")
# Generic stack-pointer-style adjust: any `subq $N, %REG` or `addq $N, %REG`.
# The CFA-tracking variable selects which register we care about per line.
X86_64_SUBQ_PATTERN = re.compile(
    r"(\s*)subq\s+\$(0x[0-9a-fA-F]+|\d+),\s*%([a-z0-9]+)", re.IGNORECASE
)
X86_64_ADDQ_PATTERN = re.compile(
    r"(\s*)addq\s+\$(0x[0-9a-fA-F]+|\d+),\s*%([a-z0-9]+)", re.IGNORECASE
)
X86_64_RET_PATTERN = re.compile(r"(\s*)retq?\s*$", re.IGNORECASE)

# Stack-alignment anchor save/restore. The pair `movq %rsp, %REG` ...
# `movq %REG, %rsp` brackets a region in which rsp moves by an
# unpredictable, data-dependent amount (e.g. `andq $-32, %rsp`). While
# inside that region we re-anchor the CFA on REG so unwinding remains
# valid regardless of how rsp shifts. A stray `mov rsp, %REG` used as a
# scratch base pointer is filtered out by requiring a matching restore
# in the same function body.
X86_64_MOV_RSP_TO_REG_PATTERN = re.compile(
    r"(\s*)movq\s+%rsp,\s*%([a-z0-9]+)\s*$", re.IGNORECASE
)
X86_64_MOV_REG_TO_RSP_PATTERN = re.compile(
    r"(\s*)movq\s+%([a-z0-9]+),\s*%rsp\s*$", re.IGNORECASE
)


def _x86_64_has_matching_rsp_restore(lines, start, reg):
    """Return True if a `movq %REG, %rsp` appears before the next ret in `lines`.

    Used to distinguish an alignment anchor save from a scratch base-pointer
    copy: only the former has a matching restore in the function body.
    """
    reg_lower = reg.lower()
    for j in range(start, len(lines)):
        candidate = lines[j].rstrip()
        m = X86_64_MOV_REG_TO_RSP_PATTERN.match(candidate)
        if m and m.group(2).lower() == reg_lower:
            return True
        if X86_64_RET_PATTERN.match(candidate):
            return False
    return False


# -----------------------------------------------------------------------------
# armv81m module-scope constants and helpers
# -----------------------------------------------------------------------------
ARMV81M_CALLEE_SAVED_GPRS = {
    "r4",
    "r5",
    "r6",
    "r7",
    "r8",
    "r9",
    "r10",
    "r11",
    "lr",
}
ARMV81M_CALLEE_SAVED_DREGS = {
    "d8",
    "d9",
    "d10",
    "d11",
    "d12",
    "d13",
    "d14",
    "d15",
}
# Register aliases: numeric form -> canonical name
ARMV81M_GPR_ALIASES = {"r14": "lr", "r15": "pc", "r13": "sp"}

ARMV81M_REG_NAME_PATTERN = re.compile(r"([a-z]+)(\d+)$")
ARMV81M_REG_RANGE_PATTERN = re.compile(r"^\s*([a-z]+)(\d+)\s*-\s*([a-z]+)(\d+)\s*$")

ARMV81M_PUSH_PATTERN = re.compile(r"(\s*)push(?:\.w)?\s*\{([^}]+)\}", re.IGNORECASE)
ARMV81M_VPUSH_PATTERN = re.compile(r"(\s*)vpush\s*\{([^}]+)\}", re.IGNORECASE)
ARMV81M_VPOP_PATTERN = re.compile(r"(\s*)vpop\s*\{([^}]+)\}", re.IGNORECASE)
ARMV81M_POP_PATTERN = re.compile(r"(\s*)pop(?:\.w)?\s*\{([^}]+)\}", re.IGNORECASE)
ARMV81M_SUB_SP_PATTERN = re.compile(
    r"(\s*)sub(?:\.w)?\s+sp,\s*(?:sp,\s*)?#(0x[0-9a-fA-F]+|\d+)", re.IGNORECASE
)
ARMV81M_ADD_SP_PATTERN = re.compile(
    r"(\s*)add(?:\.w)?\s+sp,\s*(?:sp,\s*)?#(0x[0-9a-fA-F]+|\d+)", re.IGNORECASE
)
ARMV81M_BX_LR_PATTERN = re.compile(r"(\s*)bx\s+lr\s*$", re.IGNORECASE)


def armv81m_parse_reg(s):
    """Parse a single register token, returning its canonical name
    (e.g. 'r14' -> 'lr'). Raises ValueError on unrecognised input."""
    s = s.strip().lower()
    # Named aliases
    if s in ("lr", "pc", "sp"):
        return s
    if not ARMV81M_REG_NAME_PATTERN.match(s):
        raise ValueError(f"Cannot parse register: {s}")
    return ARMV81M_GPR_ALIASES.get(s, s)


def armv81m_expand_reg_range(token):
    """Expand 'r4-r11' or 'd8-d15' into a list. Single regs returned as-is."""
    m = ARMV81M_REG_RANGE_PATTERN.match(token)
    if m:
        prefix1, lo, prefix2, hi = (
            m.group(1),
            int(m.group(2)),
            m.group(3),
            int(m.group(4)),
        )
        if prefix1 != prefix2:
            raise ValueError(f"Mismatched register prefixes in range: {token}")
        return [
            ARMV81M_GPR_ALIASES.get(f"{prefix1}{n}", f"{prefix1}{n}")
            for n in range(lo, hi + 1)
        ]
    return [armv81m_parse_reg(token)]


def armv81m_parse_reglist(reglist_str):
    """Parse a register list string, expanding ranges and normalizing aliases."""
    regs = []
    for token in reglist_str.split(","):
        regs.extend(armv81m_expand_reg_range(token))
    return regs


def add_cfi_directives(text, arch):
    lines = text.split("\n")
    result = []
    i = 0
    # x86_64: register that the CFA is currently expressed in terms of.
    # Defaults to rsp; flipped to the alignment anchor register on a
    # `mov rsp, %REG` save and back to rsp on `mov %REG, %rsp`.
    # subq/addq adjustments are reflected in CFI only when applied to
    # the current CFA register.
    x86_64_cfa_reg = "rsp"

    while i < len(lines):
        line = lines[i].rstrip()

        if arch == "aarch64":
            # Check for SIMD save pattern: stp d8,d9; stp d10,d11; stp d12,d13; stp d14,d15
            if i + 3 < len(lines):
                pattern_text = "\n".join(lines[i : i + 4])
                match = AARCH64_SIMD_SAVE_PATTERN.match(pattern_text)
                if match:
                    indent = match.group(1)
                    offsets = [match.group(j + 2) or "0" for j in range(4)]
                    for j, reg_pair in enumerate(AARCH64_SIMD_PAIRS):
                        result.append(lines[i + j].rstrip())
                        try:
                            offset_val = int(offsets[j], 0)
                            result.append(
                                f"{indent}.cfi_rel_offset d{reg_pair[0]}, 0x{offset_val:x}"
                            )
                            result.append(
                                f"{indent}.cfi_rel_offset d{reg_pair[1]}, 0x{offset_val + 8:x}"
                            )
                        except ValueError:
                            result.append(
                                f"{indent}.cfi_rel_offset d{reg_pair[0]}, {offsets[j]}"
                            )
                            result.append(
                                f"{indent}.cfi_rel_offset d{reg_pair[1]}, ({offsets[j]}+8)"
                            )
                    i += 4
                    continue

            # Check for SIMD restore pattern: ldp d8,d9; ldp d10,d11; ldp d12,d13; ldp d14,d15
            if i + 3 < len(lines):
                pattern_text = "\n".join(lines[i : i + 4])
                match = AARCH64_SIMD_RESTORE_PATTERN.match(pattern_text)
                if match:
                    indent = match.group(1)
                    for j, reg_pair in enumerate(AARCH64_SIMD_PAIRS):
                        result.append(lines[i + j].rstrip())
                        result.append(f"{indent}.cfi_restore d{reg_pair[0]}")
                        result.append(f"{indent}.cfi_restore d{reg_pair[1]}")
                    i += 4
                    continue

            # Check for GPR save pattern: stp x19,x20 through stp x29,x30
            if i + 5 < len(lines):
                pattern_text = "\n".join(lines[i : i + 6])
                match = AARCH64_GPR_SAVE_PATTERN.match(pattern_text)
                if match:
                    indent = match.group(1)
                    offsets = [match.group(j + 2) or "0" for j in range(6)]
                    for j, reg_pair in enumerate(AARCH64_GPR_PAIRS):
                        result.append(lines[i + j].rstrip())
                        try:
                            offset_val = int(offsets[j], 0)
                            result.append(
                                f"{indent}.cfi_rel_offset x{reg_pair[0]}, 0x{offset_val:x}"
                            )
                            result.append(
                                f"{indent}.cfi_rel_offset x{reg_pair[1]}, 0x{offset_val + 8:x}"
                            )
                        except ValueError:
                            result.append(
                                f"{indent}.cfi_rel_offset x{reg_pair[0]}, {offsets[j]}"
                            )
                            result.append(
                                f"{indent}.cfi_rel_offset x{reg_pair[1]}, ({offsets[j]}+8)"
                            )
                    i += 6
                    continue

            # Check for GPR restore pattern: ldp x19,x20 through ldp x29,x30
            if i + 5 < len(lines):
                pattern_text = "\n".join(lines[i : i + 6])
                match = AARCH64_GPR_RESTORE_PATTERN.match(pattern_text)
                if match:
                    indent = match.group(1)
                    for j, reg_pair in enumerate(AARCH64_GPR_PAIRS):
                        result.append(lines[i + j].rstrip())
                        result.append(f"{indent}.cfi_restore x{reg_pair[0]}")
                        result.append(f"{indent}.cfi_restore x{reg_pair[1]}")
                    i += 6
                    continue

            # Rule 7: add sp, sp, #offset -> .cfi_adjust_cfa_offset (-(offset))
            match = AARCH64_ADD_SP_PATTERN.match(line)
            if match:
                indent, offset_str = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset -{offset:#x}")
                i += 1
                continue

            # Rule 8: sub sp, sp, #offset -> .cfi_adjust_cfa_offset (offset)
            match = AARCH64_SUB_SP_PATTERN.match(line)
            if match:
                indent, offset_str = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset {offset:#x}")
                i += 1
                continue

            # Rule 2: ret -> .cfi_endproc after ret
            match = AARCH64_RET_PATTERN.match(line)
            if match:
                indent = match.group(1)
                result.append(line)
                result.append(f"{indent}.cfi_endproc")
                i += 1
                continue

        elif arch == "x86_64":
            # Check for labels and see if there's a corresponding callq
            label_match = X86_64_LABEL_PATTERN.match(line)
            if label_match:
                label = label_match.group(1)
                # Check if this label is called anywhere in the text
                if re.search(rf"\s*callq\s+{re.escape(label)}\b", text, re.IGNORECASE):
                    result.append(line)
                    result.append("        .cfi_startproc")
                    i += 1
                    continue

            # x86_64: stack-alignment anchor save. Re-anchor the CFA on
            # the destination register so unwinding survives any
            # data-dependent rsp move (e.g. `andq $-32, %rsp`) that may
            # follow. Only fires while the CFA is still on rsp - a
            # nested re-anchor would indicate malformed input.
            #
            # A `movq %rsp, %REG` may also be a scratch base-pointer
            # copy (e.g. `rep movsb` source setup) with no intent to
            # re-anchor. Distinguish by requiring a matching restore
            # `movq %REG, %rsp` later in the same function body
            # (bounded by the next ret/retq, after which cfify emits
            # `.cfi_endproc`). Without a restore, we leave the line
            # alone and keep the CFA on rsp.
            match = X86_64_MOV_RSP_TO_REG_PATTERN.match(line)
            if match and x86_64_cfa_reg == "rsp":
                indent, reg = match.group(1), match.group(2)
                if _x86_64_has_matching_rsp_restore(lines, i + 1, reg):
                    result.append(line)
                    result.append(f"{indent}.cfi_def_cfa_register %{reg}")
                    x86_64_cfa_reg = reg.lower()
                    i += 1
                    continue

            # x86_64: stack-alignment anchor restore. The CFA goes back
            # to being computed from rsp. Only fires if the source
            # register matches the current CFA anchor; an unrelated
            # `mov %rax, %rsp` is left alone.
            match = X86_64_MOV_REG_TO_RSP_PATTERN.match(line)
            if match and match.group(2).lower() == x86_64_cfa_reg:
                indent = match.group(1)
                result.append(line)
                result.append(f"{indent}.cfi_def_cfa_register %rsp")
                x86_64_cfa_reg = "rsp"
                i += 1
                continue

            # x86_64: subq $OFFSET, %REG -> .cfi_adjust_cfa_offset OFFSET
            # only when REG is the current CFA register. Adjustments to
            # other registers (or to rsp while the CFA is anchored on
            # the alignment register) are invisible to the unwinder and
            # need no CFI.
            match = X86_64_SUBQ_PATTERN.match(line)
            if match:
                indent, offset_str, reg = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                if reg.lower() == x86_64_cfa_reg:
                    result.append(f"{indent}.cfi_adjust_cfa_offset {offset:#x}")
                i += 1
                continue

            # x86_64: addq $OFFSET, %REG -> .cfi_adjust_cfa_offset -OFFSET
            match = X86_64_ADDQ_PATTERN.match(line)
            if match:
                indent, offset_str, reg = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                if reg.lower() == x86_64_cfa_reg:
                    result.append(f"{indent}.cfi_adjust_cfa_offset -{offset:#x}")
                i += 1
                continue

            # x86_64: ret/retq -> .cfi_endproc after ret
            match = X86_64_RET_PATTERN.match(line)
            if match:
                indent = match.group(1)
                result.append(line)
                result.append(f"{indent}.cfi_endproc")
                i += 1
                continue

        elif arch == "armv81m":
            # push.w {reglist} / push {reglist}
            match = ARMV81M_PUSH_PATTERN.match(line)
            if match:
                indent = match.group(1)
                regs = armv81m_parse_reglist(match.group(2))
                total_size = 4 * len(regs)
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset {total_size:#x}")
                for idx, reg in enumerate(regs):
                    if reg in ARMV81M_CALLEE_SAVED_GPRS:
                        offset = 4 * idx
                        result.append(f"{indent}.cfi_rel_offset {reg}, {offset:#x}")
                i += 1
                continue

            # vpush {d-regs}
            match = ARMV81M_VPUSH_PATTERN.match(line)
            if match:
                indent = match.group(1)
                regs = armv81m_parse_reglist(match.group(2))
                total_size = 8 * len(regs)
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset {total_size:#x}")
                for idx, reg in enumerate(regs):
                    if reg in ARMV81M_CALLEE_SAVED_DREGS:
                        offset = 8 * idx
                        result.append(f"{indent}.cfi_rel_offset {reg}, {offset:#x}")
                i += 1
                continue

            # vpop {d-regs}
            match = ARMV81M_VPOP_PATTERN.match(line)
            if match:
                indent = match.group(1)
                regs = armv81m_parse_reglist(match.group(2))
                total_size = 8 * len(regs)
                result.append(line)
                for reg in regs:
                    if reg in ARMV81M_CALLEE_SAVED_DREGS:
                        result.append(f"{indent}.cfi_restore {reg}")
                result.append(f"{indent}.cfi_adjust_cfa_offset -{total_size:#x}")
                i += 1
                continue

            # pop.w {reglist} / pop {reglist}
            match = ARMV81M_POP_PATTERN.match(line)
            if match:
                indent = match.group(1)
                regs = armv81m_parse_reglist(match.group(2))
                total_size = 4 * len(regs)
                has_pc = "pc" in regs
                result.append(line)
                for reg in regs:
                    if reg in ARMV81M_CALLEE_SAVED_GPRS:
                        result.append(f"{indent}.cfi_restore {reg}")
                    elif reg == "pc":
                        # pop into pc restores lr
                        result.append(f"{indent}.cfi_restore lr")
                result.append(f"{indent}.cfi_adjust_cfa_offset -{total_size:#x}")
                if has_pc:
                    result.append(f"{indent}.cfi_endproc")
                i += 1
                continue

            # sub.w sp, #imm / sub sp, #imm / sub sp, sp, #imm
            match = ARMV81M_SUB_SP_PATTERN.match(line)
            if match:
                indent, offset_str = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset {offset:#x}")
                i += 1
                continue

            # add.w sp, #imm / add sp, #imm / add sp, sp, #imm
            match = ARMV81M_ADD_SP_PATTERN.match(line)
            if match:
                indent, offset_str = match.groups()
                offset = (
                    int(offset_str, 16)
                    if offset_str.lower().startswith("0x")
                    else int(offset_str)
                )
                result.append(line)
                result.append(f"{indent}.cfi_adjust_cfa_offset -{offset:#x}")
                i += 1
                continue

            # bx lr — function return
            match = ARMV81M_BX_LR_PATTERN.match(line)
            if match:
                indent = match.group(1)
                result.append(line)
                result.append(f"{indent}.cfi_endproc")
                i += 1
                continue

        result.append(line)
        i += 1

    return "\n".join(result)


def main():
    parser = argparse.ArgumentParser(
        description="Add CFI directives to AArch64 assembly"
    )
    parser.add_argument("-i", "--input", help="Input file (default: stdin)")
    parser.add_argument("-o", "--output", help="Output file (default: stdout)")
    parser.add_argument(
        "--emit-cfi-proc-start",
        action="store_true",
        help="Emit .cfi_proc_start as first line",
    )
    parser.add_argument(
        "--arch",
        choices=["aarch64", "x86_64", "armv81m"],
        default="aarch64",
        help="Target architecture (default: aarch64)",
    )
    args = parser.parse_args()

    input_file = open(args.input, "r") if args.input else sys.stdin
    output_file = open(args.output, "w") if args.output else sys.stdout

    try:
        # Read all input
        text = input_file.read()

        # Add initial .cfi_startproc if requested
        if args.emit_cfi_proc_start:
            text = "        .cfi_startproc\n" + text

        # Process the text
        result = add_cfi_directives(text, args.arch)

        # Write output
        output_file.write(result)
    finally:
        if args.input:
            input_file.close()
        if args.output:
            output_file.close()


if __name__ == "__main__":
    main()