PureYAML/Sources/Parsing/EventStreamComposer.swift

extension PureYAML.Parsing {
    struct EventStreamComposer {
        let scalarParser: Parser
        let mode: EventStreamComposerMode
        // When true, the composer builds a parallel located tree alongside each value. The
        // default value path never sets this, so it stays byte-identical to before.
        var tracksLocations = false
        var anchors: [String: PureYAML.Model.Value] = [:]
        /// Each anchor's cached EXPANDED node count; see `chargeExpansion` for why.
        var anchorSizes: [String: Int] = [:]
        // The expanded-node budget remaining for this document; see `chargeExpansion`.
        var expandedNodeBudget = 5_000_000
        var containers: [EventStreamContainer] = []
        var documentValues: [PureYAML.Model.Value] = []
        var documentLocated: [PureYAML.Model.LocatedValue] = []
        var currentDocumentValue: PureYAML.Model.Value?
        var currentDocumentLocated: PureYAML.Model.LocatedValue?
        var didSeeStreamStart = false
        var didSeeStreamEnd = false

        init(
            scalarParser: Parser,
            mode: EventStreamComposerMode,
            tracksLocations: Bool = false,
        ) {
            self.scalarParser = scalarParser
            self.mode = mode
            self.tracksLocations = tracksLocations
        }

        mutating func consume(_ event: PureYAML.Parsing.Event) throws {
            switch event {
            case .streamStart:
                didSeeStreamStart = true
            case .streamEnd:
                didSeeStreamEnd = true
            case let .documentStart(_, mark):
                try startDocument(mark: mark)
            case let .documentEnd(_, mark):
                try endDocument(mark: mark)
            case let .scalar(value, anchor, tag, style, mark, endMark):
                try consumeScalar(value, anchor: anchor, tag: tag, style: style, marks: (mark, endMark))
            case let .alias(anchor, mark):
                try consumeAlias(anchor, mark: mark)
            case let .sequenceStart(anchor, _, _, mark):
                containers.append(.init(kind: .sequence, anchor: anchor, mark: mark))
            case let .sequenceEnd(mark):
                try completeSequence(mark: mark)
            case let .mappingStart(anchor, _, _, mark):
                containers.append(.init(kind: .mapping, anchor: anchor, mark: mark))
            case let .mappingEnd(mark):
                try completeMapping(mark: mark)
            }
        }

        mutating func composedValue() throws -> PureYAML.Model.Value {
            guard didSeeStreamStart, didSeeStreamEnd else {
                throw unexpectedEvent(expected: "complete YAML stream", mark: .start)
            }
            // An empty or comment-only stream has zero documents and maps to null;
            // a single document yields its value; more than one is not a single
            // document and is rejected.
            guard documentValues.count <= 1 else {
                throw unexpectedEvent(expected: "single YAML document", mark: .start)
            }
            return documentValues.first ?? .null
        }

        mutating func composedStream() throws -> [PureYAML.Stream.Document] {
            guard didSeeStreamStart, didSeeStreamEnd else {
                throw unexpectedEvent(expected: "complete YAML stream", mark: .start)
            }
            return documentValues.enumerated().map { index, value in
                PureYAML.Stream.Document(index: index, value: value)
            }
        }

        /// The single document's located tree (mirrors ``composedValue()``). Only meaningful
        /// when the composer was created with `tracksLocations: true`.
        mutating func composedLocated() throws -> PureYAML.Model.LocatedValue {
            guard didSeeStreamStart, didSeeStreamEnd else {
                throw unexpectedEvent(expected: "complete YAML stream", mark: .start)
            }
            guard documentLocated.count <= 1 else {
                throw unexpectedEvent(expected: "single YAML document", mark: .start)
            }
            return documentLocated.first ?? .scalar(.null, .init(at: .init(line: 1, column: 1)))
        }

        /// The stream's located trees, one per document (mirrors ``composedStream()``). Only
        /// meaningful when the composer was created with `tracksLocations: true`.
        mutating func composedLocatedStream() throws -> [PureYAML.Model.LocatedValue] {
            guard didSeeStreamStart, didSeeStreamEnd else {
                throw unexpectedEvent(expected: "complete YAML stream", mark: .start)
            }
            return documentLocated
        }
    }
}

private extension PureYAML.Parsing.EventStreamComposer {
    mutating func startDocument(mark: PureYAML.Parsing.Mark) throws {
        guard containers.isEmpty else {
            throw unexpectedEvent(expected: "document content end", mark: mark)
        }
        if mode == .singleDocument, !documentValues.isEmpty {
            throw PureYAML.Parsing.ParseError.unsupportedMultiDocumentStream(line: mark.line)
        }
        anchors = [:]
        anchorSizes = [:]
        expandedNodeBudget = 5_000_000
        currentDocumentValue = nil
    }

