Improve performance of enumerating constraint-bound attributes consistent across many runs (swiftlang#1226)

Author: jmschonfeld

Benchmarks/Benchmarks/AttributedString/BenchmarkAttributedString.swift

@@ -464,6 +464,43 @@ let benchmarks = {
         blackHole(hasher.finalize())
     }
 #endif
+    
+    let manyAttributesWithParagraph = {
+        var str = createManyAttributesString()
+        str.testParagraphConstrained = 2
+        return str
+    }()
+    
+    Benchmark("paragraphBoundSliceEnumeration-shortRuns") { benchmark in
+        for (value, range) in manyAttributesWithParagraph.runs[\.testParagraphConstrained] {
+            blackHole(value)
+        }
+    }
+    
+    Benchmark("paragraphBoundSliceEnumeration-shortRuns-reversed") { benchmark in
+        for (value, range) in manyAttributesWithParagraph.runs[\.testParagraphConstrained].reversed() {
+            blackHole(value)
+        }
+    }
+    
+    let longParagraphsString = {
+        var str = String(repeating: "a", count: 10000) + "\n"
+        str += String(repeating: "b", count: 10000) + "\n"
+        str += String(repeating: "c", count: 10000)
+        return AttributedString(str, attributes: AttributeContainer.testParagraphConstrained(1).testInt(2))
+    }()
+    
+    Benchmark("paragraphBoundSliceEnumeration-longRuns") { benchmark in
+        for (value, range) in longParagraphsString.runs[\.testParagraphConstrained] {
+            blackHole(value)
+        }
+    }
+    
+    Benchmark("paragraphBoundSliceEnumeration-longRuns-reversed") { benchmark in
+        for (value, range) in longParagraphsString.runs[\.testParagraphConstrained].reversed() {
+            blackHole(value)
+        }
+    }
 }
 
 

Sources/FoundationEssentials/AttributedString/AttributedString+Runs+AttributeSlices.swift

