gh-126612: Include stack effects of uops when computing maximum stack depth (#126894)

1 files changed, 37 insertions, 21 deletions

diff --git a/Tools/cases_generator/stack.py b/Tools/cases_generator/stack.py
index a954bed4df0..286f47d0cfb 100644
--- a/Tools/cases_generator/stack.py
+++ b/Tools/cases_generator/stack.py
@@ -1,8 +1,9 @@
 import re
 from analyzer import StackItem, StackEffect, Instruction, Uop, PseudoInstruction
+from collections import defaultdict
 from dataclasses import dataclass
 from cwriter import CWriter
-from typing import Iterator
+from typing import Iterator, Tuple
 
 UNUSED = {"unused"}
 
@@ -385,31 +386,46 @@ class Stack:
         self.align(other, out)
 
 
+def stacks(inst: Instruction | PseudoInstruction) -> Iterator[StackEffect]:
+    if isinstance(inst, Instruction):
+        for uop in inst.parts:
+            if isinstance(uop, Uop):
+                yield uop.stack
+    else:
+        assert isinstance(inst, PseudoInstruction)
+        yield inst.stack
+
+
+def apply_stack_effect(stack: Stack, effect: StackEffect) -> None:
+    locals: dict[str, Local] = {}
+    for var in reversed(effect.inputs):
+        _, local = stack.pop(var)
+        if var.name != "unused":
+            locals[local.name] = local
+    for var in effect.outputs:
+        if var.name in locals:
+            local = locals[var.name]
+        else:
+            local = Local.unused(var)
+        stack.push(local)
+
+
 def get_stack_effect(inst: Instruction | PseudoInstruction) -> Stack:
     stack = Stack()
+    for s in stacks(inst):
+        apply_stack_effect(stack, s)
+    return stack
 
-    def stacks(inst: Instruction | PseudoInstruction) -> Iterator[StackEffect]:
-        if isinstance(inst, Instruction):
-            for uop in inst.parts:
-                if isinstance(uop, Uop):
-                    yield uop.stack
-        else:
-            assert isinstance(inst, PseudoInstruction)
-            yield inst.stack
 
+def get_stack_effects(inst: Instruction | PseudoInstruction) -> list[Stack]:
+    """Returns a list of stack effects after each uop"""
+    result = []
+    stack = Stack()
     for s in stacks(inst):
-        locals: dict[str, Local] = {}
-        for var in reversed(s.inputs):
-            _, local = stack.pop(var)
-            if var.name != "unused":
-                locals[local.name] = local
-        for var in s.outputs:
-            if var.name in locals:
-                local = locals[var.name]
-            else:
-                local = Local.unused(var)
-            stack.push(local)
-    return stack
+        apply_stack_effect(stack, s)
+        result.append(stack.copy())
+    return result
+
 
 @dataclass
 class Storage: