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

Refactor block stack unwinding to enable finally execution.

Author: windelbouwman

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 @@ 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: