Support 'terminating return' opt in WithNode

Author: steve-s

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/generator/GeneratorWithNode.java

@@ -62,6 +62,12 @@ public GeneratorWithNode(WriteNode targetNode, StatementNode body, ExpressionNod
         this.yieldSlot = generatorInfo.nextActiveFlagIndex();
     }
 
+    @Override
+    public Object returnExecute(VirtualFrame frame) {
+        // We do not use returnExecute inside generators
+        throw new IllegalStateException();
+    }
+
     @Override
     protected Object getWithObject(VirtualFrame frame) {
         Object withObject = gen.getIterator(frame, withObjectSlot);
@@ -81,9 +87,11 @@ protected void doEnter(VirtualFrame frame, Object withObject, Object enterCallab
     }
 
     @Override
-    protected void doBody(VirtualFrame frame) {
+    protected Object doBody(VirtualFrame frame, boolean isReturn) {
+        assert !isReturn; // not supported in generators
         try {
-            super.doBody(frame);
+            super.doBody(frame, isReturn);
+            return null;
         } catch (YieldException e) {
             gen.setActive(frame, yieldSlot, true);
             throw e;

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

@@ -104,6 +104,15 @@ private PRaiseNode getRaiseNode() {
 
     @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) {
         Object withObject = getWithObject(frame);
         Object clazz = getClassNode.execute(withObject);
         Object enterCallable = enterSpecialGetter.execute(frame, clazz, withObject);
@@ -117,11 +126,12 @@ public void executeVoid(VirtualFrame frame) {
             throw getRaiseNode().raise(PythonBuiltinClassType.AttributeError, ErrorMessages.OBJ_P_HAS_NO_ATTR_S, withObject, __EXIT__);
         }
         doEnter(frame, withObject, enterCallable);
+        Object result = null;
         try {
-            doBody(frame);
+            result = doBody(frame, isReturn);
         } catch (PException exception) {
             handleException(frame, withObject, exitCallable, exception);
-            return;
+            return PNone.NONE;
         } catch (ControlFlowException e) {
             doLeave(frame, withObject, exitCallable);
             throw e;
@@ -131,9 +141,10 @@ public void executeVoid(VirtualFrame frame) {
                 throw e;
             }
             handleException(frame, withObject, exitCallable, pe);
-            return;
+            return PNone.NONE;
         }
         doLeave(frame, withObject, exitCallable);
+        return result;
     }
 
     /**
@@ -146,8 +157,8 @@ protected Object getWithObject(VirtualFrame frame) {
     /**
      * Execute the body
      */
-    protected void doBody(VirtualFrame frame) {
-        body.executeVoid(frame);
+    protected Object doBody(VirtualFrame frame, boolean isReturn) {
+        return body.genericExecute(frame, isReturn);
     }
 
     /**