    mutating func endDocument(mark: PureYAML.Parsing.Mark) throws {
        guard containers.isEmpty else {
            throw unexpectedEvent(expected: "document content end", mark: mark)
        }
        documentValues.append(currentDocumentValue ?? .null)
        currentDocumentValue = nil
        if tracksLocations {
            documentLocated.append(currentDocumentLocated ?? .scalar(.null, range(at: mark)))
            currentDocumentLocated = nil
        }
    }

    mutating func consumeScalar(
        _ value: String,
        anchor: String?,
        tag: String?,
        style: PureYAML.Parsing.ScalarStyle,
        marks: (start: PureYAML.Parsing.Mark, end: PureYAML.Parsing.Mark),
    ) throws {
        let model = try scalarParser.composeScalarValue(value, tag: tag, style: style, mark: marks.start)
        try chargeExpansion(1, mark: marks.start)
        if let anchor {
            anchors[anchor] = model
            anchorSizes[anchor] = 1
        }
        try complete(.init(
            value: model,
            mappingKey: .init(
                keyNode: .string(value),
                scalarValue: value,
                tag: scalarParser.normalizedTag(tag),
                style: style,
            ),
            mark: marks.start,
            located: tracksLocations ? .scalar(model, range(start: marks.start, end: marks.end)) : nil,
        ), cost: 1)
    }

    mutating func consumeAlias(
        _ anchor: String,
        mark: PureYAML.Parsing.Mark,
    ) throws {
        guard let value = anchors[anchor] else {
            throw PureYAML.Parsing.ParseError.undefinedAlias(
                anchor: anchor,
                line: mark.line,
                column: mark.column,
            )
        }
        // The target's own cached expanded size (itself plus every descendant, its own nested
        // aliases already resolved): an alias to a value built from other aliases charges the
        // FULL compounded cost here, in O(1), which is exactly what makes a nested bomb's true
        // size visible without ever materializing it.
        let cost = anchorSizes[anchor] ?? 1
        try chargeExpansion(cost, mark: mark)
        try complete(.init(
            value: value,
            mappingKey: .init(
                keyNode: PureYAML.Model.Key(value: value),
                scalarValue: nil,
                tag: nil,
                style: nil,
            ),
            mark: mark,
            located: tracksLocations ? .synthesized(from: value, at: range(at: mark)) : nil,
        ), cost: cost)
    }

    mutating func completeSequence(mark: PureYAML.Parsing.Mark) throws {
        guard let container = containers.popLast(), container.kind == .sequence else {
            throw unexpectedEvent(expected: "sequence start", mark: mark)
        }
        // Only the container's own wrapper is a NEW charge here; every child already charged its
        // own cost (including any alias expansion) as it completed, accumulated into
        // `logicalSize`. The cache for a later alias to THIS container needs the full total.
        try chargeExpansion(1, mark: container.mark)
        let totalCost = container.logicalSize + 1
        let value = PureYAML.Model.Value.sequence(container.values.map(\.value))
        store(value, anchor: container.anchor, cost: totalCost)
        try complete(.init(
            value: value,
            mappingKey: .init(
                keyNode: PureYAML.Model.Key(value: value),
                scalarValue: nil,
                tag: nil,
                style: nil,
            ),
            mark: container.mark,
            sequenceItems: container.values,
            located: tracksLocations
                ? .sequence(container.locatedValues, range(start: container.mark, end: mark))
                : nil,
        ), cost: totalCost)
    }

    mutating func completeMapping(mark: PureYAML.Parsing.Mark) throws {
        guard let container = containers.popLast(), container.kind == .mapping else {
            throw unexpectedEvent(expected: "mapping start", mark: mark)
        }
        guard container.pendingKey == nil else {
            throw unexpectedEvent(expected: "mapping value", mark: mark)
        }
        try chargeExpansion(1, mark: container.mark)
        let totalCost = container.logicalSize + 1
        let pairs = mergedPairs(
            from: container.mergeSources,
            localPairs: container.localPairs,
        )
        let value = PureYAML.Model.Value.mapping(.init(pairs))
        store(value, anchor: container.anchor, cost: totalCost)
        let located: PureYAML.Model.LocatedValue? = tracksLocations
            ? .mapping(
                mergedLocatedPairs(from: container.mergeLocatedSources, localPairs: container.localLocatedPairs),
                range(start: container.mark, end: mark),
            )
            : nil
        try complete(.init(
            value: value,
            mappingKey: .init(
                keyNode: PureYAML.Model.Key(value: value),
                scalarValue: nil,
                tag: nil,
                style: nil,
            ),
            mark: container.mark,
            located: located,
        ), cost: totalCost)
    }

