Merge pull request #2499 from YOCKOW/sr-10689

Author: millenomi

Sources/Foundation/DataProtocol.swift

@@ -136,34 +136,41 @@ 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 {
         guard !data.isEmpty else {
             return nil
         }
         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)
-            }
-            if matched {
-                return haystackIndex..<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)
             }
-            haystackIndex = index(after: haystackIndex)
+            position = index(after: position)
         }
         return nil
     }
@@ -173,29 +180,18 @@ extension DataProtocol {
             return nil
         }
         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
     }

Tests/Foundation/Tests/TestNSData.swift

@@ -253,6 +253,7 @@ class TestNSData: LoopbackServerTest {
             ("test_base64DecodeWithPadding1", test_base64DecodeWithPadding1),
             ("test_base64DecodeWithPadding2", test_base64DecodeWithPadding2),
             ("test_rangeOfData", test_rangeOfData),
+            ("test_sr10689_rangeOfDataProtocol", test_sr10689_rangeOfDataProtocol),
             ("test_initNSMutableData()", test_initNSMutableData),
             ("test_initNSMutableDataWithLength", test_initNSMutableDataWithLength),
             ("test_initNSMutableDataWithCapacity", test_initNSMutableDataWithCapacity),
@@ -949,6 +950,97 @@ class TestNSData: LoopbackServerTest {
 
     }
 
+    func test_sr10689_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: StaticString = #file, line: UInt = #line)
+            {
+                if let index = expectedStartIndex {
+                    let expectedRange: Range<Int> = index..<(index + fragment.count)
+                    if let someRange = range {
+                        XCTAssertEqual(data.firstRange(of: fragment, in: someRange), expectedRange, message(), file: file, line: line)
+                    } else {
+                        XCTAssertEqual(data.firstRange(of: fragment), expectedRange, message(), file: file, line: line)
+                    }
+                } else {
+                    if let someRange = range {
+                        XCTAssertNil(data.firstRange(of: fragment, in: someRange), message(), file: file, line: line)
+                    } else {
+                        XCTAssertNil(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: StaticString = #file, line: UInt = #line)
+            {
+                if let index = expectedStartIndex {
+                    let expectedRange: Range<Int> = index..<(index + fragment.count)
+                    if let someRange = range {
+                        XCTAssertEqual(data.lastRange(of: fragment, in: someRange), expectedRange, message(), file: file, line: line)
+                    } else {
+                        XCTAssertEqual(data.lastRange(of: fragment), expectedRange, message(), file: file, line: line)
+                    }
+                } else {
+                    if let someRange = range {
+                        XCTAssertNil(data.lastRange(of: fragment, in: someRange), message(), file: file, line: line)
+                    } else {
+                        XCTAssertNil(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)
+        }
+    }
+
     // Check all of the NSMutableData constructors are available.
     func test_initNSMutableData() {
         let mData = NSMutableData()