Simplify interaction between AsyncBufferSequence.Iterator and its AsyncStream

Author: Ryan Dingman

Sources/Subprocess/AsyncBufferSequence.swift

@@ -35,25 +35,25 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
     @_nonSendable
     public struct Iterator: AsyncIteratorProtocol {
         public typealias Element = Buffer
-        internal typealias Stream = AsyncThrowingStream<StreamStatus, Swift.Error>
+        internal typealias Stream = AsyncThrowingStream<Buffer, Swift.Error>
 
         private let diskIO: DiskIO
         private var buffer: [UInt8]
         private var currentPosition: Int
         private var finished: Bool
         private var streamIterator: Stream.AsyncIterator
         private let continuation: Stream.Continuation
-        private var needsNextChunk: Bool
+        private var bytesRemaining: Int
 
         internal init(diskIO: DiskIO, streamOptions: PlatformOptions.StreamOptions) {
             self.diskIO = diskIO
             self.buffer = []
             self.currentPosition = 0
             self.finished = false
-            let (stream, continuation) = AsyncThrowingStream<StreamStatus, Swift.Error>.makeStream()
+            let (stream, continuation) = AsyncThrowingStream<Buffer, Swift.Error>.makeStream()
             self.streamIterator = stream.makeAsyncIterator()
             self.continuation = continuation
-            self.needsNextChunk = true
+            self.bytesRemaining = 0
 
             #if !os(Windows)
             if let minimumBufferSize = streamOptions.minimumBufferSize {
@@ -68,28 +68,14 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
 
         public mutating func next() async throws -> Buffer? {
 
-            if needsNextChunk {
-                diskIO.readChunk(upToLength: readBufferSize, continuation: continuation)
-                needsNextChunk = false
+            if bytesRemaining <= 0 {
+                bytesRemaining = readBufferSize
+                diskIO.stream(upToLength: readBufferSize, continuation: continuation)
             }
 
-            if let status = try await streamIterator.next() {
-                switch status {
-                case .data(let data):
-                    return data
-
-                case .endOfChunk(let data):
-                    needsNextChunk = true
-                    return data
-
-                case .endOfFile:
-                    #if os(Windows)
-                    try self.diskIO.close()
-                    #else
-                    self.diskIO.close()
-                    #endif
-                    return nil
-                }
+            if let buffer = try await streamIterator.next() {
+                bytesRemaining -= buffer.count
+                return buffer
             } else {
                 #if os(Windows)
                 try self.diskIO.close()
@@ -114,17 +100,6 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
     }
 }
 
-extension AsyncBufferSequence {
-#if SubprocessSpan
-@available(SubprocessSpan, *)
-#endif
-    internal enum StreamStatus {
-        case data(AsyncBufferSequence.Buffer)
-        case endOfChunk(AsyncBufferSequence.Buffer)
-        case endOfFile
-    }
-}
-
 // MARK: - Page Size
 import _SubprocessCShims
 

Sources/Subprocess/Platforms/Subprocess+Unix.swift

@@ -415,7 +415,7 @@ extension DispatchIO {
     #if SubprocessSpan
     @available(SubprocessSpan, *)
     #endif
-    internal func readChunk(upToLength maxLength: Int, continuation: AsyncBufferSequence.Iterator.Stream.Continuation) {
+    internal func stream(upToLength maxLength: Int, continuation: AsyncBufferSequence.Iterator.Stream.Continuation) {
         self.read(
             offset: 0,
             length: maxLength,
@@ -430,26 +430,9 @@ extension DispatchIO {
             }
 
             // Treat empty data and nil as the same
-            let buffer = data.map { $0.isEmpty ? nil : $0 } ?? nil
-            let status: AsyncBufferSequence.StreamStatus
-
-            switch (buffer, done) {
-            case (.some(let data), false):
-                status = .data(AsyncBufferSequence.Buffer(data: data))
-
-            case (.some(let data), true):
-                status = .endOfChunk(AsyncBufferSequence.Buffer(data: data))
-
-            case (nil, false):
-                fatalError("Unexpectedly received no data from DispatchIO with it indicating it is not done.")
-
-            case (nil, true):
-                status = .endOfFile
-            }
-
-            continuation.yield(status)
-
-            if done {
+            if let data = data.map({ $0.isEmpty ? nil : $0 }) ?? nil {
+                continuation.yield(AsyncBufferSequence.Buffer(data: data))
+            } else if done {
                 continuation.finish()
             }
         }

Sources/Subprocess/Platforms/Subprocess+Windows.swift

@@ -1078,7 +1078,7 @@ extension FileDescriptor {
 }
 
 extension FileDescriptor {
-    internal func readChunk(upToLength maxLength: Int, continuation: AsyncBufferSequence.Iterator.Stream.Continuation) {
+    internal func stream(upToLength maxLength: Int, continuation: AsyncBufferSequence.Iterator.Stream.Continuation) {
         do {
             var totalBytesRead: Int = 0
 
@@ -1121,16 +1121,8 @@ extension FileDescriptor {
 
                 if values.count > 0 {
                     totalBytesRead += values.count
-
-                    if totalBytesRead >= maxLength {
-                        continuation.yield(.endOfChunk(AsyncBufferSequence.Buffer(data: values)))
-                        continuation.finish()
-                        return
-                    } else {
-                        continuation.yield(.data(AsyncBufferSequence.Buffer(data: values)))
-                    }
+                    continuation.yield(AsyncBufferSequence.Buffer(data: values))
                 } else {
-                    continuation.yield(.endOfFile)
                     continuation.finish()
                     return
                 }