Sapling/Sources/SaplingEditor/DocumentLineIndex.swift

import Foundation

public enum LineEndingStrategy: String, Hashable, Sendable {
    case lf
    case crlf
    case cr
    case none

    public var utf16Length: Int {
        switch self {
        case .lf, .cr:
            return 1
        case .crlf:
            return 2
        case .none:
            return 0
        }
    }
}

public struct DocumentLineBoundary: Hashable, Sendable {
    public var index: Int
    public var contentRange: NSRange
    public var lineEnding: LineEndingStrategy

    public init(index: Int, contentRange: NSRange, lineEnding: LineEndingStrategy) {
        self.index = index
        self.contentRange = contentRange
        self.lineEnding = lineEnding
    }

    public var lineEndingRange: NSRange {
        NSRange(location: contentRange.upperBound, length: lineEnding.utf16Length)
    }

    public var nextLineLocation: Int {
        contentRange.upperBound + lineEnding.utf16Length
    }
}

public struct DocumentLineIndexEdit: Hashable, Sendable {
    public var range: NSRange
    public var replacement: String

    public init(range: NSRange, replacement: String) {
        self.range = range
        self.replacement = replacement
    }

    public var replacementUTF16Length: Int {
        (replacement as NSString).length
    }

    public var insertedRangeInEditedDocument: NSRange {
        NSRange(location: range.location, length: replacementUTF16Length)
    }
}

public struct DocumentLineIndexEditResult: Hashable, Sendable {
    public var replacedLineRange: Range<Int>
    public var insertedLineRange: Range<Int>
    public var locationDelta: Int

    public var insertedLineIndexes: [Int] {
        Array(insertedLineRange)
    }
}

public struct DocumentLineIndex: Hashable, Sendable {
    public var source: String
    public var boundaries: [DocumentLineBoundary]

    public init(source: String) {
        self.source = source
        self.boundaries = Self.scanBoundaries(in: source)
    }

    public var lineCount: Int {
        boundaries.count
    }

    public func boundary(at index: Int) -> DocumentLineBoundary? {
        guard boundaries.indices.contains(index) else { return nil }
        return boundaries[index]
    }

    public func lineIndex(containing location: Int) -> Int {
        guard !boundaries.isEmpty else { return 0 }

        let clampedLocation = max(0, min(location, source.utf16.count))

        var lowerBound = 0
        var upperBound = boundaries.count - 1
        while lowerBound < upperBound {
            let midpoint = (lowerBound + upperBound) / 2
            if clampedLocation < boundaries[midpoint].nextLineLocation {
                upperBound = midpoint
            } else {
                lowerBound = midpoint + 1
            }
        }

        return boundaries[lowerBound].index
    }

    public func lineStartOffset(forLine index: Int) -> Int? {
        boundary(at: index)?.contentRange.location
    }

    public func lineContentRange(forLine index: Int) -> NSRange? {
        boundary(at: index)?.contentRange
    }

    public func lineIndexesAffected(by changedRange: NSRange, includingNeighborLines: Bool = true) -> [Int] {
        guard !boundaries.isEmpty else { return [] }

        let sourceLength = source.utf16.count
        let startLocation = max(0, min(changedRange.location, sourceLength))
        let endLocation: Int
        if changedRange.length == 0 {
            endLocation = startLocation
        } else {
            endLocation = max(startLocation, min(changedRange.upperBound - 1, sourceLength))
        }

        let startLine = lineIndex(containing: startLocation)
        let endLine = lineIndex(containing: endLocation)
        let neighborOffset = includingNeighborLines ? 1 : 0
        let lowerBound = max(0, min(startLine, endLine) - neighborOffset)
        let upperBound = min(boundaries.count - 1, max(startLine, endLine) + neighborOffset)

        return Array(lowerBound...upperBound)
    }

    public func editorLine(at index: Int, activeLineIndex: Int) -> EditorLine? {
        guard let boundary = boundary(at: index) else { return nil }
        return editorLine(for: boundary, activeLineIndex: activeLineIndex)
    }

    public func editorLines(for indexes: some Sequence<Int>, activeLineIndex: Int) -> [EditorLine] {
        indexes.compactMap { editorLine(at: $0, activeLineIndex: activeLineIndex) }
    }

    public func editorLines(activeLineIndex: Int) -> [EditorLine] {
        boundaries.map { boundary in
            editorLine(for: boundary, activeLineIndex: activeLineIndex)
        }
    }

    @discardableResult
    public mutating func replaceCharacters(in range: NSRange, with replacement: String) -> DocumentLineIndexEditResult {
        replace(DocumentLineIndexEdit(range: range, replacement: replacement))
    }

    @discardableResult
    public mutating func replace(_ edit: DocumentLineIndexEdit) -> DocumentLineIndexEditResult {
        let oldSource = source as NSString
        let oldLength = oldSource.length
        let range = clampedRange(edit.range, sourceLength: oldLength)
        let updatedSource = oldSource.replacingCharacters(in: range, with: edit.replacement)
        return replace(edit, updatedSource: updatedSource)
    }

