More efficient append(Substring), and fix use-after-free bug

Author: milseman

stdlib/public/core/SmallString.swift

@@ -243,7 +243,18 @@ extension _SmallString {
   // Appending
   @usableFromInline // testable
   @_effects(releasenone)
-  internal init?(base: _StringGuts, appending other: _StringGuts) {
+  internal init?(
+    base: __shared _StringGuts, appending other: __shared _StringGuts
+  ) {
+    self.init(
+      base: _SlicedStringGuts(base), appending: _SlicedStringGuts(other))
+  }
+
+  // Appending
+  @_effects(releasenone)
+  internal init?(
+    base: __shared _SlicedStringGuts, appending other: __shared _SlicedStringGuts
+  ) {
     guard (base.utf8Count + other.utf8Count) <= _SmallString.capacity else {
       return nil
     }

stdlib/public/core/StringGuts.swift

@@ -65,7 +65,6 @@ extension _StringGuts {
 
   @inlinable @inline(__always)
   internal init(_ storage: _StringStorage) {
-    // TODO(UTF8): We should probably store perf flags on the storage's capacity
     self.init(_StringObject(storage))
   }
 
@@ -232,53 +231,17 @@ extension _StringGuts {
   // Contents of the buffer are unspecified if nil is returned.
   @inlinable
   internal func copyUTF8(into mbp: UnsafeMutableBufferPointer<UInt8>) -> Int? {
-    let ptr = mbp.baseAddress._unsafelyUnwrappedUnchecked
-    if _fastPath(self.isFastUTF8) {
-      return self.withFastUTF8 { utf8 in
-        guard utf8.count <= mbp.count else { return nil }
-
-        let utf8Start = utf8.baseAddress._unsafelyUnwrappedUnchecked
-        ptr.initialize(from: utf8Start, count: utf8.count)
-        return utf8.count
-      }
-    }
-
-    return _foreignCopyUTF8(into: mbp)
-  }
-
-  @_effects(releasenone)
-  @usableFromInline @inline(never) // slow-path
-  internal func _foreignCopyUTF8(
-    into mbp: UnsafeMutableBufferPointer<UInt8>
-  ) -> Int? {
-    var ptr = mbp.baseAddress._unsafelyUnwrappedUnchecked
-    var numWritten = 0
-    for cu in String(self).utf8 {
-      guard numWritten < mbp.count else { return nil }
-      ptr.initialize(to: cu)
-      ptr += 1
-      numWritten += 1
-    }
-
-    return numWritten
+    // TODO(UTF8 perf): minor perf win by avoiding slicing if fast...
+    return _SlicedStringGuts(self).copyUTF8(into: mbp)
   }
 
   internal var utf8Count: Int {
     @inline(__always) get {
-      // TODO: Might be worth burning a bit for count-is-UTF-8, regardless of
-      // fast status.
-      if _fastPath(self.isFastUTF8) {
-        return self.count
-      }
-      return _foreignUTF8Count()
+      if _fastPath(self.isFastUTF8) { return count }
+      return _SlicedStringGuts(self).utf8Count
     }
   }
 
-  @_effects(releasenone)
-  @inline(never) // slow-path
-  internal func _foreignUTF8Count() -> Int {
-    return String(self).utf8.count
-  }
 }
 
 // Index
@@ -329,11 +292,25 @@ internal struct _SlicedStringGuts {
     @inline(__always) get { return _guts.isNFCFastUTF8 }
   }
 
+  @inlinable
+  internal var isASCII: Bool {
+    @inline(__always) get { return _guts.isASCII }
+  }
+
   @inlinable
   internal var isFastUTF8: Bool {
     @inline(__always) get { return _guts.isFastUTF8 }
   }
 
+  internal var utf8Count: Int {
+    @inline(__always) get {
+      if _fastPath(self.isFastUTF8) {
+        return _offsetRange.count
+      }
+      return Substring(self).utf8.count
+    }
+  }
+
   @inlinable
   internal var range: Range<String.Index> {
     @inline(__always) get {
@@ -348,6 +325,39 @@ internal struct _SlicedStringGuts {
   ) rethrows -> R {
     return try _guts.withFastUTF8(range: _offsetRange, f)
   }
-}
 
+  // Copy UTF-8 contents. Returns number written or nil if not enough space.
+  // Contents of the buffer are unspecified if nil is returned.
+  @inlinable
+  internal func copyUTF8(into mbp: UnsafeMutableBufferPointer<UInt8>) -> Int? {
+    let ptr = mbp.baseAddress._unsafelyUnwrappedUnchecked
+    if _fastPath(self.isFastUTF8) {
+      return self.withFastUTF8 { utf8 in
+        guard utf8.count <= mbp.count else { return nil }
+
+        let utf8Start = utf8.baseAddress._unsafelyUnwrappedUnchecked
+        ptr.initialize(from: utf8Start, count: utf8.count)
+        return utf8.count
+      }
+    }
+
+    return _foreignCopyUTF8(into: mbp)
+  }
 
+  @_effects(releasenone)
+  @usableFromInline @inline(never) // slow-path
+  internal func _foreignCopyUTF8(
+    into mbp: UnsafeMutableBufferPointer<UInt8>
+  ) -> Int? {
+    var ptr = mbp.baseAddress._unsafelyUnwrappedUnchecked
+    var numWritten = 0
+    for cu in Substring(self).utf8 {
+      guard numWritten < mbp.count else { return nil }
+      ptr.initialize(to: cu)
+      ptr += 1
+      numWritten += 1
+    }
+
+    return numWritten
+  }
+}

stdlib/public/core/StringGutsRangeReplaceable.swift

@@ -64,7 +64,7 @@ extension _StringGuts {
   }
 
   // Grow to accomodate at least `n` code units
-  internal mutating func grow(_ n: Int) {
+  private mutating func grow(_ n: Int) {
     defer { self._invariantCheck() }
 
     _sanityCheck(
@@ -79,7 +79,6 @@ extension _StringGuts {
           initializingFrom: $0, capacity: growthTarget, isASCII: isASCII)
       }
 
-      // TODO(UTF8): Track known ascii
       self = _StringGuts(storage)
       return
     }
@@ -88,7 +87,7 @@ extension _StringGuts {
   }
 
   @inline(never) // slow-path
-  internal mutating func _foreignGrow(_ n: Int) {
+  private mutating func _foreignGrow(_ n: Int) {
     // TODO(UTF8 perf): skip the intermediary arrays
     let selfUTF8 = Array(String(self).utf8)
     selfUTF8.withUnsafeBufferPointer {
@@ -97,54 +96,78 @@ extension _StringGuts {
     }
   }
 
-  internal mutating func append(_ other: _StringGuts) {
-    defer { self._invariantCheck() }
-
-    // Try to form a small string if possible
-    if !hasNativeStorage {
-      if let smol = _SmallString(base: self, appending: other) {
-        self = _StringGuts(smol)
-        return
-      }
+  // Ensure unique native storage with sufficient capacity for the following
+  // append.
+  private mutating func prepareForAppendInPlace(
+    otherUTF8Count otherCount: Int
+  ) {
+    defer {
+      _sanityCheck(self.uniqueNativeUnusedCapacity != nil,
+        "growth should produce uniqueness")
+      _sanityCheck(self.uniqueNativeUnusedCapacity! >= self.count + otherCount,
+        "growth should produce enough capacity")
     }
 
     // See if we can accomodate without growing or copying. If we have
     // sufficient capacity, we do not need to grow, and we can skip the copy if
     // unique. Otherwise, growth is required.
-    let otherUTF8Count = other.utf8Count
     let sufficientCapacity: Bool
-    if let unused = self.nativeUnusedCapacity, unused >= otherUTF8Count {
+    if let unused = self.nativeUnusedCapacity, unused >= otherCount {
       sufficientCapacity = true
     } else {
       sufficientCapacity = false
     }
-    if !self.isUniqueNative || !sufficientCapacity {
-      let totalCount = self.utf8Count + otherUTF8Count
-
-      // Non-unique storage: just make a copy of the appropriate size, otherwise
-      // grow like an array.
-      let growthTarget: Int
-      if sufficientCapacity {
-        growthTarget = totalCount
-      } else {
-        growthTarget = Swift.max(
-          totalCount, _growArrayCapacity(nativeCapacity ?? 0))
+    if self.isUniqueNative && sufficientCapacity {
+      return
+    }
+
+    let totalCount = self.utf8Count + otherCount
+
+    // Non-unique storage: just make a copy of the appropriate size, otherwise
+    // grow like an array.
+    let growthTarget: Int
+    if sufficientCapacity {
+      growthTarget = totalCount
+    } else {
+      growthTarget = Swift.max(
+        totalCount, _growArrayCapacity(nativeCapacity ?? 0))
+    }
+    self.grow(growthTarget)
+  }
+
+  internal mutating func append(_ other: _StringGuts) {
+    // TODO(UTF8 perf): Minor potential perf win to elevating smol fast-path
+    // prior to slicing.
+    append(_SlicedStringGuts(other))
+  }
+
+  internal mutating func append(_ slicedOther: _SlicedStringGuts) {
+    defer { self._invariantCheck() }
+
+    // Try to form a small string if possible
+    if !hasNativeStorage {
+      if let smol = _SmallString(
+        base: _SlicedStringGuts(self), appending: slicedOther
+      ) {
+        self = _StringGuts(smol)
+        return
       }
-      self.grow(growthTarget)
     }
 
-    _sanityCheck(self.uniqueNativeUnusedCapacity != nil,
-      "growth should produce uniqueness")
+    prepareForAppendInPlace(otherUTF8Count: slicedOther.utf8Count)
 
-    if other.isFastUTF8 {
-      let otherIsASCII = other.isASCII
-      other.withFastUTF8 { self.appendInPlace($0, isASCII: otherIsASCII) }
+    if slicedOther.isFastUTF8 {
+      let otherIsASCII = slicedOther.isASCII
+      slicedOther.withFastUTF8 { otherUTF8 in
+        self.appendInPlace(otherUTF8, isASCII: otherIsASCII)
+      }
       return
     }
-    _foreignAppendInPlace(other)
+
+    _foreignAppendInPlace(slicedOther)
   }
 
-  internal mutating func appendInPlace(
+  private mutating func appendInPlace(
     _ other: UnsafeBufferPointer<UInt8>, isASCII: Bool
   ) {
     self._object.nativeStorage.appendInPlace(other, isASCII: isASCII)
@@ -155,11 +178,11 @@ extension _StringGuts {
   }
 
   @inline(never) // slow-path
-  internal mutating func _foreignAppendInPlace(_ other: _StringGuts) {
+  private mutating func _foreignAppendInPlace(_ other: _SlicedStringGuts) {
     _sanityCheck(!other.isFastUTF8)
     _sanityCheck(self.uniqueNativeUnusedCapacity != nil)
 
-    var iter = String(other).utf8.makeIterator()
+    var iter = Substring(other).utf8.makeIterator()
     self._object.nativeStorage.appendInPlace(&iter, isASCII: other.isASCII)
 
     // We re-initialize from the modified storage to pick up new count, flags,

stdlib/public/core/StringRangeReplaceableCollection.swift

@@ -152,8 +152,7 @@ extension String: RangeReplaceableCollection {
   }
 
   public mutating func append(contentsOf newElements: Substring) {
-    // TODO(UTF8 perf): This is a horribly slow means...
-    self.append(newElements._ephemeralString)
+    self._guts.append(newElements._slicedGuts)
   }
 
   /// Appends the characters in the given sequence to the string.

stdlib/public/core/StringStorage.swift

@@ -157,19 +157,6 @@ private func determineCodeUnitCapacity(_ desiredCapacity: Int) -> Int {
 
 // Creation
 extension _StringStorage {
-  @_effects(releasenone)
-  @nonobjc
-  private static func create(
-    capacity: Int, countAndFlags: CountAndFlags
-  ) -> _StringStorage {
-    _sanityCheck(capacity >= countAndFlags.count)
-
-    let realCapacity = determineCodeUnitCapacity(capacity)
-    _sanityCheck(realCapacity > capacity)
-    return _StringStorage.create(
-      realCodeUnitCapacity: realCapacity, countAndFlags: countAndFlags)
-  }
-
   @inline(never) // rdar://problem/44542202
   @_effects(releasenone)
   @nonobjc
@@ -192,6 +179,19 @@ extension _StringStorage {
     return storage
   }
 
+  @_effects(releasenone)
+  @nonobjc
+  private static func create(
+    capacity: Int, countAndFlags: CountAndFlags
+  ) -> _StringStorage {
+    _sanityCheck(capacity >= countAndFlags.count)
+
+    let realCapacity = determineCodeUnitCapacity(capacity)
+    _sanityCheck(realCapacity > capacity)
+    return _StringStorage.create(
+      realCodeUnitCapacity: realCapacity, countAndFlags: countAndFlags)
+  }
+
   @_effects(releasenone)
   @nonobjc
   internal static func create(

stdlib/public/core/Substring.swift

@@ -105,6 +105,11 @@ public struct Substring {
     _invariantCheck()
   }
 
+  @inline(__always)
+  internal init(_ slice: _SlicedStringGuts) {
+    self.init(String(slice._guts)[slice.range])
+  }
+
   /// Creates an empty substring.
   @inlinable @inline(__always)
   public init() {