Merge pull request #1995 from itaiferber/data-underestimated-sequence-initializer-fix

Fix off-by-one when initializing Data with discontiguous, underestimated Sequences

Author: itaiferber

Foundation/Data.swift

@@ -2085,9 +2085,8 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
                 // ... and append the rest byte-wise, buffering through an InlineData.
                 var buffer = InlineData()
                 while let element = iter.next() {
-                    if buffer.count < buffer.capacity {
-                        buffer.append(byte: element)
-                    } else {
+                    buffer.append(byte: element)
+                    if buffer.count == buffer.capacity {
                         buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
                         buffer.count = 0
                     }
@@ -2379,9 +2378,8 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
             // ... and append the rest byte-wise, buffering through an InlineData.
             var buffer = InlineData()
             while let element = iter.next() {
-                if buffer.count < buffer.capacity {
-                    buffer.append(byte: element)
-                } else {
+                buffer.append(byte: element)
+                if buffer.count == buffer.capacity {
                     buffer.withUnsafeBytes { _representation.append(contentsOf: $0) }
                     buffer.count = 0
                 }

TestFoundation/TestNSData.swift

@@ -521,6 +521,8 @@ class TestNSData: LoopbackServerTest {
             ("test_copyBytes2", test_copyBytes2),
             ("test_sliceOfSliceViaRangeExpression", test_sliceOfSliceViaRangeExpression),
             ("test_appendingSlices", test_appendingSlices),
+            ("test_appendingNonContiguousSequence_exactCount", test_appendingNonContiguousSequence_exactCount),
+            ("test_appendingNonContiguousSequence_underestimatedCount", test_appendingNonContiguousSequence_underestimatedCount),
             ("test_sequenceInitializers", test_sequenceInitializers),
             ("test_reversedDataInit", test_reversedDataInit),
             ("test_replaceSubrangeReferencingMutable", test_replaceSubrangeReferencingMutable),
@@ -4391,6 +4393,55 @@ extension TestNSData {
         d2.append(slice)
         XCTAssertEqual(Data([1]), slice)
     }
+    
+    // This test uses `repeatElement` to produce a sequence -- the produced sequence reports its actual count as its `.underestimatedCount`.
+    func test_appendingNonContiguousSequence_exactCount() {
+        var d = Data()
+        
+        // d should go from .empty representation to .inline.
+        // Appending a small enough sequence to fit in .inline should actually be copied.
+        d.append(contentsOf: 0x00...0x01)
+        expectEqual(Data([0x00, 0x01]), d)
+        
+        // Appending another small sequence should similarly still work.
+        d.append(contentsOf: 0x02...0x02)
+        expectEqual(Data([0x00, 0x01, 0x02]), d)
+        
+        // If we append a sequence of elements larger than a single InlineData, the internal append here should buffer.
+        // We want to make sure that buffering in this way does not accidentally drop trailing elements on the floor.
+        d.append(contentsOf: 0x03...0x2F)
+        expectEqual(Data([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                          0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
+                          0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+                          0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F,
+                          0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27,
+                          0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F]), d)
+    }
+    
+    // This test is like test_appendingNonContiguousSequence_exactCount but uses a sequence which reports 0 for its `.underestimatedCount`.
+    // This attempts to hit the worst-case scenario of `Data.append<S>(_:)` -- a discontiguous sequence of unknown length.
+    func test_appendingNonContiguousSequence_underestimatedCount() {
+        var d = Data()
+        
+        // d should go from .empty representation to .inline.
+        // Appending a small enough sequence to fit in .inline should actually be copied.
+        d.append(contentsOf: (0x00...0x01).makeIterator()) // `.makeIterator()` produces a sequence whose `.underestimatedCount` is 0.
+        expectEqual(Data([0x00, 0x01]), d)
+        
+        // Appending another small sequence should similarly still work.
+        d.append(contentsOf: (0x02...0x02).makeIterator()) // `.makeIterator()` produces a sequence whose `.underestimatedCount` is 0.
+        expectEqual(Data([0x00, 0x01, 0x02]), d)
+        
+        // If we append a sequence of elements larger than a single InlineData, the internal append here should buffer.
+        // We want to make sure that buffering in this way does not accidentally drop trailing elements on the floor.
+        d.append(contentsOf: (0x03...0x2F).makeIterator()) // `.makeIterator()` produces a sequence whose `.underestimatedCount` is 0.
+        expectEqual(Data([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+                          0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
+                          0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+                          0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F,
+                          0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27,
+                          0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F]), d)
+    }
 
     func test_sequenceInitializers() {
         let seq = repeatElement(UInt8(0x02), count: 3) // ensure we fall into the sequence case