Update Buffer to be based on DispatchData.Slice

Author: iCharlesHu

Sources/Subprocess/AsyncBufferSequence.swift

@@ -37,22 +37,23 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
         public typealias Element = Buffer
 
         private let diskIO: DiskIO
-        private var buffer: [UInt8]
-        private var currentPosition: Int
-        private var finished: Bool
+        private var buffer: [Buffer]
 
         internal init(diskIO: DiskIO) {
             self.diskIO = diskIO
             self.buffer = []
-            self.currentPosition = 0
-            self.finished = false
         }
 
-        public func next() async throws -> Buffer? {
-            let data = try await self.diskIO.readChunk(
+        public mutating func next() async throws -> Buffer? {
+            // If we have more left in buffer, use that
+            guard self.buffer.isEmpty else {
+                return self.buffer.removeFirst()
+            }
+            // Read more data
+            let data = try await self.diskIO.read(
                 upToLength: readBufferSize
             )
-            if data == nil {
+            guard let data = data else {
                 // We finished reading. Close the file descriptor now
                 #if os(Windows)
                 try self.diskIO.close()
@@ -61,7 +62,15 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
                 #endif
                 return nil
             }
-            return data
+            let createdBuffers = Buffer.createFrom(data)
+            // Most (all?) cases there should be only one buffer
+            // because DispatchData are motsly contiguous
+            if _fastPath(createdBuffers.count == 1) {
+                // No need to push to the stack
+                return createdBuffers[0]
+            }
+            self.buffer = createdBuffers
+            return self.buffer.removeFirst()
         }
     }
 
@@ -144,7 +153,10 @@ extension AsyncBufferSequence {
                         capacity: buffer.data.count
                     )
                     defer { temporary.deallocate() }
-                    let actualBytesCopied = buffer.data.copyBytes(to: temporary)
+                    let actualBytesCopied = buffer.data.copyBytes(
+                        to: temporary,
+                        count: buffer.data.count
+                    )
 
                     // Calculate how many CodePoint elements we have
                     let elementCount = actualBytesCopied / MemoryLayout<Encoding.CodeUnit>.stride

Sources/Subprocess/Buffer.swift

@@ -9,7 +9,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#if canImport(Darwin) || canImport(Glibc) || canImport(Android) || canImport(Musl)
 @preconcurrency internal import Dispatch
+#endif
 
 #if SubprocessSpan
 @available(SubprocessSpan, *)
@@ -18,16 +20,32 @@ extension AsyncBufferSequence {
     /// A immutable collection of bytes
     public struct Buffer: Sendable {
         #if os(Windows)
-        internal var data: [UInt8]
+        internal let data: [UInt8]
 
         internal init(data: [UInt8]) {
             self.data = data
         }
+
+        internal static func createFrom(_ data: [UInt8]) -> [Buffer] {
+            return [.init(data: data)]
+        }
         #else
-        internal var data: DispatchData
+        // We need to keep the backingData alive while Slice is alive
+        internal let backingData: DispatchData
+        internal let data: DispatchData.Slice
 
-        internal init(data: DispatchData) {
+        internal init(data: DispatchData.Slice, backingData: DispatchData) {
             self.data = data
+            self.backingData = backingData
+        }
+
+        internal static func createFrom(_ data: DispatchData) -> [Buffer] {
+            let slices = data.slices
+            // In most (all?) cases data should only have one slice
+            if _fastPath(slices.count == 1) {
+                return [.init(data: slices[0], backingData: data)]
+            }
+            return slices.map{ .init(data: $0, backingData: data) }
         }
         #endif
     }
@@ -64,70 +82,20 @@ extension AsyncBufferSequence.Buffer {
     public func withUnsafeBytes<ResultType>(
         _ body: (UnsafeRawBufferPointer) throws -> ResultType
     ) rethrows -> ResultType {
-        #if os(Windows)
         return try self.data.withUnsafeBytes(body)
-        #else
-        // Although DispatchData was designed to be uncontiguous, in practice
-        // we found that almost all DispatchData are contiguous. Therefore
-        // we can access this body in O(1) most of the time.
-        return try self.data.withUnsafeBytes { ptr in
-            let bytes = UnsafeRawBufferPointer(start: ptr, count: self.data.count)
-            return try body(bytes)
-        }
-        #endif
     }
 
     #if SubprocessSpan
     // Access the storge backing this Buffer
     public var bytes: RawSpan {
         @lifetime(borrow self)
         borrowing get {
-            var backing: SpanBacking?
-            #if os(Windows)
-            self.data.withUnsafeBufferPointer {
-                backing = .pointer($0)
-            }
-            #else
-            self.data.enumerateBytes { buffer, byteIndex, stop in
-                if _fastPath(backing == nil) {
-                    // In practice, almost all `DispatchData` is contiguous
-                    backing = .pointer(buffer)
-                } else {
-                    // This DispatchData is not contiguous. We need to copy
-                    // the bytes out
-                    let contents = Array(buffer)
-                    switch backing! {
-                    case .pointer(let ptr):
-                        // Convert the ptr to array
-                        let existing = Array(ptr)
-                        backing = .array(existing + contents)
-                    case .array(let array):
-                        backing = .array(array + contents)
-                    }
-                }
-            }
-            #endif
-            guard let backing = backing else {
-                let empty = UnsafeRawBufferPointer(start: nil, count: 0)
-                let span = RawSpan(_unsafeBytes: empty)
-                return _overrideLifetime(of: span, to: self)
-            }
-            switch backing {
-            case .pointer(let ptr):
-                let span = RawSpan(_unsafeElements: ptr)
-                return _overrideLifetime(of: span, to: self)
-            case .array(let array):
-                let span = array.span.bytes
-                return _overrideLifetime(of: span, to: self)
-            }
+            let ptr = self.data.withUnsafeBytes { $0 }
+            let bytes = RawSpan(_unsafeBytes: ptr)
+            return _overrideLifetime(of: bytes, to: self)
         }
     }
     #endif  // SubprocessSpan
-
-    private enum SpanBacking {
-        case pointer(UnsafeBufferPointer<UInt8>)
-        case array([UInt8])
-    }
 }
 
 // MARK: - Hashable, Equatable
@@ -144,51 +112,53 @@ extension AsyncBufferSequence.Buffer: Equatable, Hashable {
 
     public func hash(into hasher: inout Hasher) {
         self.data.withUnsafeBytes { ptr in
-            let bytes = UnsafeRawBufferPointer(
-                start: ptr,
-                count: self.data.count
-            )
-            hasher.combine(bytes: bytes)
+            hasher.combine(bytes: ptr)
         }
     }
     #endif
 }
 
-// MARK: - Initializers
-#if SubprocessSpan
-@available(SubprocessSpan, *)
-#endif
-extension String {
-    /// Create a String with the given encoding from `Buffer`.
-    /// - Parameters:
-    ///   - buffer: the buffer to copy from
-    ///   - encoding: the encoding to encode Self with
-    public init?<Encoding: _UnicodeEncoding>(buffer: AsyncBufferSequence.Buffer, as encoding: Encoding.Type) {
-#if os(Windows)
-        let source = buffer.data.map { Encoding.CodeUnit($0) }
-        self = String(decoding: source, as: encoding)
-#else
-        self = buffer.withUnsafeBytes { ptr in
-            return String(
-                decoding: ptr.bindMemory(to: Encoding.CodeUnit.self).lazy.map { $0 },
-                as: encoding
-            )
+// MARK: - DispatchData.Block
+#if canImport(Darwin) || canImport(Glibc) || canImport(Android) || canImport(Musl)
+extension DispatchData {
+    /// Unfortunitely `DispatchData.Region` is not available on Linux, hence our own wrapper
+    internal struct Slice: @unchecked Sendable, RandomAccessCollection {
+        typealias Element = UInt8
+
+        internal let bytes: UnsafeBufferPointer<UInt8>
+
+        internal var startIndex: Int { self.bytes.startIndex }
+        internal var endIndex: Int { self.bytes.endIndex }
+
+        internal init(bytes: UnsafeBufferPointer<UInt8>) {
+            self.bytes = bytes
+        }
+
+        internal func withUnsafeBytes<ResultType>(_ body: (UnsafeRawBufferPointer) throws -> ResultType) rethrows -> ResultType {
+            return try body(UnsafeRawBufferPointer(self.bytes))
+        }
+
+        @discardableResult
+        internal func copyBytes<DestinationType>(
+            to ptr: UnsafeMutableBufferPointer<DestinationType>, count: Int
+        ) -> Int {
+            self.bytes.copyBytes(to: ptr, count: count)
+        }
+
+        subscript(position: Int) -> UInt8 {
+            _read {
+                yield self.bytes[position]
+            }
         }
-#endif
     }
-}
 
-#if SubprocessSpan
-@available(SubprocessSpan, *)
-#endif
-extension Array where Element == UInt8 {
-    /// Create an Array from `Buffer`
-    /// - Parameter buffer: the buffer to copy from
-    public init(buffer: AsyncBufferSequence.Buffer) {
-#if os(Windows)
-        self = buffer.data
-#else
-        self = Array(buffer.data)
-#endif
+    internal var slices: [Slice] {
+        var slices = [Slice]()
+        enumerateBytes { (bytes, index, stop) in
+            slices.append(Slice(bytes: bytes))
+        }
+        return slices
     }
 }
+
+#endif

Sources/Subprocess/IO/Output.swift

@@ -161,29 +161,17 @@ public struct BytesOutput: OutputProtocol {
     internal func captureOutput(
         from diskIO: consuming TrackedPlatformDiskIO?
     ) async throws -> [UInt8] {
-        var diskIOBox: TrackedPlatformDiskIO? = consume diskIO
-        return try await withCheckedThrowingContinuation { continuation in
-            let _diskIO = diskIOBox.take()
-            guard let _diskIO = _diskIO else {
-                // Show not happen due to type system constraints
-                fatalError("Trying to capture output without file descriptor")
-            }
-            _diskIO.readUntilEOF(upToLength: self.maxSize) { result in
-                switch result {
-                case .success(let data):
-                    // FIXME: remove workaround for
-                    // rdar://143992296
-                    // https://github.com/swiftlang/swift-subprocess/issues/3
-                    #if os(Windows)
-                    continuation.resume(returning: data)
-                    #else
-                    continuation.resume(returning: data.array())
-                    #endif
-                case .failure(let error):
-                    continuation.resume(throwing: error)
-                }
-            }
+        guard let diskIO = diskIO else {
+            // Show not happen due to type system constraints
+            fatalError("Trying to capture output without file descriptor")
         }
+        #if os(Windows)
+        let result = try await diskIO.fileDescriptor.read(upToLength: self.maxSize)
+        #else
+        let result = try await diskIO.dispatchIO.read(upToLength: self.maxSize)
+        #endif
+
+        return result?.array() ?? []
     }
 
     #if SubprocessSpan
@@ -315,34 +303,23 @@ extension OutputProtocol {
         if let bytesOutput = self as? BytesOutput {
             return try await bytesOutput.captureOutput(from: diskIO) as! Self.OutputType
         }
-        var diskIOBox: TrackedPlatformDiskIO? = consume diskIO
-        return try await withCheckedThrowingContinuation { continuation in
-            if OutputType.self == Void.self {
-                continuation.resume(returning: () as! OutputType)
-                return
-            }
-            guard let _diskIO = diskIOBox.take() else {
-                // Show not happen due to type system constraints
-                fatalError("Trying to capture output without file descriptor")
-            }
 
-            _diskIO.readUntilEOF(upToLength: self.maxSize) { result in
-                do {
-                    switch result {
-                    case .success(let data):
-                        // FIXME: remove workaround for
-                        // rdar://143992296
-                        // https://github.com/swiftlang/swift-subprocess/issues/3
-                        let output = try self.output(from: data)
-                        continuation.resume(returning: output)
-                    case .failure(let error):
-                        continuation.resume(throwing: error)
-                    }
-                } catch {
-                    continuation.resume(throwing: error)
-                }
-            }
+        if OutputType.self == Void.self {
+            return () as! OutputType
+        }
+
+        guard let diskIO = diskIO else {
+            // Show not happen due to type system constraints
+            fatalError("Trying to capture output without file descriptor")
         }
+
+        #if os(Windows)
+        let result = try await diskIO.fileDescriptor.read(upToLength: self.maxSize)
+        #else
+        let result = try await diskIO.dispatchIO.read(upToLength: self.maxSize)
+        #endif
+
+        return try self.output(from: result ?? .empty)
     }
 }