From 1e71e90dc0c13bb6d7542a9bb54f8287fa274c2f Mon Sep 17 00:00:00 2001
From: Kevin Sweeney <kevin_sweeney@apple.com>
Date: Wed, 9 Oct 2019 19:01:57 +0100
Subject: [PATCH] [5.1] [SR-7749] Poor performance of
 String.replacingOccurrences(of:with:) in corelibs-foundation

This supercedes #1620.

(cherry picked from commit 4cd59de72296e109af4977fc7e502de59e20dcd7)
---
 Foundation/NSString.swift         | 44 ++++++++++++++++++-----
 TestFoundation/TestNSString.swift | 58 +++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 8 deletions(-)

diff --git a/Foundation/NSString.swift b/Foundation/NSString.swift
index 65bec45c97..bd1216b024 100644
--- a/Foundation/NSString.swift
+++ b/Foundation/NSString.swift
@@ -1452,6 +1452,18 @@ extension NSMutableString {
         }
         return 0
     }
+
+    private static func makeFindResultsRangeIterator(findResults: CFArray, count: Int, backwards: Bool) -> AnyIterator<NSRange> {
+        var index = 0
+        return AnyIterator<NSRange>() { () -> NSRange? in
+            defer { index += 1 }
+            if index < count {
+                let rangePtr = CFArrayGetValueAtIndex(findResults, backwards ? count - index - 1 : index)
+                return NSRange(rangePtr!.load(as: CFRange.self))
+            }
+            return nil
+        }
+    }
     
     public func replaceOccurrences(of target: String, with replacement: String, options: CompareOptions = [], range searchRange: NSRange) -> Int {
         let backwards = options.contains(.backwards)
@@ -1462,19 +1474,35 @@ extension NSMutableString {
         if options.contains(.regularExpression) {
             return _replaceOccurrencesOfRegularExpressionPattern(target, withTemplate:replacement, options:options, range: searchRange)
         }
-        
 
-        if let findResults = CFStringCreateArrayWithFindResults(kCFAllocatorSystemDefault, _cfObject, target._cfObject, CFRange(searchRange), options._cfValue(true)) {
-            let numOccurrences = CFArrayGetCount(findResults)
-            for cnt in 0..<numOccurrences {
-                let rangePtr = CFArrayGetValueAtIndex(findResults, backwards ? cnt : numOccurrences - cnt - 1)
-                replaceCharacters(in: NSRange(rangePtr!.load(as: CFRange.self)), with: replacement)
+        guard let findResults = CFStringCreateArrayWithFindResults(kCFAllocatorSystemDefault, _cfObject, target._cfObject, CFRange(searchRange), options._cfValue(true)) else {
+             return 0
+        }
+        let numOccurrences = CFArrayGetCount(findResults)
+
+        let rangeIterator = NSMutableString.makeFindResultsRangeIterator(findResults: findResults, count: numOccurrences, backwards: backwards)
+
+        guard type(of: self) == NSMutableString.self else {
+            // If we're dealing with non NSMutableString, mutations must go through `replaceCharacters` (documented behavior)
+            for range in rangeIterator {
+                replaceCharacters(in: range, with: replacement)
             }
+
             return numOccurrences
-        } else {
-            return 0
         }
 
+        var newStorage = Substring()
+        var sourceStringCurrentIndex = _storage.startIndex
+        for range in rangeIterator {
+            let matchStartIndex = String.Index(utf16Offset: range.location, in: _storage)
+            let matchEndIndex = String.Index(utf16Offset: range.location + range.length, in: _storage)
+            newStorage += _storage[sourceStringCurrentIndex ..< matchStartIndex]
+            newStorage += replacement
+            sourceStringCurrentIndex = matchEndIndex
+        }
+        newStorage += _storage[sourceStringCurrentIndex ..< _storage.endIndex]
+        _storage = String(newStorage)
+        return numOccurrences
     }
     
     public func applyTransform(_ transform: String, reverse: Bool, range: NSRange, updatedRange resultingRange: NSRangePointer?) -> Bool {
diff --git a/TestFoundation/TestNSString.swift b/TestFoundation/TestNSString.swift
index 70372b7a6a..0dd4cf326a 100755
--- a/TestFoundation/TestNSString.swift
+++ b/TestFoundation/TestNSString.swift
@@ -1413,6 +1413,64 @@ extension TestNSString {
         XCTAssertEqual(str4.replacingOccurrences(of: "\n\r", with: " "), "Hello\r\rworld.")
     }
 
+    func test_replacingOccurrencesInSubclass() {
+        class TestMutableString: NSMutableString {
+            private var wrapped: NSMutableString
+            var replaceCharactersCount: Int = 0
+
+            override var length: Int {
+                return wrapped.length
+            }
+
+            override func character(at index: Int) -> unichar {
+                return wrapped.character(at: index)
+            }
+
+            override func replaceCharacters(in range: NSRange, with aString: String) {
+                defer { replaceCharactersCount += 1 }
+                wrapped.replaceCharacters(in: range, with: aString)
+            }
+
+            override func mutableCopy(with zone: NSZone? = nil) -> Any {
+                return wrapped.mutableCopy()
+            }
+
+            required init(stringLiteral value: StaticString) {
+                wrapped = .init(stringLiteral: value)
+                super.init(stringLiteral: value)
+            }
+
+            required init(capacity: Int) {
+                fatalError("init(capacity:) has not been implemented")
+            }
+
+            required init(string aString: String) {
+                fatalError("init(string:) has not been implemented")
+            }
+
+            required convenience init?(coder aDecoder: NSCoder) {
+                fatalError("init(coder:) has not been implemented")
+            }
+
+            required init(characters: UnsafePointer<unichar>, length: Int) {
+                fatalError("init(characters:length:) has not been implemented")
+            }
+
+            required convenience init(extendedGraphemeClusterLiteral value: StaticString) {
+                fatalError("init(extendedGraphemeClusterLiteral:) has not been implemented")
+            }
+
+            required convenience init(unicodeScalarLiteral value: StaticString) {
+                fatalError("init(unicodeScalarLiteral:) has not been implemented")
+            }
+        }
+
+        let testString = TestMutableString(stringLiteral: "ababab")
+        XCTAssertEqual(testString.replacingOccurrences(of: "ab", with: "xx"), "xxxxxx")
+        XCTAssertEqual(testString.replaceCharactersCount, 3)
+    }
+
+
     func test_fileSystemRepresentation() {
         let name = "☃" as NSString
         let result = name.fileSystemRepresentation