[stdlib] Fix Substring.UnicodeScalarView.replaceSubrange

Author: lorentey

stdlib/public/core/StringGutsRangeReplaceable.swift

@@ -342,6 +342,60 @@ extension _StringGuts {
     return Range(_uncheckedBounds: (i, j))
   }
 
+  // - Returns: The encoded offset range of the replaced contents in the result.
+  @discardableResult
+  internal mutating func replaceSubrange<C>(
+    _ bounds: Range<Index>,
+    with newElements: C
+  ) -> Range<Int>
+  where C: Collection, C.Iterator.Element == UnicodeScalar {
+    if isUniqueNative {
+      if let repl = newElements as? String.UnicodeScalarView {
+        if repl._guts.isFastUTF8 {
+          return repl._guts.withFastUTF8 {
+            uniqueNativeReplaceSubrange(
+              bounds, with: $0, isASCII: repl._guts.isASCII)
+          }
+        }
+      } else if let repl = newElements as? Substring.UnicodeScalarView {
+        if repl._wholeGuts.isFastUTF8 {
+          return repl._wholeGuts.withFastUTF8(range: repl._offsetRange) {
+            uniqueNativeReplaceSubrange(
+              bounds, with: $0, isASCII: repl._wholeGuts.isASCII)
+          }
+        }
+      }
+      if #available(SwiftStdlib 5.1, *) {
+        return uniqueNativeReplaceSubrange(
+          bounds, with: newElements.lazy.flatMap { $0.utf8 })
+      } else {
+        // FIXME: The stdlib should not have a deployment target this ancient.
+        let c = newElements.reduce(0) { $0 + UTF8.width($1) }
+        var utf8: [UInt8] = []
+        utf8.reserveCapacity(c)
+        utf8 = newElements.reduce(into: utf8) { utf8, next in
+          next.withUTF8CodeUnits { utf8.append(contentsOf: $0) }
+        }
+        return uniqueNativeReplaceSubrange(bounds, with: utf8)
+      }
+    }
+
+    var result = String.UnicodeScalarView()
+    // FIXME: It should be okay to get rid of excess capacity
+    // here. rdar://problem/45635432
+    if let capacity = self.nativeCapacity {
+      result.reserveCapacity(capacity)
+    }
+    let selfStr = String.UnicodeScalarView(self)
+    result.append(contentsOf: selfStr[..<bounds.lowerBound])
+    let i = result._guts.count
+    result.append(contentsOf: newElements)
+    let j = result._guts.count
+    result.append(contentsOf: selfStr[bounds.upperBound...])
+    self = result._guts
+    return Range(_uncheckedBounds: (i, j))
+  }
+
   // - Returns: The encoded offset range of the replaced contents in the result.
   internal mutating func uniqueNativeReplaceSubrange(
     _ bounds: Range<Index>,
@@ -386,5 +440,116 @@ extension _StringGuts {
     self = _StringGuts(_object.nativeStorage)
     return Range(_uncheckedBounds: (start, start + replCount))
   }
+
+  /// Run `body` to mutate the given `subrange` of this string within
+  /// `startIndex ..< endIndex`, then update `startIndex` and `endIndex` to be
+  /// valid positions in the resulting string, addressing the same (logical)
+  /// locations as in the original string.
+  ///
+  /// This is used by both `Substring` and `Substring.UnicodeScalarView` to
+  /// implement their `replaceSubrange` methods.
+  ///
+  /// - Parameter subrange: A scalar-aligned offset range in this string.
+  /// - Parameter startIndex: The start index of the substring that performs
+  ///   this operation.
+  /// - Parameter endIndex: The end index of the substring that performs this
+  ///   operations.
+  /// - Parameter body: The mutation operation to execute on `self`. The
+  ///   returned offset range must correspond to `subrange` in the resulting
+  ///   string.
+  internal mutating func mutateSubrangeInSubstring(
+    subrange: Range<Index>,
+    startIndex: inout Index,
+    endIndex: inout Index,
+    with body: (inout _StringGuts) -> Range<Int>
+  ) {
+    _internalInvariant(
+      subrange.lowerBound >= startIndex && subrange.upperBound <= endIndex)
+
+    if _slowPath(isKnownUTF16) {
+      // UTF-16 (i.e., foreign) string. The mutation will convert this to the
+      // native UTF-8 encoding, so we need to do some extra work to preserve our
+      // bounds.
+      let utf8StartOffset = String(self).utf8.distance(
+        from: self.startIndex, to: startIndex)
+      let oldUTF8Count = String(self).utf8.distance(
+        from: startIndex, to: endIndex)
+
+      let oldUTF8SubrangeCount = String(self).utf8.distance(
+        from: subrange.lowerBound, to: subrange.upperBound)
+
+      let newUTF8Subrange = body(&self)
+      _internalInvariant(!isKnownUTF16)
+
+      let newUTF8Count =
+        oldUTF8Count + newUTF8Subrange.count - oldUTF8SubrangeCount
+
+      // Get the character stride in the entire string, not just the substring.
+      // (Characters in a substring may end beyond the bounds of it.)
+      let newStride = _opaqueCharacterStride(
+        startingAt: utf8StartOffset, in: utf8StartOffset ..< count)
+
+      startIndex = String.Index(
+        encodedOffset: utf8StartOffset,
+        transcodedOffset: 0,
+        characterStride: newStride)._scalarAligned._knownUTF8
+      if isOnGraphemeClusterBoundary(startIndex) {
+        startIndex = startIndex._characterAligned
+      }
+
+      endIndex = String.Index(
+        encodedOffset: utf8StartOffset + newUTF8Count,
+        transcodedOffset: 0)._scalarAligned._knownUTF8
+      return
+    }
+
+    // UTF-8 string.
+
+    let oldRange = subrange._encodedOffsetRange
+    let newRange = body(&self)
+
+    let oldBounds = Range(
+      _uncheckedBounds: (startIndex._encodedOffset, endIndex._encodedOffset))
+    let newBounds = Range(_uncheckedBounds: (
+        oldBounds.lowerBound,
+        oldBounds.upperBound &+ newRange.count &- oldRange.count))
+
+    // Update `startIndex` if necessary. The replacement may have invalidated
+    // its cached character stride and character alignment flag, but not its
+    // stored offset, encoding, or scalar alignment.
+    //
+    // We are exploiting the fact that mutating the string _after_ the scalar
+    // following the end of the character at `startIndex` cannot possibly change
+    // the length of that character. (This is true because `index(after:)` never
+    // needs to look ahead by more than one Unicode scalar.)
+    let oldStride = startIndex.characterStride ?? 0
+    if oldRange.lowerBound <= oldBounds.lowerBound &+ oldStride {
+      // Get the character stride in the entire string, not just the substring.
+      // (Characters in a substring may end beyond the bounds of it.)
+      let newStride = _opaqueCharacterStride(
+        startingAt: newBounds.lowerBound,
+        in: newBounds.lowerBound ..< self.count)
+      var newStart = String.Index(
+        encodedOffset: newBounds.lowerBound,
+        characterStride: newStride
+      )._scalarAligned._knownUTF8
+
+      // Preserve character alignment flag if possible
+      if startIndex._isCharacterAligned,
+        (oldRange.lowerBound > oldBounds.lowerBound ||
+         isOnGraphemeClusterBoundary(newStart)) {
+        newStart = newStart._characterAligned
+      }
+
+      startIndex = newStart
+    }
+
+    // Update `endIndex`.
+    if newBounds.upperBound != endIndex._encodedOffset {
+      endIndex = Index(
+        _encodedOffset: newBounds.upperBound
+      )._scalarAligned._knownUTF8
+    }
+  }
 }
 

