Merge pull request #1321 from RustPython/try-return-finally

Refactor block stack unwinding to enable finally execution.

bytecode/src/bytecode.rs

@@ -178,6 +178,26 @@ pub enum Instruction {
         start: Label,
         end: Label,
     },
+
+    /// Setup a finally handler, which will be called whenever one of this events occurs:
+    /// - the block is popped
+    /// - the function returns
+    /// - an exception is returned
+    SetupFinally {
+        handler: Label,
+    },
+
+    /// Enter a finally block, without returning, excepting, just because we are there.
+    EnterFinally,
+
+    /// Marker bytecode for the end of a finally sequence.
+    /// When this bytecode is executed, the eval loop does one of those things:
+    /// - Continue at a certain bytecode position
+    /// - Propagate the exception
+    /// - Return from a function
+    /// - Do nothing at all, just continue
+    EndFinally,
+
     SetupExcept {
         handler: Label,
     },
@@ -475,7 +495,10 @@ impl Instruction {
             YieldValue => w!(YieldValue),
             YieldFrom => w!(YieldFrom),
             SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]),
-            SetupExcept { handler } => w!(SetupExcept, handler),
+            SetupExcept { handler } => w!(SetupExcept, label_map[handler]),
+            SetupFinally { handler } => w!(SetupFinally, label_map[handler]),
+            EnterFinally => w!(EnterFinally),
+            EndFinally => w!(EndFinally),
             SetupWith { end } => w!(SetupWith, end),
             CleanupWith { end } => w!(CleanupWith, end),
             PopBlock => w!(PopBlock),

compiler/src/compile.rs

