Merge pull request #31601 from lorentey/data-append

[Foundation] Data.append: Switch to using resetBytes

Author: lorentey

stdlib/public/Darwin/Foundation/Data.swift

@@ -1415,65 +1415,6 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
             }
         }
 
-        @inlinable
-        @_alwaysEmitIntoClient
-        internal mutating func _truncateOrZeroExtend(toCount newCount: Int) {
-            switch self {
-            case .empty:
-                if newCount == 0 {
-                    return
-                } else if InlineData.canStore(count: newCount) {
-                    self = .inline(InlineData(count: newCount))
-                } else if InlineSlice.canStore(count: newCount) {
-                    self = .slice(InlineSlice(count: newCount))
-                } else {
-                    self = .large(LargeSlice(count: newCount))
-                }
-            case .inline(var inline):
-                if newCount == 0 {
-                    self = .empty
-                } else if InlineData.canStore(count: newCount) {
-                    guard inline.count != newCount else { return }
-                    inline.count = newCount
-                    self = .inline(inline)
-                } else if InlineSlice.canStore(count: newCount) {
-                    var slice = InlineSlice(inline)
-                    slice.count = newCount
-                    self = .slice(slice)
-                } else {
-                    var slice = LargeSlice(inline)
-                    slice.count = newCount
-                    self = .large(slice)
-                }
-            case .slice(var slice):
-                if newCount == 0 && slice.startIndex == 0 {
-                    self = .empty
-                } else if slice.startIndex == 0 && InlineData.canStore(count: newCount) {
-                    self = .inline(InlineData(slice, count: newCount))
-                } else if InlineSlice.canStore(count: newCount + slice.startIndex) {
-                    guard slice.count != newCount else { return }
-                    self = .empty // TODO: remove this when mgottesman lands optimizations
-                    slice.count = newCount
-                    self = .slice(slice)
-                } else {
-                    var newSlice = LargeSlice(slice)
-                    newSlice.count = newCount
-                    self = .large(newSlice)
-                }
-            case .large(var slice):
-                if newCount == 0 && slice.startIndex == 0 {
-                    self = .empty
-                } else if slice.startIndex == 0 && InlineData.canStore(count: newCount) {
-                    self = .inline(InlineData(slice, count: newCount))
-                } else {
-                    guard slice.count != newCount else { return }
-                    self = .empty // TODO: remove this when mgottesman lands optimizations
-                    slice.count = newCount
-                    self = .large(slice)
-                }
-            }
-        }
-
         @inlinable // This is @inlinable as reasonably small.
         var count: Int {
             get {
@@ -1485,7 +1426,69 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
                 }
             }
             set(newValue) {
-                _truncateOrZeroExtend(toCount: newValue)
+                // HACK: The definition of this inline function takes an inout reference to self, giving the optimizer a unique referencing guarantee.
+                //       This allows us to avoid excessive retain-release traffic around modifying enum values, and inlining the function then avoids the additional frame.
+                @inline(__always)
+                func apply(_ representation: inout _Representation, _ newValue: Int) -> _Representation? {
+                    switch representation {
+                    case .empty:
+                        if newValue == 0 {
+                            return nil
+                        } else if InlineData.canStore(count: newValue) {
+                            return .inline(InlineData(count: newValue))
+                        } else if InlineSlice.canStore(count: newValue) {
+                            return .slice(InlineSlice(count: newValue))
+                        } else {
+                            return .large(LargeSlice(count: newValue))
+                        }
+                    case .inline(var inline):
+                        if newValue == 0 {
+                            return .empty
+                        } else if InlineData.canStore(count: newValue) {
+                            guard inline.count != newValue else { return nil }
+                            inline.count = newValue
+                            return .inline(inline)
+                        } else if InlineSlice.canStore(count: newValue) {
+                            var slice = InlineSlice(inline)
+                            slice.count = newValue
+                            return .slice(slice)
+                        } else {
+                            var slice = LargeSlice(inline)
+                            slice.count = newValue
+                            return .large(slice)
+                        }
+                    case .slice(var slice):
+                        if newValue == 0 && slice.startIndex == 0 {
+                            return .empty
+                        } else if slice.startIndex == 0 && InlineData.canStore(count: newValue) {
+                            return .inline(InlineData(slice, count: newValue))
+                        } else if InlineSlice.canStore(count: newValue + slice.startIndex) {
+                            guard slice.count != newValue else { return nil }
+                            representation = .empty // TODO: remove this when mgottesman lands optimizations
+                            slice.count = newValue
+                            return .slice(slice)
+                        } else {
+                            var newSlice = LargeSlice(slice)
+                            newSlice.count = newValue
+                            return .large(newSlice)
+                        }
+                    case .large(var slice):
+                        if newValue == 0 && slice.startIndex == 0 {
+                            return .empty
+                        } else if slice.startIndex == 0 && InlineData.canStore(count: newValue) {
+                            return .inline(InlineData(slice, count: newValue))
+                        } else {
+                            guard slice.count != newValue else { return nil}
+                            representation = .empty // TODO: remove this when mgottesman lands optimizations
+                            slice.count = newValue
+                            return .large(slice)
+                        }
+                    }
+                }
+
+                if let rep = apply(&self, newValue) {
+                    self = rep
+                }
             }
         }
 
