Optimize allocations in replaceCharacters

SignalCoreKit/src/String+OWS.swift

@@ -91,24 +91,37 @@ public extension String {
         return String(safePrefix(upperBoundCharCount))
     }
 
-    func replaceCharacters(characterSet: CharacterSet, replacement: String) -> String {
-        guard var range = self.rangeOfCharacter(from: characterSet) else {
-            return self
+    func replaceCharacters(
+        characterSet: CharacterSet,
+        replacement: String
+    ) -> String {
+        let endIndex = self.endIndex
+        var startIndex = self.startIndex
+
+        // Build up a list of ranges that need to be replaced
+        var ranges = [Range<String.Index>]()
+        while startIndex < endIndex, let range = self.rangeOfCharacter(from: characterSet, options: [], range: startIndex..<endIndex) {
+            ranges.append(range)
+            startIndex = range.upperBound
         }
 
+        // Don't do any allocation for unchanged strings
+        guard ranges.count > 0 else { return self }
+
+        // Create the result string and set up a capacity close to the final string
         var result = ""
-        var remaining = self[...]
-        while true {
-            result += remaining[..<range.lowerBound]
+        result.reserveCapacity(self.count)
+
+        // Iterate through the ranges, appending the string between the last
+        // match and the next, and then appending the replacement string
+        var currentIndex = self.startIndex
+        for range in ranges {
+            result += self[currentIndex..<range.lowerBound]
             result += replacement
-            remaining = remaining[range.upperBound...]
-            guard let nextRange = remaining.rangeOfCharacter(from: characterSet) else {
-                result += remaining
-                break
-            }
-            range = nextRange
+            currentIndex = range.upperBound
         }
-
+        // Add the remainder of the string
+        result += self[currentIndex..<endIndex]
         return result
     }
 

SignalCoreKitTests/src/StringSanitizerTests.swift

@@ -58,3 +58,77 @@ class StringSanitizerTests: XCTestCase {
         XCTAssertEqual(sanitizer.sanitized, expected)
     }
 }
+
+class StringReplacementTests: XCTestCase {
+
+    func testEquivalent() {
+
+        let testCases: [String: String] = [
+            "": "",
+            " ": "",
+            "         ": "",
+            "a": "a",
+            "abcd": "abcd",
+            " abcd ": "abcd",
+            "abcd ": "abcd",
+            " abcd": "abcd",
+            "ab cd": "abcd",
+            "ab  1 cd ": "ab1cd",
+            "ab            cd ": "abcd"
+        ]
+
+        for key in testCases.keys {
+            let expectedResult = testCases[key]
+
+            let result = key.replaceCharacters(characterSet: .whitespacesAndNewlines, replacement: "")
+
+            XCTAssertEqual(result, expectedResult)
+        }
+    }
+
+    func testEquivalent2() {
+
+        let testCases: [String: String] = [
+            "": "",
+            "abcd": "abcd",
+            " abcd ": "X abcdX ",
+            "abcd ": "abcdX ",
+            " abcd": "X abcd",
+            "ab cd": "abX cd",
+            "ab  1 cd ": "abX X 1X cdX "
+        ]
+
+        for key in testCases.keys {
+            let expectedResult = testCases[key]
+
+            let result = key.replaceCharacters(characterSet: .whitespacesAndNewlines, replacement: "X ")
+
+            XCTAssertEqual(result, expectedResult)
+        }
+    }
+
+     func testEquivalent3() {
+
+        let testCases: [String: String] = [
+            "": "",
+            "abcd": "",
+            " abcd ": "  ",
+            "abcd ": " ",
+            " abcd": " ",
+            "ab cd": " ",
+            "ab  1 cd ": "  1  ",
+            "ab  1 ZcdX ": "  1 ZX "
+        ]
+
+        for key in testCases.keys {
+            let expectedResult = testCases[key]
+
+            var characterSetUnion = CharacterSet.punctuationCharacters
+            characterSetUnion.formUnion(.lowercaseLetters)
+
+            let result = key.replaceCharacters(characterSet: characterSetUnion, replacement: "")
+
+            XCTAssertEqual(result, expectedResult)
+        }
+    }
+}