@@ -708,12 +708,20 @@ impl<O: OutputStream> Compiler<O> {
         &mut self,
         body: &[ast::Statement],
         handlers: &[ast::ExceptHandler],
-        orelse: &Option<Vec<ast::Statement>>,
-        finalbody: &Option<Vec<ast::Statement>>,
+        orelse: &Option<ast::Suite>,
+        finalbody: &Option<ast::Suite>,
     ) -> Result<(), CompileError> {
         let mut handler_label = self.new_label();
-        let finally_label = self.new_label();
+        let finally_handler_label = self.new_label();
         let else_label = self.new_label();
+
+        // Setup a finally block if we have a finally statement.
+        if finalbody.is_some() {
+            self.emit(Instruction::SetupFinally {
+                handler: finally_handler_label,
+            });
+        }
+
         // try:
         self.emit(Instruction::SetupExcept {
             handler: handler_label,
@@ -765,26 +773,28 @@ impl<O: OutputStream> Compiler<O> {
             // Handler code:
             self.compile_statements(&handler.body)?;
             self.emit(Instruction::PopException);
+
+            if finalbody.is_some() {
+                self.emit(Instruction::PopBlock); // pop finally block
+                                                  // We enter the finally block, without exception.
+                self.emit(Instruction::EnterFinally);
+            }
+
             self.emit(Instruction::Jump {
-                target: finally_label,
+                target: finally_handler_label,
             });
 
             // Emit a new label for the next handler
             self.set_label(handler_label);
             handler_label = self.new_label();
         }
+
         self.emit(Instruction::Jump {
             target: handler_label,
         });
         self.set_label(handler_label);
         // If code flows here, we have an unhandled exception,
-        // emit finally code and raise again!
-        // Duplicate finally code here:
-        // TODO: this bytecode is now duplicate, could this be
-        // improved?
-        if let Some(statements) = finalbody {
-            self.compile_statements(statements)?;
-        }
+        // raise the exception again!
         self.emit(Instruction::Raise { argc: 0 });
 
         // We successfully ran the try block:
@@ -794,12 +804,20 @@ impl<O: OutputStream> Compiler<O> {
             self.compile_statements(statements)?;
         }
 
+        if finalbody.is_some() {
+            self.emit(Instruction::PopBlock); // pop finally block
+
+            // We enter the finally block, without return / exception.
+            self.emit(Instruction::EnterFinally);
+        }
+
         // finally:
-        self.set_label(finally_label);
+        self.set_label(finally_handler_label);
         if let Some(statements) = finalbody {
             self.compile_statements(statements)?;
+            self.emit(Instruction::EndFinally);
         }
-        // unimplemented!();
+
         Ok(())
     }
 

tests/snippets/try_exceptions.py

@@ -187,6 +187,19 @@ with assertRaises(ZeroDivisionError):
             pass
         raise
 
+# try-return-finally behavior:
+l = []
+def foo():
+    try:
+        return 33
+    finally:
+        l.append(1337)
+
+r = foo()
+assert r == 33
+assert l == [1337]
+
+
 # Regression https://github.com/RustPython/RustPython/issues/867
 for _ in [1, 2]:
     try:

vm/src/frame.rs

@@ -43,6 +43,14 @@ enum BlockType {
     TryExcept {
         handler: bytecode::Label,
     },
+    Finally {
+        handler: bytecode::Label,
+    },
+
+    /// Active finally sequence
+    FinallyHandler {
+        reason: Option<UnwindReason>,
+    },
     With {
         end: bytecode::Label,
         context_manager: PyObjectRef,
@@ -52,6 +60,25 @@ enum BlockType {
 
 pub type FrameRef = PyRef<Frame>;
 
+/// The reason why we might be unwinding a block.
+/// This could be return of function, exception being
+/// raised, a break or continue being hit, etc..
+#[derive(Clone, Debug)]
+enum UnwindReason {
+    /// We are returning a value from a return statement.
+    Returning { value: PyObjectRef },
+
+    /// We hit an exception, so unwind any try-except and finally blocks.
+    Raising { exception: PyObjectRef },
+
+    // NoWorries,
+    /// We are unwinding blocks, since we hit break
+    Break,
+
+    /// We are unwinding blocks since we hit a continue statements.
+    Continue,
+}
+
 pub struct Frame {
     pub code: bytecode::CodeObject,
     // We need 1 stack per frame
@@ -137,9 +164,12 @@ impl Frame {
                         vm.ctx.new_str(run_obj_name.clone()),
                     ]);
                     objlist::PyListRef::try_from_object(vm, traceback)?.append(raise_location, vm);
-                    match self.unwind_exception(vm, exception) {
-                        None => {}
-                        Some(exception) => {
+                    match self.unwind_blocks(vm, UnwindReason::Raising { exception }) {
+                        Ok(None) => {}
+                        Ok(Some(result)) => {
+                            break Ok(result);
+                        }
+                        Err(exception) => {
                             // TODO: append line number to traceback?
                             // traceback.append();
                             break Err(exception);
@@ -151,9 +181,10 @@ impl Frame {
     }
 
     pub fn throw(&self, vm: &VirtualMachine, exception: PyObjectRef) -> PyResult<ExecutionResult> {
-        match self.unwind_exception(vm, exception) {
-            None => self.run(vm),
-            Some(exception) => Err(exception),
+        match self.unwind_blocks(vm, UnwindReason::Raising { exception }) {
+            Ok(None) => self.run(vm),
+            Ok(Some(result)) => Ok(result),
+            Err(exception) => Err(exception),
         }
     }
 
@@ -345,11 +376,7 @@ impl Frame {
             bytecode::Instruction::CompareOperation { ref op } => self.execute_compare(vm, op),
             bytecode::Instruction::ReturnValue => {
                 let value = self.pop_value();
-                if let Some(exc) = self.unwind_blocks(vm) {
-                    Err(exc)
-                } else {
-                    Ok(Some(ExecutionResult::Return(value)))
-                }
+                self.unwind_blocks(vm, UnwindReason::Returning { value })
             }
             bytecode::Instruction::YieldValue => {
                 let value = self.pop_value();
@@ -382,6 +409,31 @@ impl Frame {
                 self.push_block(BlockType::TryExcept { handler: *handler });
                 Ok(None)
             }
+            bytecode::Instruction::SetupFinally { handler } => {
+                self.push_block(BlockType::Finally { handler: *handler });
+                Ok(None)
+            }
+            bytecode::Instruction::EnterFinally => {
+                self.push_block(BlockType::FinallyHandler { reason: None });
+                Ok(None)
+            }
+            bytecode::Instruction::EndFinally => {
+                // Pop the finally handler from the stack, and recall
+                // what was the reason we were in this finally clause.
+                let block = self.pop_block();
+
+                if let BlockType::FinallyHandler { reason } = block.typ {
+                    if let Some(reason) = reason {
+                        self.unwind_blocks(vm, reason)
+                    } else {
+                        Ok(None)
+                    }
+                } else {
+                    panic!(
+                        "Block type must be finally handler when reaching EndFinally instruction!"
+                    );
+                }
+            }
             bytecode::Instruction::SetupWith { end } => {
                 let context_manager = self.pop_value();
                 // Call enter:
@@ -394,14 +446,14 @@ impl Frame {
                 Ok(None)
             }
             bytecode::Instruction::CleanupWith { end: end1 } => {
-                let block = self.pop_block().unwrap();
+                let block = self.pop_block();
                 if let BlockType::With {
                     end: end2,
                     context_manager,
                 } = &block.typ
                 {
                     debug_assert!(end1 == end2);
-                    self.call_context_manager_exit_no_exception(vm, &context_manager)?;
+                    self.call_context_manager_exit(vm, &context_manager, None)?;
                 } else {
                     unreachable!("Block stack is incorrect, expected a with block");
                 }
@@ -409,7 +461,7 @@ impl Frame {
                 Ok(None)
             }
             bytecode::Instruction::PopBlock => {
-                self.pop_block().expect("no pop to block");
+                self.pop_block();
                 Ok(None)
             }
             bytecode::Instruction::GetIter => {
@@ -567,29 +619,12 @@ impl Frame {
                 Err(exception)
             }
 
-            bytecode::Instruction::Break => {
-                let block = self.unwind_loop(vm);
-                if let BlockType::Loop { end, .. } = block.typ {
-                    self.pop_block();
-                    self.jump(end);
-                } else {
-                    unreachable!()
-                }
-                Ok(None)
-            }
+            bytecode::Instruction::Break => self.unwind_blocks(vm, UnwindReason::Break),
             bytecode::Instruction::Pass => {
                 // Ah, this is nice, just relax!
                 Ok(None)
             }
-            bytecode::Instruction::Continue => {
-                let block = self.unwind_loop(vm);
-                if let BlockType::Loop { start, .. } = block.typ {
-                    self.jump(start);
-                } else {
-                    unreachable!();
-                }
-                Ok(None)
-            }
+            bytecode::Instruction::Continue => self.unwind_blocks(vm, UnwindReason::Continue),
             bytecode::Instruction::PrintExpr => {
                 let expr = self.pop_value();
                 if !expr.is(&vm.get_none()) {
@@ -675,7 +710,7 @@ impl Frame {
                 Ok(None)
             }
             bytecode::Instruction::PopException {} => {
-                let block = self.pop_block().unwrap(); // this asserts that the block is_some.
+                let block = self.pop_block();
                 if let BlockType::ExceptHandler = block.typ {
                     vm.pop_exception().expect("Should have exception in stack");
                     Ok(None)
@@ -758,134 +793,138 @@ impl Frame {
         Ok(None)
     }
 
-    // Unwind all blocks:
+    /// Unwind blocks.
+    /// The reason for unwinding gives a hint on what to do when
+    /// unwinding a block.
+    /// Optionally returns an exception.
     #[cfg_attr(feature = "flame-it", flame("Frame"))]
-    fn unwind_blocks(&self, vm: &VirtualMachine) -> Option<PyObjectRef> {
-        while let Some(block) = self.pop_block() {
+    fn unwind_blocks(&self, vm: &VirtualMachine, mut reason: UnwindReason) -> FrameResult {
+        // First unwind all existing blocks on the block stack:
+        while let Some(block) = self.current_block() {
             match block.typ {
-                BlockType::Loop { .. } => {}
-                BlockType::TryExcept { .. } => {
-                    // TODO: execute finally handler
-                }
-                BlockType::With {
-                    context_manager, ..
-                } => {
-                    match self.call_context_manager_exit_no_exception(vm, &context_manager) {
-                        Ok(..) => {}
-                        Err(exc) => {
-                            // __exit__ went wrong,
-                            return Some(exc);
-                        }
+                BlockType::Loop { start, end } => match &reason {
+                    UnwindReason::Break => {
+                        self.pop_block();
+                        self.jump(end);
+                        return Ok(None);
                     }
-                }
-                BlockType::ExceptHandler => {
-                    vm.pop_exception().expect("Should have exception in stack");
-                }
-            }
-        }
-
-        None
-    }
-
-    #[cfg_attr(feature = "flame-it", flame("Frame"))]
-    fn unwind_loop(&self, vm: &VirtualMachine) -> Block {
-        loop {
-            let block = self.current_block().expect("not in a loop");
-            match block.typ {
-                BlockType::Loop { .. } => break block,
-                BlockType::TryExcept { .. } => {
-                    // TODO: execute finally handler
-                }
-                BlockType::With {
-                    context_manager, ..
-                } => match self.call_context_manager_exit_no_exception(vm, &context_manager) {
-                    Ok(..) => {}
-                    Err(exc) => {
-                        panic!("Exception in with __exit__ {:?}", exc);
+                    UnwindReason::Continue => {
+                        self.jump(start);
+                        return Ok(None);
+                    }
+                    _ => {
+                        self.pop_block();
                     }
                 },
-                BlockType::ExceptHandler => {
-                    vm.pop_exception().expect("Should have exception in stack");
-                }
-            }
-
-            self.pop_block();
-        }
-    }
-
-    #[cfg_attr(feature = "flame-it", flame("Frame"))]
-    fn unwind_exception(&self, vm: &VirtualMachine, exc: PyObjectRef) -> Option<PyObjectRef> {
-        // unwind block stack on exception and find any handlers:
-        while let Some(block) = self.pop_block() {
-            match block.typ {
-                BlockType::TryExcept { handler } => {
-                    self.push_block(BlockType::ExceptHandler {});
-                    self.push_value(exc.clone());
-                    vm.push_exception(exc);
+                BlockType::Finally { handler } => {
+                    self.pop_block();
+                    self.push_block(BlockType::FinallyHandler {
+                        reason: Some(reason.clone()),
+                    });
                     self.jump(handler);
-                    return None;
+                    return Ok(None);
+                }
+                BlockType::TryExcept { handler } => {
+                    self.pop_block();
+                    if let UnwindReason::Raising { exception } = &reason {
+                        self.push_block(BlockType::ExceptHandler {});
+                        self.push_value(exception.clone());
+                        vm.push_exception(exception.clone());
+                        self.jump(handler);
+                        return Ok(None);
+                    }
                 }
                 BlockType::With {
-                    end,
                     context_manager,
+                    end,
                 } => {
-                    match self.call_context_manager_exit(vm, &context_manager, exc.clone()) {
-                        Ok(exit_result_obj) => {
-                            match objbool::boolval(vm, exit_result_obj) {
-                                // If __exit__ method returned True, suppress the exception and continue execution.
-                                Ok(suppress_exception) => {
-                                    if suppress_exception {
-                                        self.jump(end);
-                                        return None;
-                                    } else {
-                                        // go on with the stack unwinding.
+                    self.pop_block();
+                    match &reason {
+                        UnwindReason::Raising { exception } => {
+                            match self.call_context_manager_exit(
+                                vm,
+                                &context_manager,
+                                Some(exception.clone()),
+                            ) {
+                                Ok(exit_result_obj) => {
+                                    match objbool::boolval(vm, exit_result_obj) {
+                                        // If __exit__ method returned True, suppress the exception and continue execution.
+                                        Ok(suppress_exception) => {
+                                            if suppress_exception {
+                                                self.jump(end);
+                                                return Ok(None);
+                                            } else {
+                                                // go on with the stack unwinding.
+                                            }
+                                        }
+                                        Err(exit_exc) => {
+                                            reason = UnwindReason::Raising {
+                                                exception: exit_exc,
+                                            };
+                                            // return Err(exit_exc);
+                                        }
                                     }
                                 }
                                 Err(exit_exc) => {
-                                    return Some(exit_exc);
+                                    // TODO: what about original exception?
+                                    reason = UnwindReason::Raising {
+                                        exception: exit_exc,
+                                    };
+                                    // return Err(exit_exc);
                                 }
                             }
                         }
-                        Err(exit_exc) => {
-                            // TODO: what about original exception?
-                            return Some(exit_exc);
+                        _ => {
+                            match self.call_context_manager_exit(vm, &context_manager, None) {
+                                Ok(..) => {}
+                                Err(exit_exc) => {
+                                    // __exit__ went wrong,
+                                    reason = UnwindReason::Raising {
+                                        exception: exit_exc,
+                                    };
+                                    // return Err(exc);
+                                }
+                            }
                         }
                     }
                 }
-                BlockType::Loop { .. } => {}
+                BlockType::FinallyHandler { .. } => {
+                    self.pop_block();
+                }
                 BlockType::ExceptHandler => {
+                    self.pop_block();
                     vm.pop_exception().expect("Should have exception in stack");
                 }
             }
         }
-        Some(exc)
-    }
 
-    fn call_context_manager_exit_no_exception(
-        &self,
-        vm: &VirtualMachine,
-        context_manager: &PyObjectRef,
-    ) -> PyResult {
-        // TODO: do we want to put the exit call on the stack?
-        // TODO: what happens when we got an error during execution of __exit__?
-        vm.call_method(
-            context_manager,
-            "__exit__",
-            vec![vm.ctx.none(), vm.ctx.none(), vm.ctx.none()],
-        )
+        // We do not have any more blocks to unwind. Inspect the reason we are here:
+        match reason {
+            UnwindReason::Raising { exception } => Err(exception),
+            UnwindReason::Returning { value } => Ok(Some(ExecutionResult::Return(value))),
+            UnwindReason::Break | UnwindReason::Continue => {
+                panic!("Internal error: break or continue must occur within a loop block.")
+            } // UnwindReason::NoWorries => Ok(None),
+        }
     }
 
     fn call_context_manager_exit(
         &self,
         vm: &VirtualMachine,
         context_manager: &PyObjectRef,
-        exc: PyObjectRef,
+        exc: Option<PyObjectRef>,
     ) -> PyResult {
         // TODO: do we want to put the exit call on the stack?
         // TODO: what happens when we got an error during execution of __exit__?
-        let exc_type = exc.class().into_object();
-        let exc_val = exc.clone();
-        let exc_tb = vm.ctx.none(); // TODO: retrieve traceback?
+        let (exc_type, exc_val, exc_tb) = if let Some(exc) = exc {
+            let exc_type = exc.class().into_object();
+            let exc_val = exc.clone();
+            let exc_tb = vm.ctx.none(); // TODO: retrieve traceback?
+            (exc_type, exc_val, exc_tb)
+        } else {
+            (vm.ctx.none(), vm.ctx.none(), vm.ctx.none())
+        };
+
         vm.call_method(context_manager, "__exit__", vec![exc_type, exc_val, exc_tb])
     }
 
@@ -964,7 +1003,7 @@ impl Frame {
         let target_pc = self.code.label_map[&label];
         #[cfg(feature = "vm-tracing-logging")]
         trace!("jump from {:?} to {:?}", self.lasti, target_pc);
-        *self.lasti.borrow_mut() = target_pc;
+        self.lasti.replace(target_pc);
     }
 
     fn execute_make_function(
@@ -1177,10 +1216,14 @@ impl Frame {
         });
     }
 
-    fn pop_block(&self) -> Option<Block> {
-        let block = self.blocks.borrow_mut().pop()?;
+    fn pop_block(&self) -> Block {
+        let block = self
+            .blocks
+            .borrow_mut()
+            .pop()
+            .expect("No more blocks to pop!");
         self.stack.borrow_mut().truncate(block.level);
-        Some(block)
+        block
     }
 
     fn current_block(&self) -> Option<Block> {