Merge pull request #28639 from YOCKOW/sr-10689-cherry-pick

SR-10689: Fix bugs of DataProtocol's firstRange(of:in:)/lastRange(of:in:).

Author: Kyle Macomber

stdlib/public/Darwin/Foundation/DataProtocol.swift

@@ -142,60 +142,56 @@ extension DataProtocol {
         return self.copyBytes(to: UnsafeMutableRawBufferPointer(start: ptr.baseAddress, count: ptr.count * MemoryLayout<DestinationType>.stride), from: range)
     }
 
+    private func matches<D: DataProtocol>(_ data: D, from index: Index) -> Bool {
+        var haystackIndex = index
+        var needleIndex = data.startIndex
+        
+        while true {
+            guard self[haystackIndex] == data[needleIndex] else { return false }
+            
+            haystackIndex = self.index(after: haystackIndex)
+            needleIndex = data.index(after: needleIndex)
+            if needleIndex == data.endIndex {
+                // i.e. needle is found.
+                return true
+            } else if haystackIndex == endIndex {
+                return false
+            }
+        }
+    }
+    
     public func firstRange<D: DataProtocol, R: RangeExpression>(of data: D, in range: R) -> Range<Index>? where R.Bound == Index {
         let r = range.relative(to: self)
-        let rangeCount = distance(from: r.lowerBound, to: r.upperBound)
-        if rangeCount < data.count {
+        let length = data.count
+        
+        if length == 0 || length > distance(from: r.lowerBound, to: r.upperBound) {
             return nil
         }
-        var haystackIndex = r.lowerBound
-        let haystackEnd = index(r.upperBound, offsetBy: -data.count)
-        while haystackIndex < haystackEnd {
-            var compareIndex = haystackIndex
-            var needleIndex = data.startIndex
-            let needleEnd = data.endIndex
-            var matched = true
-            while compareIndex < haystackEnd && needleIndex < needleEnd {
-                if self[compareIndex] != data[needleIndex] {
-                    matched = false
-                    break
-                }
-                needleIndex = data.index(after: needleIndex)
-                compareIndex = index(after: compareIndex)
+        
+        var position = r.lowerBound
+        while position < r.upperBound && distance(from: position, to: r.upperBound) >= length {
+            if matches(data, from: position) {
+                return position..<index(position, offsetBy: length)
             }
-            if matched {
-                return haystackIndex..<compareIndex
-            }
-            haystackIndex = index(after: haystackIndex)
+            position = index(after: position)
         }
         return nil
     }
-
+    
     public func lastRange<D: DataProtocol, R: RangeExpression>(of data: D, in range: R) -> Range<Index>? where R.Bound == Index {
         let r = range.relative(to: self)
-        let rangeCount = distance(from: r.lowerBound, to: r.upperBound)
-        if rangeCount < data.count {
+        let length = data.count
+        
+        if length == 0 || length > distance(from: r.lowerBound, to: r.upperBound) {
             return nil
         }
-        var haystackIndex = r.upperBound
-        let haystackStart = index(r.lowerBound, offsetBy: data.count)
-        while haystackIndex > haystackStart {
-            var compareIndex = haystackIndex
-            var needleIndex = data.endIndex
-            let needleStart = data.startIndex
-            var matched = true
-            while compareIndex > haystackStart && needleIndex > needleStart {
-                if self[compareIndex] != data[needleIndex] {
-                    matched = false
-                    break
-                }
-                needleIndex = data.index(before: needleIndex)
-                compareIndex = index(before: compareIndex)
-            }
-            if matched {
-                return compareIndex..<haystackIndex
+        
+        var position = index(r.upperBound, offsetBy: -length)
+        while position >= r.lowerBound {
+            if matches(data, from: position) {
+                return position..<index(position, offsetBy: length)
             }
-            haystackIndex = index(before: haystackIndex)
+            position = index(before: position)
         }
         return nil
     }

test/stdlib/TestData.swift

@@ -1077,6 +1077,99 @@ class TestData : TestDataSuper {
         expectEqual(slice4[0], 8)
     }
 
+    func test_rangeOfDataProtocol() {
+        // https://bugs.swift.org/browse/SR-10689
+        
+        let base = Data([0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03,
+                         0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03])
+        let subdata = base[10..<13] // [0x02, 0x03, 0x00]
+        let oneByte = base[14..<15] // [0x02]
+        
+        do { // firstRange(of:in:)
+            func assertFirstRange(_ data: Data, _ fragment: Data, range: ClosedRange<Int>? = nil,
+                                  expectedStartIndex: Int?,
+                                  _ message: @autoclosure () -> String = "",
+                                  file: String = #file, line: UInt = #line) {
+                if let index = expectedStartIndex {
+                    let expectedRange: Range<Int> = index..<(index + fragment.count)
+                    if let someRange = range {
+                        expectEqual(data.firstRange(of: fragment, in: someRange), expectedRange, message(), file: file, line: line)
+                    } else {
+                        expectEqual(data.firstRange(of: fragment), expectedRange, message(), file: file, line: line)
+                    }
+                } else {
+                    if let someRange = range {
+                        expectNil(data.firstRange(of: fragment, in: someRange), message(), file: file, line: line)
+                    } else {
+                        expectNil(data.firstRange(of: fragment), message(), file: file, line: line)
+                    }
+                }
+            }
+            
+            assertFirstRange(base, base, expectedStartIndex: base.startIndex)
+            assertFirstRange(base, subdata, expectedStartIndex: 2)
+            assertFirstRange(base, oneByte, expectedStartIndex: 2)
+            
+            assertFirstRange(subdata, base, expectedStartIndex: nil)
+            assertFirstRange(subdata, subdata, expectedStartIndex: subdata.startIndex)
+            assertFirstRange(subdata, oneByte, expectedStartIndex: subdata.startIndex)
+            
+            assertFirstRange(oneByte, base, expectedStartIndex: nil)
+            assertFirstRange(oneByte, subdata, expectedStartIndex: nil)
+            assertFirstRange(oneByte, oneByte, expectedStartIndex: oneByte.startIndex)
+            
+            assertFirstRange(base, subdata, range: 1...14, expectedStartIndex: 2)
+            assertFirstRange(base, subdata, range: 6...8, expectedStartIndex: 6)
+            assertFirstRange(base, subdata, range: 8...10, expectedStartIndex: nil)
+            
+            assertFirstRange(base, oneByte, range: 1...14, expectedStartIndex: 2)
+            assertFirstRange(base, oneByte, range: 6...6, expectedStartIndex: 6)
+            assertFirstRange(base, oneByte, range: 8...9, expectedStartIndex: nil)
+        }
+        
+        do { // lastRange(of:in:)
+            func assertLastRange(_ data: Data, _ fragment: Data, range: ClosedRange<Int>? = nil,
+                                 expectedStartIndex: Int?,
+                                 _ message: @autoclosure () -> String = "",
+                                 file: String = #file, line: UInt = #line) {
+                if let index = expectedStartIndex {
+                    let expectedRange: Range<Int> = index..<(index + fragment.count)
+                    if let someRange = range {
+                        expectEqual(data.lastRange(of: fragment, in: someRange), expectedRange, message(), file: file, line: line)
+                    } else {
+                        expectEqual(data.lastRange(of: fragment), expectedRange, message(), file: file, line: line)
+                    }
+                } else {
+                  if let someRange = range {
+                      expectNil(data.lastRange(of: fragment, in: someRange), message(), file: file, line: line)
+                  } else {
+                      expectNil(data.lastRange(of: fragment), message(), file: file, line: line)
+                  }
+                }
+            }
+            
+            assertLastRange(base, base, expectedStartIndex: base.startIndex)
+            assertLastRange(base, subdata, expectedStartIndex: 10)
+            assertLastRange(base, oneByte, expectedStartIndex: 14)
+            
+            assertLastRange(subdata, base, expectedStartIndex: nil)
+            assertLastRange(subdata, subdata, expectedStartIndex: subdata.startIndex)
+            assertLastRange(subdata, oneByte, expectedStartIndex: subdata.startIndex)
+            
+            assertLastRange(oneByte, base, expectedStartIndex: nil)
+            assertLastRange(oneByte, subdata, expectedStartIndex: nil)
+            assertLastRange(oneByte, oneByte, expectedStartIndex: oneByte.startIndex)
+            
+            assertLastRange(base, subdata, range: 1...14, expectedStartIndex: 10)
+            assertLastRange(base, subdata, range: 6...8, expectedStartIndex: 6)
+            assertLastRange(base, subdata, range: 8...10, expectedStartIndex: nil)
+            
+            assertLastRange(base, oneByte, range: 1...14, expectedStartIndex: 14)
+            assertLastRange(base, oneByte, range: 6...6, expectedStartIndex: 6)
+            assertLastRange(base, oneByte, range: 8...9, expectedStartIndex: nil)
+        }
+    }
+
     func test_sliceAppending() {
         // https://bugs.swift.org/browse/SR-4473
         var fooData = Data()
@@ -4022,6 +4115,7 @@ DataTests.test("test_noCopyBehavior") { TestData().test_noCopyBehavior() }
 DataTests.test("test_doubleDeallocation") { TestData().test_doubleDeallocation() }
 DataTests.test("test_repeatingValueInitialization") { TestData().test_repeatingValueInitialization() }
 DataTests.test("test_rangeZoo") { TestData().test_rangeZoo() }
+DataTests.test("test_rangeOfDataProtocol") { TestData().test_rangeOfDataProtocol() }
 DataTests.test("test_sliceAppending") { TestData().test_sliceAppending() }
 DataTests.test("test_replaceSubrange") { TestData().test_replaceSubrange() }
 DataTests.test("test_sliceWithUnsafeBytes") { TestData().test_sliceWithUnsafeBytes() }