Fix async comprehension exception handler range and SEND stack depth

youknowone · youknowone · commit 41860a93a3ec · 2026-02-01T12:16:31.000+09:00
- Move PopBlock before compile_store in async comprehension so
  STORE_SUBSCR is not covered by the END_ASYNC_FOR exception handler
- Add SEND jump target depth = depth - 1 in stackdepth calculation
- Remove expectedFailure from three now-passing test_coroutines tests
diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
@@ -1364,7 +1364,6 @@ async def foo():
             self.fail('invalid asynchronous context manager did not fail')
 
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: "'async with' received an object from __aexit__ that does not implement __await__: int" does not match "'range_iterator' object is not callable"
     def test_with_8(self):
         CNT = 0
 
@@ -2170,7 +2169,6 @@ async def func(): pass
                              f"coroutine {coro_repr}")
             self.assertIn("was never awaited", str(cm.unraisable.exc_value))
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: StopAsyncIteration not raised
     def test_for_assign_raising_stop_async_iteration(self):
         class BadTarget:
             def __setitem__(self, key, value):
@@ -2204,7 +2202,6 @@ async def run_gen():
             return 'end'
         self.assertEqual(run_async(run_gen()), ([], 'end'))
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: StopAsyncIteration not raised
     def test_for_assign_raising_stop_async_iteration_2(self):
         class BadIterable:
             def __iter__(self):
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -7414,9 +7414,11 @@ impl Compiler {
                 )?;
                 self.emit_load_const(ConstantData::None);
                 self.compile_yield_from_sequence(true)?;
-                self.compile_store(&generator.target)?;
+                // POP_BLOCK before store: only __anext__/yield_from are
+                // protected by SetupFinally targeting END_ASYNC_FOR.
                 emit!(self, PseudoInstruction::PopBlock);
                 self.pop_fblock(FBlockType::AsyncComprehensionGenerator);
+                self.compile_store(&generator.target)?;
             } else {
                 emit!(
                     self,
diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs
@@ -837,8 +837,15 @@ impl CodeInfo {
                         }
                         stackdepth_push(&mut stack, &mut start_depths, ins.target, handler_depth);
                     } else {
-                        let target_depth = depth.checked_add_signed(effect).ok_or({
-                            if effect < 0 {
+                        // SEND's jump path has different stack effect than
+                        // normal path: jump replaces receiver (depth - 1),
+                        // normal pushes above receiver (depth + 0).
+                        let jump_effect = match instr.real() {
+                            Some(Instruction::Send { .. }) => -1i32,
+                            _ => effect,
+                        };
+                        let target_depth = depth.checked_add_signed(jump_effect).ok_or({
+                            if jump_effect < 0 {
                                 InternalError::StackUnderflow
                             } else {
                                 InternalError::StackOverflow
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -937,9 +937,9 @@ impl ExecutingFrame<'_> {
                 Ok(None)
             }
             Instruction::EndAsyncFor => {
-                // END_ASYNC_FOR pops (awaitable, exc) from stack
-                // Stack: [awaitable, exc] -> []
-                // exception_unwind pushes exception to stack before jumping to handler
+                // Pops (awaitable, exc) from stack.
+                // If exc is StopAsyncIteration, clears it (normal loop end).
+                // Otherwise re-raises.
                 let exc = self.pop_value();
                 let _awaitable = self.pop_value();
 
@@ -1979,22 +1979,20 @@ impl ExecutingFrame<'_> {
                 Ok(Some(ExecutionResult::Yield(value)))
             }
             Instruction::Send { target } => {
-                // Stack: (receiver, value) -> (receiver, retval)
-                // On StopIteration: replace value with stop value and jump to target
+                // (receiver, v -- receiver, retval)
+                // Pops v, sends it to receiver. On yield, pushes retval
+                // (so stack = [..., receiver, retval]). On return/StopIteration,
+                // also pushes retval and jumps to END_SEND which will pop receiver.
                 let exit_label = target.get(arg);
                 let val = self.pop_value();
                 let receiver = self.top_value();
 
                 match self._send(receiver, val, vm)? {
                     PyIterReturn::Return(value) => {
-                        // Value yielded, push it back for YIELD_VALUE
-                        // Stack: (receiver, retval)
                         self.push_value(value);
                         Ok(None)
                     }
                     PyIterReturn::StopIteration(value) => {
-                        // StopIteration: replace top with stop value, jump to exit
-                        // Stack: (receiver, value) - receiver stays, v replaced
                         let value = vm.unwrap_or_none(value);
                         self.push_value(value);
                         self.jump(exit_label);
