PureYAMLStreaming/Sources/PureYAMLStreaming/ResumableDocumentScanner.swift

import PureYAML

public extension PureYAMLStreaming {
    /// Resumable document-boundary scanner driven by `PureYAML.Streaming.Scanner`.
    ///
    /// Chunks are tokenized by PureYAML's resumable token scanner (proven
    /// byte-for-byte identical to batch scanning across the yaml-test-suite
    /// corpus), and the buffered source text is cut immediately before each
    /// `documentStart` token, or before the first directive token that follows
    /// an explicit `...` document end. Unlike the previous line-based
    /// heuristic, a `---` that the grammar does not treat as a document start
    /// (for example inside a quoted scalar) can never misfire as a boundary,
    /// directives stay attached to the document they introduce, and a chunk
    /// boundary may fall anywhere, including mid-token.
    ///
    /// Guarantees:
    /// - Lossless: concatenating every emitted segment (from `consume(_:)` and
    ///   `finish()`, in order) reproduces the fed input exactly.
    /// - Boundary-sound: no YAML document ever spans two segments, so parsing
    ///   the segments one at a time yields the same documents as batch parsing
    ///   the whole stream.
    /// - Bounded: after each cut the token scanner is restarted on the
    ///   remaining buffer (a document boundary fully resets scanner state), so
    ///   retained state is bounded by the current document, not the stream.
    struct ResumableDocumentScanner: Sendable {
        private var scanner = PureYAML.Streaming.Scanner()
        private var buffer: [Unicode.Scalar] = []
        private var scannerFedCount = 0
        private var segmentHasContent = false
        private var inDirectivePreamble = false
        private var afterDocumentEnd = false
        private var isFinal = false

        public init() {}

        /// Appends a chunk of YAML source and returns every document source
        /// completed by it. A chunk may end anywhere, including mid-token.
        public mutating func consume(_ chunk: String) throws -> [String] {
            buffer.append(contentsOf: chunk.unicodeScalars)
            feedScanner()
            return try drainTokens()
        }

        /// Like `consume(_:)`, but wraps each completed document source in a
        /// ``StreamingToken``.
        public mutating func consumeToken(_ chunk: String) throws -> [StreamingToken] {
            try consume(chunk).map(StreamingToken.documentSource)
        }

        /// Marks the end of input and returns the remaining document sources:
        /// any boundary that only end-of-input could confirm (for example a
        /// trailing `---` line), followed by the final buffered document.
        public mutating func finish() throws -> [String] {
            isFinal = true
            feedScanner()
            scanner.finish()
            var documents = try drainTokens()
            if !buffer.isEmpty {
                documents.append(String(String.UnicodeScalarView(buffer)))
            }
            self = ResumableDocumentScanner()
            return documents
        }

        /// Feeds the token scanner every buffered scalar it has not seen yet,
        /// except a trailing CR while more input may still arrive. Whether
        /// that CR is a lone break or the first half of a CR LF pair is
        /// undecidable until the next scalar (or end of input) is known, and
        /// PureYAML's persisted plain-scalar suspension counts a pair split
        /// exactly there as two line breaks, shifting every later token mark
        /// down a line (observed at PureYAML commit 2fcf546). Withholding the
        /// CR keeps every break the scanner sees whole, so the token line
        /// numbers this type cuts by always match batch scanning.
        private mutating func feedScanner() {
            var target = buffer.count
            if !isFinal, buffer.last == "\r" {
                target -= 1
            }
            guard target > scannerFedCount else {
                return
            }
            scanner.feed(String(String.UnicodeScalarView(buffer[scannerFedCount ..< target])))
            scannerFedCount = target
        }

        private mutating func drainTokens() throws -> [String] {
            var documents: [String] = []
            drain: while true {
                switch try scanner.next() {
                case let .token(token):
                    handle(token, documents: &documents)
                case .needMoreInput, .done:
                    break drain
                }
            }
            return documents
        }

        private mutating func handle(_ token: PureYAML.Streaming.Token, documents: inout [String]) {
            switch token.kind {
            case .streamStart, .streamEnd:
                return
            case .documentStart:
                if segmentHasContent, !inDirectivePreamble {
                    // The restarted scanner re-emits this token from the new
                    // buffer start; that re-emission sets the segment flags.
                    documents.append(cutSegment(beforeLine: token.range.start.line))
                    return
                }
                inDirectivePreamble = false
                afterDocumentEnd = false
                segmentHasContent = true
            case .versionDirective, .tagDirective, .reservedDirective:
                // A directive is a legal document boundary only after an
                // explicit `...`. Without one, batch parsing rejects the
                // stream, so the directive must stay in the current segment
                // for the per-segment parse to reject it the same way.
                if segmentHasContent, !inDirectivePreamble, afterDocumentEnd {
                    documents.append(cutSegment(beforeLine: token.range.start.line))
                    return
                }
                inDirectivePreamble = true
                afterDocumentEnd = false
                segmentHasContent = true
            case .documentEnd:
                inDirectivePreamble = false
                afterDocumentEnd = true
                segmentHasContent = true
            default:
                inDirectivePreamble = false
                afterDocumentEnd = false
                segmentHasContent = true
            }
        }

        /// Cuts the buffer at the start of `line` (0-based, relative to the
        /// current scanner instance, whose input always begins at `buffer[0]`),
        /// emits the prefix, and restarts the token scanner on the remainder so
        /// that no state older than the current document is retained.
        ///
        /// Boundary tokens (`---` and directives) always start at column 0, so
        /// the cut offset is that line's first scalar. Line counting mirrors
        /// the scanner's reader exactly: CR LF is one break, and CR, LF, NEL,
        /// LINE SEPARATOR, and PARAGRAPH SEPARATOR each end a line.
        private mutating func cutSegment(beforeLine line: Int) -> String {
            var offset = 0
            var linesSeen = 0
            while linesSeen < line, offset < buffer.count {
                if buffer[offset] == "\r", offset + 1 < buffer.count, buffer[offset + 1] == "\n" {
                    offset += 2
                    linesSeen += 1
                } else if Self.isBreak(buffer[offset]) {
                    offset += 1
                    linesSeen += 1
                } else {
                    offset += 1
                }
            }

            let segment = String(String.UnicodeScalarView(buffer[..<offset]))
            buffer.removeFirst(offset)
            segmentHasContent = false
            inDirectivePreamble = false
            afterDocumentEnd = false
            scanner = PureYAML.Streaming.Scanner()
            scannerFedCount = 0
            feedScanner()
            if isFinal {
                scanner.finish()
            }
            return segment
        }

        /// The scanner reader's break set (libyaml's `IS_BREAK`).
        private static func isBreak(_ scalar: Unicode.Scalar) -> Bool {
            switch scalar {
            case "\r", "\n", "\u{85}", "\u{2028}", "\u{2029}": true
            default: false
            }
        }
    }
}