gh-132967: Add type attributes to the optimizer DSL

Lib/test/test_generated_cases.py

@@ -59,14 +59,14 @@ class TestEffects(unittest.TestCase):
     def test_effect_sizes(self):
         stack = Stack()
         inputs = [
-            x := StackItem("x", None, "1"),
-            y := StackItem("y", None, "oparg"),
-            z := StackItem("z", None, "oparg*2"),
+            x := StackItem("x", None, "1", []),
+            y := StackItem("y", None, "oparg", []),
+            z := StackItem("z", None, "oparg*2", []),
         ]
         outputs = [
-            StackItem("x", None, "1"),
-            StackItem("b", None, "oparg*4"),
-            StackItem("c", None, "1"),
+            StackItem("x", None, "1", []),
+            StackItem("b", None, "oparg*4", []),
+            StackItem("c", None, "1", []),
         ]
         null = CWriter.null()
         stack.pop(z, null)
@@ -2253,5 +2253,28 @@ def test_validate_uop_unused_size_mismatch(self):
                                     "Inputs must have equal sizes"):
             self.run_cases_test(input, input2, output)
 
+    def test_uop_type_attribute_input(self):
+        input = """
+        op(OP, (foo -- )) {
+        }
+        """
+        input2 = """
+        op(OP, (type(&PyLong_Type) foo -- )) {
+            (void)foo;
+        }
+        """
+        output = """
+        case OP: {
+            JitOptSymbol *foo;
+            foo = stack_pointer[-1];
+            (void)foo;
+            assert(sym_matches_type(foo, &PyLong_Type));
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/Core_and_Builtins/2025-05-11-06-25-26.gh-issue-132967.AwR3m5.rst

@@ -0,0 +1 @@
+Add support for attributes to the bytecodes DSL.

Tools/cases_generator/analyzer.py

@@ -5,7 +5,7 @@
 import re
 from typing import Optional, Callable
 
-from parser import Stmt, SimpleStmt, BlockStmt, IfStmt, WhileStmt
+from parser import Stmt, SimpleStmt, BlockStmt, IfStmt, WhileStmt, StackAttribute
 
 @dataclass
 class EscapingCall:
@@ -137,13 +137,14 @@ class StackItem:
     name: str
     type: str | None
     size: str
+    attributes: list[StackAttribute]
     peek: bool = False
     used: bool = False
 
     def __str__(self) -> str:
         size = f"[{self.size}]" if self.size else ""
         type = "" if self.type is None else f"{self.type} "
-        return f"{type}{self.name}{size} {self.peek}"
+        return f"{self.attributes} {type}{self.name}{size} {self.peek}"
 
     def is_array(self) -> bool:
         return self.size != ""
@@ -345,7 +346,7 @@ def override_error(
 def convert_stack_item(
     item: parser.StackEffect, replace_op_arg_1: str | None
 ) -> StackItem:
-    return StackItem(item.name, item.type, item.size)
+    return StackItem(item.name, item.type, item.size, item.attributes)
 
 def check_unused(stack: list[StackItem], input_names: dict[str, lexer.Token]) -> None:
     "Unused items cannot be on the stack above used, non-peek items"

Tools/cases_generator/optimizer_generator.py

@@ -24,6 +24,7 @@
 from typing import TextIO
 from lexer import Token
 from stack import Local, Stack, StackError, Storage
+from parser import TYPE
 
 DEFAULT_OUTPUT = ROOT / "Python/optimizer_cases.c.h"
 DEFAULT_ABSTRACT_INPUT = (ROOT / "Python/optimizer_bytecodes.c").absolute().as_posix()
@@ -147,6 +148,14 @@ def goto_label(self, goto: Token, label: Token, storage: Storage) -> None:
         self.out.emit(goto)
         self.out.emit(label)
 
+def get_type(item: StackItem) -> str | None:
+    for attribute in item.attributes:
+        if attribute.ident == TYPE:
+            return attribute.expr
+    return None
+
+
+
 def write_uop(
     override: Uop | None,
     uop: Uop,
@@ -182,7 +191,12 @@ def write_uop(
             for var in storage.inputs:  # type: ignore[possibly-undefined]
                 var.in_local = False
             _, storage = emitter.emit_tokens(override, storage, None, False)
+            # Emit type effects.
             out.start_line()
+            for input_ in override.stack.inputs:
+                typ = get_type(input_)
+                if typ is not None:
+                    emitter.emit(f"assert(sym_matches_type({input_.name}, {typ}));\n")
             storage.flush(out)
         else:
             emit_default(out, uop, stack)

Tools/cases_generator/parser.py

@@ -18,6 +18,8 @@
     WhileStmt,
     BlockStmt,
     MacroIfStmt,
+    StackAttribute,
+    TYPE,
 )
 
 import pprint

Tools/cases_generator/parsing.py

@@ -244,15 +244,24 @@ def accept(self, visitor: Visitor) -> None:
 
     __hash__ = object.__hash__
 
+@dataclass
+class StackAttribute(Node):
+    ident: str
+    expr: str
+
+    def __repr__(self) -> str:
+        return f"{self.ident}({self.expr})"
+
 @dataclass
 class StackEffect(Node):
     name: str = field(compare=False)  # __eq__ only uses type, cond, size
+    attributes: list[StackAttribute] = field(compare=False)
     type: str = ""  # Optional `:type`
     size: str = ""  # Optional `[size]`
     # Note: size cannot be combined with type or cond
 
     def __repr__(self) -> str:
-        items = [self.name, self.type, self.size]
+        items = [self.attributes, self.name, self.type, self.size]
         while items and items[-1] == "":
             del items[-1]
         return f"StackEffect({', '.join(repr(item) for item in items)})"
@@ -274,6 +283,14 @@ class OpName(Node):
     name: str
 
 
+TYPE = "type"
+# We have to do this at the parsing stage and not
+# lexing stage as we want to allow this to be used as
+# a normal identifier in C code.
+STACK_ATTRIBUTES = {
+    TYPE,
+}
+
 InputEffect = StackEffect | CacheEffect
 OutputEffect = StackEffect
 UOp = OpName | CacheEffect
@@ -458,10 +475,27 @@ def cache_effect(self) -> CacheEffect | None:
                     return CacheEffect(tkn.text, size)
         return None
 
+    def stack_attributes(self) -> list[StackAttribute]:
+        # IDENTIFIER '(' expression ')'
+        res = []
+        while tkn := self.expect(lx.IDENTIFIER):
+            if self.expect(lx.LPAREN):
+                if tkn.text not in STACK_ATTRIBUTES:
+                    raise self.make_syntax_error(f"Stack attribute {tkn.text} is not recognized.")
+                expr = self.expression()
+                assert expr is not None
+                self.require(lx.RPAREN)
+                res.append(StackAttribute(tkn.text.strip(), expr.size.strip()))
+            else:
+                self.backup()
+                break
+        return res
+
     @contextual
     def stack_effect(self) -> StackEffect | None:
-        # IDENTIFIER [':' IDENTIFIER [TIMES]] ['if' '(' expression ')']
-        # | IDENTIFIER '[' expression ']'
+        # stack_attributes IDENTIFIER [':' IDENTIFIER [TIMES]] ['if' '(' expression ')']
+        # | stack_attributes IDENTIFIER '[' expression ']'
+        stack_attributes = self.stack_attributes()
         if tkn := self.expect(lx.IDENTIFIER):
             type_text = ""
             if self.expect(lx.COLON):
@@ -476,7 +510,7 @@ def stack_effect(self) -> StackEffect | None:
                     raise self.make_syntax_error("Expected expression")
                 self.require(lx.RBRACKET)
                 size_text = size.text.strip()
-            return StackEffect(tkn.text, type_text, size_text)
+            return StackEffect(tkn.text, stack_attributes, type_text, size_text)
         return None
 
     @contextual

Tools/cases_generator/stack.py

@@ -168,7 +168,7 @@ def from_memory(defn: StackItem, offset: PointerOffset) -> "Local":
 
     @staticmethod
     def register(name: str) -> "Local":
-        item = StackItem(name, None, "", False, True)
+        item = StackItem(name, None, "", [], False, True)
         return Local(item, None, True)
 
     def kill(self) -> None: