CreateStorageFromIteratorNode: improve fast paths for builtin iterators

Author: steve-s

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/SequenceStorageNodes.java

@@ -40,15 +40,22 @@
  */
 package com.oracle.graal.python.builtins.objects.iterator;
 
+import static com.oracle.graal.python.builtins.PythonBuiltinClassType.PIterator;
 import static com.oracle.graal.python.builtins.PythonBuiltinClassType.TypeError;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__LENGTH_HINT__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__LEN__;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.PNotImplemented;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodes.GetSequenceStorageNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
+import com.oracle.graal.python.builtins.objects.iterator.IteratorNodesFactory.GetInternalIteratorSequenceStorageNodeGen;
+import com.oracle.graal.python.builtins.objects.list.PList;
 import com.oracle.graal.python.builtins.objects.str.PString;
+import com.oracle.graal.python.builtins.objects.tuple.PTuple;
 import com.oracle.graal.python.builtins.objects.type.SpecialMethodSlot;
 import com.oracle.graal.python.lib.PyIndexCheckNode;
 import com.oracle.graal.python.lib.PyNumberAsSizeNode;
@@ -74,6 +81,7 @@
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.GenerateNodeFactory;
 import com.oracle.truffle.api.dsl.GenerateUncached;
 import com.oracle.truffle.api.dsl.ImportStatic;
@@ -85,8 +93,6 @@
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.profiles.ConditionProfile;
 import com.oracle.truffle.api.profiles.LoopConditionProfile;
-import java.util.ArrayList;
-import java.util.List;
 
 public abstract class IteratorNodes {
 
@@ -181,6 +187,102 @@ int length(VirtualFrame frame, Object iterable,
         }
     }
 
+    @ImportStatic(PGuards.class)
+    public abstract static class GetInternalIteratorSequenceStorage extends Node {
+        public static GetInternalIteratorSequenceStorage getUncached() {
+            return GetInternalIteratorSequenceStorageNodeGen.getUncached();
+        }
+
+        /**
+         * The argument must be a builtin iterator, which points to the first element of the
+         * internal sequence storage. Returns {@code null} if the sequence storage is not available
+         * or if the iterator is not pointing to the first item in the storage.
+         */
+        public final SequenceStorage execute(PBuiltinIterator iterator) {
+            assert GetClassNode.getUncached().execute(iterator) == PIterator;
+            assert iterator.index == 0;
+            return executeInternal(iterator);
+        }
+
+        protected abstract SequenceStorage executeInternal(PBuiltinIterator iterator);
+
+        @Specialization(guards = "isList(it.sequence)")
+        static SequenceStorage doSequenceList(PSequenceIterator it) {
+            return ((PList) it.sequence).getSequenceStorage();
+        }
+
+        @Specialization
+        static SequenceStorage doSequenceLong(PLongSequenceIterator it) {
+            return it.sequence;
+        }
+
+        @Specialization
+        static SequenceStorage doSequenceDouble(PDoubleSequenceIterator it) {
+            return it.sequence;
+        }
+
+        @Specialization
+        static SequenceStorage doSequenceObj(PObjectSequenceIterator it) {
+            return it.sequence;
+        }
+
+        @Specialization
+        static SequenceStorage doSequenceIntSeq(PIntegerSequenceIterator it) {
+            return it.sequence;
+        }
+
+        @Specialization(guards = "isPTuple(it.sequence)")
+        static SequenceStorage doSequenceTuple(PSequenceIterator it) {
+            return ((PTuple) it.sequence).getSequenceStorage();
+        }
+
+        @Fallback
+        static SequenceStorage doOthers(PBuiltinIterator it) {
+            return null;
+        }
+    }
+
+    @ImportStatic(PGuards.class)
+    public abstract static class BuiltinIteratorLengthHint extends Node {
+        @Child GetInternalIteratorSequenceStorage getSeqStorage = GetInternalIteratorSequenceStorageNodeGen.create();
+        private final ConditionProfile noStorageProfile = ConditionProfile.createBinaryProfile();
+
+        /**
+         * The argument must be a builtin iterator. Returns {@code -1} if the length hint is not
+         * available.
+         */
+        public final int execute(PBuiltinIterator iterator) {
+            assert GetClassNode.getUncached().execute(iterator) == PIterator;
+            SequenceStorage result = getSeqStorage.execute(iterator);
+            if (noStorageProfile.profile(result != null)) {
+                return result.length();
+            }
+            return executeInternal(iterator);
+        }
+
+        protected abstract int executeInternal(PBuiltinIterator iterator);
+
+        @Specialization
+        static int doString(PStringIterator it) {
+            return it.value.length();
+        }
+
+        @Specialization
+        static int doSequenceArr(PArrayIterator it) {
+            return it.array.getLength();
+        }
+
+        @Specialization
+        static int doSequenceIntRange(PIntRangeIterator it) {
+            return it.getLength();
+        }
+
+        @Fallback
+        static int doOthers(PBuiltinIterator it) {
+            return -1;
+        }
+    }
+
     @GenerateUncached
     public abstract static class IsIteratorObjectNode extends Node {
 

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

@@ -67,7 +67,7 @@ protected void increaseCapacityExact(int newCapacity) {
 
     @Override
     public SequenceStorage copy() {
-        return new BoolSequenceStorage(Arrays.copyOf(values, length));
+        return new BoolSequenceStorage(PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/BoolSequenceStorage.java

@@ -75,7 +75,7 @@ protected void increaseCapacityExact(int newCapacity) {
 
     @Override
     public SequenceStorage copy() {
-        return new ByteSequenceStorage(Arrays.copyOf(values, length));
+        return new ByteSequenceStorage(PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/ByteSequenceStorage.java

@@ -69,7 +69,7 @@ protected void increaseCapacityExact(int newCapacity) {
 
     @Override
     public SequenceStorage copy() {
-        return new DoubleSequenceStorage(Arrays.copyOf(values, length));
+        return new DoubleSequenceStorage(PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/DoubleSequenceStorage.java

@@ -69,7 +69,7 @@ protected void increaseCapacityExact(int newCapacity) {
 
     @Override
     public SequenceStorage copy() {
-        return new IntSequenceStorage(Arrays.copyOf(values, length));
+        return new IntSequenceStorage(PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/IntSequenceStorage.java

@@ -71,7 +71,7 @@ protected void increaseCapacityExact(int newCapacity) {
 
     @Override
     public SequenceStorage copy() {
-        return new LongSequenceStorage(Arrays.copyOf(values, length));
+        return new LongSequenceStorage(PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/LongSequenceStorage.java

@@ -143,7 +143,7 @@ public String getClassName() {
 
     @Override
     public SequenceStorage copy() {
-        return new MroSequenceStorage(getClassName(), Arrays.copyOf(values, length));
+        return new MroSequenceStorage(getClassName(), PythonUtils.arrayCopyOf(values, length));
     }
 
     @Override

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/MroSequenceStorage.java

@@ -112,12 +112,12 @@ public Object[] getInternalArray() {
 
     @Override
     public Object[] getCopyOfInternalArray() {
-        return Arrays.copyOf(values, length);
+        return PythonUtils.arrayCopyOf(values, length);
     }
 
     @Override
     public void increaseCapacityExactWithCopy(int newCapacity) {
-        values = Arrays.copyOf(values, newCapacity);
+        values = PythonUtils.arrayCopyOf(values, newCapacity);
         capacity = values.length;
     }