    mutating func complete(_ node: PureYAML.Parsing.EventStreamComposedNode, cost: Int) throws {
        guard !containers.isEmpty else {
            currentDocumentValue = node.value
            if tracksLocations {
                currentDocumentLocated = node.located
            }
            return
        }

        switch containers[containers.count - 1].kind {
        case .sequence:
            containers[containers.count - 1].values.append(node)
            containers[containers.count - 1].logicalSize += cost
            if tracksLocations {
                containers[containers.count - 1].locatedValues.append(locatedNode(node))
            }
        case .mapping:
            try completeMappingNode(node)
            containers[containers.count - 1].logicalSize += cost
        }
    }

    mutating func completeMappingNode(_ node: PureYAML.Parsing.EventStreamComposedNode) throws {
        let index = containers.count - 1
        guard let key = containers[index].pendingKey else {
            containers[index].pendingKey = node.mappingKey
            if tracksLocations {
                containers[index].pendingKeyNode = node
            }
            return
        }

        if key.isMergeKey {
            try containers[index].mergeSources.append(contentsOf: mergeSources(
                from: node,
            ))
            if tracksLocations {
                containers[index].mergeLocatedSources.append(contentsOf: locatedMergeSources(from: node))
            }
        } else {
            containers[index].localPairs.append(.init(keyNode: key.keyNode, value: node.value))
            if tracksLocations {
                let keyNode = containers[index].pendingKeyNode
                containers[index].localLocatedPairs.append(.init(
                    keyNode: key.keyNode,
                    keyRange: keyNode.flatMap(\.located)?.range ?? locatedNode(node).range,
                    value: locatedNode(node),
                ))
            }
        }
        containers[index].pendingKey = nil
        if tracksLocations {
            containers[index].pendingKeyNode = nil
        }
    }

    mutating func store(
        _ value: PureYAML.Model.Value,
        anchor: String?,
        cost: Int,
    ) {
        guard let anchor else {
            return
        }
        anchors[anchor] = value
        anchorSizes[anchor] = cost
    }

    func mergeSources(
        from node: PureYAML.Parsing.EventStreamComposedNode,
    ) throws -> [[PureYAML.Model.Pair]] {
        switch node.value {
        case let .mapping(mapping):
            return [mapping.pairs]
        case let .sequence(values):
            guard !node.sequenceItems.isEmpty else {
                return try values.map { value in
                    guard case let .mapping(mapping) = value else {
                        throw invalidMergeValue(mark: node.mark)
                    }
                    return mapping.pairs
                }
            }
            return try values.enumerated().map { index, value in
                let mark = node.sequenceItems[index].mark
                guard case let .mapping(mapping) = value else {
                    throw invalidMergeValue(mark: mark)
                }
                return mapping.pairs
            }
        case .bool, .double, .int, .null, .string:
            throw invalidMergeValue(mark: node.mark)
        }
    }

    func mergedPairs(
        from mergeSources: [[PureYAML.Model.Pair]],
        localPairs: [PureYAML.Model.Pair],
    ) -> [PureYAML.Model.Pair] {
        guard !mergeSources.isEmpty else {
            return localPairs
        }

        let localKeys = Set(localPairs.map(\.keyNode))
        var inheritedKeys = Set<PureYAML.Model.Key>()
        var inheritedPairs: [PureYAML.Model.Pair] = []
        for source in mergeSources {
            var sourceKeys = Set<PureYAML.Model.Key>()
            for pair in source {
                guard !localKeys.contains(pair.keyNode),
                      !inheritedKeys.contains(pair.keyNode)
                else {
                    continue
                }
                inheritedPairs.append(pair)
                sourceKeys.insert(pair.keyNode)
            }
            inheritedKeys.formUnion(sourceKeys)
        }
        return inheritedPairs + localPairs
    }

    func invalidMergeValue(
        mark: PureYAML.Parsing.Mark,
    ) -> PureYAML.Parsing.ParseError {
        PureYAML.Parsing.ParseError.invalidMergeValue(
            line: mark.line,
            column: mark.column,
        )
    }

    func unexpectedEvent(
        expected: String,
        mark: PureYAML.Parsing.Mark,
    ) -> PureYAML.Parsing.ParseError {
        PureYAML.Parsing.ParseError.unexpectedEvent(
            expected: expected,
            actual: "streaming event",
            line: mark.line,
            column: mark.column,
        )
    }
}