Autogen opcodes metadata (#7983)

tools/opcode_metadata/opcodes.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import collections
+import dataclasses
+import re
+import typing
+import warnings
+
+import utils
+from cpython import SKIP_PROPERTIES, Family, Properties, get_analysis, get_stack_effect
+from utils import SKIP_OVERRIDE, Override, OverrideConfs, StackEffect, to_pascal_case
+
+if typing.TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class OpcodeInfo:
+    enum_name: str
+    size: str
+    opcodes: tuple[Opcode, ...]
+
+    @property
+    def deopts(self) -> dict[str, list[str]]:
+        analysis = get_analysis()
+        names = {opcode.rust_name for opcode in self}
+
+        res = collections.defaultdict(list)
+        for family in analysis.families.values():
+            family_name = to_pascal_case(family.name)
+            if family_name not in names:
+                continue
+
+            for member in family.members:
+                member_name = to_pascal_case(member.name)
+                if member.name == family_name:
+                    continue
+
+                res[family_name].append(member_name)
+
+        return dict(res)
+
+    def __iter__(self):
+        yield from self.opcodes
+
+    @classmethod
+    def iter_infos(
+        cls, text: str, override_confs: OverrideConfs
+    ) -> Iterable[typing.Self]:
+        for block_match in re.finditer(
+            r"define_opcodes!\s*\((.+?)\);", text, re.DOTALL
+        ):
+            block = block_match.group(1).strip()
+
+            size = re.search(r"#\[repr\((\w+)\)\]", block).group(1)
+            enum_name = re.search(
+                r"#\[repr\(\w+\)\]\s*pub\s+enum\s+(\w+)\s*;", block
+            ).group(1)
+
+            second_enum_match = re.search(r"pub\s+enum\s+(\w+)\s*\{", block, re.DOTALL)
+            entries = utils.extract_enum_body(block, second_enum_match.end() - 1)
+
+            opcodes = tuple(sorted(iter_opcodes(entries, override_confs)))
+
+            yield cls(enum_name, size, opcodes)
+
+
+def iter_opcodes(text: str, override_confs: OverrideConfs) -> Iterable[Opcode]:
+    analysis = get_analysis()
+    # Split on commas that are followed by a newline + an uppercase letter (new entry)
+    entries = map(str.strip, re.split(r",\s*\n\s*(?=[A-Z])", text))
+    for entry in entries:
+        if not entry:
+            continue
+
+        opcode = Opcode.from_str(entry)
+
+        rust_name = opcode.rust_name
+        override = override_confs.get(rust_name, SKIP_OVERRIDE)
+
+        cpython_name = opcode.cpython_name
+
+        kwargs = {}
+        if instr := analysis.instructions.get(cpython_name):
+            kwargs["properties"] = instr.properties
+            kwargs["family"] = getattr(instr, "family", None)
+            kwargs["cache_entry"] = getattr(instr, "size", -1)
+
+            stack = get_stack_effect(instr)
+
+            popped = (-stack.base_offset).to_c()
+            pushed = (stack.logical_sp - stack.base_offset).to_c()
+            kwargs["stack_effect"] = StackEffect(popped=popped, pushed=pushed)
+        elif override == SKIP_OVERRIDE:
+            warnings.warn(
+                f"Could not get instruction metadata for {rust_name}"
+                " from CPython or override conf"
+            )
+
+        yield dataclasses.replace(opcode, override=override, **kwargs)
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class Opcode:
+    rust_name: str
+    id: int
+    have_argument: bool = False
+    cache_entry: int = 0
+    stack_effect: StackEffect | None = None
+    properties: Properties = dataclasses.field(default_factory=lambda: SKIP_PROPERTIES)
+    family: Family | None = None
+    override: Override = dataclasses.field(default_factory=Override)
+
+    @property
+    def is_instrumented(self) -> bool:
+        if (res := self.override.is_instrumented) is not None:
+            return res
+
+        return self.cpython_name.startswith("INSTRUMENTED_")
+
+    @property
+    def cpython_name(self):
+        return utils.to_upper_snake_case(self.rust_name)
+
+    @property
+    def cpy_popped(self) -> str | None:
+        return getattr(self.stack_effect, "popped", None)
+
+    @property
+    def cpy_pushed(self) -> str | None:
+        return getattr(self.stack_effect, "pushed", None)
+
+    @property
+    def stack_effect_popped(self) -> str:
+        ove_popped = self.override.stack_effect.popped
+
+        if (ove_popped is None) and (self.cpy_popped is None):
+            raise ValueError(f"{self.rust_name} is missing popped stack_effect")
+
+        return ove_popped or self.cpy_popped
+
+    @property
+    def stack_effect_pushed(self) -> str:
+        ove_pushed = self.override.stack_effect.pushed
+
+        if (ove_pushed is None) and (self.cpy_pushed is None):
+            raise ValueError(f"{self.rust_name} is missing pushed stack_effect")
+
+        return ove_pushed or self.cpy_pushed
+
+    @classmethod
+    def from_str(cls, entry: str) -> typing.Self:
+        rust_name = re.match(r"(\w+)", entry).group(1)
+        id_num = re.findall(r"= (\d+)", entry)[0]
+        have_argument = "Arg<" in entry
+        return cls(rust_name, int(id_num), have_argument=have_argument)
+
+    def __lt__(self, other: typing.Self) -> bool:
+        return self.id < other.id