RustPython/scripts/generate_opcode_metadata.py

"""
Generate Lib/_opcode_metadata.py for RustPython bytecode.

This file generates opcode metadata that is compatible with CPython 3.13.
"""

import itertools
import pathlib
import re
import typing

ROOT = pathlib.Path(__file__).parents[1]
BYTECODE_FILE = (
    ROOT / "crates" / "compiler-core" / "src" / "bytecode" / "instruction.rs"
)
OPCODE_METADATA_FILE = ROOT / "Lib" / "_opcode_metadata.py"


# Opcodes that needs to be first, regardless of their opcode ID.
PRIORITY_OPMAP = {
    "CACHE",
    "RESERVED",
    "RESUME",
    "INSTRUMENTED_LINE",
    "ENTER_EXECUTOR",
}


def to_snake_case(s: str) -> str:
    res = re.sub(r"(?<=[a-z0-9])([A-Z])", r"_\1", s)
    return re.sub(r"(\D)(\d+)$", r"\1_\2", res).upper()


class Opcode(typing.NamedTuple):
    rust_name: str
    cpython_name: str
    id: int
    have_oparg: bool

    @property
    def is_instrumented(self):
        return self.cpython_name.startswith("INSTRUMENTED_")

    @classmethod
    def from_str(cls, text: str):
        # Split on commas that are followed by a newline + an uppercase letter (new entry)
        entries = re.split(r",\s*\n\s*(?=[A-Z])", text)
        for entry in entries:
            entry = entry.strip()
            if not entry:
                continue
            have_oparg = "Arg<" in entry  # Hacky but works
            rust_name = re.match(r"(\w+)", entry).group(1)
            id_num, cpython_name = re.search(r'\((\d+),\s*"([^"]+)"\)', entry).groups()
            yield cls(
                rust_name=rust_name,
                cpython_name=cpython_name,
                id=int(id_num),
                have_oparg=have_oparg,
            )

    def __lt__(self, other: typing.Self) -> bool:
        sprio, oprio = (
            opcode.cpython_name not in PRIORITY_OPMAP for opcode in (self, other)
        )
        return (sprio, self.id) < (oprio, other.id)


def extract_enum_body(text: str, name: str) -> str:
    # Find the start of the enum block
    start_match = re.search(rf"enum\s+{name}\s*\{{", text)
    if not start_match:
        return None

    # Manually track brace depth from that point
    depth = 0
    start = start_match.end() - 1  # position of opening '{'
    for i, ch in enumerate(text[start:], start):
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                # Return only the inner content (excluding outer braces)
                return text[start + 1 : i]


def build_deopts(contents: str) -> dict[str, list[str]]:
    raw_body = re.search(
        r"fn deopt\(self\) -> Option<Self>(.*)", contents, re.DOTALL
    ).group(1)
    body = "\n".join(
        itertools.takewhile(
            lambda l: not l.startswith("_ =>"),  # Take until reaching fallback
            filter(
                lambda l: (
                    not l.startswith(
                        ("//", "Some(match")
                    )  # Skip comments or start of match
                ),
                map(str.strip, raw_body.splitlines()),
            ),
        )
    ).removeprefix("{")

    depth = 0
    arms = []
    buf = []
    for char in body:
        if char == "{":
            depth += 1
        elif char == "}":
            depth -= 1

        if depth == 0 and (char in ("}", ",")):
            arm = "".join(buf).strip()
            arms.append(arm)
            buf = []
        else:
            buf.append(char)

    # last arm
    arms.append("".join(buf))
    arms = [arm for arm in arms if arm]

    deopts = {}
    for arm in arms:
        *specialized, deopt = map(to_snake_case, re.findall(r"Self::(\w*)\b", arm))
        deopts[deopt] = specialized

    return deopts


contents = BYTECODE_FILE.read_text(encoding="utf-8")

deopts = build_deopts(contents)

enum_body = "\n".join(
    extract_enum_body(contents, enum_name)
    for enum_name in ("Instruction", "PseudoInstruction")
)
opcodes = list(Opcode.from_str(enum_body))

have_oparg = min(opcode.id for opcode in opcodes if opcode.have_oparg) - 1
min_instrumented = min(opcode.id for opcode in opcodes if opcode.is_instrumented)

# Generate the output file
output = """# This file is generated by scripts/generate_opcode_metadata.py
# for RustPython bytecode format (CPython 3.14 compatible opcode numbers).
# Do not edit!
"""

output += "\n_specializations = {\n"

for key, lst in deopts.items():
    output += f'    "{key}": [\n'
    for item in lst:
        output += f'        "{item}",\n'
    output += "    ],\n"

output += "}\n"

specialized = set(itertools.chain.from_iterable(deopts.values()))
output += "\n_specialized_opmap = {\n"
for opcode in sorted(opcodes, key=lambda op: op.cpython_name):
    cpython_name = opcode.cpython_name
    if cpython_name not in specialized:
        continue

    output += f"    '{cpython_name}': {opcode.id},\n"

output += "}\n"

output += "\nopmap = {\n"

for opcode in sorted(opcodes):
    cpython_name = opcode.cpython_name
    if cpython_name in specialized:
        continue

    output += f"    '{cpython_name}': {opcode.id},\n"

output += "}\n"

output += f"""
HAVE_ARGUMENT = {have_oparg}
MIN_INSTRUMENTED_OPCODE = {min_instrumented}
"""

OPCODE_METADATA_FILE.write_text(output, encoding="utf-8")