Implement execution of finally block. Fixes RustPython#1306.

windelbouwman · windelbouwman · commit 2ca6e5d399ed · 2019-08-27T21:16:59.000+02:00
diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
@@ -187,11 +187,15 @@ pub enum Instruction {
         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 {
@@ -493,6 +497,7 @@ impl Instruction {
             SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]),
             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),
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -712,13 +712,13 @@ impl<O: OutputStream> Compiler<O> {
         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_label,
+                handler: finally_handler_label,
             });
         }
 
@@ -773,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:
@@ -802,8 +804,15 @@ 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);
diff --git a/tests/snippets/try_exceptions.py b/tests/snippets/try_exceptions.py
@@ -187,6 +187,19 @@ def f():
             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:
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -46,6 +46,11 @@ enum BlockType {
     Finally {
         handler: bytecode::Label,
     },
+
+    /// Active finally sequence
+    FinallyHandler {
+        reason: Option<UnwindReason>,
+    },
     With {
         end: bytecode::Label,
         context_manager: PyObjectRef,
@@ -58,6 +63,7 @@ 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 },
@@ -407,10 +413,27 @@ impl Frame {
                 self.push_block(BlockType::Finally { handler: *handler });
                 Ok(None)
             }
-            bytecode::Instruction::EndFinally => {
-                let _block = self.pop_block().unwrap();
+            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:
@@ -423,7 +446,7 @@ 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,
@@ -438,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 => {
@@ -687,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)
@@ -781,7 +804,7 @@ impl Frame {
             match block.typ {
                 BlockType::Loop { start, end } => match &reason {
                     UnwindReason::Break => {
-                        self.pop_block().unwrap();
+                        self.pop_block();
                         self.jump(end);
                         return Ok(None);
                     }
@@ -790,33 +813,32 @@ impl Frame {
                         return Ok(None);
                     }
                     _ => {
-                        self.pop_block().unwrap();
+                        self.pop_block();
                     }
                 },
-                BlockType::Finally { .. } => {
-                    self.pop_block().unwrap();
-                    // TODO!
+                BlockType::Finally { handler } => {
+                    self.pop_block();
+                    self.push_block(BlockType::FinallyHandler {
+                        reason: Some(reason.clone()),
+                    });
+                    self.jump(handler);
+                    return Ok(None);
                 }
                 BlockType::TryExcept { handler } => {
-                    self.pop_block().unwrap();
-                    match &reason {
-                        UnwindReason::Raising { exception } => {
-                            self.push_block(BlockType::ExceptHandler {});
-                            self.push_value(exception.clone());
-                            vm.push_exception(exception.clone());
-                            self.jump(handler);
-                            return Ok(None);
-                        }
-                        _ => {
-                            // No worries!
-                        }
+                    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 {
                     context_manager,
                     end,
                 } => {
-                    self.pop_block().unwrap();
+                    self.pop_block();
                     match &reason {
                         UnwindReason::Raising { exception } => {
                             match self.call_context_manager_exit(
@@ -866,8 +888,11 @@ impl Frame {
                         }
                     }
                 }
+                BlockType::FinallyHandler { .. } => {
+                    self.pop_block();
+                }
                 BlockType::ExceptHandler => {
-                    self.pop_block().unwrap();
+                    self.pop_block();
                     vm.pop_exception().expect("Should have exception in stack");
                 }
             }
@@ -881,8 +906,6 @@ impl Frame {
                 panic!("Internal error: break or continue must occur within a loop block.")
             } // UnwindReason::NoWorries => Ok(None),
         }
-
-        // None
     }
 
     fn call_context_manager_exit(
@@ -1193,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> {
