Do not assume storage type of bytes object

Author: fangerer

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/io/BytesIOBuiltins.java

@@ -97,6 +97,7 @@
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
 import com.oracle.graal.python.builtins.objects.common.HashingStorageLibrary;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.EnsureCapacityNode;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetCapacityNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetInternalArrayNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetInternalByteArrayNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetInternalObjectArrayNode;
@@ -684,24 +685,27 @@ protected static boolean shouldCopy(PBytesIO self) {
             return self.getStringSize() <= 1 || self.getExports() > 0;
         }
 
-        protected static boolean shouldUnshare(PBytesIO self) {
-            return self.getStringSize() != self.getBufCapacity();
+        protected static boolean shouldUnshare(GetCapacityNode getCapacityNode, PBytesIO self) {
+            int capacity = getCapacityNode.execute(self.getBuf().getSequenceStorage());
+            return self.getStringSize() != capacity;
         }
 
         @Specialization(guards = {"self.hasBuf()", "shouldCopy(self)"})
-        Object copy(PBytesIO self,
+        Object doCopy(PBytesIO self,
                         @Cached GetInternalByteArrayNode getBytes) {
             byte[] buf = getBytes.execute(self.getBuf().getSequenceStorage());
             return factory().createBytes(PythonUtils.arrayCopyOf(buf, self.getStringSize()));
         }
 
-        @Specialization(guards = {"self.hasBuf()", "!shouldCopy(self)", "!shouldUnshare(self)"})
-        static Object doit(PBytesIO self) {
+        @Specialization(guards = {"self.hasBuf()", "!shouldCopy(self)", "!shouldUnshare(getCapacityNode, self)"}, limit = "1")
+        static Object doShare(PBytesIO self,
+                        @SuppressWarnings("unused") @Shared("getCapacityNode") @Cached GetCapacityNode getCapacityNode) {
             return self.getBuf();
         }
 
-        @Specialization(guards = {"self.hasBuf()", "!shouldCopy(self)", "shouldUnshare(self)"})
-        Object unshare(PBytesIO self,
+        @Specialization(guards = {"self.hasBuf()", "!shouldCopy(self)", "shouldUnshare(getCapacityNode, self)"}, limit = "1")
+        Object doUnshare(PBytesIO self,
+                        @SuppressWarnings("unused") @Shared("getCapacityNode") @Cached GetCapacityNode getCapacityNode,
                         @Cached GetInternalArrayNode internalArray) {
             // if (SHARED_BUF(self))
             unshareBuffer(self, self.getStringSize(), internalArray, factory());

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/io/PBytesIO.java

@@ -43,10 +43,9 @@
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
 import com.oracle.graal.python.builtins.objects.memoryview.BufferLifecycleManager;
 import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
-import com.oracle.graal.python.runtime.sequence.storage.BasicSequenceStorage;
 import com.oracle.truffle.api.object.Shape;
 
-public class PBytesIO extends PythonBuiltinObject {
+public final class PBytesIO extends PythonBuiltinObject {
     private PBytes buf;
     private int pos;
     private int stringSize;
@@ -89,11 +88,6 @@ public void setStringSize(int size) {
         this.stringSize = size;
     }
 
-    public int getBufCapacity() {
-        // Casting is safe as we only create/replace buf internally.
-        return ((BasicSequenceStorage) buf.getSequenceStorage()).capacity();
-    }
-
     public int getExports() {
         return exports.getExports().get();
     }

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

@@ -3066,6 +3066,27 @@ private static NativeSequenceStorage reallocNativeSequenceStorage(NativeSequence
         }
     }
 
+    @GenerateUncached
+    public abstract static class GetCapacityNode extends Node {
+
+        public abstract int execute(SequenceStorage s);
+
+        @Specialization
+        static int doBasicSequenceStorage(BasicSequenceStorage s) {
+            return s.capacity();
+        }
+
+        @Specialization
+        static int doNativeSequenceStorage(NativeSequenceStorage s) {
+            return s.getCapacity();
+        }
+
+        @Specialization
+        static int doBasicSequenceStorage(@SuppressWarnings("unused") EmptySequenceStorage s) {
+            return 0;
+        }
+    }
+
     @GenerateUncached
     @ImportStatic(SequenceStorageBaseNode.class)
     public abstract static class GetInternalArrayNode extends Node {