further simplifications in generator nodes

Author: lukasstadler

graalpython/com.oracle.graal.python.test/src/tests/test_generator-if-and-loop.py

@@ -40,3 +40,22 @@ def foo(arg):
 
     assert foo([1,2,3,4]) == [1,2,3,4]
     assert foo("hello") == ["h", "e", "l", "l", "o"]
+
+def gen_effect(n):
+    if effect(n) == 5:
+        yield 1
+        yield 2
+    else:
+        yield 3
+        yield 4
+
+effect_count = 0
+def effect(n):
+    global effect_count
+    effect_count += 1
+    return n
+
+def test_sideeffect():
+    l = list(gen_effect(5))
+    assert l == [1,2]
+    assert effect_count == 1

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

@@ -75,7 +75,6 @@
 import com.oracle.graal.python.nodes.function.ClassBodyRootNode;
 import com.oracle.graal.python.nodes.function.FunctionRootNode;
 import com.oracle.graal.python.nodes.generator.DictConcatNode;
-import com.oracle.graal.python.nodes.generator.ListAppendNode;
 import com.oracle.graal.python.nodes.generator.YieldNode;
 import com.oracle.graal.python.nodes.generator.YieldResumeNode;
 import com.oracle.graal.python.nodes.literal.BooleanLiteralNode;
@@ -281,11 +280,6 @@ public PNode createSetLiteral(Set<PNode> values) {
         return new SetLiteralNode(convertedValues);
     }
 
-    public PNode createListAppend(FrameSlot frameSlot, PNode right) {
-        PNode readList = createReadLocal(frameSlot);
-        return ListAppendNode.create(readList, right);
-    }
-
     public PNode createUnaryOperation(String string, PNode operand) {
         switch (string) {
             case "+":
@@ -523,10 +517,6 @@ public PNode createDictionaryConcat(PNode... dictNodes) {
         return DictConcatNode.create(dictNodes);
     }
 
-    public PNode createListAppend(PNode leftNode, PNode rightNode) {
-        return ListAppendNode.create(leftNode, rightNode);
-    }
-
     public PNode callBuiltin(String string, PNode argument) {
         return PythonCallNode.create(getBuiltin(string), new PNode[]{argument}, new PNode[0], EmptyNode.create(), EmptyNode.create());
     }

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

@@ -32,7 +32,7 @@ public abstract class LoopNode extends StatementNode {
 
     public abstract PNode getBody();
 
-    protected final void reportLoopCount(@SuppressWarnings("unused") int value) {
-        // ignore for the time being
+    protected final void reportLoopCount(int value) {
+        com.oracle.truffle.api.nodes.LoopNode.reportLoopCount(this, value);
     }
 }

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

@@ -27,18 +27,39 @@
 
 import java.util.List;
 
+import com.oracle.graal.python.builtins.objects.PNone;
 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;
     }
 
@@ -58,24 +79,30 @@ public GeneratorBlockNode insertNodesBefore(PNode insertBefore, List<PNode> inse
     @ExplodeLoop
     @Override
     public Object execute(VirtualFrame frame) {
-        Object result = null;
-
-        for (int i = 0; i < statements.length; i++) {
-            final int currentIndex = gen.getIndex(frame, indexSlot);
-
-            if (i < currentIndex) {
-                continue;
+        int startIndex = gen.getIndex(frame, indexSlot);
+        int i = 0;
+        int nextIndex = 0;
+        try {
+            Object result = PNone.NONE;
+            for (i = 0; i < statements.length; i++) {
+                if (i >= startIndex) {
+                    result = statements[i].execute(frame);
+                }
+            }
+            return result;
+        } catch (YieldException e) {
+            seenYield.enter();
+            if (i < isYield.length && isYield[i]) {
+                // continue after the yield
+                nextIndex = i + 1;
+            } else {
+                nextIndex = i;
+            }
+            throw e;
+        } finally {
+            if (needsUpdateProfile.profile(nextIndex != startIndex)) {
+                gen.setIndex(frame, indexSlot, nextIndex);
             }
-
-            result = statements[i].execute(frame);
-            gen.setIndex(frame, indexSlot, currentIndex + 1);
         }
-
-        reset(frame);
-        return result;
-    }
-
-    public void reset(VirtualFrame frame) {
-        gen.setIndex(frame, indexSlot, 0);
     }
 }

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

@@ -25,9 +25,5 @@
  */
 package com.oracle.graal.python.nodes.generator;
 
-import com.oracle.truffle.api.frame.VirtualFrame;
-
 public interface GeneratorControlNode {
-
-    void reset(VirtualFrame frame);
 }