stdlib/public/core/Substring.swift

@@ -470,7 +470,6 @@ extension Substring: StringProtocol {
     // Note: SE-0180 requires us to use `subrange` bounds even if they aren't
     // `Character` aligned. (We still have to round things down to the nearest
     // scalar boundary, though, or we may generate ill-formed encodings.)
-    defer { _invariantCheck() }
     let subrange = _validateScalarRange(subrange)
 
     // Replacing the range is easy -- we can just reuse `String`'s
@@ -488,85 +487,14 @@ extension Substring: StringProtocol {
     //      merged with the Character preceding/following the replaced range.
     //
     // The best way to avoid problems in these cases is to lower index
-    // calculations to Unicode scalars (or below). In this implementation, we
-    // are measuring things in UTF-8 code units, for efficiency.
-
-    if _slowPath(_wholeGuts.isKnownUTF16) {
-      // UTF-16 (i.e., foreign) string. The mutation will convert this to the
-      // native UTF-8 encoding, so we need to do some extra work to preserve our
-      // bounds.
-      let utf8StartOffset = _slice._base.utf8.distance(
-        from: _slice._base.startIndex,
-        to: _slice._startIndex)
-      let oldUTF8Count = self.utf8.count
-
-      let oldSubrangeCount = self.utf8.distance(
-        from: subrange.lowerBound, to: subrange.upperBound)
-
-      let newUTF8Subrange = _slice._base._guts.replaceSubrange(
-        subrange, with: newElements)
-      _internalInvariant(!_wholeGuts.isKnownUTF16)
-
-      let newUTF8Count = oldUTF8Count + newUTF8Subrange.count - oldSubrangeCount
-
-      // Get the character stride in the entire string, not just the substring.
-      // (Characters in a substring may end beyond the bounds of it.)
-      let newStride = _wholeGuts._opaqueCharacterStride(
-        startingAt: utf8StartOffset,
-        in: utf8StartOffset ..< _wholeGuts.count)
-
-      _slice._startIndex = String.Index(
-        encodedOffset: utf8StartOffset,
-        transcodedOffset: 0,
-        characterStride: newStride)._scalarAligned._knownUTF8
-      _slice._endIndex = String.Index(
-        encodedOffset: utf8StartOffset + newUTF8Count,
-        transcodedOffset: 0)._scalarAligned._knownUTF8
-      return
-    }
-
-    // UTF-8 string.
-
-    let oldRange = Range(_uncheckedBounds: (
-        subrange.lowerBound._encodedOffset, subrange.upperBound._encodedOffset))
-
-    let newRange = _slice._base._guts.replaceSubrange(
-      subrange, with: newElements)
-
-    let newOffsetBounds = Range(_uncheckedBounds: (
-        startIndex._encodedOffset,
-        endIndex._encodedOffset &+ newRange.count &- oldRange.count))
-
-    // Update `startIndex` if necessary. The replacement may have invalidated
-    // its cached character stride, but not its stored offset.
-    //
-    // We are exploiting the fact that mutating the string _after_ the scalar
-    // following the end of the character at `startIndex` cannot possibly change
-    // the length of that character. (This is true because `index(after:)` never
-    // needs to look ahead by more than one Unicode scalar.)
-    if
-      let stride = startIndex.characterStride,
-      oldRange.lowerBound <= startIndex._encodedOffset &+ stride
-    {
-      // Get the character stride in the entire string, not just the substring.
-      // (Characters in a substring may end beyond the bounds of it.)
-      let newStride = _wholeGuts._opaqueCharacterStride(
-        startingAt: newOffsetBounds.lowerBound,
-        in: newOffsetBounds.lowerBound ..< _wholeGuts.count)
-      _slice._startIndex = String.Index(
-        encodedOffset: startIndex._encodedOffset,
-        characterStride: newStride
-      )._scalarAligned._knownUTF8
-    }
-
-    // Update endIndex.
-    if newOffsetBounds.upperBound != endIndex._encodedOffset {
-      _slice._endIndex = Index(
-        _encodedOffset: newOffsetBounds.upperBound
-      )._scalarAligned._knownUTF8
-    }
+    // calculations to Unicode scalars (or below).
+    _slice._base._guts.mutateSubrangeInSubstring(
+      subrange: subrange,
+      startIndex: &_slice._startIndex,
+      endIndex: &_slice._endIndex,
+      with: { $0.replaceSubrange(subrange, with: newElements) })
 
-    // TODO(lorentey): Mark new bounds character aligned if possible
+    _invariantCheck()
   }
 
   /// Creates a string from the given Unicode code units in the specified
@@ -1209,17 +1137,32 @@ extension String {
   }
 }
 
