keep flag for YieldNode instead of special handling in GeneratorBlockNode

Author: lukasstadler

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/NodeFactory.java

@@ -76,7 +76,6 @@
 import com.oracle.graal.python.nodes.function.FunctionRootNode;
 import com.oracle.graal.python.nodes.generator.DictConcatNode;
 import com.oracle.graal.python.nodes.generator.YieldNode;
-import com.oracle.graal.python.nodes.generator.YieldResumeNode;
 import com.oracle.graal.python.nodes.literal.BooleanLiteralNode;
 import com.oracle.graal.python.nodes.literal.BuiltinsLiteralNode;
 import com.oracle.graal.python.nodes.literal.BytesLiteralNode;
@@ -213,7 +212,7 @@ public StatementNode createBreakTarget(PNode forNode) {
     }
 
     public PNode createYield(PNode right, FrameSlot returnSlot) {
-        return createBlock(new YieldNode(createWriteLocal(right, returnSlot)), new YieldResumeNode());
+        return new YieldNode(createWriteLocal(right, returnSlot));
     }
 
     public PNode createIntegerLiteral(int value) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/control/BaseBlockNode.java

@@ -39,8 +39,6 @@ protected BaseBlockNode(PNode[] statements) {
         assert statements.length > 0;
     }
 
-    public abstract BaseBlockNode insertNodesBefore(PNode insertBefore, List<PNode> insertees);
-
     protected PNode[] insertStatementsBefore(PNode insertBefore, List<PNode> insertees) {
         int insertAt = -1;
         for (int i = 0; i < statements.length; i++) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/control/BlockNode.java

@@ -25,11 +25,8 @@
  */
 package com.oracle.graal.python.nodes.control;
 
-import java.util.List;
-
 import com.oracle.graal.python.nodes.EmptyNode;
 import com.oracle.graal.python.nodes.PNode;
-import com.oracle.graal.python.nodes.generator.YieldNode;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.nodes.UnexpectedResultException;
@@ -46,17 +43,12 @@ public static PNode create(PNode... statements) {
         if (length == 0) {
             return EmptyNode.create();
         } else if (length == 1) {
-            return statements[0] instanceof YieldNode ? new BlockNode(statements) : statements[0];
+            return statements[0];
         } else {
             return new BlockNode(statements);
         }
     }
 
-    @Override
-    public BaseBlockNode insertNodesBefore(PNode insertBefore, List<PNode> insertees) {
-        return new BlockNode(insertStatementsBefore(insertBefore, insertees));
-    }
-
     @ExplodeLoop
     private void executeFirst(VirtualFrame frame) {
         for (int i = 0; i < statements.length - 1; i++) {

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

@@ -31,35 +31,21 @@
 import com.oracle.graal.python.nodes.PNode;
 import com.oracle.graal.python.nodes.control.BaseBlockNode;
 import com.oracle.graal.python.runtime.exception.YieldException;
-import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.profiles.BranchProfile;
 import com.oracle.truffle.api.profiles.ConditionProfile;
 
 public final class GeneratorBlockNode extends BaseBlockNode implements GeneratorControlNode {
 
-    private static final boolean[] EMPTY = new boolean[0];
-
     @Child private GeneratorAccessNode gen = GeneratorAccessNode.create();
 
     private final ConditionProfile needsUpdateProfile = ConditionProfile.createBinaryProfile();
     private final BranchProfile seenYield = BranchProfile.create();
-    @CompilationFinal(dimensions = 1) private final boolean[] isYield;
     private final int indexSlot;
 
     public GeneratorBlockNode(PNode[] statements, int indexSlot) {
         super(statements);
-        boolean[] yields = EMPTY;
-        for (int i = 0; i < statements.length; i++) {
-            if (statements[i] instanceof YieldNode) {
-                if (yields.length == 0) {
-                    yields = new boolean[statements.length];
-                }
-                yields[i] = true;
-            }
-        }
-        this.isYield = yields;
         this.indexSlot = indexSlot;
     }
 
@@ -71,7 +57,6 @@ public int getIndexSlot() {
         return indexSlot;
     }
 
-    @Override
     public GeneratorBlockNode insertNodesBefore(PNode insertBefore, List<PNode> insertees) {
         return new GeneratorBlockNode(insertStatementsBefore(insertBefore, insertees), getIndexSlot());
     }
@@ -92,12 +77,7 @@ public Object execute(VirtualFrame frame) {
             return result;
         } catch (YieldException e) {
             seenYield.enter();
-            if (i < isYield.length && isYield[i]) {
-                // continue after the yield
-                nextIndex = i + 1;
-            } else {
-                nextIndex = i;
-            }
+            nextIndex = i;
             throw e;
         } finally {
             if (needsUpdateProfile.profile(nextIndex != startIndex)) {

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

@@ -25,14 +25,26 @@
  */
 package com.oracle.graal.python.nodes.generator;
 
+import com.oracle.graal.python.builtins.objects.PNone;
+import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.nodes.PNode;
 import com.oracle.graal.python.nodes.statement.StatementNode;
+import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.exception.YieldException;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.profiles.BranchProfile;
 
 public class YieldNode extends StatementNode implements GeneratorControlNode {
 
     @Child private PNode right;
+    @Child private GeneratorAccessNode access = GeneratorAccessNode.create();
+
+    @CompilationFinal private int flagSlot;
+
+    private final BranchProfile gotException = BranchProfile.create();
+    private final BranchProfile gotValue = BranchProfile.create();
+    private final BranchProfile gotNothing = BranchProfile.create();
 
     public YieldNode(PNode right) {
         this.right = right;
@@ -44,7 +56,28 @@ public PNode getRhs() {
 
     @Override
     public Object execute(VirtualFrame frame) {
-        right.execute(frame);
-        throw YieldException.INSTANCE;
+        if (access.isActive(frame, flagSlot)) {
+            // yield is active -> resume
+            access.setActive(frame, flagSlot, false);
+            Object specialArgument = PArguments.getSpecialArgument(frame);
+            if (specialArgument == null) {
+                gotNothing.enter();
+                return PNone.NONE;
+            } else if (specialArgument instanceof PException) {
+                gotException.enter();
+                throw (PException) specialArgument;
+            } else {
+                gotValue.enter();
+                return specialArgument;
+            }
+        } else {
+            right.executeVoid(frame);
+            access.setActive(frame, flagSlot, true);
+            throw YieldException.INSTANCE;
+        }
+    }
+
+    public void setFlagSlot(int slot) {
+        this.flagSlot = slot;
     }
 }

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

@@ -124,6 +124,7 @@ public RootCallTarget translate() {
 
     private void replaceYield(YieldNode yield) {
         PNode current = yield;
+        yield.setFlagSlot(nextActiveFlagSlot());
 
         if (yield.getParent() instanceof GeneratorReturnTargetNode) {
             // if this yield is the only thing in the body, we introduce a block