Skip to content

[stdlib] Fix bug in small string uninitialized init #36667

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.

Sign up for GitHub

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 1 commit into from
Apr 2, 2021
Merged

[stdlib] Fix bug in small string uninitialized init #36667

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
18 changes: 15 additions & 3 deletions stdlib/public/core/SmallString.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,13 @@ extension _SmallString {
internal func _invariantCheck() {
_internalInvariant(count <= _SmallString.capacity)
_internalInvariant(isASCII == computeIsASCII())

// No bits should be set between the last code unit and the discriminator
var copy = self
withUnsafeBytes(of: &copy._storage) {
_internalInvariant(
$0[count..<_SmallString.capacity].allSatisfy { $0 == 0 })
}
}
#endif // INTERNAL_CHECKS_ENABLED

Expand Down Expand Up @@ -213,14 +220,19 @@ extension _SmallString {
}

// Overwrite stored code units, including uninitialized. `f` should return the
// new count.
// new count. This will re-establish the invariant after `f` that all bits
// between the last code unit and the discriminator are unset.
@inline(__always)
internal mutating func withMutableCapacity(
fileprivate mutating func withMutableCapacity(
_ f: (UnsafeMutableRawBufferPointer) throws -> Int
) rethrows {
let len = try withUnsafeMutableBytes(of: &self._storage) {
(rawBufPtr: UnsafeMutableRawBufferPointer) -> Int in
return try f(rawBufPtr)
let len = try f(rawBufPtr)
UnsafeMutableRawBufferPointer(
rebasing: rawBufPtr[len...]
).initializeMemory(as: UInt8.self, repeating: 0)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@atrick is this a proper use of initializeMemory? It's basically a convenient way to mask off our unused bytes and make sure they're set to 0.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@milseman Now that I'm looking at this again, I think this needs to rebind the memory; otherwise, after initializing memory as UInt8, portions of the memory are bound simultaneously as RawBitPattern.self and UInt8.self 😕

let count = self.count
let len = try withUnsafeMutableBytes(of: &self._storage) {
  (rawBufPtr: UnsafeMutableRawBufferPointer) -> Int in
  let len = try f(rawBufPtr)
  rawBufPtr.bindMemory(to: UInt8.self, capacity: count)
  defer {
    _ = rawBufPtr.bindMemory(to: RawBitPattern.self, capacity: 1)
  }
  UnsafeMutableRawBufferPointer(
    rebasing: rawBufPtr[len...]
  ).initializeMemory(as: UInt8.self, repeating: 0)
  return len
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see any API for initializing raw bytes without using a typed pointer. You can do the silly rebinding thing above or just write your own raw initialize that doesn't care about the in-memory type

var nextPtr = self
    for _ in 0..<count {
      Builtin.initialize(repeatedValue, nextPtr._rawValue)
      nextPtr += MemoryLayout<T>.stride
    }

return len
}
if len == 0 {
self = _SmallString()
Expand Down
24 changes: 24 additions & 0 deletions test/stdlib/StringCreate.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,3 +80,27 @@ if #available(macOS 10.16, iOS 14.0, watchOS 7.0, tvOS 14.0, *) {
expectTrue(empty.isEmpty)
}
}

if #available(macOS 10.16, iOS 14.0, watchOS 7.0, tvOS 14.0, *) {
StringCreateTests.test("Small string unsafeUninitializedCapacity") {
let str1 = "42"
let str2 = String(42)
expectEqual(str1, str2)

let str3 = String(unsafeUninitializedCapacity: 2) {
$0[0] = UInt8(ascii: "4")
$0[1] = UInt8(ascii: "2")
return 2
}
expectEqual(str1, str3)

// Write into uninitialized space
let str4 = String(unsafeUninitializedCapacity: 3) {
$0[0] = UInt8(ascii: "4")
$0[1] = UInt8(ascii: "2")
$0[2] = UInt8(ascii: "X")
return 2
}
expectEqual(str1, str4)
}
}