Fix with __exit__ error handling

youknowone · youknowone · commit 462f54f90693 · 2024-04-25T15:17:50.000+09:00
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -925,15 +925,15 @@ impl ExecutingFrame<'_> {
                     _ => None,
                 };
 
-                let exit = self.pop_value();
+                let exit = self.top_value();
 
                 let args = if let Some(exc) = exc {
                     vm.split_exception(exc)
                 } else {
                     (vm.ctx.none(), vm.ctx.none(), vm.ctx.none())
                 };
                 let exit_res = exit.call(args, vm)?;
-                self.push_value(exit_res);
+                self.replace_top(exit_res);
 
                 Ok(None)
             }
@@ -994,7 +994,7 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::GetANext => {
                 let aiter = self.top_value();
-                let awaitable = vm.call_special_method(&aiter, identifier!(vm, __anext__), ())?;
+                let awaitable = vm.call_special_method(aiter, identifier!(vm, __anext__), ())?;
                 let awaitable = if awaitable.payload_is::<PyCoroutine>() {
                     awaitable
                 } else {
@@ -1954,6 +1954,13 @@ impl ExecutingFrame<'_> {
         self.state.stack.drain(stack_len - count..)
     }
 
+    #[inline]
+    fn replace_top(&mut self, mut top: PyObjectRef) -> PyObjectRef {
+        let last = self.state.stack.last_mut().unwrap();
+        std::mem::swap(&mut top, last);
+        top
+    }
+
     #[inline]
     #[track_caller] // not a real track_caller but top_value is not very useful
     fn top_value(&self) -> &PyObject {