@@ -2385,16 +2388,17 @@ public struct Data : ReferenceConvertible, Equatable, Hashable, RandomAccessColl
         // Copy as much as we can in one shot.
         let underestimatedCount = elements.underestimatedCount
         let originalCount = _representation.count
-        _representation._truncateOrZeroExtend(toCount: originalCount + underestimatedCount)
+        resetBytes(in: self.endIndex ..< self.endIndex + underestimatedCount)
         var (iter, copiedCount): (S.Iterator, Int) = _representation.withUnsafeMutableBytes { buffer in
+            assert(buffer.count == originalCount + underestimatedCount)
             let start = buffer.baseAddress!.assumingMemoryBound(to: UInt8.self) + originalCount
             let b = UnsafeMutableBufferPointer(start: start, count: buffer.count - originalCount)
             return elements._copyContents(initializing: b)
         }
         guard copiedCount == underestimatedCount else {
             // We can't trap here. We have to allow an underfilled buffer
             // to emulate the previous implementation.
-            _representation._truncateOrZeroExtend(toCount: originalCount + copiedCount)
+            _representation.replaceSubrange(startIndex + originalCount + copiedCount ..< endIndex, with: nil, count: 0)
             return
         }
 

test/stdlib/TestData.swift

@@ -3883,6 +3883,99 @@ class TestData : TestDataSuper {
             }
         }
     }
+
+    // This is a (potentially invalid) sequence that produces a configurable number of 42s and has a freely customizable `underestimatedCount`.
+    struct TestSequence: Sequence {
+        typealias Element = UInt8
+        struct Iterator: IteratorProtocol {
+            var _remaining: Int
+            init(_ count: Int) {
+                _remaining = count
+            }
+            mutating func next() -> UInt8? {
+                guard _remaining > 0 else { return nil }
+                _remaining -= 1
+                return 42
+            }
+        }
+        let underestimatedCount: Int
+        let count: Int
+
+        func makeIterator() -> Iterator {
+            return Iterator(count)
+        }
+    }
+
+    func test_init_TestSequence() {
+        // Underestimated count
+        do {
+            let d = Data(TestSequence(underestimatedCount: 0, count: 10))
+            expectEqual(10, d.count)
+            expectEqual(Array(repeating: 42 as UInt8, count: 10), Array(d))
+        }
+
+        // Very underestimated count (to exercise realloc path)
+        do {
+            let d = Data(TestSequence(underestimatedCount: 0, count: 1000))
+            expectEqual(1000, d.count)
+            expectEqual(Array(repeating: 42 as UInt8, count: 1000), Array(d))
+        }
+
+        // Exact count
+        do {
+            let d = Data(TestSequence(underestimatedCount: 10, count: 10))
+            expectEqual(10, d.count)
+            expectEqual(Array(repeating: 42 as UInt8, count: 10), Array(d))
+        }
+
+        // Overestimated count. This is an illegal case, so trapping would be fine.
+        // However, for compatibility with the implementation in Swift 5, Data
+        // handles this case by simply truncating itself to the actual size.
+        do {
+            let d = Data(TestSequence(underestimatedCount: 20, count: 10))
+            expectEqual(10, d.count)
+            expectEqual(Array(repeating: 42 as UInt8, count: 10), Array(d))
+        }
+    }
+
+    func test_append_TestSequence() {
+        let base = Data(Array(repeating: 23 as UInt8, count: 10))
+
+        // Underestimated count
+        do {
+            var d = base
+            d.append(contentsOf: TestSequence(underestimatedCount: 0, count: 10))
+            expectEqual(20, d.count)
+            expectEqual(Array(base) + Array(repeating: 42 as UInt8, count: 10),
+                           Array(d))
+        }
+
+        // Very underestimated count (to exercise realloc path)
+        do {
+            var d = base
+            d.append(contentsOf: TestSequence(underestimatedCount: 0, count: 1000))
+            expectEqual(1010, d.count)
+            expectEqual(Array(base) + Array(repeating: 42 as UInt8, count: 1000), Array(d))
+        }
+
+        // Exact count
+        do {
+            var d = base
+            d.append(contentsOf: TestSequence(underestimatedCount: 10, count: 10))
+            expectEqual(20, d.count)
+            expectEqual(Array(base) + Array(repeating: 42 as UInt8, count: 10), Array(d))
+        }
+
+        // Overestimated count. This is an illegal case, so trapping would be fine.
+        // However, for compatibility with the implementation in Swift 5, Data
+        // handles this case by simply truncating itself to the actual size.
+        do {
+            var d = base
+            d.append(contentsOf: TestSequence(underestimatedCount: 20, count: 10))
+            expectEqual(20, d.count)
+            expectEqual(Array(base) + Array(repeating: 42 as UInt8, count: 10), Array(d))
+        }
+    }
 }
 
 #if !FOUNDATION_XCTEST
@@ -4211,6 +4304,8 @@ if #available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) {
 }
 DataTests.test("test_increaseCount") { TestData().test_increaseCount() }
 DataTests.test("test_decreaseCount") { TestData().test_decreaseCount() }
+DataTests.test("test_increaseCount") { TestData().test_init_TestSequence() }
+DataTests.test("test_decreaseCount") { TestData().test_append_TestSequence() }
 
 
 // XCTest does not have a crash detection, whereas lit does