gh-104909: Implement conditional stack effects for macros

Tools/cases_generator/generate_cases.py

@@ -174,14 +174,12 @@ def block(self, head: str):
 
     def stack_adjust(
         self,
-        diff: int,
         input_effects: list[StackEffect],
         output_effects: list[StackEffect],
     ):
-        # TODO: Get rid of 'diff' parameter
         shrink, isym = list_effect_size(input_effects)
         grow, osym = list_effect_size(output_effects)
-        diff += grow - shrink
+        diff = grow - shrink
         if isym and isym != osym:
             self.emit(f"STACK_SHRINK({isym});")
         if diff < 0:
@@ -191,14 +189,20 @@ def stack_adjust(
         if osym and osym != isym:
             self.emit(f"STACK_GROW({osym});")
 
-    def declare(self, dst: StackEffect, src: StackEffect | None):
+    def declare(
+            self, dst: StackEffect,
+            src: StackEffect | None,
+            *,
+            # Don't initialize from dst.cond; used for conditional *outputs*.
+            init_null: bool = True,
+        ):
         if dst.name == UNUSED:
             return
         typ = f"{dst.type}" if dst.type else "PyObject *"
         if src:
             cast = self.cast(dst, src)
             init = f" = {cast}{src.name}"
-        elif dst.cond:
+        elif init_null and dst.cond:
             init = " = NULL"
         else:
             init = ""
@@ -338,7 +342,6 @@ def write(self, out: Formatter) -> None:
 
         # Write net stack growth/shrinkage
         out.stack_adjust(
-            0,
             [ieff for ieff in self.input_effects],
             [oeff for oeff in self.output_effects],
         )
@@ -491,7 +494,7 @@ class OverriddenInstructionPlaceHolder:
     name: str
 
 
-AnyInstruction = Instruction | MacroInstruction
+AnyInstruction = Instruction | MacroInstruction | PseudoInstruction
 INSTR_FMT_PREFIX = "INSTR_FMT_"
 
 
@@ -530,6 +533,7 @@ def error(self, msg: str, node: parser.Node) -> None:
     macros: dict[str, parser.Macro]
     macro_instrs: dict[str, MacroInstruction]
     families: dict[str, parser.Family]
+    pseudos: dict[str, parser.Pseudo]
     pseudo_instrs: dict[str, PseudoInstruction]
 
     def parse(self) -> None:
@@ -587,7 +591,7 @@ def parse_file(self, filename: str, instrs_idx: dict[str, int]) -> None:
 
         # Parse from start
         psr.setpos(start)
-        thing: parser.InstDef | parser.Macro | parser.Family | None
+        thing: parser.InstDef | parser.Macro | parser.Pseudo | parser.Family | None
         thing_first_token = psr.peek()
         while thing := psr.definition():
             match thing:
@@ -820,9 +824,14 @@ def stack_analysis(
 
         Ignore cache effects.
 
-        Return the list of variable names and the initial stack pointer.
+        Return the list of variables (as StackEffects) and the initial stack pointer.
         """
         lowest = current = highest = 0
+        conditions: dict[int, str] = {}  # Indexed by 'current'.
+        last_instr: Instruction | None = None
+        for thing in components:
+            if isinstance(thing, Instruction):
+                last_instr = thing
         for thing in components:
             match thing:
                 case Instruction() as instr:
@@ -835,9 +844,24 @@ def stack_analysis(
                             "which are not supported in macro instructions",
                             instr.inst,  # TODO: Pass name+location of macro
                         )
+                    if any(eff.cond for eff in instr.input_effects):
+                        self.error(
+                            f"Instruction {instr.name!r} has conditional input stack effect, "
+                            "which are not supported in macro instructions",
+                            instr.inst,  # TODO: Pass name+location of macro
+                        )
+                    if any(eff.cond for eff in instr.output_effects) and instr is not last_instr:
+                        self.error(
+                            f"Instruction {instr.name!r} has conditional output stack effect, "
+                            "but is not the last instruction in a macro",
+                            instr.inst,  # TODO: Pass name+location of macro
+                        )
                     current -= len(instr.input_effects)
                     lowest = min(lowest, current)
-                    current += len(instr.output_effects)
+                    for eff in instr.output_effects:
+                        if eff.cond:
+                            conditions[current] = eff.cond
+                        current += 1
                     highest = max(highest, current)
                 case parser.CacheEffect():
                     pass
@@ -846,9 +870,9 @@ def stack_analysis(
         # At this point, 'current' is the net stack effect,
         # and 'lowest' and 'highest' are the extremes.
         # Note that 'lowest' may be negative.
-        # TODO: Reverse the numbering.
         stack = [
-            StackEffect(f"_tmp_{i+1}", "") for i in reversed(range(highest - lowest))
+            StackEffect(f"_tmp_{i}", "", conditions.get(highest - i, ""))
+            for i in reversed(range(1, highest - lowest + 1))
         ]
         return stack, -lowest
 
@@ -880,15 +904,17 @@ def effect_str(effects: list[StackEffect]) -> str:
                 low = 0
                 sp = 0
                 high = 0
+                pushed_symbolic: list[str] = []
                 for comp in parts:
                     for effect in comp.instr.input_effects:
                         assert not effect.cond, effect
                         assert not effect.size, effect
                         sp -= 1
                         low = min(low, sp)
                     for effect in comp.instr.output_effects:
-                        assert not effect.cond, effect
                         assert not effect.size, effect
+                        if effect.cond:
+                            pushed_symbolic.append(maybe_parenthesize(f"{maybe_parenthesize(effect.cond)} ? 1 : 0"))
                         sp += 1
                         high = max(sp, high)
                 if high != max(0, sp):
@@ -898,9 +924,10 @@ def effect_str(effects: list[StackEffect]) -> str:
                     # calculations to use the micro ops.
                     self.error("Macro has virtual stack growth", thing)
                 popped = str(-low)
-                pushed = str(sp - low)
+                pushed_symbolic.append(str(sp - low - len(pushed_symbolic)))
+                pushed = " + ".join(pushed_symbolic)
             case parser.Pseudo():
-                instr = self.pseudos[thing.name]
+                instr = self.pseudo_instrs[thing.name]
                 popped = pushed = None
                 # Calculate stack effect, and check that it's the the same
                 # for all targets.
@@ -1168,12 +1195,17 @@ def wrap_macro(self, mac: MacroInstruction):
                 src = None
                 if i < mac.initial_sp:
                     src = StackEffect(f"stack_pointer[-{mac.initial_sp - i}]", "")
-                self.out.declare(var, src)
+                self.out.declare(var, src, init_null=False)
 
             yield
 
-            # TODO: Use slices of mac.stack instead of numeric values
-            self.out.stack_adjust(mac.final_sp - mac.initial_sp, [], [])
+            # The input effects should have no conditionals.
+            # Only the output effects do (for now).
+            ieffects = [
+                StackEffect(eff.name, eff.type) if eff.cond else eff
+                for eff in mac.stack[:mac.initial_sp]
+            ]
+            self.out.stack_adjust(ieffects, mac.stack[:mac.final_sp])
 
             for i, var in enumerate(reversed(mac.stack[: mac.final_sp]), 1):
                 dst = StackEffect(f"stack_pointer[-{i}]", "")

Tools/cases_generator/test_generator.py

@@ -432,3 +432,38 @@ def test_cond_effect():
         }
     """
     run_cases_test(input, output)
+
+def test_macro_cond_effect():
+    input = """
+        op(A, (left, right --)) {
+            # Body of A
+        }
+        op(B, (-- extra if (oparg), res)) {
+            # Body of B
+        }
+        macro(M) = A + B;
+    """
+    output = """
+        TARGET(M) {
+            PyObject *_tmp_1 = stack_pointer[-1];
+            PyObject *_tmp_2 = stack_pointer[-2];
+            {
+                PyObject *right = _tmp_1;
+                PyObject *left = _tmp_2;
+                # Body of A
+            }
+            {
+                PyObject *extra = NULL;
+                PyObject *res;
+                # Body of B
+                if (oparg) { _tmp_2 = extra; }
+                _tmp_1 = res;
+            }
+            STACK_SHRINK(1);
+            STACK_GROW((oparg ? 1 : 0));
+            stack_pointer[-1] = _tmp_1;
+            if (oparg) { stack_pointer[-2] = _tmp_2; }
+            DISPATCH();
+        }
+    """
+    run_cases_test(input, output)