From c1e9647107c854439a9864b6ec4f6784aeb94ed5 Mon Sep 17 00:00:00 2001
From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
Date: Wed, 29 May 2024 10:47:56 +0100
Subject: gh-119689: generate stack effect metadata for pseudo instructions
 (#119691)

---
 Tools/cases_generator/stack.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

(limited to 'Tools/cases_generator/stack.py')

diff --git a/Tools/cases_generator/stack.py b/Tools/cases_generator/stack.py
index 5aecac39aef..7f07a6805b1 100644
--- a/Tools/cases_generator/stack.py
+++ b/Tools/cases_generator/stack.py
@@ -1,7 +1,8 @@
 import re
-from analyzer import StackItem, Instruction, Uop
+from analyzer import StackItem, StackEffect, Instruction, Uop, PseudoInstruction
 from dataclasses import dataclass
 from cwriter import CWriter
+from typing import Iterator
 
 UNUSED = {"unused"}
 
@@ -208,13 +209,20 @@ class Stack:
         return f"/* Variables: {[v.name for v in self.variables]}. Base offset: {self.base_offset.to_c()}. Top offset: {self.top_offset.to_c()} */"
 
 
-def get_stack_effect(inst: Instruction) -> Stack:
+def get_stack_effect(inst: Instruction | PseudoInstruction) -> Stack:
     stack = Stack()
-    for uop in inst.parts:
-        if not isinstance(uop, Uop):
-            continue
-        for var in reversed(uop.stack.inputs):
+    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
+
+    for s in stacks(inst):
+        for var in reversed(s.inputs):
             stack.pop(var)
-        for i, var in enumerate(uop.stack.outputs):
+        for var in s.outputs:
             stack.push(var)
     return stack
-- 
cgit v1.2.3