RustPython · youknowone · Aug 10, 2022
diff --git a/bytecode/src/lib.rs b/bytecode/src/lib.rs
@@ -73,7 +73,7 @@ impl Constant for ConstantData {
 }
 
 /// A Constant Bag
-pub trait ConstantBag: Sized + Copy {
+pub trait ConstantBag: Sized + std::marker::Copy {
     type Constant: Constant;
     fn make_constant<C: Constant>(&self, constant: BorrowedConstant<C>) -> Self::Constant;
     fn make_name(&self, name: &str) -> <Self::Constant as Constant>::Name;
@@ -245,6 +245,9 @@ pub enum Instruction {
     Rotate3,
     Duplicate,
     Duplicate2,
+    Copy {
+        nth: u32,
+    },
     GetIter,
     Continue {
         target: Label,
@@ -331,8 +334,10 @@ pub enum Instruction {
     SetupWith {
         end: Label,
     },
+    WithExceptStart,
     WithCleanupStart,
     WithCleanupFinish,
+    PushExcInfo,
     PopBlock,
     Raise {
         kind: RaiseKind,
@@ -973,6 +978,7 @@ impl Instruction {
             Rotate2 | Rotate3 => 0,
             Duplicate => 1,
             Duplicate2 => 2,
+            Copy { .. } => 1,
             GetIter => 0,
             Continue { .. } => 0,
             Break { .. } => 0,
@@ -1022,14 +1028,12 @@ impl Instruction {
                 }
             }
             SetupWith { .. } => {
-                if jump {
-                    0
-                } else {
-                    1
-                }
+                1
             }
+            WithExceptStart => 1,
             WithCleanupStart => 0,
             WithCleanupFinish => -1,
+            PushExcInfo => 1,
             PopBlock => 0,
             Raise { kind } => -(*kind as u8 as i32),
             BuildString { size }
@@ -1160,6 +1164,7 @@ impl Instruction {
             Rotate3 => w!(Rotate3),
             Duplicate => w!(Duplicate),
             Duplicate2 => w!(Duplicate2),
+            Copy { nth } => w!(Copy, nth),
             GetIter => w!(GetIter),
             Continue { target } => w!(Continue, target),
             Break { target } => w!(Break, target),
@@ -1187,8 +1192,10 @@ impl Instruction {
             EnterFinally => w!(EnterFinally),
             EndFinally => w!(EndFinally),
             SetupWith { end } => w!(SetupWith, end),
+            WithExceptStart => w!(WithExceptStart),
             WithCleanupStart => w!(WithCleanupStart),
             WithCleanupFinish => w!(WithCleanupFinish),
+            PushExcInfo => w!(PushExcInfo),
             BeforeAsyncWith => w!(BeforeAsyncWith),
             SetupAsyncWith { end } => w!(SetupAsyncWith, end),
             PopBlock => w!(PopBlock),

diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -1387,6 +1387,43 @@ impl Compiler {
         Ok(())
     }
 
+    fn compile_call_exit_with_nones(&mut self) {
+        for _ in 0..3 {
+            self.emit_constant(ConstantData::None);
+        }
+        self.emit(Instruction::CallFunctionPositional { nargs: 3 });
+    }
+
+    fn compile_pop_except_and_reraise(&mut self) {
+        // Stack contents
+        // [exc_info, last_i, exc]            COPY        3
+        // [exc_info, last_i, exc, exc_info]  POP_EXCEPT
+        // [exc_info, last_i, exc]            RERAISE      1
+        // (exception_unwind clears the stack)
+
+        self.emit(Instruction::Copy { nth: 3 });
+        self.emit(Instruction::PopException);
+        self.emit(Instruction::Raise {
+            kind: bytecode::RaiseKind::Reraise,
+        });
+    }
+
+    fn compile_with_except_finish(&mut self, cleanup: bytecode::Label) {
+        let exit = self.new_block();
+        self.emit(Instruction::JumpIfTrue { target: exit });
+        self.emit(Instruction::Raise {
+            kind: bytecode::RaiseKind::Reraise,
+        });
+        self.switch_to_block(cleanup);
+        self.compile_pop_except_and_reraise();
+        self.switch_to_block(exit);
+        self.emit(Instruction::Pop);
+        self.emit(Instruction::PopBlock);
+        self.emit(Instruction::PopException);
+        self.emit(Instruction::Pop);
+        self.emit(Instruction::Pop);
+    }
+
     fn compile_with(
         &mut self,
         items: &[ast::Withitem],
@@ -1401,6 +1438,7 @@ impl Compiler {
             return Err(self.error(CompileErrorType::EmptyWithItems));
         };
 
+        let block = self.new_block();
         let final_block = {
             let final_block = self.new_block();
             self.compile_expression(&item.context_expr)?;
@@ -1416,6 +1454,7 @@ impl Compiler {
                 self.emit(Instruction::SetupWith { end: final_block });
             }
 
+            self.switch_to_block(block);
             match &item.optional_vars {
                 Some(var) => {
                     self.set_source_location(var.location);
@@ -1443,18 +1482,29 @@ impl Compiler {
         self.set_source_location(with_location);
         self.emit(Instruction::PopBlock);
 
-        self.emit(Instruction::EnterFinally);
+        // End of body; start the cleanup.
+
+        let exit_block = self.new_block();
+
+        self.compile_call_exit_with_nones();
+        self.emit(Instruction::Pop);
+        self.emit(Instruction::Jump { target: exit_block });
 
         self.switch_to_block(final_block);
-        self.emit(Instruction::WithCleanupStart);
+
+        let cleanup_block = self.new_block();
+        // self.emit(Instruction::SetupCleanup { target: cleanup_block });
+        self.emit(Instruction::PushExcInfo);
+        self.emit(Instruction::WithExceptStart);
 
         if is_async {
             self.emit(Instruction::GetAwaitable);
             self.emit_constant(ConstantData::None);
             self.emit(Instruction::YieldFrom);
         }
 
-        self.emit(Instruction::WithCleanupFinish);
+        self.compile_with_except_finish(cleanup_block);
+        self.switch_to_block(exit_block);
 
         Ok(())
     }

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -4,8 +4,8 @@ use crate::{
         asyncgenerator::PyAsyncGenWrappedValue,
         function::{PyCell, PyCellRef, PyFunction},
         tuple::{PyTuple, PyTupleTyped},
-        PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator, PyList, PySet,
-        PySlice, PyStr, PyStrInterned, PyStrRef, PyTraceback, PyType,
+        PyBaseException, PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator,
+        PyInt, PyList, PySet, PySlice, PyStr, PyStrInterned, PyStrRef, PyTraceback, PyType,
     },
     bytecode,
     convert::{IntoObject, ToPyResult},
@@ -44,6 +44,9 @@ enum BlockType {
     Finally {
         handler: bytecode::Label,
     },
+    With {
+        handler: bytecode::Label,
+    },
 
     /// Active finally sequence
     FinallyHandler {
@@ -627,6 +630,11 @@ impl ExecutingFrame<'_> {
                 self.push_value(top);
                 Ok(None)
             }
+            bytecode::Instruction::Copy { nth } => {
+                let obj = self.nth_value(*nth).to_owned();
+                self.push_value(obj);
+                Ok(None)
+            }
             // splitting the instructions like this offloads the cost of "dynamic" dispatch (on the
             // amount to rotate) to the opcode dispatcher, and generates optimized code for the
             // concrete cases we actually have
@@ -773,6 +781,23 @@ impl ExecutingFrame<'_> {
                     );
                 }
             }
+            bytecode::Instruction::PushExcInfo => {
+                let value = self.pop_value();
+
+                let exc_val = vm
+                    .pop_exception()
+                    .map_or_else(|| vm.ctx.none(), |e| e.into());
+                self.push_value(exc_val);
+
+                self.push_value(value.clone());
+                // assert(PyExceptionInstance_Check(value));
+                vm.push_exception(Some(
+                    value
+                        .downcast::<PyBaseException>()
+                        .expect("value must be exception"),
+                ));
+                Ok(None)
+            }
             bytecode::Instruction::SetupWith { end } => {
                 let context_manager = self.pop_value();
                 let enter_res = vm.call_special_method(
@@ -782,7 +807,7 @@ impl ExecutingFrame<'_> {
                 )?;
                 let exit = context_manager.get_attr(identifier!(vm, __exit__), vm)?;
                 self.push_value(exit);
-                self.push_block(BlockType::Finally { handler: *end });
+                self.push_block(BlockType::With { handler: *end });
                 self.push_value(enter_res);
                 Ok(None)
             }
@@ -802,6 +827,30 @@ impl ExecutingFrame<'_> {
                 self.push_value(enter_res);
                 Ok(None)
             }
+            bytecode::Instruction::WithExceptStart => {
+                // At the top of the stack are 4 values:
+                // - TOP = exc_info()
+                // - SECOND = previous exception
+                // - THIRD: lasti of exception in exc_info()
+                // - FOURTH: the context.__exit__ bound method
+                // We call FOURTH(type(TOP), TOP, GetTraceback(TOP)).
+                // Then we push the __exit__ return value.
+
+                let val = self.top_value();
+                assert!(self.peek(3).payload_is::<PyInt>());
+                let exit_func = self.peek(4);
+
+                let (exc, val, tb) = vm.split_exception(
+                    val.to_owned()
+                        .downcast::<PyBaseException>()
+                        .expect("top value must be exception"),
+                );
+
+                let res = vm.invoke(exit_func, (exc, val, tb))?;
+                self.push_value(res);
+
+                Ok(None)
+            }
             bytecode::Instruction::WithCleanupStart => {
                 let block = self.current_block().unwrap();
                 let reason = match block.typ {
@@ -1161,6 +1210,24 @@ impl ExecutingFrame<'_> {
                         return Ok(None);
                     }
                 }
+                BlockType::With { handler } => {
+                    self.pop_block();
+                    if let UnwindReason::Raising { exception } = &reason {
+                        let exit_func = self.pop_value();
+
+                        let prev_exc = vm.current_exception();
+                        vm.set_exception(Some(exception.clone()));
+                        self.push_block(BlockType::FinallyHandler {
+                            reason: Some(reason),
+                            prev_exc,
+                        });
+                        vm.invoke(&exit_func, (vm.ctx.none(), vm.ctx.none(), vm.ctx.none()))?;
+
+                        self.jump(handler);
+
+                        return Ok(None);
+                    }
+                }
                 BlockType::FinallyHandler { prev_exc, .. }
                 | BlockType::ExceptHandler { prev_exc } => {
                     self.pop_block();
@@ -1810,6 +1877,14 @@ impl ExecutingFrame<'_> {
         }
     }
 
+    #[inline]
+    fn top_value(&self) -> &PyObjectRef {
+        match self.state.stack.last() {
+            Some(x) => x,
+            None => self.fatal("tried to get top value but there was nothing on the stack"),
+        }
+    }
+
     fn pop_multiple(&mut self, count: usize) -> crate::common::boxvec::Drain<PyObjectRef> {
         let stack_len = self.state.stack.len();
         self.state.stack.drain(stack_len - count..)
@@ -1828,12 +1903,20 @@ impl ExecutingFrame<'_> {
         }
     }
 
+    // deprecate this in favor of `peek`
     #[inline]
     fn nth_value(&self, depth: u32) -> &PyObject {
         let stack = &self.state.stack;
         &stack[stack.len() - depth as usize - 1]
     }
 
+    /// depth starts with 1
+    #[inline]
+    fn peek(&self, depth: u32) -> &PyObject {
+        let stack = &self.state.stack;
+        &stack[stack.len() - depth as usize]
+    }
+
     #[cold]
     #[inline(never)]
     fn fatal(&self, msg: &'static str) -> ! {

diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -227,6 +227,13 @@ impl VirtualMachine {
         stdlib::builtins::make_module(self, self.builtins.clone().into());
         stdlib::sys::init_module(self, self.sys_module.as_ref(), self.builtins.as_ref());
 
+        self.compile(
+            include_str!("../../../with0.py"),
+            rustpython_compiler_core::compile::Mode::Exec,
+            "_frozen_importlib".to_owned(),
+        )
+        .unwrap();
+
         let mut inner_init = || -> PyResult<()> {
             #[cfg(not(target_arch = "wasm32"))]
             import::import_builtin(self, "_signal")?;