Repository: apple/swift-nio
File: Sources/NIOCore/NIODecodedAsyncSequence.swift (Swift)

Code before:
            case readingFromBuffer
            case readLastChunkFromBuffer
            case finishedDecoding
        }

        @usableFromInline
        var baseIterator: AsyncSeq.AsyncIterator
        @usableFromInline
        var processor: NIOSingleStepByteToMessageProcessor<Decoder>
        @usableFromInline
        var state: State

        @inlinable
        init(base: NIODecodedAsyncSequence) {
            self.baseIterator = base.asyncSequence.makeAsyncIterator()
            self.processor = NIOSingleStepByteToMessageProcessor(
                base.decoder,
                maximumBufferSize: base.maximumBufferSize
            )
            self.state = .readingFromBuffer
        }

        /// Retrieve the next element from the ``DecodeSequence``.
        @inlinable
        public mutating func next() async throws -> Element? {
            if self.state == .finishedDecoding { return nil }

            let readLastChunkFromBuffer = self.state == .readLastChunkFromBuffer

            // Decode from the buffer if possible
            let (decoded, ended) = try self.processor.decodeNext(
                decodeMode: readLastChunkFromBuffer ? .last : .normal,
                seenEOF: readLastChunkFromBuffer
            )

            if ended {
                // We expect `decodeNext()` to only return `ended == true` only if we've notified it
                // that we've read the last chunk from the buffer.
                assert(readLastChunkFromBuffer)
                self.state = .finishedDecoding
                return decoded
            }

            // If we've managed to decode a message, return it
            if let decoded {
                return decoded
            }

            // If `ended == false` and if `readLastChunkFromBuffer == true` then we must have manged
            // to decode a message. Otherwise something is wrong in `decodeNext()`.
            // We precondition here to avoid infinite recursion.

Review comment (suggestion):
Can you exhaustively switch over the `state` each time? It makes the code easier to reason about and is generally a rule we follow through the NIO repos. It also makes it easier to reason about state transitions as you can make the transition within the body of the `case` that you're switching over.

Code after:
/// ```swift
/// let myDecoder = MyNIOSingleStepByteToMessageDecoder()
/// let baseSequence = MyAsyncSequence<ByteBuffer>(...)
/// let decodedSequence = baseSequence.decode(using: myDecoder)
///
/// for try await element in decodedSequence {
///     print("Decoded an element!", element)
/// }
/// ```
@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
public struct NIODecodedAsyncSequence<
    Base: AsyncSequence,
    Decoder: NIOSingleStepByteToMessageDecoder
> where Base.Element == ByteBuffer {
    @usableFromInline
    var asyncSequence: Base
    @usableFromInline
    var decoder: Decoder
    @usableFromInline
    var maximumBufferSize: Int?

    @inlinable
    init(asyncSequence: Base, decoder: Decoder, maximumBufferSize: Int? = nil) {
        self.asyncSequence = asyncSequence
        self.decoder = decoder
        self.maximumBufferSize = maximumBufferSize
    }
}

@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
extension NIODecodedAsyncSequence: AsyncSequence {
    public typealias Element = Decoder.InboundOut

    /// Create an ``AsyncIterator`` for this ``NIODecodedAsyncSequence``.
    @inlinable
    public func makeAsyncIterator() -> AsyncIterator {
        AsyncIterator(base: self)
    }

    /// An ``AsyncIterator`` over a ``NIODecodedAsyncSequence``.
    public struct AsyncIterator: AsyncIteratorProtocol {
        @usableFromInline
        enum State: Sendable {
            case readingFromBuffer
            case readLastChunkFromBuffer
            case finishedDecoding
        }

        @usableFromInline
        var baseIterator: Base.AsyncIterator
        @usableFromInline