Merge pull request #82097 from glessard/rdar152615664-issue82024-62-with-small-string-spans

[6.2, stdlib] fix small-string UTF8Span support

Author: glessard

stdlib/private/StdlibUnittest/StdlibUnittest.swift

@@ -752,11 +752,14 @@ public func expectNil<T>(_ value: T?,
 }
 
 @discardableResult
-public func expectNotNil<T>(_ value: T?,
+@lifetime(copy value)
+public func expectNotNil<T: ~Copyable & ~Escapable>(
+  _ value: consuming T?,
   _ message: @autoclosure () -> String = "",
   stackTrace: SourceLocStack = SourceLocStack(),
   showFrame: Bool = true,
-  file: String = #file, line: UInt = #line) -> T? {
+  file: String = #file, line: UInt = #line
+) -> T? {
   if value == nil {
     expectationFailure("expected optional to be non-nil", trace: message(),
       stackTrace: stackTrace.pushIf(showFrame, file: file, line: line))

stdlib/public/core/StringGuts.swift

@@ -17,6 +17,7 @@ import SwiftShims
 // functionality and guidance for efficiently working with Strings.
 //
 @frozen
+@_addressableForDependencies
 public // SPI(corelibs-foundation)
 struct _StringGuts: @unchecked Sendable {
   @usableFromInline

stdlib/public/core/StringUTF8View.swift

@@ -89,7 +89,6 @@ extension String {
   ///     print(String(s1.utf8.prefix(15))!)
   ///     // Prints "They call me 'B"
   @frozen
-  @_addressableForDependencies
   public struct UTF8View: Sendable {
     @usableFromInline
     internal var _guts: _StringGuts

stdlib/public/core/Substring.swift

@@ -630,7 +630,6 @@ extension Substring: LosslessStringConvertible {
 
 extension Substring {
   @frozen
-  @_addressableForDependencies
   public struct UTF8View: Sendable {
     @usableFromInline
     internal var _slice: Slice<String.UTF8View>

stdlib/public/core/UTF8Span.swift

@@ -201,36 +201,112 @@ extension String {
   }
 
   @available(SwiftStdlib 6.2, *)
-  public var utf8Span: UTF8Span {
+  private var _span: Span<UTF8.CodeUnit> {
     @lifetime(borrow self)
     borrowing get {
-      let isKnownASCII = _guts.isASCII
-      let utf8 = self.utf8
-      let span = utf8.span
-      let result = unsafe UTF8Span(
-        unchecked: span,
-        isKnownASCII: isKnownASCII)
-      return unsafe _overrideLifetime(result, borrowing: self)
+#if _runtime(_ObjC)
+      // handle non-UTF8 Objective-C bridging cases here
+      if !_guts.isFastUTF8, _guts._object.hasObjCBridgeableObject {
+        let storage = _guts._getOrAllocateAssociatedStorage()
+        let (start, count) = unsafe (storage.start, storage.count)
+        let span = unsafe Span(_unsafeStart: start, count: count)
+        return unsafe _overrideLifetime(span, borrowing: self)
+      }
+#endif
+      let count = _guts.count
+      if _guts.isSmall {
+        let a = Builtin.addressOfBorrow(self)
+        let address = unsafe UnsafePointer<UTF8.CodeUnit>(a)
+        let span = unsafe Span(_unsafeStart: address, count: count)
+        return unsafe _overrideLifetime(span, borrowing: self)
+      }
+      let isFastUTF8 = _guts.isFastUTF8
+      _precondition(isFastUTF8, "String must be contiguous UTF8")
+      let buffer = unsafe _guts._object.fastUTF8
+      let span = unsafe Span(_unsafeElements: buffer)
+      return unsafe _overrideLifetime(span, borrowing: self)
     }
   }
-}
 
-extension Substring {
+  /// A UTF8span over the code units that make up this string.
+  ///
+  /// - Note: In the case of bridged UTF16 String instances (on Apple
+  /// platforms,) this property transcodes the code units the first time
+  /// it is called. The transcoded buffer is cached, and subsequent calls
+  /// to `span` can reuse the buffer.
+  ///
+  ///  Returns: a `UTF8Span` over the code units of this String.
+  ///
+  ///  Complexity: O(1) for native UTF8 Strings,
+  ///    amortized O(1) for bridged UTF16 Strings.
   @available(SwiftStdlib 6.2, *)
   public var utf8Span: UTF8Span {
     @lifetime(borrow self)
     borrowing get {
-      let isKnownASCII = base._guts.isASCII
-      let utf8 = self.utf8
-      let span = utf8.span
-      let result = unsafe UTF8Span(
-        unchecked: span,
-        isKnownASCII: isKnownASCII)
-      return unsafe _overrideLifetime(result, borrowing: self)
+      unsafe UTF8Span(unchecked: _span, isKnownASCII: _guts.isASCII)
     }
   }
 }
 
+extension Substring {
 
+  @available(SwiftStdlib 6.2, *)
+  private var _span: Span<UTF8.CodeUnit> {
+    @lifetime(borrow self)
+    borrowing get {
+#if _runtime(_ObjC)
+      // handle non-UTF8 Objective-C bridging cases here
+      if !_wholeGuts.isFastUTF8, _wholeGuts._object.hasObjCBridgeableObject {
+        let base: String.UTF8View = _slice._base.utf8
+        let first = base._foreignDistance(from: base.startIndex, to: startIndex)
+        let count = base._foreignDistance(from: startIndex, to: endIndex)
+        let span = base.span._extracting(first..<(first &+ count))
+        return unsafe _overrideLifetime(span, borrowing: self)
+      }
+#endif
+      let first = _slice._startIndex._encodedOffset
+      let end = _slice._endIndex._encodedOffset
+      if _wholeGuts.isSmall {
+        let a = Builtin.addressOfBorrow(self)
+        let offset = first &+ (2 &* MemoryLayout<String.Index>.stride)
+        let start = unsafe UnsafePointer<UTF8.CodeUnit>(a).advanced(by: offset)
+        let span = unsafe Span(_unsafeStart: start, count: end &- first)
+        return unsafe _overrideLifetime(span, borrowing: self)
+      }
+      let isFastUTF8 = _wholeGuts.isFastUTF8
+      _precondition(isFastUTF8, "Substring must be contiguous UTF8")
+      var span = unsafe Span(_unsafeElements: _wholeGuts._object.fastUTF8)
+      span = span._extracting(first..<end)
+      return unsafe _overrideLifetime(span, borrowing: self)
+    }
+  }
 
-
+  /// A UTF8Span over the code units that make up this substring.
+  ///
+  /// - Note: In the case of bridged UTF16 String instances (on Apple
+  /// platforms,) this property needs to transcode the code units every time
+  /// it is called.
+  /// For example, if `string` has the bridged UTF16 representation,
+  ///     for word in string.split(separator: " ") {
+  ///         useSpan(word.span)
+  ///     }
+  ///  is accidentally quadratic because of this issue. A workaround is to
+  ///  explicitly convert the string into its native UTF8 representation:
+  ///      var nativeString = consume string
+  ///      nativeString.makeContiguousUTF8()
+  ///      for word in nativeString.split(separator: " ") {
+  ///          useSpan(word.span)
+  ///      }
+  ///  This second option has linear time complexity, as expected.
+  ///
+  ///  Returns: a `UTF8Span` over the code units of this Substring.
+  ///
+  ///  Complexity: O(1) for native UTF8 Strings, O(n) for bridged UTF16 Strings.
+  @available(SwiftStdlib 6.2, *)
+  public var utf8Span: UTF8Span {
+    @lifetime(borrow self)
+    borrowing get {
+      unsafe UTF8Span(unchecked: _span, isKnownASCII: base._guts.isASCII)
+    }
+  }
+}

test/SILOptimizer/lifetime_dependence/semantics.swift

@@ -417,7 +417,7 @@ public func testTrivialInoutBorrow(p: inout UnsafePointer<Int>) -> Span<Int> {
 
 private let immortalInt = 0
 
-private let immortalString = ""
+private let immortalStrings: [String] = []
 
 @lifetime(immortal)
 func testImmortalInt() -> Span<Int> {
@@ -427,10 +427,10 @@ func testImmortalInt() -> Span<Int> {
 }
 
 @lifetime(immortal)
-func testImmortalString() -> Span<String> {
-  let nilBasedBuffer = UnsafeBufferPointer<String>(start: nil, count: 0)
+func testImmortalStrings() -> Span<[String]> {
+  let nilBasedBuffer = UnsafeBufferPointer<[String]>(start: nil, count: 0)
   let span = Span(base: nilBasedBuffer.baseAddress, count: nilBasedBuffer.count)
-  return _overrideLifetime(span, borrowing: immortalString)
+  return _overrideLifetime(span, borrowing: immortalStrings)
 }
 
 let ptr = UnsafePointer<Int>(bitPattern: 1)!

test/api-digester/Outputs/stability-stdlib-source-base.swift.expected

@@ -363,8 +363,6 @@ Func ContiguousArray.withUnsafeMutableBufferPointer(_:) is now without rethrows
 
 // Adoption of @_addressableForDependencies
 Struct CollectionOfOne is now with @_addressableForDependencies
-Struct String.UTF8View is now with @_addressableForDependencies
-Struct Substring.UTF8View is now with @_addressableForDependencies
 
 Protocol CodingKey has added inherited protocol SendableMetatype
 Protocol Error has added inherited protocol SendableMetatype

test/api-digester/stability-stdlib-abi-without-asserts.test

@@ -825,8 +825,7 @@ Func _SliceBuffer.withUnsafeMutableBufferPointer(_:) has mangled name changing f
 Struct String.Index has added a conformance to an existing protocol CustomDebugStringConvertible
 
 Struct CollectionOfOne is now with @_addressableForDependencies
-Struct String.UTF8View is now with @_addressableForDependencies
-Struct Substring.UTF8View is now with @_addressableForDependencies
+Struct _StringGuts is now with @_addressableForDependencies
 
 Enum _SwiftifyInfo is a new API without '@available'
 Enum _SwiftifyExpr is a new API without '@available'

test/stdlib/OptionalGeneralizations.swift

@@ -92,3 +92,23 @@ suite.test("Initializer references") {
     expectTrue(r != nil)
   }
 }
+
+suite.test("expectNotNil()") {
+  func opt1<T: ~Copyable>(_ t: consuming T) -> T? { Optional.some(t) }
+  _ = expectNotNil(opt1(TrivialStruct()))
+  _ = expectNotNil(opt1(NoncopyableStruct()))
+  _ = expectNotNil(opt1(RegularClass()))
+#if $NonescapableTypes
+  @lifetime(copy t)
+  func opt2<T: ~Copyable & ~Escapable>(_ t: consuming T) -> T? { t }
+
+  let ne = NonescapableStruct()
+  _ = expectNotNil(opt2(ne))
+
+  let ncne = NoncopyableNonescapableStruct()
+  _ = expectNotNil(opt2(ncne))
+
+  let nent = NonescapableNontrivialStruct()
+  _ = expectNotNil(opt2(nent))
+#endif
+}

test/stdlib/Span/StringUTF8SpanProperty.swift

@@ -85,3 +85,45 @@ suite.test("Span from Large Native String's Substring")
     expectEqual(span[i], u[i])
   }
 }
+
+suite.test("Span from String.utf8Span")
+.require(.stdlib_6_2).code {
+  guard #available(SwiftStdlib 6.2, *) else { return }
+
+  let s = String(200)
+  let utf8span = s.utf8Span
+  let span1 = utf8span.span
+  let utf8view = s.utf8
+  let span2 = utf8view.span
+  expectEqual(span1.count, span2.count)
+  for (i,j) in zip(span1.indices, span2.indices) {
+    expectEqual(span1[i], span2[j])
+  }
+}
+
+suite.test("UTF8Span from Span")
+.require(.stdlib_6_2).code {
+  guard #available(SwiftStdlib 6.2, *) else { return }
+
+  let s = String(200).utf8
+  let span1 = s.span
+  guard let utf8 = expectNotNil(try? UTF8Span(validating: span1)) else { return }
+
+  let span2 = utf8.span
+  expectTrue(span1.isIdentical(to: span2))
+}
+
+suite.test("Span from Substring.utf8Span")
+.require(.stdlib_6_2).code {
+  guard #available(SwiftStdlib 6.2, *) else { return }
+
+  let s = String(22000).dropFirst().dropLast()
+  let utf8span = s.utf8Span
+  let span1 = utf8span.span
+  let utf8view = s.utf8
+  let span2 = utf8view.span
+  expectEqual(span1.count, span2.count)
+  for (i,j) in zip(span1.indices, span2.indices) {
+    expectEqual(span1[i], span2[j])
+  }
+}