Performance improvements of reading and writing through StreamBox (#1732)

* Performance improvements of reading and writing through StreamBox

* Optimize StreamBox.ReadInto

Author: BCSharp

Src/IronPython.Modules/nt.cs

@@ -1802,13 +1802,12 @@ public static PythonTuple waitpid(int pid, int options) {
 
         public static int write(CodeContext/*!*/ context, int fd, [NotNone] IBufferProtocol data) {
             try {
+                using var buffer = data.GetBuffer();
                 PythonContext pythonContext = context.LanguageContext;
                 var streams = pythonContext.FileManager.GetStreams(fd);
-                using var buffer = data.GetBuffer();
-                var bytes = buffer.AsReadOnlySpan();
                 if (!streams.WriteStream.CanWrite) throw PythonOps.OSError(9, "Bad file descriptor");
 
-                return streams.Write(bytes);
+                return streams.Write(buffer);
             } catch (Exception e) {
                 throw ToPythonException(e);
             }

Src/IronPython/Modules/_fileio.cs

@@ -351,14 +351,7 @@ public BigInteger readinto([NotNone] IBufferProtocol buffer) {
 
                 _checkClosed();
 
-                var span = pythonBuffer.AsSpan();
-                for (int i = 0; i < span.Length; i++) {
-                    int b = _streams.ReadStream.ReadByte();
-                    if (b == -1) return i;
-                    span[i] = (byte)b;
-                }
-
-                return span.Length;
+                return _streams.ReadInto(pythonBuffer);
             }
 
             public override BigInteger readinto(CodeContext/*!*/ context, object buf) {
@@ -444,10 +437,9 @@ public override bool writable(CodeContext/*!*/ context) {
             public override BigInteger write(CodeContext/*!*/ context, object b) {
                 var bufferProtocol = Converter.Convert<IBufferProtocol>(b);
                 using var buffer = bufferProtocol.GetBuffer();
-                var bytes = buffer.AsReadOnlySpan();
 
                 EnsureWritable();
-                return _streams.Write(bytes);
+                return _streams.Write(buffer);
             }
 
             #endregion

Src/IronPython/Runtime/BufferProtocol.cs

@@ -308,23 +308,56 @@ public static void CopyTo(this IPythonBuffer buffer, Span<byte> dest) {
             }
         }
 
+        /// <summary>
+        /// Obtain the underlying array, if possible.
+        /// The returned array is unsafe because it should not be written to.
+        /// </summary>
         internal static byte[]? AsUnsafeArray(this IPythonBuffer buffer) {
             if (!buffer.IsCContiguous())
                 return null;
 
-            if (buffer.Object is Bytes b)
-                return b.UnsafeByteArray;
-
-            if (buffer.Object is Memory<byte> mem) {
-                if (MemoryMarshal.TryGetArray(mem, out ArraySegment<byte> seg) && seg.Array != null && seg.Offset == 0 && seg.Count == seg.Array.Length)
+            ReadOnlySpan<byte> bufdata = buffer.AsReadOnlySpan();
+            if (buffer.Object is Bytes b) {
+                if (b.UnsafeByteArray.AsSpan() == bufdata)
+                    return b.UnsafeByteArray;
+            } else if (buffer.Object is ByteArray ba) {
+                byte[] arrdata = ba.UnsafeByteList.Data;
+                if (arrdata.AsSpan() == bufdata)
+                    return arrdata;
+            } else if (buffer.Object is Memory<byte> mem) {
+                if (MemoryMarshal.TryGetArray(mem, out ArraySegment<byte> seg) && seg.Array is not null && seg.Array.AsSpan() == bufdata)
                     return seg.Array;
             } else if (buffer.Object is ReadOnlyMemory<byte> rom) {
-                if (MemoryMarshal.TryGetArray(rom, out ArraySegment<byte> seg) && seg.Array != null && seg.Offset == 0 && seg.Count == seg.Array.Length)
+                if (MemoryMarshal.TryGetArray(rom, out ArraySegment<byte> seg) && seg.Array is not null && seg.Array.AsSpan() == bufdata)
+                    return seg.Array;
+            }
+
+            return null;
+        }
+
+        /// <summary>
+        /// Obtain the underlying writable array, if possible.
+        /// The returned array is unsafe because it can be longer than the buffer.
+        /// </summary>
+        internal static byte[]? AsUnsafeWritableArray(this IPythonBuffer buffer) {
+            if (!buffer.IsCContiguous() || buffer.IsReadOnly)
+                return null;
+
+            Span<byte> bufdata = buffer.AsSpan();
+            if (buffer.Object is ByteArray ba) {
+                byte[] arrdata = ba.UnsafeByteList.Data;
+                if (UseSameMemory(arrdata, bufdata))
+                    return arrdata;
+            } else if (buffer.Object is Memory<byte> mem) {
+                if (MemoryMarshal.TryGetArray(mem, out ArraySegment<byte> seg) && seg.Array is not null && UseSameMemory(seg.Array, bufdata))
                     return seg.Array;
             }
 
             return null;
         }
+
+        private static bool UseSameMemory(byte[] arr, ReadOnlySpan<byte> span)
+            => arr.Length >= span.Length && arr.AsSpan(0, span.Length) == span;
     }
 
     public ref struct BufferBytesEnumerator {

Src/IronPython/Runtime/PythonFileManager.cs

@@ -109,17 +109,44 @@ public byte[] Read(int count) {
             return buffer;
         }
 
-        public int Write(ReadOnlySpan<byte> bytes) {
+        public int ReadInto(IPythonBuffer buffer) {
 #if NETCOREAPP
+            return _readStream.Read(buffer.AsSpan());
+#else
+            byte[]? bytes = buffer.AsUnsafeWritableArray();
+            if (bytes is not null) {
+                return _readStream.Read(bytes, 0, buffer.NumBytes());
+            }
+
+            const int chunkSize = 0x400; // 1 KiB
+            bytes = new byte[chunkSize];
+            var span = buffer.AsSpan();
+            for (int pos = 0; pos < span.Length; pos += chunkSize) {
+                int toRead = Math.Min(chunkSize, span.Length - pos);
+                int hasRead = _readStream.Read(bytes, 0, toRead);
+                bytes.AsSpan(0, hasRead).CopyTo(span.Slice(pos));
+                if (hasRead < toRead) return pos + hasRead;
+            }
+            return span.Length;
+#endif
+        }
+
+        public int Write(IPythonBuffer buffer) {
+            int count;
+#if NETCOREAPP
+            ReadOnlySpan<byte> bytes = buffer.AsReadOnlySpan();
+            count = bytes.Length;
             _writeStream.Write(bytes);
 #else
-            _writeStream.Write(bytes.ToArray(), 0, bytes.Length);
+            byte[] bytes = buffer.AsUnsafeArray() ?? buffer.AsUnsafeWritableArray() ?? buffer.ToArray();
+            count = buffer.NumBytes();
+            _writeStream.Write(bytes, 0, count);
 #endif
             _writeStream.Flush(); // IO at this level is not supposed to buffer so we need to call Flush.
             if (!IsSingleStream) {
                 _readStream.Seek(_writeStream.Position, SeekOrigin.Begin);
             }
-            return bytes.Length;
+            return count;
         }
 
         public void Flush() {