Ensure stack is consistent at end of for loops

Author: markshannon

Lib/test/test_generated_cases.py

@@ -2080,6 +2080,42 @@ def test_recording_uop_between_real_uops_rejected(self):
                                     "non-recording, non-specializing uops"):
             self.run_cases_test(input, "")
 
+    def test_escaping_in_loop(self):
+        input = """
+        inst(TEST_LOOP_SPILL, (a -- a)) {
+            int n = (int)oparg;
+            for (int i = 0; i < n; i++) {
+                escaping_inside_loop(a);
+            }
+            escaping_after_loop(a);
+        }
+        """
+        output = """
+        TARGET(TEST_LOOP_SPILL) {
+            #if _Py_TAIL_CALL_INTERP
+            int opcode = TEST_LOOP_SPILL;
+            (void)(opcode);
+            #endif
+            frame->instr_ptr = next_instr;
+            next_instr += 1;
+            INSTRUCTION_STATS(TEST_LOOP_SPILL);
+            _PyStackRef a;
+            a = stack_pointer[-1];
+            int n = (int)oparg;
+            for (int i = 0; i < n; i++) {
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                _PyFrame_StackPointerValidate(frame);
+                escaping_inside_loop(a);
+                _PyFrame_StackPointerInvalidate(frame);
+            }
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            _PyFrame_StackPointerValidate(frame);
+            escaping_after_loop(a);
+            _PyFrame_StackPointerInvalidate(frame);
+            DISPATCH();
+        }
+        """
+        self.run_cases_test(input, output)
 
 class TestRecorderTableGeneration(unittest.TestCase):
 

Tools/cases_generator/generators_common.py

@@ -712,7 +712,10 @@ def emit_ForStmt(
         self.out.emit(stmt.for_)
         for tkn in stmt.header:
             self.out.emit(tkn)
-        return self._emit_stmt(stmt.body, uop, storage, inst)
+        reachable, brace, body_storage = self._emit_stmt(stmt.body, uop, storage.copy(), inst)
+        body_storage.merge(storage, self.out)
+        self.out.emit(brace)
+        return reachable, None, storage
 
     def emit_WhileStmt(
         self,