Add PSequence specializations for array constructor and extend

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/ArrayModuleBuiltins.java

@@ -42,6 +42,8 @@
 import com.oracle.graal.python.builtins.objects.array.PArray;
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
 import com.oracle.graal.python.builtins.objects.bytes.PBytesLike;
+import com.oracle.graal.python.builtins.objects.common.SequenceNodes;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.function.PKeyword;
 import com.oracle.graal.python.builtins.objects.module.PythonModule;
 import com.oracle.graal.python.builtins.objects.object.PythonObjectLibrary;
@@ -60,6 +62,8 @@
 import com.oracle.graal.python.runtime.PythonCore;
 import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
+import com.oracle.graal.python.runtime.sequence.PSequence;
+import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
 import com.oracle.graal.python.util.BufferFormat;
 import com.oracle.graal.python.util.OverflowException;
 import com.oracle.graal.python.util.PythonUtils;
@@ -188,7 +192,43 @@ PArray arrayWithBytesInitializer(VirtualFrame frame, Object cls, String typeCode
                 return array;
             }
 
-            // TODO impl for PSequence and PArray or use lenght_hint
+            @Specialization
+            PArray arrayArrayInitializer(VirtualFrame frame, Object cls, String typeCode, PArray initializer,
+                            @Cached ArrayNodes.PutValueNode putValueNode,
+                            @Cached ArrayNodes.GetValueNode getValueNode) {
+                BufferFormat format = getFormatChecked(typeCode);
+                try {
+                    PArray array = getFactory().createArray(cls, typeCode, format, initializer.getLength());
+                    for (int i = 0; i < initializer.getLength(); i++) {
+                        putValueNode.execute(frame, array, i, getValueNode.execute(initializer, i));
+                    }
+                    return array;
+                } catch (OverflowException e) {
+                    CompilerDirectives.transferToInterpreterAndInvalidate();
+                    throw raise(MemoryError);
+                }
+            }
+
+            @Specialization
+            PArray arraySequenceInitializer(VirtualFrame frame, Object cls, String typeCode, PSequence initializer,
+                            @Cached ArrayNodes.PutValueNode putValueNode,
+                            @Cached SequenceNodes.GetSequenceStorageNode getSequenceStorageNode,
+                            @Cached SequenceStorageNodes.LenNode lenNode,
+                            @Cached SequenceStorageNodes.GetItemScalarNode getItemNode) {
+                BufferFormat format = getFormatChecked(typeCode);
+                SequenceStorage storage = getSequenceStorageNode.execute(initializer);
+                int lenght = lenNode.execute(storage);
+                try {
+                    PArray array = getFactory().createArray(cls, typeCode, format, lenght);
+                    for (int i = 0; i < lenght; i++) {
+                        putValueNode.execute(frame, array, i, getItemNode.execute(storage, i));
+                    }
+                    return array;
+                } catch (OverflowException e) {
+                    CompilerDirectives.transferToInterpreterAndInvalidate();
+                    throw raise(MemoryError);
+                }
+            }
 
             @Specialization(guards = "!isBytes(initializer)", limit = "3")
             PArray arrayIteratorInitializer(VirtualFrame frame, Object cls, String typeCode, Object initializer,

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/array/ArrayBuiltins.java

@@ -89,6 +89,7 @@
 import com.oracle.graal.python.nodes.subscript.SliceLiteralNode;
 import com.oracle.graal.python.nodes.util.CastToJavaStringNode;
 import com.oracle.graal.python.runtime.exception.PException;
+import com.oracle.graal.python.runtime.sequence.PSequence;
 import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
 import com.oracle.graal.python.util.BufferFormat;
 import com.oracle.graal.python.util.OverflowException;
@@ -697,12 +698,43 @@ Object extend(PArray self, PArray value) {
             }
         }
 
-        @Specialization(guards = "self.getFormat() != value.getFormat()")
-        @SuppressWarnings("unused")
-        Object error(PArray self, PArray value) {
-            // CPython allows extending an array with an arbitrary iterable. Except a differently
-            // formatted array. Weird
-            throw raise(TypeError, "can only extend with array of same kind");
+        @Specialization
+        Object extend(VirtualFrame frame, PArray self, PSequence value,
+                        @Cached ArrayNodes.PutValueNode putValueNode,
+                        @Cached SequenceNodes.GetSequenceStorageNode getSequenceStorageNode,
+                        @Cached SequenceStorageNodes.LenNode lenNode,
+                        @Cached SequenceStorageNodes.GetItemScalarNode getItemNode) {
+            SequenceStorage storage = getSequenceStorageNode.execute(value);
+            int storageLenght = lenNode.execute(storage);
+            boolean capacityEnsured = false;
+            try {
+                self.ensureCapacity(PythonUtils.addExact(self.getLength(), storageLenght));
+                capacityEnsured = true;
+            } catch (OverflowException e) {
+                CompilerDirectives.transferToInterpreterAndInvalidate();
+                // Let it fail later, so that it fails in the same state as the generic
+                // specialization
+            }
+            int lenght = self.getLength();
+            for (int i = 0; i < storageLenght; i++) {
+                // The whole extend is not atomic, just individual inserts are. That's the same as
+                // in CPython
+                if (capacityEnsured) {
+                    lenght++;
+                } else {
+                    try {
+                        lenght = PythonUtils.addExact(lenght, 1);
+                        self.ensureCapacity(lenght);
+                    } catch (OverflowException e) {
+                        CompilerDirectives.transferToInterpreterAndInvalidate();
+                        throw raise(MemoryError);
+                    }
+                }
+                putValueNode.execute(frame, self, lenght - 1, getItemNode.execute(storage, i));
+                self.setLenght(lenght);
+            }
+
+            return PNone.NONE;
         }
 
         @Specialization(guards = "!isArray(value)", limit = "3")
@@ -736,6 +768,14 @@ Object extend(VirtualFrame frame, PArray self, Object value,
 
             return PNone.NONE;
         }
+
+        @Specialization(guards = "self.getFormat() != value.getFormat()")
+        @SuppressWarnings("unused")
+        Object error(PArray self, PArray value) {
+            // CPython allows extending an array with an arbitrary iterable. Except a differently
+            // formatted array. Weird
+            throw raise(TypeError, "can only extend with array of same kind");
+        }
     }
 
     @Builtin(name = "insert", minNumOfPositionalArgs = 3, numOfPositionalOnlyArgs = 3, parameterNames = {"$self", "index", "value"})