Fix: properly convert RHS when writing slice into byte array

Author: steve-s

graalpython/com.oracle.graal.python.test/src/tests/test_bytes.py

@@ -239,6 +239,13 @@ def test_setslice():
     b[5:10] = [4, 5, 6, 7, 8]
     assert b == bytearray(b"hello\x04\x05\x06\x07\x08hello")
 
+    # assign list with integers backed by generic storage
+    b = bytearray(b"hellohellohello")
+    rhs = ['hello', 5, 6, 7, 8]
+    rhs[0] = 4
+    b[5:10] = rhs
+    assert b == bytearray(b"hello\x04\x05\x06\x07\x08hello")
+
     # assign range
     b = bytearray(b"hellohellohello")
     b[5:10] = range(5)

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/ByteArrayBuiltins.java

@@ -56,6 +56,7 @@
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.function.PKeyword;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.ByteArrayGeneralizationNode;
 import com.oracle.graal.python.builtins.objects.iterator.IteratorNodes;
 import com.oracle.graal.python.builtins.objects.memoryview.PMemoryView;
 import com.oracle.graal.python.builtins.objects.object.PythonObjectLibrary;
@@ -232,11 +233,15 @@ Object error(Object self, Object idx, Object value) {
         }
 
         protected SequenceStorageNodes.SetItemNode createSetItem() {
+            // Note the error message should never be reached, because the storage should always be
+            // writeable and so SetItemScalarNode should always have a specialization for it and
+            // inside that specialization the conversion of RHS may fail and produce Python level
+            // ValueError
             return SequenceStorageNodes.SetItemNode.create(NormalizeIndexNode.forBytearray(), "an integer is required");
         }
 
         protected SequenceStorageNodes.SetItemNode createSetSlice() {
-            return SequenceStorageNodes.SetItemNode.create(NormalizeIndexNode.forBytearray(), "can assign only bytes, buffers, or iterables of ints in range(0, 256)");
+            return SequenceStorageNodes.SetItemNode.create(NormalizeIndexNode.forBytearray(), ByteArrayGeneralizationNode.CACHED_SUPPLIER);
         }
 
         protected static boolean isMemoryView(Object value) {

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

@@ -1140,12 +1140,12 @@ protected SequenceStorage doSliceSequence(SequenceStorage storage, PSlice slice,
             int len = lenNode.execute(storage);
             SliceInfo info = adjustIndices.execute(len, unadjusted);
             try {
-                setItemSliceNode.execute(storage, info, sequence);
+                setItemSliceNode.execute(storage, info, sequence, true);
                 return storage;
             } catch (SequenceStoreException e) {
                 generalizeProfile.enter();
                 SequenceStorage generalized = generalizeStore(storage, e.getIndicationValue());
-                setItemSliceNode.execute(generalized, info, sequence);
+                setItemSliceNode.execute(generalized, info, sequence, false);
                 return generalized;
             }
         }
@@ -1167,12 +1167,12 @@ protected SequenceStorage doSliceGeneric(SequenceStorage storage, PSlice slice,
             // must not use iterable again. It could have sice-effects.
             PList values = constructListNode.execute(iterable);
             try {
-                setItemSliceNode.execute(storage, info, values);
+                setItemSliceNode.execute(storage, info, values, true);
                 return storage;
             } catch (SequenceStoreException e) {
                 generalizeProfile.enter();
                 SequenceStorage generalized = generalizeStore(storage, e.getIndicationValue());
-                setItemSliceNode.execute(generalized, info, values);
+                setItemSliceNode.execute(generalized, info, values, false);
                 return generalized;
             }
         }
@@ -1321,21 +1321,25 @@ public static SetItemScalarNode create() {
     @ImportStatic({ListStorageType.class, SequenceStorageBaseNode.class})
     public abstract static class SetItemSliceNode extends Node {
 
-        public abstract void execute(SequenceStorage s, SliceInfo info, Object iterable);
+        public abstract void execute(SequenceStorage s, SliceInfo info, Object iterable, boolean canGeneralize);
+
+        public final void execute(SequenceStorage s, SliceInfo info, Object iterable) {
+            execute(s, info, iterable, true);
+        }
 
         @Specialization(guards = "hasStorage(seq)")
-        static void doStorage(SequenceStorage s, SliceInfo info, PSequence seq,
+        static void doStorage(SequenceStorage s, SliceInfo info, PSequence seq, boolean canGeneralize,
                         @Shared("setStorageSliceNode") @Cached SetStorageSliceNode setStorageSliceNode,
                         @Cached GetSequenceStorageNode getSequenceStorageNode) {
-            setStorageSliceNode.execute(s, info, getSequenceStorageNode.execute(seq));
+            setStorageSliceNode.execute(s, info, getSequenceStorageNode.execute(seq), canGeneralize);
         }
 
         @Specialization
-        static void doGeneric(SequenceStorage s, SliceInfo info, Object iterable,
+        static void doGeneric(SequenceStorage s, SliceInfo info, Object iterable, boolean canGeneralize,
                         @Shared("setStorageSliceNode") @Cached SetStorageSliceNode setStorageSliceNode,
                         @Cached ListNodes.ConstructListNode constructListNode) {
             PList list = constructListNode.execute(iterable);
-            setStorageSliceNode.execute(s, info, list.getSequenceStorage());
+            setStorageSliceNode.execute(s, info, list.getSequenceStorage(), canGeneralize);
         }
 
         public static SetItemSliceNode create() {
@@ -1416,10 +1420,10 @@ protected static boolean isBasicSequenceStorage(Object o) {
     @ImportStatic(SequenceStorageBaseNode.class)
     abstract static class SetStorageSliceNode extends Node {
 
-        public abstract void execute(SequenceStorage s, SliceInfo info, SequenceStorage iterable);
+        public abstract void execute(SequenceStorage s, SliceInfo info, SequenceStorage iterable, boolean canGeneralize);
 
         @Specialization(limit = "MAX_ARRAY_STORAGES", guards = {"self.getClass() == cachedClass", "self.getClass() == sequence.getClass()", "replacesWholeSequence(cachedClass, self, info)"})
-        static void doWholeSequence(BasicSequenceStorage self, @SuppressWarnings("unused") SliceInfo info, BasicSequenceStorage sequence,
+        static void doWholeSequence(BasicSequenceStorage self, @SuppressWarnings("unused") SliceInfo info, BasicSequenceStorage sequence, @SuppressWarnings("unused") boolean canGeneralize,
                         @Cached("self.getClass()") Class<? extends BasicSequenceStorage> cachedClass) {
             BasicSequenceStorage selfProfiled = cachedClass.cast(self);
             BasicSequenceStorage otherProfiled = cachedClass.cast(sequence);
@@ -1428,8 +1432,8 @@ static void doWholeSequence(BasicSequenceStorage self, @SuppressWarnings("unused
             selfProfiled.minimizeCapacity();
         }
 
-        @Specialization(guards = {"isDataTypeCompatibleNode.execute(self, values)", "sinfo.step == 1"})
-        static void singleStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage values,
+        @Specialization(guards = {"!canGeneralize || isDataTypeCompatibleNode.execute(self, values)", "sinfo.step == 1"})
+        static void singleStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage values, @SuppressWarnings("unused") boolean canGeneralize,
                         @Cached @SuppressWarnings("unused") IsDataTypeCompatibleNode isDataTypeCompatibleNode,
                         @Cached LenNode lenNode,
                         @Cached SetLenNode setLenNode,
@@ -1456,8 +1460,8 @@ static void singleStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage va
                             memoryError, negGrowth, posGrowth, raiseNode);
         }
 
-        @Specialization(guards = {"isDataTypeCompatibleNode.execute(self, values)", "sinfo.step != 1"})
-        static void multiStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage values,
+        @Specialization(guards = {"!canGeneralize || isDataTypeCompatibleNode.execute(self, values)", "sinfo.step != 1"})
+        static void multiStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage values, @SuppressWarnings("unused") boolean canGeneralize,
                         @Cached @SuppressWarnings("unused") IsDataTypeCompatibleNode isDataTypeCompatibleNode,
                         @Cached ConditionProfile wrongLength,
                         @Cached ConditionProfile deleteSlice,
@@ -1490,8 +1494,8 @@ static void multiStep(SequenceStorage self, SliceInfo sinfo, SequenceStorage val
             }
         }
 
-        @Specialization(guards = "!isAssignCompatibleNode.execute(self, sequence)")
-        static void doError(@SuppressWarnings("unused") SequenceStorage self, @SuppressWarnings("unused") SliceInfo info, SequenceStorage sequence,
+        @Specialization(guards = {"canGeneralize", "!isAssignCompatibleNode.execute(self, sequence)"})
+        static void doError(@SuppressWarnings("unused") SequenceStorage self, @SuppressWarnings("unused") SliceInfo info, SequenceStorage sequence, @SuppressWarnings("unused") boolean canGeneralize,
                         @Cached @SuppressWarnings("unused") IsAssignCompatibleNode isAssignCompatibleNode) {
             throw new SequenceStoreException(sequence.getIndicativeValue());
         }
@@ -2765,6 +2769,13 @@ static int doGeneric(VirtualFrame frame, SequenceStorage left, Object item,
         }
     }
 
+    /**
+     * Generalization node must convert given storage to a storage that is able to be written any
+     * number of any valid elements. I.e., there must be a specialization handling that storage type
+     * in {@link SetItemScalarNode}. Note: it is possible that the RHS of the write may be invalid
+     * element, e.g., large integer when the storage is bytes array storage, but in such case the
+     * {@link SetItemScalarNode} will correctly raise Python level {@code ValueError}.
+     */
     public abstract static class GeneralizationNode extends Node {
         public abstract SequenceStorage execute(SequenceStorage toGeneralize, Object indicationValue);
 
@@ -2836,6 +2847,49 @@ public static NoGeneralizationCustomMessageNode create(String msg) {
         }
     }
 
+    /**
+     * Byte array is specific that it supports being written using slice from other iterables as
+     * long as the individual elements written can be converted to bytes. Arrays from the array
+     * module do not support this, they can be written only individual elements or slice from other
+     * array of the same type.
+     *
+     * This node works with the assumption that all storages of byte arrays support writing of any
+     * number of bytes. There is no actual generalization of the storage, but instead we tell the
+     * caller that it should try again with the same storage (in the second try, it should try to
+     * write whatever is in RHS no matter of the type of the storage of RHS).
+     *
+     * This is only limitation of this node. Shall we ever want to support byte arrays that can be
+     * backed by different types of storage, we'd only need to change this node to accommodate for
+     * that and return the most generic storage of those.
+     */
+    public static class ByteArrayGeneralizationNode extends GeneralizationNode {
+        public static ByteArrayGeneralizationNode UNCACHED = new ByteArrayGeneralizationNode();
+
+        public static final GenNodeSupplier SUPPLIER = new GenNodeSupplier() {
+            @Override
+            public GeneralizationNode getUncached() {
+                return UNCACHED;
+            }
+
+            @Override
+            public GeneralizationNode create() {
+                return new ByteArrayGeneralizationNode();
+            }
+        };
+
+        public static final Supplier<GeneralizationNode> CACHED_SUPPLIER = new Supplier<GeneralizationNode>() {
+            @Override
+            public GeneralizationNode get() {
+                return new ByteArrayGeneralizationNode();
+            }
+        };
+
+        @Override
+        public SequenceStorage execute(SequenceStorage toGeneralize, @SuppressWarnings("unused") Object indicationValue) {
+            return toGeneralize;
+        }
+    }
+
     /**
      * Implements list generalization rules; previously in 'SequenceStroage.generalizeFor'.
      */

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/util/CastToByteNode.java

@@ -147,10 +147,15 @@ protected byte doBoolean(boolean value) {
         return value ? (byte) 1 : (byte) 0;
     }
 
-    @Specialization(guards = "coerce")
+    @Specialization
     protected byte doBytes(VirtualFrame frame, PBytesLike value,
                     @Cached("create()") SequenceStorageNodes.GetItemNode getItemNode) {
-        return doIntOvf(getItemNode.executeInt(frame, value.getSequenceStorage(), 0));
+        // Workaround GR-26346
+        if (coerce) {
+            return doIntOvf(getItemNode.executeInt(frame, value.getSequenceStorage(), 0));
+        } else {
+            return doGeneric(value);
+        }
     }
 
     @Specialization(guards = "plib.isForeignObject(value)", limit = "1")