Skip to content

[Foundation] Data.append: Switch to using resetBytes #31601

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 7, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 66 additions & 62 deletions stdlib/public/Darwin/Foundation/Data.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
}
}
}

Expand Down Expand Up @@ -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
}

Expand Down
95 changes: 95 additions & 0 deletions test/stdlib/TestData.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -4208,6 +4301,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
Expand Down