GH-98831: Implement super-instruction generation (#99084)

Co-authored-by: C.A.M. Gerlach <CAM.Gerlach@Gerlach.CAM>

2 files changed, 66 insertions, 14 deletions

diff --git a/Tools/cases_generator/generate_cases.py b/Tools/cases_generator/generate_cases.py
index ec2481b..c5a436d 100644
--- a/Tools/cases_generator/generate_cases.py
+++ b/Tools/cases_generator/generate_cases.py
@@ -11,7 +11,7 @@ import re
 import sys
 
 import parser
-from parser import InstDef
+from parser import InstDef  # TODO: Use parser.InstDef
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument("-i", "--input", type=str, default="Python/bytecodes.c")
@@ -29,19 +29,24 @@ def eopen(filename: str, mode: str = "r"):
     return open(filename, mode)
 
 
-def parse_cases(src: str, filename: str|None = None) -> tuple[list[InstDef], list[parser.Family]]:
+def parse_cases(
+    src: str, filename: str|None = None
+) -> tuple[list[InstDef], list[parser.Super], list[parser.Family]]:
     psr = parser.Parser(src, filename=filename)
     instrs: list[InstDef] = []
+    supers: list[parser.Super] = []
     families: list[parser.Family] = []
     while not psr.eof():
         if inst := psr.inst_def():
             assert inst.block
-            instrs.append(InstDef(inst.name, inst.inputs, inst.outputs, inst.block))
+            instrs.append(inst)
+        elif sup := psr.super_def():
+            supers.append(sup)
         elif fam := psr.family_def():
             families.append(fam)
         else:
             raise psr.make_syntax_error(f"Unexpected token")
-    return instrs, families
+    return instrs, supers, families
 
 
 def always_exits(block: parser.Block) -> bool:
@@ -62,7 +67,7 @@ def always_exits(block: parser.Block) -> bool:
     return line.startswith(("goto ", "return ", "DISPATCH", "GO_TO_", "Py_UNREACHABLE()"))
 
 
-def write_cases(f: io.TextIOBase, instrs: list[InstDef]):
+def write_cases(f: io.TextIOBase, instrs: list[InstDef], supers: list[parser.Super]):
     predictions = set()
     for inst in instrs:
         for target in re.findall(r"(?:PREDICT|GO_TO_INSTRUCTION)\((\w+)\)", inst.block.text):
@@ -70,8 +75,10 @@ def write_cases(f: io.TextIOBase, instrs: list[InstDef]):
     indent = "        "
     f.write(f"// This file is generated by {os.path.relpath(__file__)}\n")
     f.write("// Do not edit!\n")
+    instr_index: dict[str, InstDef] = {}
     for instr in instrs:
         assert isinstance(instr, InstDef)
+        instr_index[instr.name] = instr
         f.write(f"\n{indent}TARGET({instr.name}) {{\n")
         if instr.name in predictions:
             f.write(f"{indent}    PREDICTED({instr.name});\n")
@@ -102,6 +109,22 @@ def write_cases(f: io.TextIOBase, instrs: list[InstDef]):
         # Write trailing '}'
         f.write(f"{indent}}}\n")
 
+    for sup in supers:
+        assert isinstance(sup, parser.Super)
+        components = [instr_index[name] for name in sup.ops]
+        f.write(f"\n{indent}TARGET({sup.name}) {{\n")
+        for i, instr in enumerate(components):
+            if i > 0:
+                f.write(f"{indent}    NEXTOPARG();\n")
+                f.write(f"{indent}    next_instr++;\n")
+            text = instr.block.to_text(-4)
+            textlines = text.splitlines(True)
+            textlines = [line for line in textlines if not line.strip().startswith("PREDICTED(")]
+            text = "".join(textlines)
+            f.write(f"{indent}    {text.strip()}\n")
+        f.write(f"{indent}    DISPATCH();\n")
+        f.write(f"{indent}}}\n")
+
 
 def main():
     args = arg_parser.parse_args()
@@ -110,21 +133,22 @@ def main():
     begin = srclines.index("// BEGIN BYTECODES //")
     end = srclines.index("// END BYTECODES //")
     src = "\n".join(srclines[begin+1 : end])
-    instrs, families = parse_cases(src, filename=args.input)
-    ninstrs = nfamilies = 0
+    instrs, supers, families = parse_cases(src, filename=args.input)
+    ninstrs = nsupers = nfamilies = 0
     if not args.quiet:
         ninstrs = len(instrs)
+        nsupers = len(supers)
         nfamilies = len(families)
         print(
-            f"Read {ninstrs} instructions "
+            f"Read {ninstrs} instructions, {nsupers} supers, "
             f"and {nfamilies} families from {args.input}",
             file=sys.stderr,
         )
     with eopen(args.output, "w") as f:
-        write_cases(f, instrs)
+        write_cases(f, instrs, supers)
     if not args.quiet:
         print(
-            f"Wrote {ninstrs} instructions to {args.output}",
+            f"Wrote {ninstrs + nsupers} instructions to {args.output}",
             file=sys.stderr,
         )
 
diff --git a/Tools/cases_generator/parser.py b/Tools/cases_generator/parser.py
index d5e4de2..f603bc6 100644
--- a/Tools/cases_generator/parser.py
+++ b/Tools/cases_generator/parser.py
@@ -39,13 +39,16 @@ class Node:
 
     @property
     def text(self) -> str:
+        return self.to_text()
+
+    def to_text(self, dedent: int = 0) -> str:
         context = self.context
         if not context:
             return ""
         tokens = context.owner.tokens
         begin = context.begin
         end = context.end
-        return lx.to_text(tokens[begin:end])
+        return lx.to_text(tokens[begin:end], dedent)
 
 
 @dataclass
@@ -62,6 +65,12 @@ class InstDef(Node):
 
 
 @dataclass
+class Super(Node):
+    name: str
+    ops: list[str]
+
+
+@dataclass
 class Family(Node):
     name: str
     members: list[str]
@@ -156,17 +165,36 @@ class Parser(PLexer):
         return self.input()  # TODO: They're not quite the same.
 
     @contextual
-    def family_def(self) -> Family | None:
+    def super_def(self) -> Super | None:
+        if (tkn := self.expect(lx.IDENTIFIER)) and tkn.text == "super":
+            if self.expect(lx.LPAREN):
+                if (tkn := self.expect(lx.IDENTIFIER)):
+                    if self.expect(lx.RPAREN):
+                        if self.expect(lx.EQUALS):
+                            if ops := self.ops():
+                                res = Super(tkn.text, ops)
+                                return res
+
+    def ops(self) -> list[str] | None:
         here = self.getpos()
+        if tkn := self.expect(lx.IDENTIFIER):
+            ops = [tkn.text]
+            while self.expect(lx.PLUS):
+                if tkn := self.require(lx.IDENTIFIER):
+                    ops.append(tkn.text)
+            self.require(lx.SEMI)
+            return ops
+
+    @contextual
+    def family_def(self) -> Family | None:
         if (tkn := self.expect(lx.IDENTIFIER)) and tkn.text == "family":
             if self.expect(lx.LPAREN):
                 if (tkn := self.expect(lx.IDENTIFIER)):
-                    name = tkn.text
                     if self.expect(lx.RPAREN):
                         if self.expect(lx.EQUALS):
                             if members := self.members():
                                 if self.expect(lx.SEMI):
-                                    return Family(name, members)
+                                    return Family(tkn.text, members)
         return None
 
     def members(self):