swift-foundation/Sources/FoundationEssentials/AttributedString/DiscontiguousAttributedSubstring.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2024 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#if FOUNDATION_FRAMEWORK
@_spi(Unstable) internal import CollectionsInternal
#elseif canImport(_RopeModule)
internal import _RopeModule
#elseif canImport(_FoundationCollections)
internal import _FoundationCollections
#endif

@dynamicMemberLookup
@available(FoundationPreview 6.2, *)
public struct DiscontiguousAttributedSubstring: Sendable {
    /// The guts of the base attributed string.
    internal var _guts: AttributedString.Guts
    
    internal var _indices: RangeSet<BigString.Index>
    
    internal var _identity: Int = 0
    
    internal init(_ guts: AttributedString.Guts, in indices: RangeSet<BigString.Index>) {
        self._guts = guts
        // Forcibly resolve bounds and round them down to nearest scalar boundary.
        var ranges = Array(indices.ranges)
        for i in ranges.indices {
            let slice = _guts.string.unicodeScalars[ranges[i]]
            ranges[i] = Range(uncheckedBounds: (slice.startIndex, slice.endIndex))
        }
        self._indices = RangeSet(ranges)
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring {
    public var base: AttributedString {
        return AttributedString(_guts)
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring : CustomStringConvertible {
    public var description: String {
        _guts.description(in: _indices)
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring : Equatable {
    public static func == (lhs: Self, rhs: Self) -> Bool {
        AttributedString.Guts._characterwiseIsEqual(lhs.runs, to: rhs.runs)
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring : Hashable {
    public func hash(into hasher: inout Hasher) {
        for range in _indices.ranges {
            _guts.characterwiseHash(in: range, into: &hasher)
        }
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring : AttributedStringAttributeMutation {
    internal mutating func ensureUniqueReference() {
        // Note: slices should never discard the data outside their bounds, so we must make a
        // copy of the entire base string here.
        //
        // (Discarding out-of-range data would change index values, interfere with "in-place"
        // mutations of slices via the subscript accessors, and it would confuse the semantics of
        // the `base` property.)
        if !isKnownUniquelyReferenced(&_guts) {
            _guts = _guts.copy()
        }
    }
    
    public mutating func setAttributes(_ attributes: AttributeContainer) {
        ensureUniqueReference()
        for range in _indices.ranges {
            _guts.setAttributes(attributes.storage, in: range)
        }
    }
    
    public mutating func mergeAttributes(_ attributes: AttributeContainer, mergePolicy:  AttributedString.AttributeMergePolicy = .keepNew) {
        ensureUniqueReference()
        for range in _indices.ranges {
            _guts.mergeAttributes(attributes, in: range, mergePolicy: mergePolicy)
        }
    }
    
    public mutating func replaceAttributes(_ attributes: AttributeContainer, with others: AttributeContainer) {
        guard attributes != others else {
            return
        }
        ensureUniqueReference()
        let hasConstrainedAttributes = attributes.storage.hasConstrainedAttributes || others.storage.hasConstrainedAttributes
        var fixupRanges = [Range<Int>]()
        for range in _indices.ranges {
            _guts.runs(in: range._utf8OffsetRange).updateEach(
                when: { $0.matches(attributes.storage) },
                with: { runAttributes, utf8Range in
                    for key in attributes.storage.keys {
                        runAttributes[key] = nil
                    }
                    runAttributes.mergeIn(others)
                    if hasConstrainedAttributes {
                        fixupRanges.append(utf8Range)
                    }
                })
        }
        for range in fixupRanges {
            // FIXME: Collect boundary constraints.
            _guts.enforceAttributeConstraintsAfterMutation(in: range, type: .attributes)
        }
    }
    
    public subscript(bounds: some RangeExpression<AttributedString.Index>) -> DiscontiguousAttributedSubstring {
        let characterView = AttributedString.CharacterView(_guts)
        let bounds = bounds.relative(to: characterView)._bstringRange
        if let first = _indices.ranges.first, let last = _indices.ranges.last, first.lowerBound <= bounds.lowerBound, last.upperBound >= bounds.upperBound {
            return DiscontiguousAttributedSubstring(_guts, in: _indices.intersection(RangeSet(bounds)))
        }
        preconditionFailure("Attributed string index range \(bounds) is out of bounds")
    }
    
    public subscript(bounds: RangeSet<AttributedString.Index>) -> DiscontiguousAttributedSubstring {
        let bounds = bounds._bstringIndices
        if bounds.ranges.isEmpty {
            return DiscontiguousAttributedSubstring(_guts, in: bounds)
        } else if let first = _indices.ranges.first,
                  let last = _indices.ranges.last,
                  let firstBounds = bounds.ranges.first,
                  let lastBounds = bounds.ranges.last,
                  first.lowerBound <= firstBounds.lowerBound,
                  last.upperBound >= lastBounds.upperBound {
            return DiscontiguousAttributedSubstring(_guts, in: _indices.intersection(bounds))
        }
        preconditionFailure("Attributed string index range \(bounds) is out of bounds")
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring {
    public var characters: DiscontiguousSlice<AttributedString.CharacterView> {
        AttributedString.CharacterView(_guts)[_indices._attributedStringIndices(version: _guts.version)]
    }
    
    public var unicodeScalars: DiscontiguousSlice<AttributedString.UnicodeScalarView> {
        AttributedString.UnicodeScalarView(_guts)[_indices._attributedStringIndices(version: _guts.version)]
    }
    
    public var runs: AttributedString.Runs {
        AttributedString.Runs(_guts, in: _indices)
    }
}

@available(FoundationPreview 6.2, *)
extension DiscontiguousAttributedSubstring {
    public subscript<K: AttributedStringKey>(_: K.Type) -> K.Value? where K.Value : Sendable {
        get {
            var result: AttributedString._AttributeValue?
            for range in _indices.ranges {
                guard let value = _guts.getUniformValue(in: range, key: K.self) else {
                    return nil
                }
                if let previous = result, previous != value {
                    return nil
                }
                result = value
            }
            return result?.rawValue(as: K.self)
        }
        set {
            ensureUniqueReference()
            if let v = newValue {
                for range in _indices.ranges {
                    _guts.setAttributeValue(v, forKey: K.self, in: range)
                }
            } else {
                for range in _indices.ranges {
                    _guts.removeAttributeValue(forKey: K.self, in: range)
                }
            }
        }
    }
    
    @inlinable // Trivial implementation, allows callers to optimize away the keypath allocation
    public subscript<K: AttributedStringKey>(
        dynamicMember keyPath: KeyPath<AttributeDynamicLookup, K>
    ) -> K.Value? where K.Value : Sendable {
        get { self[K.self] }
        set { self[K.self] = newValue }
    }
    
    public subscript<S: AttributeScope>(
        dynamicMember keyPath: KeyPath<AttributeScopes, S.Type>
    ) -> ScopedAttributeContainer<S> {
        get {
            var attributes = AttributedString._AttributeStorage()
            var first = true
            for range in _indices.ranges {
                let value = _guts.getUniformValues(in: range)
                guard !first else {
                    attributes = value
                    first = false
                    continue
                }
                attributes = attributes.filterWithoutInvalidatingDependents {
                    guard let value = value[$0.key] else { return false }
                    return value == $0.value
                }
                if attributes.isEmpty {
                    break
                }
            }
            return ScopedAttributeContainer(attributes)
            
        }
        _modify {
            ensureUniqueReference()
            var container = ScopedAttributeContainer<S>()
            defer {
                if let removedKey = container.removedKey {
                    for range in _indices.ranges {
                        _guts.removeAttributeValue(forKey: removedKey, in: range)
                    }
                } else {
                    for range in _indices.ranges {
                        _guts.mergeAttributes(AttributeContainer(container.storage), in: range)
                    }
                }
            }
            yield &container
        }
    }
}

@available(FoundationPreview 6.2, *)
extension AttributedString {
    public init(_ substring: DiscontiguousAttributedSubstring) {
        let created = AttributedString.Guts()
        for range in substring._indices.ranges {
            created.replaceSubrange(
                created.string.endIndex ..< created.string.endIndex,
                with: AttributedSubstring(substring._guts, in: range)
            )
        }
        self.init(created)
    }
}

@available(FoundationPreview 6.2, *)
extension AttributedStringProtocol {
    public subscript(_ indices: RangeSet<AttributedString.Index>) -> DiscontiguousAttributedSubstring {
        let range = Range(uncheckedBounds: (startIndex, endIndex))._bstringRange
        let newIndices = indices._bstringIndices.intersection(RangeSet(range))
        return DiscontiguousAttributedSubstring(__guts, in: newIndices)
    }
}

@available(FoundationPreview 6.2, *)
extension AttributedString {
    public subscript(_ indices: RangeSet<AttributedString.Index>) -> DiscontiguousAttributedSubstring {
        get {
            let range = Range(uncheckedBounds: (startIndex, endIndex))._bstringRange
            let newIndices = indices._bstringIndices.intersection(RangeSet(range))
            return DiscontiguousAttributedSubstring(_guts, in: newIndices)
        }
        _modify {
            ensureUniqueReference()
            let range = Range(uncheckedBounds: (startIndex, endIndex))._bstringRange
            let newIndices = indices._bstringIndices.intersection(RangeSet(range))
            var view = DiscontiguousAttributedSubstring(_guts, in: newIndices)
            let ident = Self._nextModifyIdentity
            view._identity = ident
            _guts = Guts() // Preserve uniqueness of view
            defer {
                if view._identity != ident {
                    fatalError("Mutating a DiscontiguousAttributedSubstring by replacing it with another from a different source is unsupported")
                }
                _guts = view._guts
            }
            yield &view
        }
        set {
            // The behavior of this function can be semantically confusing - we are required to have a setter in order to allow mutations such as attrStr[rangeSet].foregroundColor = .green, but wholesale replacements like attrStr[rangeSet] = otherAttrStr[otherRangeSet] then become possible
            // The general principle taken here is that the behavior of this function should be defined such that the _modify behavior is equivalent to a get followed by a mutation and then a set
            // Therefore, this function must interpolate the sliced contents of the newValue into the discontiguous chunks referenced by indices. This makes the behavior unclear when newValue has more values than what fit into indices, and therefore this causes a precondition failure.
            ensureUniqueReference()
            let other = AttributedString(newValue)
            var idxInOther = other.unicodeScalars.endIndex
            for range in indices.ranges.lazy.reversed() {
                let unicodeScalarLength = _guts.string.unicodeScalars.distance(from: range.lowerBound._value, to: range.upperBound._value)
                guard let startIdx = other.unicodeScalars.index(idxInOther, offsetBy: -unicodeScalarLength, limitedBy: other.unicodeScalars.startIndex) else {
                    preconditionFailure("Cannot set a DiscontiguousAttributedSubstring on a discontiguous slice of a different size")
                }
                let content = other[startIdx ..< idxInOther]
                _guts.replaceSubrange(range._bstringRange, with: content)
                idxInOther = startIdx
            }
            precondition(idxInOther == other.unicodeScalars.startIndex, "Cannot set a DiscontiguousAttributedSubstring on a discontiguous slice of a different size")
        }
    }
    
    public mutating func removeSubranges(_ subranges: RangeSet<Index>) {
        ensureUniqueReference()
        for range in subranges.ranges.lazy.reversed() {
            _guts.replaceSubrange(range._bstringRange, with: AttributedString())
        }
    }
}