-// FIXME: The other String views should be RangeReplaceable too.
 extension Substring.UnicodeScalarView: RangeReplaceableCollection {
   @inlinable
   public init() { _slice = Slice.init() }
 
   public mutating func replaceSubrange<C: Collection>(
     _ subrange: Range<Index>, with replacement: C
   ) where C.Element == Element {
-    // TODO(lorentey): Don't forward to slice
     let subrange = _wholeGuts.validateScalarRange(subrange, in: _bounds)
-    _slice.replaceSubrange(subrange, with: replacement)
+
+    // Replacing the range is easy -- we can just reuse `String`'s
+    // implementation. However, we must also update `startIndex` and `endIndex`
+    // to keep them valid & pointing to the same positions, which is somewhat
+    // tricky.
+    //
+    // In Swift <=5.6, this used to forward to `Slice.replaceSubrange`, which
+    // (incorrectly) assumes that indices before the replaced subrange are
+    // preserved after the mutation. (This isn't true for strings, esp. when the
+    // original value is UTF-16 encoded.)
+
+    _slice._base._guts.mutateSubrangeInSubstring(
+      subrange: subrange,
+      startIndex: &_slice._startIndex,
+      endIndex: &_slice._endIndex,
+      with: { $0.replaceSubrange(subrange, with: replacement) })
+
+    _invariantCheck()
   }
 }