Exhaust primitive iterators in 'ForNode'.

Author: fangerer

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/iterator/IteratorBuiltins.java

@@ -85,10 +85,10 @@ public Object next(PArrayIterator self,
 
         @Specialization
         public int next(PIntegerSequenceIterator self) {
-            if (!self.stopIterationReached && self.index < self.sequence.length()) {
+            if (!self.isExhausted() && self.index < self.sequence.length()) {
                 return self.sequence.getIntItemNormalized(self.index++);
             }
-            self.stopIterationReached = true;
+            self.setExhausted();
             throw raise(StopIteration);
         }
 
@@ -114,19 +114,19 @@ public int next(PRangeReverseIterator self) {
 
         @Specialization
         public double next(PDoubleSequenceIterator self) {
-            if (!self.stopIterationReached && self.index < self.sequence.length()) {
+            if (!self.isExhausted() && self.index < self.sequence.length()) {
                 return self.sequence.getDoubleItemNormalized(self.index++);
             }
-            self.stopIterationReached = true;
+            self.setExhausted();
             throw raise(StopIteration);
         }
 
         @Specialization
         public long next(PLongSequenceIterator self) {
-            if (!self.stopIterationReached && self.index < self.sequence.length()) {
+            if (!self.isExhausted() && self.index < self.sequence.length()) {
                 return self.sequence.getLongItemNormalized(self.index++);
             }
-            self.stopIterationReached = true;
+            self.setExhausted();
             throw raise(StopIteration);
         }
 
@@ -157,10 +157,10 @@ public Object next(PSequenceIterator self,
                         @Cached("createClassProfile()") ValueProfile sequenceProfile,
                         @Cached("createGetItem()") SequenceStorageNodes.GetItemNode getItemNode) {
             PSequence sequence = sequenceProfile.profile(self.getPSequence());
-            if (!self.stopIterationReached && self.index < sequence.len()) {
+            if (!self.isExhausted() && self.index < sequence.len()) {
                 return getItemNode.execute(sequence.getSequenceStorage(), self.index++);
             }
-            self.stopIterationReached = true;
+            self.setExhausted();
             throw raise(StopIteration);
         }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/iterator/PBuiltinIterator.java

@@ -33,10 +33,20 @@
  */
 public abstract class PBuiltinIterator extends PythonBuiltinObject {
 
+    private boolean exhausted = false;
+
     public PBuiltinIterator(PythonClass clazz) {
         super(clazz);
     }
 
+    public void setExhausted() {
+        exhausted = true;
+    }
+
+    public boolean isExhausted() {
+        return exhausted;
+    }
+
     @Override
     public String toString() {
         return "<iterator object at " + hashCode() + ">";

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

@@ -95,6 +95,7 @@ protected boolean doIntegerIterator(VirtualFrame frame, PIntegerIterator iterato
                     @Cached("iterator.getClass()") Class<? extends PIntegerIterator> clazz) {
         PIntegerIterator profiledIterator = clazz.cast(iterator);
         if (!profiledIterator.hasNext()) {
+            profiledIterator.setExhausted();
             return false;
         }
         ((WriteNode) target).doWrite(frame, profiledIterator.next());
@@ -106,6 +107,7 @@ protected boolean doLongIterator(VirtualFrame frame, PLongIterator iterator,
                     @Cached("iterator.getClass()") Class<? extends PLongIterator> clazz) {
         PLongIterator profiledIterator = clazz.cast(iterator);
         if (!profiledIterator.hasNext()) {
+            profiledIterator.setExhausted();
             return false;
         }
         ((WriteNode) target).doWrite(frame, profiledIterator.next());
@@ -117,6 +119,7 @@ protected boolean doDoubleIterator(VirtualFrame frame, PDoubleIterator iterator,
                     @Cached("iterator.getClass()") Class<? extends PDoubleIterator> clazz) {
         PDoubleIterator profiledIterator = clazz.cast(iterator);
         if (!profiledIterator.hasNext()) {
+            profiledIterator.setExhausted();
             return false;
         }
         ((WriteNode) target).doWrite(frame, profiledIterator.next());