Support 'terminating return' opt in TryExceptNode

Author: steve-s

graalpython/com.oracle.graal.python.test/src/tests/test_return.py

@@ -51,3 +51,58 @@ def method(v):
 
     assert method(10) is None
     assert method("value") is None
+
+def test_return_in_finally():
+    def override_explicit():
+        try:
+            return 0
+        finally:
+            return 42
+
+    def override_implicit():
+        try:
+            pass
+        finally:
+            return 'correct'
+
+    assert override_explicit() == 42
+    assert override_implicit() == 'correct'
+
+def test_return_in_try():
+    def basic(arg):
+        try:
+            if arg:
+                raise ValueError()
+            return 1 # can be 'terminating'
+        except ValueError as e:
+            return 2 # can be 'terminating'
+
+    assert basic(True) == 2
+    assert basic(False) == 1
+
+    def with_orelse(arg):
+        try:
+            if arg:
+                return 'try'
+        except:
+            pass
+        else:
+            return 'else'
+
+    assert with_orelse(True) == 'try'
+    assert with_orelse(False) == 'else'
+
+
+def test_return_in_try_finally():
+    def foo(arg):
+        try:
+            if arg:
+                raise ValueError()
+            return 'try'
+        except ValueError as e:
+            return 'except'
+        finally:
+            return 'finally'
+
+    assert foo(True) == 'finally'
+    assert foo(False) == 'finally'

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/statement/StatementNode.java

@@ -53,6 +53,16 @@ public Object returnExecute(VirtualFrame frame) {
         return PNone.NONE;
     }
 
+    // Helper method
+    protected final Object genericExecute(VirtualFrame frame, boolean isReturnExecute) {
+        if (isReturnExecute) {
+            return returnExecute(frame);
+        } else {
+            executeVoid(frame);
+            return null;
+        }
+    }
+
     public void markAsTryBlock() {
         isTryBlock = true;
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/statement/TryExceptNode.java

@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 
 import com.oracle.graal.python.PythonLanguage;
+import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.exception.PBaseException;
 import com.oracle.graal.python.builtins.objects.object.PythonObjectLibrary;
 import com.oracle.graal.python.nodes.PNode;
@@ -86,24 +87,39 @@ public TryExceptNode(StatementNode body, ExceptNode[] exceptNodes, StatementNode
 
     @Override
     public void executeVoid(VirtualFrame frame) {
+        executeImpl(frame, false);
+    }
+
+    @Override
+    public Object returnExecute(VirtualFrame frame) {
+        return executeImpl(frame, true);
+    }
+
+    private Object executeImpl(VirtualFrame frame, boolean isReturn) {
         // The following statement is a no-op, but it helps graal to optimize the exception handler
         // by moving the cast to PException to the beginning
         saveExceptionState(frame);
 
         try {
-            body.executeVoid(frame);
+            if (isReturn && orelse == null) {
+                return body.returnExecute(frame);
+            } else {
+                body.executeVoid(frame);
+            }
         } catch (PException ex) {
             if (!catchPException(frame, ex)) {
                 throw ex;
             }
-            return;
+            // To keep it simple do not run the exception handlers with "terminating return" opt
+            // If we reach here, no explicit return could have happened as it would throw
+            return PNone.NONE;
         } catch (AbstractTruffleException e) {
             if (excLib == null) {
                 CompilerDirectives.transferToInterpreterAndInvalidate();
                 excLib = insert(InteropLibrary.getFactory().createDispatched(3));
             }
             if (excLib.isException(e) && catchTruffleException(frame, e)) {
-                return;
+                return PNone.NONE;
             }
             throw e;
         } catch (ControlFlowException e) {
@@ -113,14 +129,14 @@ public void executeVoid(VirtualFrame frame) {
             if (pe != null) {
                 boolean handled = catchPException(frame, pe);
                 if (handled) {
-                    return;
+                    return PNone.NONE;
                 } else {
                     throw pe.getExceptionForReraise();
                 }
             }
             throw e;
         }
-        orelse.executeVoid(frame);
+        return orelse.genericExecute(frame, isReturn);
     }
 
     @ExplodeLoop(kind = LoopExplosionKind.FULL_EXPLODE_UNTIL_RETURN)