[GR-22536] Fix list comprehensions raising StopIteration when empty due to condition

PullRequest: graalpython/894

Author: msimacek

graalpython/com.oracle.graal.python.test/src/com/oracle/graal/python/test/generator/BuiltinIntrinsificationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017, 2018, Oracle and/or its affiliates.
+ * Copyright (c) 2017, 2020, Oracle and/or its affiliates.
  * Copyright (c) 2014, Regents of the University of California
  *
  * All rights reserved.
@@ -38,4 +38,18 @@ public void simpleListComp() {
                         "print(ll)";
         assertPrints("[0, 1, 2, 3, 4]\n", source);
     }
+
+    @Test
+    public void simpleListCompCondition() {
+        String source = "ll = list(i for i in range(5) if i % 2 == 0)\n" + //
+                        "print(ll)";
+        assertPrints("[0, 2, 4]\n", source);
+    }
+
+    @Test
+    public void emptyListComp() {
+        String source = "ll = list(i for i in range(5) if i == 6)\n" + //
+                        "print(ll)";
+        assertPrints("[]\n", source);
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/list/ListBuiltins.java

@@ -112,6 +112,7 @@
 import com.oracle.truffle.api.dsl.TypeSystemReference;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.nodes.UnexpectedResultException;
 import com.oracle.truffle.api.profiles.ConditionProfile;
 
 @CoreFunctions(extendClasses = PythonBuiltinClassType.PList)
@@ -217,36 +218,50 @@ PNone listRange(PList list, PRange range) {
             return PNone.NONE;
         }
 
-        @Specialization(guards = "iterable.isPRangeIterator()")
+        @Specialization(guards = "iterable.isPRangeIterator()", rewriteOn = UnexpectedResultException.class)
         PNone listPGenerator(VirtualFrame frame, PList list, PGenerator iterable,
                         @Cached GetIteratorNode getIteratorNode,
                         @Cached SequenceStorageNodes.AppendNode appendNode,
                         @Cached GetNextNode getNextNode,
                         @Cached IsBuiltinClassProfile errorProfile,
                         @Cached("createBinaryProfile()") ConditionProfile stepProfile,
-                        @Cached("createBinaryProfile()") ConditionProfile positiveRangeProfile) {
+                        @Cached("createBinaryProfile()") ConditionProfile positiveRangeProfile) throws UnexpectedResultException {
             clearStorage(list);
             Object iterObj = getIteratorNode.executeWith(frame, iterable);
             SequenceStorage storage = EmptySequenceStorage.INSTANCE;
 
             PRangeIterator range = (PRangeIterator) iterable.getIterator();
-            final int len = range.getLength(stepProfile, positiveRangeProfile);
-            if (len > 0) {
-                Object value = getNextNode.execute(frame, iterObj);
-                storage = SequenceStorageFactory.createStorage(value, len);
-                storage = appendNode.execute(storage, value, ListGeneralizationNode.SUPPLIER);
-                while (true) {
-                    try {
-                        storage = appendNode.execute(storage, getNextNode.execute(frame, iterObj), ListGeneralizationNode.SUPPLIER);
-                    } catch (PException e) {
-                        e.expectStopIteration(errorProfile);
-                        break;
+            final int estimatedMaxLen = range.getLength(stepProfile, positiveRangeProfile);
+            int realLen = 0;
+            if (estimatedMaxLen > 0) {
+                Object value = null;
+                try {
+                    value = getNextNode.execute(frame, iterObj);
+                    realLen++;
+                } catch (PException e) {
+                    e.expectStopIteration(errorProfile);
+                }
+                if (value != null) {
+                    storage = SequenceStorageFactory.createStorage(value, estimatedMaxLen);
+                    storage = appendNode.execute(storage, value, ListGeneralizationNode.SUPPLIER);
+                    while (true) {
+                        try {
+                            storage = appendNode.execute(storage, getNextNode.execute(frame, iterObj), ListGeneralizationNode.SUPPLIER);
+                            realLen++;
+                        } catch (PException e) {
+                            e.expectStopIteration(errorProfile);
+                            break;
+                        }
                     }
                 }
             }
 
             list.setSequenceStorage(storage);
-            return PNone.NONE;
+            if (realLen == estimatedMaxLen) {
+                return PNone.NONE;
+            } else {
+                throw new UnexpectedResultException(PNone.NONE);
+            }
         }
 
         @Specialization(guards = {"!isNoValue(iterable)", "!isString(iterable)"})