Fix code gen for RESUME and YIELD_VALUE

Summary: `RESUME` now has a mask that can be applied, and 'YIELD_VALUE` no longer gets set the length of the exception stack.

Reviewed By: martindemello

Differential Revision: D80820416

fbshipit-source-id: 462a4e5a672cbd770ba1f57229c46b40982549b0

Author: DinoV

cinderx/PythonLib/cinderx/compiler/pyassem.py

@@ -9,7 +9,7 @@
 from collections import defaultdict
 from contextlib import contextmanager, redirect_stdout
 from dataclasses import dataclass
-from enum import IntEnum
+from enum import IntEnum, IntFlag
 
 try:
     # pyre-ignore[21]: No _inline_cache_entries
@@ -54,6 +54,16 @@
 MAX_COPY_SIZE = 4
 
 
+class ResumeOparg(IntFlag):
+    ScopeEntry = 0
+    Yield = 1
+    YieldFrom = 2
+    Await = 3
+
+    LocationMask = 0x03
+    Depth1Mask = 0x04
+
+
 def sign(a: float) -> float:
     if not isinstance(a, float):
         raise TypeError(f"Must be a real number, not {type(a)}")
@@ -2700,6 +2710,66 @@ def make_explicit_jump_block(self) -> Block:
         res.num_predecessors = 1
         return res
 
+    def label_exception_targets(self) -> None:
+        def push_todo_block(block: Block) -> None:
+            todo_stack.append(block)
+            visited.add(block)
+
+        todo_stack: list[Block] = [self.entry]
+        visited: set[Block] = {self.entry}
+        except_stack = []
+        self.entry.except_stack = except_stack
+
+        while todo_stack:
+            block = todo_stack.pop()
+            assert block in visited
+            except_stack = block.except_stack
+            block.except_stack = []
+            handler = except_stack[-1] if except_stack else None
+            last_yield_except_depth = -1
+            for instr in block.insts:
+                if instr.opname in SETUP_OPCODES:
+                    target = instr.target
+                    assert target, instr
+                    if target not in visited:
+                        # Copy the current except stack into the target's except stack
+                        target.except_stack = list(except_stack)
+                        push_todo_block(target)
+                    handler = self.push_except_block(except_stack, instr)
+                elif instr.opname == "POP_BLOCK":
+                    except_stack.pop()
+                    handler = except_stack[-1] if except_stack else None
+                    instr.set_to_nop()
+                elif instr.is_jump(self.opcode) and instr.opname != "END_ASYNC_FOR":
+                    instr.exc_handler = handler
+                    if instr.target not in visited:
+                        target = instr.target
+                        assert target
+                        if block.has_fallthrough:
+                            # Copy the current except stack into the block's except stack
+                            target.except_stack = list(except_stack)
+                        else:
+                            # Move the current except stack to the block and start a new one
+                            target.except_stack = except_stack
+                            except_stack = []
+                        push_todo_block(target)
+                elif instr.opname == "YIELD_VALUE":
+                    last_yield_except_depth = len(except_stack)
+                    instr.exc_handler = handler
+                elif instr.opname == "RESUME":
+                    instr.exc_handler = handler
+                    if instr.oparg != ResumeOparg.ScopeEntry:
+                        if last_yield_except_depth == 1:
+                            instr.ioparg |= ResumeOparg.Depth1Mask
+                        last_yield_except_depth = -1
+                else:
+                    instr.exc_handler = handler
+
+            if block.has_fallthrough and block.next and block.next not in visited:
+                assert except_stack is not None
+                block.next.except_stack = except_stack
+                push_todo_block(block.next)
+
     _const_opcodes: set[str] = set(PyFlowGraph312._const_opcodes) | {"LOAD_SMALL_INT"}
 
 

cinderx/PythonLib/cinderx/compiler/pycodegen.py

@@ -61,6 +61,7 @@
         PyFlowGraph314,
         PyFlowGraphCinder310,
         PyFlowGraphCinder312,
+        ResumeOparg,
         SrcLocation,
     )
     from .symbols import (
@@ -2748,13 +2749,6 @@ def __init__(
         self.node = node
 
 
-class ResumeOparg(IntEnum):
-    ScopeEntry = 0
-    Yield = 1
-    YieldFrom = 2
-    Await = 3
-
-
 class CodeGenerator310(CodeGenerator):
     flow_graph: type[PyFlowGraph] = PyFlowGraph310
     _SymbolVisitor = SymbolVisitor310
@@ -4309,6 +4303,9 @@ def emit_print(self) -> None:
         self.set_no_pos()
         self.emit("POP_TOP")
 
+    def emit_yield_value_for_yield_from(self) -> None:
+        self.emit("YIELD_VALUE")
+
     def emit_yield_from(self, await_: bool = False) -> None:
         send = self.newBlock("send")
         fail = self.newBlock("fail")
@@ -4321,7 +4318,7 @@ def emit_yield_from(self, await_: bool = False) -> None:
 
         # Setup try/except to handle StopIteration
         self.emit("SETUP_FINALLY", fail)
-        self.emit("YIELD_VALUE")
+        self.emit_yield_value_for_yield_from()
         self.emit_noline("POP_BLOCK")
         self.emit_resume(ResumeOparg.Await if await_ else ResumeOparg.YieldFrom)
         self.emit("JUMP_NO_INTERRUPT", send)
@@ -6376,6 +6373,9 @@ def visitWhile(self, node: ast.While) -> None:
 
         self.nextBlock(after)
 
+    def emit_yield_value_for_yield_from(self) -> None:
+        self.emit("YIELD_VALUE", 1)
+
 
 class CinderCodeGenerator310(CinderCodeGenBase, CodeGenerator310):
     flow_graph = PyFlowGraphCinder310