    @discardableResult
    public mutating func replace(_ edit: DocumentLineIndexEdit, updatedSource: String) -> DocumentLineIndexEditResult {
        let oldSource = source as NSString
        let oldLength = oldSource.length
        let range = clampedRange(edit.range, sourceLength: oldLength)
        let replacement = edit.replacement
        let replacementLength = (replacement as NSString).length
        let locationDelta = replacementLength - range.length

        let lowerProbe = max(0, range.location - 1)
        let upperProbe = min(oldLength, range.upperBound + 1)
        let lowerLineIndex = max(0, lineIndex(containing: lowerProbe) - 1)
        let upperLineIndex = min(boundaries.count - 1, lineIndex(containing: upperProbe) + 1)
        let replacedLineRange = lowerLineIndex..<(upperLineIndex + 1)
        let scanStart = boundaries[lowerLineIndex].contentRange.location
        let oldScanEnd = boundaries[upperLineIndex].nextLineLocation

        source = updatedSource

        let newSource = source as NSString
        let newScanEnd = max(scanStart, min(newSource.length, oldScanEnd + locationDelta))
        let scannedBoundaries = Self.scanBoundaries(
            in: newSource,
            range: NSRange(location: scanStart, length: newScanEnd - scanStart),
            startingIndex: lowerLineIndex,
            includeTrailingBlankLine: newScanEnd == newSource.length
        )

        boundaries.replaceSubrange(replacedLineRange, with: scannedBoundaries)

        let insertedLineRange = lowerLineIndex..<(lowerLineIndex + scannedBoundaries.count)
        let indexDelta = scannedBoundaries.count - replacedLineRange.count
        let followingStart = insertedLineRange.upperBound
        if followingStart < boundaries.count {
            for boundaryIndex in followingStart..<boundaries.count {
                boundaries[boundaryIndex].index += indexDelta
                boundaries[boundaryIndex].contentRange.location += locationDelta
            }
        }

        return DocumentLineIndexEditResult(
            replacedLineRange: replacedLineRange,
            insertedLineRange: insertedLineRange,
            locationDelta: locationDelta
        )
    }

    private func clampedRange(_ range: NSRange, sourceLength: Int) -> NSRange {
        let location = max(0, min(range.location, sourceLength))
        return NSRange(
            location: location,
            length: max(0, min(range.length, sourceLength - location))
        )
    }

    private func editorLine(for boundary: DocumentLineBoundary, activeLineIndex: Int) -> EditorLine {
        let nsSource = source as NSString
        return EditorLine(
            index: boundary.index,
            source: nsSource.substring(with: boundary.contentRange),
            range: boundary.contentRange,
            mode: boundary.index == activeLineIndex ? .source : .rendered
        )
    }

    private static func scanBoundaries(in source: String) -> [DocumentLineBoundary] {
        let nsSource = source as NSString
        return scanBoundaries(
            in: nsSource,
            range: NSRange(location: 0, length: nsSource.length),
            startingIndex: 0,
            includeTrailingBlankLine: true
        )
    }

    private static func scanBoundaries(
        in source: NSString,
        range: NSRange,
        startingIndex: Int,
        includeTrailingBlankLine: Bool
    ) -> [DocumentLineBoundary] {
        let sourceLength = range.upperBound
        guard sourceLength > 0 else {
            return [
                DocumentLineBoundary(
                    index: startingIndex,
                    contentRange: NSRange(location: range.location, length: 0),
                    lineEnding: .none
                )
            ]
        }

        var boundaries: [DocumentLineBoundary] = []
        var lineStart = range.location
        var lineIndex = startingIndex
        var endedWithLineEnding = false

        while lineStart < sourceLength {
            var cursor = lineStart
            while cursor < sourceLength,
                  !isLineEndingStart(source.character(at: cursor)) {
                cursor += 1
            }

            let contentRange = NSRange(location: lineStart, length: cursor - lineStart)
            if cursor < sourceLength {
                let lineEnding = lineEndingStrategy(at: cursor, in: source, scanEnd: sourceLength)
                boundaries.append(DocumentLineBoundary(
                    index: lineIndex,
                    contentRange: contentRange,
                    lineEnding: lineEnding
                ))
                cursor += lineEnding.utf16Length
                lineStart = cursor
                endedWithLineEnding = cursor == sourceLength
            } else {
                boundaries.append(DocumentLineBoundary(
                    index: lineIndex,
                    contentRange: contentRange,
                    lineEnding: .none
                ))
                lineStart = sourceLength
                endedWithLineEnding = false
            }

            lineIndex += 1
        }

        if endedWithLineEnding, includeTrailingBlankLine {
            boundaries.append(DocumentLineBoundary(
                index: lineIndex,
                contentRange: NSRange(location: sourceLength, length: 0),
                lineEnding: .none
            ))
        }

        return boundaries
    }

    private static func lineEndingStrategy(at location: Int, in source: NSString, scanEnd: Int? = nil) -> LineEndingStrategy {
        let character = source.character(at: location)
        if character == carriageReturnUTF16 {
            let nextLocation = location + 1
            let upperBound = scanEnd ?? source.length
            if nextLocation < upperBound,
               source.character(at: nextLocation) == lineFeedUTF16 {
                return .crlf
            }
            return .cr
        }
        return .lf
    }

    private static func isLineEndingStart(_ character: unichar) -> Bool {
        character == lineFeedUTF16 || character == carriageReturnUTF16
    }

    private static let lineFeedUTF16: unichar = 10
    private static let carriageReturnUTF16: unichar = 13
}

private extension NSRange {
    var upperBound: Int {
        location + length
    }
}