[stdlib] String.UnicodeScalarView: Review index validation

Author: lorentey

stdlib/public/core/StringGuts.swift

@@ -367,12 +367,14 @@ extension _StringGuts {
   /// not set the flags that this method relies on. However, false positives
   /// cannot happen: if this method detects a mismatch, then it is guaranteed to
   /// be a real one.
+  @_alwaysEmitIntoClient
   @inline(__always)
   internal func ensureMatchingEncoding(_ i: String.Index) -> String.Index {
     if _fastPath(!isForeign && i._canBeUTF8) { return i }
     return _slowEnsureMatchingEncoding(i)
   }
 
+  @_alwaysEmitIntoClient
   @inline(never)
   internal func _slowEnsureMatchingEncoding(_ i: String.Index) -> String.Index {
     _internalInvariant(isForeign || !i._canBeUTF8)
@@ -428,6 +430,7 @@ extension _StringGuts {
   /// - has an encoding that matches this string,
   /// - is within the bounds of this string, and
   /// - is aligned on a scalar boundary.
+  @_alwaysEmitIntoClient
   internal func validateScalarIndex(_ i: String.Index) -> String.Index {
     let i = ensureMatchingEncoding(i)
     _precondition(i._encodedOffset < count, "String index is out of bounds")
@@ -441,6 +444,7 @@ extension _StringGuts {
   /// - has an encoding that matches this string,
   /// - is within `start ..< end`, and
   /// - is aligned on a scalar boundary.
+  @_alwaysEmitIntoClient
   internal func validateScalarIndex(
     _ i: String.Index,
     from start: String.Index,
@@ -460,6 +464,7 @@ extension _StringGuts {
   /// - has an encoding that matches this string,
   /// - is within the bounds of this string (including the `endIndex`), and
   /// - is aligned on a scalar boundary.
+  @_alwaysEmitIntoClient
   internal func validateInclusiveScalarIndex(
     _ i: String.Index
   ) -> String.Index {

stdlib/public/core/StringUnicodeScalarView.swift

@@ -106,42 +106,38 @@ extension String.UnicodeScalarView: BidirectionalCollection {
   /// - Precondition: The next location exists.
   @inlinable @inline(__always)
   public func index(after i: Index) -> Index {
-    // TODO(String performance): isASCII fast-path
-
-    // TODO(lorentey): Review index validation
-    _precondition(i < endIndex, "String index is out of bounds")
-    let i = _guts.scalarAlign(i)
+    let i = _guts.validateScalarIndex(i)
+    return _uncheckedIndex(after: i)
+  }
 
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  internal func _uncheckedIndex(after i: Index) -> Index {
+    // TODO(String performance): isASCII fast-path
     if _fastPath(_guts.isFastUTF8) {
       let len = _guts.fastUTF8ScalarLength(startingAt: i._encodedOffset)
       return i.encoded(offsetBy: len)._scalarAligned._knownUTF8
     }
-
     return _foreignIndex(after: i)
   }
 
-  @_alwaysEmitIntoClient // Swift 5.1 bug fix
-  public func distance(from start: Index, to end: Index) -> Int {
-    // TODO(lorentey): Review index validation
-    return _distance(from: _guts.scalarAlign(start), to: _guts.scalarAlign(end))
-  }
-
   /// Returns the previous consecutive location before `i`.
   ///
   /// - Precondition: The previous location exists.
   @inlinable @inline(__always)
   public func index(before i: Index) -> Index {
-    // TODO(lorentey): Review index validation
-    // TODO(String performance): isASCII fast-path
-
-    // Note: bounds checking in `index(before:)` is tricky as scalar aligning an
-    // index may need to access storage, but it may also move it closer towards
-    // the `startIndex`. Therefore, we must check against the `endIndex` before
-    // aligning, but we need to delay the `i > startIndex` check until after.
-    _precondition(i <= endIndex, "String index is out of bounds")
-    let i = _guts.scalarAlign(i)
+    let i = _guts.validateInclusiveScalarIndex(i)
+    // Note: Aligning an index may move it closer towards the `startIndex`, so
+    // the `i > startIndex` check needs to come after rounding.
     _precondition(i > startIndex, "String index is out of bounds")
 
+    return _uncheckedIndex(before: i)
+  }
+
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  internal func _uncheckedIndex(before i: Index) -> Index {
+    // TODO(String performance): isASCII fast-path
     if _fastPath(_guts.isFastUTF8) {
       let len = _guts.withFastUTF8 { utf8 in
         _utf8ScalarLength(utf8, endingAt: i._encodedOffset)
@@ -171,11 +167,80 @@ extension String.UnicodeScalarView: BidirectionalCollection {
   ///   must be less than the view's end index.
   @inlinable @inline(__always)
   public subscript(position: Index) -> Unicode.Scalar {
-    // TODO(lorentey): Review index validation
-    String(_guts)._boundsCheck(position)
-    let i = _guts.scalarAlign(position)
+    let i = _guts.validateScalarIndex(position)
     return _guts.errorCorrectedScalar(startingAt: i._encodedOffset).0
   }
+
+  @_alwaysEmitIntoClient // Swift 5.1 bug fix
+  public func distance(from start: Index, to end: Index) -> Int {
+    let start = _guts.validateInclusiveScalarIndex(start)
+    let end = _guts.validateInclusiveScalarIndex(end)
+
+    var i = start
+    var count = 0
+    if i < end {
+      while i < end {
+        count += 1
+        i = _uncheckedIndex(after: i)
+      }
+    }
+    else if i > end {
+      while i > end {
+        count -= 1
+        i = _uncheckedIndex(before: i)
+      }
+    }
+    return count
+  }
+
+  @_alwaysEmitIntoClient
+  public func index(_ i: Index, offsetBy distance: Int) -> Index {
+    var i = _guts.validateInclusiveScalarIndex(i)
+
+    if distance >= 0 {
+      for _ in stride(from: 0, to: distance, by: 1) {
+        _precondition(i._encodedOffset < _guts.count, "String index is out of bounds")
+        i = _uncheckedIndex(after: i)
+      }
+    } else {
+      for _ in stride(from: 0, to: distance, by: -1) {
+        _precondition(i._encodedOffset > 0, "String index is out of bounds")
+        i = _uncheckedIndex(before: i)
+      }
+    }
+    return _guts.markEncoding(i)
+  }
+
+  @_alwaysEmitIntoClient
+  public func index(
+    _ i: Index, offsetBy distance: Int, limitedBy limit: Index
+  ) -> Index? {
+    // Note: `limit` is intentionally not scalar aligned to ensure our behavior
+    // exactly matches the documentation above. We do need to ensure it has a
+    // matching encoding, though. The same goes for `start`, which is used to
+    // determine whether the limit applies at all.
+    let limit = _guts.ensureMatchingEncoding(limit)
+    let start = _guts.ensureMatchingEncoding(i)
+
+    var i = _guts.validateInclusiveScalarIndex(i)
+
+    if distance >= 0 {
+      for _ in stride(from: 0, to: distance, by: 1) {
+        guard limit < start || i < limit else { return nil }
+        _precondition(i._encodedOffset < _guts.count, "String index is out of bounds")
+        i = _uncheckedIndex(after: i)
+      }
+      guard limit < start || i <= limit else { return nil }
+    } else {
+      for _ in stride(from: 0, to: distance, by: -1) {
+        guard limit > start || i > limit else { return nil }
+        _precondition(i._encodedOffset > 0, "String index is out of bounds")
+        i = _uncheckedIndex(before: i)
+      }
+      guard limit > start || i >= limit else { return nil }
+    }
+    return _guts.markEncoding(i)
+  }
 }
 
 extension String.UnicodeScalarView {
@@ -318,9 +383,8 @@ extension String.UnicodeScalarView: RangeReplaceableCollection {
     _ bounds: Range<Index>,
     with newElements: C
   ) where C: Collection, C.Element == Unicode.Scalar {
-    // TODO(lorentey): Review index validation
     // TODO(String performance): Skip extra String and Array allocation
-
+    let bounds = _guts.validateScalarRange(bounds)
     let utf8Replacement = newElements.flatMap { String($0).utf8 }
     let replacement = utf8Replacement.withUnsafeBufferPointer {
       return String._uncheckedFromUTF8($0)
@@ -423,9 +487,8 @@ extension String.UnicodeScalarView {
 
   @available(swift, introduced: 4)
   public subscript(r: Range<Index>) -> String.UnicodeScalarView.SubSequence {
-    // TODO(lorentey): Review index validation
-    _failEarlyRangeCheck(r, bounds: startIndex..<endIndex)
-    return String.UnicodeScalarView.SubSequence(self, _bounds: r)
+    let r = _guts.validateScalarRange(r)
+    return SubSequence(_unchecked: self, bounds: r)
   }
 }
 

stdlib/public/core/Substring.swift

@@ -1018,14 +1018,24 @@ extension Substring {
     internal var _slice: Slice<String.UnicodeScalarView>
 
     /// Creates an instance that slices `base` at `_bounds`.
-    @inlinable
+    internal init(
+      _unchecked base: String.UnicodeScalarView, bounds: Range<Index>
+    ) {
+      _slice = Slice(base: base, bounds: bounds)
+    }
+
+    /// Creates an instance that slices `base` at `_bounds`.
+    @usableFromInline // This used to be inlinable before 5.7
+    @available(*, deprecated) // Use `init(_unchecked:)` in new code.
     internal init(_ base: String.UnicodeScalarView, _bounds: Range<Index>) {
       let start = base._guts.scalarAlign(_bounds.lowerBound)
       let end = base._guts.scalarAlign(_bounds.upperBound)
-      _slice = Slice(
-        base: String(base._guts).unicodeScalars,
-        bounds: Range(_uncheckedBounds: (start, end)))
+      _slice = Slice(base: base, bounds: Range(_uncheckedBounds: (start, end)))
     }
+
+    @_alwaysEmitIntoClient
+    @inline(__always)
+    internal var _wholeGuts: _StringGuts { _slice._base._guts }
   }
 }
 
@@ -1038,87 +1048,77 @@ extension Substring.UnicodeScalarView: BidirectionalCollection {
   //
   // Plumb slice operations through
   //
-  @inlinable
-  public var startIndex: Index { return _slice.startIndex }
+  @inlinable @inline(__always)
+  public var startIndex: Index { _slice._startIndex }
 
-  @inlinable
-  public var endIndex: Index { return _slice.endIndex }
+  @inlinable @inline(__always)
+  public var endIndex: Index { _slice._endIndex }
 
   @inlinable
   public subscript(index: Index) -> Element {
-    // TODO(lorentey): Review index validation
-    return _slice[index]
+    let index = _wholeGuts.validateScalarIndex(
+      index, from: startIndex, to: endIndex)
+    return _wholeGuts.errorCorrectedScalar(startingAt: index._encodedOffset).0
   }
 
   @inlinable
   public var indices: Indices {
-    // TODO(lorentey): Review index validation
     return _slice.indices
   }
 
   @inlinable
   public func index(after i: Index) -> Index {
-    // TODO(lorentey): Review index validation
-    return _slice.index(after: i)
+    _slice._base.index(after: i)
   }
 
   @inlinable
   public func formIndex(after i: inout Index) {
-    // TODO(lorentey): Review index validation
-    _slice.formIndex(after: &i)
+    _slice._base.formIndex(after: &i)
   }
 
   @inlinable
   public func index(_ i: Index, offsetBy n: Int) -> Index {
-    // TODO(lorentey): Review index validation
-    return _slice.index(i, offsetBy: n)
+    _slice._base.index(i, offsetBy: n)
   }
 
   @inlinable
   public func index(
     _ i: Index, offsetBy n: Int, limitedBy limit: Index
   ) -> Index? {
-    // TODO(lorentey): Review index validation
-    return _slice.index(i, offsetBy: n, limitedBy: limit)
+    _slice._base.index(i, offsetBy: n, limitedBy: limit)
   }
 
   @inlinable
   public func distance(from start: Index, to end: Index) -> Int {
-    // TODO(lorentey): Review index validation
-    return _slice.distance(from: start, to: end)
+    _slice._base.distance(from: start, to: end)
   }
 
   @inlinable
   public func _failEarlyRangeCheck(_ index: Index, bounds: Range<Index>) {
-    // TODO(lorentey): Review index validation
-    _slice._failEarlyRangeCheck(index, bounds: bounds)
+    _slice._base._failEarlyRangeCheck(index, bounds: bounds)
   }
 
   @inlinable
   public func _failEarlyRangeCheck(
     _ range: Range<Index>, bounds: Range<Index>
   ) {
-    // TODO(lorentey): Review index validation
-    _slice._failEarlyRangeCheck(range, bounds: bounds)
+    _slice._base._failEarlyRangeCheck(range, bounds: bounds)
   }
 
   @inlinable
   public func index(before i: Index) -> Index {
-    // TODO(lorentey): Review index validation
-    return _slice.index(before: i)
+    _slice._base.index(before: i)
   }
 
   @inlinable
   public func formIndex(before i: inout Index) {
-    // TODO(lorentey): Review index validation
-    _slice.formIndex(before: &i)
+    _slice._base.formIndex(before: &i)
   }
 
-  @inlinable
   public subscript(r: Range<Index>) -> Substring.UnicodeScalarView {
-    // TODO(lorentey): Review index validation
-    _failEarlyRangeCheck(r, bounds: startIndex..<endIndex)
-    return Substring.UnicodeScalarView(_slice._base, _bounds: r)
+    // Note: This used to be inlinable until Swift 5.7
+    let r = _wholeGuts.validateScalarRange(r, from: startIndex, to: endIndex)
+    return Substring.UnicodeScalarView(_unchecked: _slice._base, bounds: r)
   }
 }