PureYAML/Sources/Parsing/LibYAML.Scanner.Streaming.swift

extension PureYAML.Parsing.LibYAML.Scanner {
    /// The result of one `scanStep()` call.
    enum ScanOutcome: Equatable {
        /// A token is safe to release to the caller now.
        case token(PureYAML.Parsing.LibYAML.Token)
        /// The buffered input was exhausted mid-fetch; append more input (via
        /// `reader.append(_:)`) and call `scanStep()` again. The scanner's state is
        /// unchanged from before this call, as if it had not been called.
        case needMoreInput
        /// STREAM-END has been released and no further tokens remain.
        case done
    }

    /// Resumable, one-token-at-a-time counterpart to `scanAll()`. Where `scanAll()`
    /// assumes the full input is already buffered, `scanStep()` may be driven across
    /// a `Reader` still receiving chunks (`isFinal == false`): a fetch that runs off
    /// the buffered end is rolled back in full (as if it had not been attempted)
    /// rather than mistaken for true end of input.
    ///
    /// A token already queued is not always safe to hand out immediately: libyaml's
    /// simple-key mechanism can still splice a KEY (and a BLOCK-MAPPING-START) in
    /// front of it once a later `:` is found (`saveSimpleKey`/`insert(_:atTokenNumber:)`
    /// in `LibYAML.Scanner.swift`/`LibYAML.IndentEngine.swift`). Port of libyaml's
    /// `yaml_parser_fetch_more_tokens` release condition: the token at the queue
    /// front (position `tokensParsed`) is safe to release only when no currently
    /// possible simple key (at any flow level) is still recorded at that exact
    /// position; such a key could still resolve into a splice targeting that slot.
    mutating func scanStep() throws -> ScanOutcome {
        while true {
            if isFrontSafeToRelease, let token = indentEngine.dequeue() {
                return .token(token)
            }
            if streamEndProduced {
                return .done
            }
            // A plain scalar in progress resumes from exactly where it stopped,
            // with no checkpoint at all: `runPlainScalarLoop` never leaves `reader`
            // or `state` in a state that needs rolling back (see
            // `PlainScalarSuspended`), so there is nothing to restore if it runs
            // dry again.
            switch try resumeSuspendedPlainScalarIfNeeded() {
            case .notSuspended:
                break
            case .tokenReady:
                continue
            case .stillDry:
                return .needMoreInput
            }
            if let outcome = try fetchNextTokenCheckpointed() {
                return outcome
            }
        }
    }

    /// A dry run can only ever be reported while `!reader.isFinal` (see
    /// `Reader.peek(_:)`), so once the caller has appended every chunk and
    /// called `markFinal()`, no fetch can ever need rolling back. Skipping
    /// the snapshot in that case matters: it is the entire cost difference
    /// between `scanStep()` and `scanAll()` once the whole input is already
    /// buffered (measured in StreamingScannerPerformanceTests), since neither
    /// reads nor mutations inside `fetchNextToken()` become any cheaper by
    /// being wrapped in an unused checkpoint. Returns `nil` to mean "keep
    /// looping in `scanStep()`", matching a fetch that produced no directly
    /// returnable outcome (a non-plain-scalar token was merely enqueued, or a
    /// document/directive token dispatched with no scalar involved).
    private mutating func fetchNextTokenCheckpointed() throws -> ScanOutcome? {
        guard !reader.isFinal else {
            try fetchNextToken()
            return nil
        }
        let checkpoint = self
        do {
            try fetchNextToken()
        } catch let suspended as PlainScalarSuspended {
            // `fetchPlainScalar()` calls `saveSimpleKey()` and sets
            // `simpleKeyAllowed = false` before ever calling `scanPlainScalar()`,
            // and a plain-scalar-only resume has no way to undo those: it is
            // only safe to keep them applied and resume past this suspend if
            // real progress (content) exists to preserve. `phase == .start`
            // with an empty `state` means the very first check inside the loop
            // (before `consumePlainContent` ever ran, e.g. `documentMarkerAhead()`'s
            // own lookahead went dry) suspended: nothing would be saved by
            // resuming instead of retrying from scratch, and retrying from
            // scratch is what correctly re-derives whether this token should
            // have been dispatched as a plain scalar at all. `dispatchToken()`'s
            // own `isDocumentMarker` check, called before `fetchPlainScalar()`
            // is ever reached, can suffer the exact same dry-lookahead problem
            // `runPlainScalarLoop` guards against; only a full retry re-checks it
            // with more data, since nothing here can undo a wrong dispatch.
            if suspended.phase == .start, suspended.state.string.isEmpty {
                self = checkpoint
                return .needMoreInput
            }
            suspendedPlainScalar = suspended
            return .needMoreInput
        } catch {
            if reader.ranDry {
                self = checkpoint
                return .needMoreInput
            }
            throw error
        }
        if reader.ranDry {
            self = checkpoint
            return .needMoreInput
        }
        return nil
    }

    private var isFrontSafeToRelease: Bool {
        guard !indentEngine.tokens.isEmpty else {
            return false
        }
        return !simpleKeys.contains { $0.possible && $0.tokenNumber == indentEngine.tokensParsed }
    }

    private enum ResumeOutcome {
        /// No plain scalar was suspended; `scanStep()` should fall through
        /// to its normal fetch path.
        case notSuspended
        /// The suspended scalar was resumed to completion and its token was
        /// enqueued; `scanStep()` should loop back to the top and try to
        /// release it, not fall through to a fresh fetch.
        case tokenReady
        /// The resume itself ran dry again; `scanStep()` should report
        /// `.needMoreInput`.
        case stillDry
    }

    /// Extracted out of `scanStep()` (which was tripping the cyclomatic
    /// complexity limit) rather than inlined, this also fixes a real O(N^2)
    /// bug the inline version had: extracting `state` and immediately
    /// overwriting `suspended.state` with a fresh, unrelated value before
    /// calling `resumePlainScalar` is required, not optional. A plain
    /// `resumePlainScalar(suspended)` call, passing the whole suspended
    /// value, would keep a reference to it alive (as an ordinary borrowed
    /// value parameter) for this call's entire duration regardless of what
    /// the callee does internally, leaving `state`'s string buffer with two
    /// live owners throughout `runPlainScalarLoop` and forcing copy-on-write
    /// to fully copy the accumulated scalar on every single append.
    /// Overwriting `suspended.state` here, before the call, drops that
    /// reference for good, so `state` is the sole owner by the time
    /// `resumePlainScalar` starts mutating it (see
    /// `StreamingScannerPerformanceTests.test_measureLongPlainScalar`, which
    /// caught this at a measured 29x versus the ~3-4x this fix gives).
    private mutating func resumeSuspendedPlainScalarIfNeeded() throws -> ResumeOutcome {
        guard var suspended = suspendedPlainScalar else {
            return .notSuspended
        }
        suspendedPlainScalar = nil
        var state = suspended.state
        suspended.state = PlainState(endMark: state.endMark)
        let indent = suspended.indent
        let startMark = suspended.startMark
        let phase = suspended.phase
        do {
            let token = try resumePlainScalar(&state, indent: indent, startMark: startMark, phase: phase)
            indentEngine.enqueue(token)
            return .tokenReady
        } catch let resuspended as PlainScalarSuspended {
            suspendedPlainScalar = resuspended
            return .stillDry
        }
    }
}