@@ -24,13 +24,13 @@ extension AttributedString.Runs {
 
         let runs: Runs
         let _names: [String]
-        let _constraints: [AttributeRunBoundaries]
+        let _constraints: Set<AttributeRunBoundaries?>
 
         init(runs: Runs) {
             self.runs = runs
             // FIXME: ☠️ Get these from a proper cache in runs._guts.
             _names = [T.name]
-            _constraints = T._constraintsInvolved
+            _constraints = [T.runBoundaries]
         }
 
         public struct Iterator: IteratorProtocol, Sendable {
@@ -162,13 +162,13 @@ extension AttributedString.Runs {
 
         let runs : Runs
         let _names: [String]
-        let _constraints: [AttributeRunBoundaries]
+        let _constraints: Set<AttributeRunBoundaries?>
 
         init(runs: Runs) {
             self.runs = runs
             // FIXME: ☠️ Get these from a proper cache in runs._guts.
             _names = [T.name, U.name]
-            _constraints = Array(_contents: T.runBoundaries, U.runBoundaries)
+            _constraints = [T.runBoundaries, U.runBoundaries]
         }
 
         public struct Iterator: IteratorProtocol, Sendable {
@@ -320,13 +320,13 @@ extension AttributedString.Runs {
 
         let runs : Runs
         let _names: [String]
-        let _constraints: [AttributeRunBoundaries]
+        let _constraints: Set<AttributeRunBoundaries?>
 
         init(runs: Runs) {
             self.runs = runs
             // FIXME: ☠️ Get these from a proper cache in runs._guts.
             _names = [T.name, U.name, V.name]
-            _constraints = Array(_contents: T.runBoundaries, U.runBoundaries, V.runBoundaries)
+            _constraints = [T.runBoundaries, U.runBoundaries, V.runBoundaries]
         }
 
         public struct Iterator: IteratorProtocol, Sendable {
@@ -490,14 +490,13 @@ extension AttributedString.Runs {
 
         let runs : Runs
         let _names: [String]
-        let _constraints: [AttributeRunBoundaries]
+        let _constraints: Set<AttributeRunBoundaries?>
 
         init(runs: Runs) {
             self.runs = runs
             // FIXME: ☠️ Get these from a proper cache in runs._guts.
             _names = [T.name, U.name, V.name, W.name]
-            _constraints = Array(
-                _contents: T.runBoundaries, U.runBoundaries, V.runBoundaries, W.runBoundaries)
+            _constraints = [T.runBoundaries, U.runBoundaries, V.runBoundaries, W.runBoundaries]
         }
 
         public struct Iterator: IteratorProtocol, Sendable {
@@ -675,18 +674,19 @@ extension AttributedString.Runs {
 
         let runs : Runs
         let _names: [String]
-        let _constraints: [AttributeRunBoundaries]
+        let _constraints: Set<AttributeRunBoundaries?>
 
         init(runs: Runs) {
             self.runs = runs
             // FIXME: ☠️ Get these from a proper cache in runs._guts.
             _names = [T.name, U.name, V.name, W.name]
-            _constraints = Array(
-                _contents: T.runBoundaries,
+            _constraints = [
+                T.runBoundaries,
                 U.runBoundaries,
                 V.runBoundaries,
                 W.runBoundaries,
-                X.runBoundaries)
+                X.runBoundaries
+            ]
         }
 
         public struct Iterator: IteratorProtocol, Sendable {
@@ -955,75 +955,3 @@ extension AttributedString.Runs {
 }
 
 #endif // FOUNDATION_FRAMEWORK
-
-extension RangeReplaceableCollection {
-    internal init(_contents item1: Element?) {
-        self.init()
-        if let item1 { self.append(item1) }
-    }
-
-    internal init(_contents item1: Element?, _ item2: Element?) {
-        self.init()
-        var c = 0
-        if item1 != nil { c &+= 1 }
-        if item2 != nil { c &+= 1 }
-        guard c > 0 else { return }
-        self.reserveCapacity(c)
-        if let item1 { self.append(item1) }
-        if let item2 { self.append(item2) }
-    }
-
-    internal init(_contents item1: Element?, _ item2: Element?, _ item3: Element?) {
-        self.init()
-        var c = 0
-        if item1 != nil { c &+= 1 }
-        if item2 != nil { c &+= 1 }
-        if item3 != nil { c &+= 1 }
-        guard c > 0 else { return }
-        self.reserveCapacity(c)
-        if let item1 { self.append(item1) }
-        if let item2 { self.append(item2) }
-        if let item3 { self.append(item3) }
-    }
-
-    internal init(
-        _contents item1: Element?, _ item2: Element?, _ item3: Element?, _ item4: Element?
-    ) {
-        self.init()
-        var c = 0
-        if item1 != nil { c &+= 1 }
-        if item2 != nil { c &+= 1 }
-        if item3 != nil { c &+= 1 }
-        if item4 != nil { c &+= 1 }
-        guard c > 0 else { return }
-        self.reserveCapacity(c)
-        if let item1 { self.append(item1) }
-        if let item2 { self.append(item2) }
-        if let item3 { self.append(item3) }
-        if let item4 { self.append(item4) }
-    }
-
-    internal init(
-        _contents item1: Element?,
-        _ item2: Element?,
-        _ item3: Element?,
-        _ item4: Element?,
-        _ item5: Element?
-    ) {
-        self.init()
-        var c = 0
-        if item1 != nil { c &+= 1 }
-        if item2 != nil { c &+= 1 }
-        if item3 != nil { c &+= 1 }
-        if item4 != nil { c &+= 1 }
-        if item5 != nil { c &+= 1 }
-        guard c > 0 else { return }
-        self.reserveCapacity(c)
-        if let item1 { self.append(item1) }
-        if let item2 { self.append(item2) }
-        if let item3 { self.append(item3) }
-        if let item4 { self.append(item4) }
-        if let item5 { self.append(item5) }
-    }
-}
-

Sources/FoundationEssentials/AttributedString/AttributedString+Runs.swift

@@ -481,19 +481,36 @@ extension AttributedString.Runs {
     internal func _slicedRunBoundary(
         after i: AttributedString.Index,
         attributeNames: [String],
-        constraints: [AttributeRunBoundaries],
+        constraints: Set<AttributeRunBoundaries?>,
         endOfCurrent: Bool
     ) -> AttributedString.Index {
         precondition(
             self._strBounds.contains(i._value),
             "AttributedString index is out of bounds")
         precondition(!attributeNames.isEmpty)
         let r = _guts.findRun(at: i._value)
+        let currentRangeIdx = _strBounds.rangeIdx(containing: i._value)
+        let currentRange = _strBounds.ranges[currentRangeIdx]
+        
+        guard constraints.count != 1 || constraints.contains(nil) else {
+            // We have a single constraint and attributes are guaranteed to be consistent between constraint boundaries
+            // This means that we will not break until the next constraint boundary, so we don't need to enumerate the actual run contents
+            let constraintBreak = _guts.string._firstConstraintBreak(in: i._value ..< currentRange.upperBound, with: constraints)
+            if !endOfCurrent && constraintBreak == currentRange.upperBound {
+                // No constraint break, return the next subrange start or the end index
+                if currentRangeIdx == _strBounds.ranges.count - 1 {
+                    return .init(currentRange.upperBound, version: _guts.version)
+                } else {
+                    return .init(_strBounds.ranges[currentRangeIdx + 1].lowerBound, version: _guts.version)
+                }
+            } else {
+                return .init(constraintBreak, version: _guts.version)
+            }
+        }
+        
         let endRun = _lastOfMatchingRuns(with: r.runIndex, comparing: attributeNames)
         let utf8End = endRun.utf8Offset + _guts.runs[endRun].length
         let strIndexEnd = _guts.string.utf8.index(r.start, offsetBy: utf8End - r.start.utf8Offset)
-        let currentRangeIdx = _strBounds.rangeIdx(containing: i._value)
-        let currentRange = _strBounds.ranges[currentRangeIdx]
         if strIndexEnd < currentRange.upperBound {
             // The coalesced run ends within the current range, so just look for the next break in the coalesced run
             return .init(_guts.string._firstConstraintBreak(in: i._value ..< strIndexEnd, with: constraints), version: _guts.version)
@@ -520,7 +537,7 @@ extension AttributedString.Runs {
     internal func _slicedRunBoundary(
         before i: AttributedString.Index,
         attributeNames: [String],
-        constraints: [AttributeRunBoundaries],
+        constraints: Set<AttributeRunBoundaries?>,
         endOfPrevious: Bool
     ) -> AttributedString.Index {
         precondition(
@@ -545,6 +562,11 @@ extension AttributedString.Runs {
             currentStringIdx = currentRange.upperBound
             if endOfPrevious { return .init(currentStringIdx, version: _guts.version) }
         }
+        
+        guard constraints.count != 1 || constraints.contains(nil) else {
+            return .init(_guts.string._lastConstraintBreak(in: currentRange.lowerBound ..< currentStringIdx, with: constraints), version: _guts.version)
+        }
+        
         let beforeStringIdx = _guts.string.utf8.index(before: currentStringIdx)
         let r = _guts.runs.index(atUTF8Offset: beforeStringIdx.utf8Offset)
         let startRun = _firstOfMatchingRuns(with: r.index, comparing: attributeNames)
@@ -561,7 +583,7 @@ extension AttributedString.Runs {
     internal func _slicedRunBoundary(
         roundingDown i: AttributedString.Index,
         attributeNames: [String],
-        constraints: [AttributeRunBoundaries]
+        constraints: Set<AttributeRunBoundaries?>
     ) -> (index: AttributedString.Index, runIndex: AttributedString._InternalRuns.Index) {
         precondition(
             _strBounds.contains(i._value) || i._value == endIndex._stringIndex,
@@ -571,8 +593,22 @@ extension AttributedString.Runs {
         if r.runIndex.offset == endIndex._runOffset {
             return (i, r.runIndex)
         }
-        let startRun = _firstOfMatchingRuns(with: r.runIndex, comparing: attributeNames)
         let currentRange = _strBounds.ranges[_strBounds.rangeIdx(containing: i._value)]
+        
+        guard constraints.count != 1 || constraints.contains(nil) else {
+            let nextIndex = _guts.string.unicodeScalars.index(after: i._value)
+            let constraintBreak = _guts.string._lastConstraintBreak(in: currentRange.lowerBound ..< nextIndex, with: constraints)
+            var runIdx = r.runIndex
+            while runIdx.utf8Offset > constraintBreak.utf8Offset {
+                _guts.runs.formIndex(before: &runIdx)
+            }
+            return (
+                .init(constraintBreak, version: _guts.version),
+                runIdx
+            )
+        }
+        
+        let startRun = _firstOfMatchingRuns(with: r.runIndex, comparing: attributeNames)
         let stringStart = Swift.max(
             _guts.string.utf8.index(r.start, offsetBy: startRun.utf8Offset - r.start.utf8Offset),
             currentRange.lowerBound)
@@ -587,9 +623,9 @@ extension AttributedString.Runs {
 extension BigString {
     internal func _firstConstraintBreak(
         in range: Range<Index>,
-        with constraints: [AttributedString.AttributeRunBoundaries]
+        with constraints: Set<AttributedString.AttributeRunBoundaries?>
     ) -> Index {
-        guard !constraints.isEmpty, !range.isEmpty else { return range.upperBound }
+        guard !range.isEmpty else { return range.upperBound }
 
         var r = range
         if
@@ -602,7 +638,7 @@ extension BigString {
         if constraints._containsScalarConstraint {
             // Note: we need to slice runs on matching scalars even if they don't carry
             // the attributes we're looking for.
-            let scalars: [UnicodeScalar] = constraints.compactMap { $0._constrainedScalar }
+            let scalars: [UnicodeScalar] = constraints.compactMap { $0?._constrainedScalar }
             if let firstBreak = self.unicodeScalars[r]._findFirstScalarBoundary(for: scalars) {
                 r = r.lowerBound ..< firstBreak
             }
@@ -613,9 +649,9 @@ extension BigString {
 
     internal func _lastConstraintBreak(
         in range: Range<Index>,
-        with constraints: [AttributedString.AttributeRunBoundaries]
+        with constraints: Set<AttributedString.AttributeRunBoundaries?>
     ) -> Index {
-        guard !constraints.isEmpty, !range.isEmpty else { return range.lowerBound }
+        guard !range.isEmpty else { return range.lowerBound }
 
         var r = range
         if
@@ -628,7 +664,7 @@ extension BigString {
         if constraints._containsScalarConstraint {
             // Note: we need to slice runs on matching scalars even if they don't carry
             // the attributes we're looking for.
-            let scalars: [UnicodeScalar] = constraints.compactMap { $0._constrainedScalar }
+            let scalars: [UnicodeScalar] = constraints.compactMap { $0?._constrainedScalar }
             if let lastBreak = self.unicodeScalars[r]._findLastScalarBoundary(for: scalars) {
                 r = lastBreak ..< r.upperBound
             }

Sources/FoundationEssentials/AttributedString/AttributedStringAttribute.swift

@@ -123,14 +123,6 @@ extension AttributedStringKey {
     public static var invalidationConditions : Set<AttributedString.AttributeInvalidationCondition>? { nil }
 }
 
-extension AttributedStringKey {
-    // FIXME: ☠️ Allocating an Array here is not a good idea.
-    static var _constraintsInvolved: [AttributedString.AttributeRunBoundaries] {
-        guard let rb = runBoundaries else { return [] }
-        return [rb]
-    }
-}
-
 // MARK: Attribute Scopes
 
 @dynamicMemberLookup @frozen

Sources/FoundationEssentials/AttributedString/AttributedStringAttributeConstrainingBehavior.swift

@@ -83,6 +83,13 @@ extension Collection where Element == AttributedString.AttributeRunBoundaries {
     }
 }
 
+@available(macOS 13, iOS 16, tvOS 16, watchOS 9, *)
+extension Collection where Element == AttributedString.AttributeRunBoundaries? {
+    var _containsScalarConstraint: Bool {
+        self.contains { $0?._isScalarConstrained ?? false }
+    }
+}
+
 @available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)
 extension AttributedString.Guts {