Merge pull request #16733 from milseman/the_incredible_shrinking_index

[String] Shrink and simplify String.Index

Author: milseman

stdlib/public/SDK/Foundation/ExtraStringAPIs.swift

@@ -16,7 +16,7 @@ extension String.UTF16View.Index {
   @available(swift, obsoleted: 4.0)
   public init(_ offset: Int) {
     _precondition(offset >= 0, "Negative UTF16 index offset not allowed")
-    self.init(_offset: offset)
+    self.init(encodedOffset: offset)
   }
 
   @available(swift, deprecated: 3.2)

stdlib/public/core/StringCharacterView.swift

@@ -221,18 +221,14 @@ extension String._CharacterView : BidirectionalCollection {
   /// this view's `_baseOffset`.
   @inlinable // FIXME(sil-serialize-all)
   internal func _toBaseIndex(_ index: Index) -> Index {
-    return Index(
-      encodedOffset: index.encodedOffset - _baseOffset,
-      index._cache)
+    return String.Index(from: index, adjustingEncodedOffsetBy: -_baseOffset)
   }
 
   /// Translates an index in the underlying base string into a view index using
   /// this view's `_baseOffset`.
   @inlinable // FIXME(sil-serialize-all)
   internal func _toViewIndex(_ index: Index) -> Index {
-    return Index(
-      encodedOffset: index.encodedOffset + _baseOffset,
-      index._cache)
+    return String.Index(from: index, adjustingEncodedOffsetBy: _baseOffset)
   }
 
   /// The position of the first character in a nonempty character view.

stdlib/public/core/StringComparable.swift

@@ -27,6 +27,14 @@ extension _StringGuts {
     if left._bitwiseEqualTo(right) {
       return true
     }
+    if left._isSmall && right._isSmall {
+      // TODO: Ensure normality when adding UTF-8 support
+      _sanityCheck(left._isASCIIOrSmallASCII && right._isASCIIOrSmallASCII,
+        "Need to ensure normality")
+
+      // Equal small strings should be bitwise equal if ASCII
+      return false
+    }
     return compare(left, to: right) == 0
   }
 
@@ -50,6 +58,10 @@ extension _StringGuts {
     if left._bitwiseEqualTo(right) {
       return false
     }
+    if left._isSmall && right._isSmall {
+      // Small strings compare lexicographically if ASCII
+      return left._smallUTF8String._compare(right._smallUTF8String) == .less
+    }
     return compare(left, to: right) == -1
   }
 
@@ -72,14 +84,14 @@ extension _StringGuts {
   ) -> Int {
     defer { _fixLifetime(left) }
     defer { _fixLifetime(right) }
-    
+
     if left.isASCII && right.isASCII {
       let leftASCII = left._unmanagedASCIIView[leftRange]
       let rightASCII = right._unmanagedASCIIView[rightRange]
       let result = leftASCII.compareASCII(to: rightASCII)
       return result
     }
-    
+
     let leftBits = left.rawBits
     let rightBits = right.rawBits
 
@@ -92,14 +104,14 @@ extension _StringGuts {
   ) -> Int {
     defer { _fixLifetime(left) }
     defer { _fixLifetime(right) }
-    
+
     if left.isASCII && right.isASCII {
       let leftASCII = left._unmanagedASCIIView
       let rightASCII = right._unmanagedASCIIView
       let result = leftASCII.compareASCII(to: rightASCII)
       return result
     }
-    
+
     let leftBits = left.rawBits
     let rightBits = right.rawBits
 

stdlib/public/core/StringComparison.swift

@@ -120,25 +120,19 @@ func _compareUnicode(
 // TODO: coalesce many of these into a protocol to simplify the code
 
 extension _SmallUTF8String {
+  @inlinable
   func _compare(_ other: _SmallUTF8String) -> _Ordering {
 #if arch(i386) || arch(arm)
     _conditionallyUnreachable()
 #else
-    if _fastPath(self.isASCII && other.isASCII) {
-      // TODO: fast in-register comparison
-      return self.withUnmanagedASCII { selfView in
-        return other.withUnmanagedASCII { otherView in
-          return _Ordering(signedNotation: selfView.compareASCII(to: otherView))
-        }
-      }
-    }
-
-    // TODO: fast in-register comparison
-    return self.withUnmanagedUTF16 { selfView in
-      return other.withUnmanagedUTF16 { otherView in
-        return selfView._compare(otherView)
-      }
+    // TODO: Ensure normality when adding UTF-8 support
+    _sanityCheck(self.isASCII && other.isASCII, "Need to ensure normality")
+    if self._storage == other._storage { return .equal }
+    for i in 0..<Swift.min(self.count, other.count) {
+      if self[i] < other[i] { return .less }
+      if self[i] > other[i] { return .greater }
     }
+    return self.count < other.count ? .less : .greater
 #endif // 64-bit
   }
   func _compare(_contiguous other: _StringGuts) -> _Ordering {

stdlib/public/core/StringGraphemeBreaking.swift

@@ -18,29 +18,10 @@ internal var _CR: UInt8 { return 0x0d }
 @inlinable // FIXME(sil-serialize-all)
 internal var _LF: UInt8 { return 0x0a }
 
-extension String.Index {
-  @inlinable // FIXME(sil-serialize-all)
-  internal init(encodedOffset: Int, characterStride stride: Int) {
-    if _slowPath(stride == 0 || stride > UInt16.max) {
-      // Don't store a 0 stride for the endIndex
-      // or a truncated stride for an overlong grapheme cluster.
-      self.init(encodedOffset: encodedOffset)
-      return
-    }
-    self.init(
-      encodedOffset: encodedOffset,
-      .character(stride: UInt16(truncatingIfNeeded: stride)))
-  }
-}
-
 extension _StringVariant {
   @inlinable
   internal func _stride(at i: String.Index) -> Int {
-    if case .character(let stride) = i._cache {
-      // TODO: should _fastPath the case somehow
-      _sanityCheck(stride > 0)
-      return Int(stride)
-    }
+    if let stride = i.characterStride { return stride }
     return characterStride(atOffset: i.encodedOffset)
   }
 

stdlib/public/core/StringIndex.swift

@@ -13,54 +13,57 @@ extension String {
   /// A position of a character or code unit in a string.
   @_fixed_layout // FIXME(sil-serialize-all)
   public struct Index {
+    internal typealias _UTF8Buffer = UTF8.EncodedScalar
+
     @usableFromInline // FIXME(sil-serialize-all)
-    internal var _compoundOffset : UInt64
+    internal var _compoundOffset: UInt64
+
     @usableFromInline
-    internal var _cache: _Cache
+    internal var _utf8Buffer = _UTF8Buffer()
 
-    internal typealias _UTF8Buffer = _ValidUTF8Buffer<UInt64>
-    @_frozen // FIXME(sil-serialize-all)
     @usableFromInline
-    internal enum _Cache {
-    case utf16
-    case utf8(buffer: _UTF8Buffer)
-    case character(stride: UInt16)
-    case unicodeScalar(value: Unicode.Scalar)
-    }
+    internal var _graphemeStrideCache: UInt16 = 0
   }
 }
 
 /// Convenience accessors
-extension String.Index._Cache {
-  @inlinable // FIXME(sil-serialize-all)
-  internal var utf16: Void? {
-    if case .utf16 = self { return () } else { return nil }
-  }
+extension String.Index {
   @inlinable // FIXME(sil-serialize-all)
-  internal var utf8: String.Index._UTF8Buffer? {
-    if case .utf8(let r) = self { return r } else { return nil }
+  internal var utf8Buffer: String.Index._UTF8Buffer? {
+    guard !_utf8Buffer.isEmpty else { return nil }
+    return _utf8Buffer
   }
+
   @inlinable // FIXME(sil-serialize-all)
-  internal var character: UInt16? {
-    if case .character(let r) = self { return r } else { return nil }
+  internal var characterStride: Int? {
+    guard _graphemeStrideCache > 0 else { return nil }
+    return Int(truncatingIfNeeded: _graphemeStrideCache)
   }
+
+  // TODO: Probably worth carving a bit for, or maybe a isSubScalar bit...
   @inlinable // FIXME(sil-serialize-all)
-  internal var unicodeScalar: UnicodeScalar? {
-    if case .unicodeScalar(let r) = self { return r } else { return nil }
+  internal var isUTF8: Bool {
+    return self.utf8Buffer != nil || self.transcodedOffset > 0
   }
 }
 
 extension String.Index : Equatable {
+  // A combined code unit and transcoded offset, for comparison purposes
+  @inlinable // FIXME(sil-serialize-all)
+  internal var _orderingValue: UInt64 {
+    return _compoundOffset
+  }
+
   @inlinable // FIXME(sil-serialize-all)
   public static func == (lhs: String.Index, rhs: String.Index) -> Bool {
-    return lhs._compoundOffset == rhs._compoundOffset
+    return lhs._orderingValue == rhs._orderingValue
   }
 }
 
 extension String.Index : Comparable {
   @inlinable // FIXME(sil-serialize-all)
   public static func < (lhs: String.Index, rhs: String.Index) -> Bool {
-    return lhs._compoundOffset < rhs._compoundOffset
+    return lhs._orderingValue < rhs._orderingValue
   }
 }
 
@@ -72,56 +75,67 @@ extension String.Index : Hashable {
   ///   of this instance.
   @inlinable // FIXME(sil-serialize-all)
   public func hash(into hasher: inout Hasher) {
-    hasher.combine(_compoundOffset)
+    hasher.combine(_orderingValue)
   }
 }
 
 extension String.Index {
-  internal typealias _Self = String.Index
-  
+  @inline(__always)
+  @inlinable
+  internal init(encodedOffset: Int, transcodedOffset: Int) {
+    let cuOffset = UInt64(truncatingIfNeeded: encodedOffset)
+    _sanityCheck(
+      cuOffset & 0xFFFF_0000_0000_0000 == 0, "String length capped at 48bits")
+    let transOffset = UInt64(truncatingIfNeeded: transcodedOffset)
+    _sanityCheck(transOffset <= 4, "UTF-8 max transcoding is 4 code units")
+
+    self._compoundOffset = cuOffset &<< 2 | transOffset
+  }
+
+  @inline(__always)
+  @inlinable
+  internal init(from other: String.Index, adjustingEncodedOffsetBy adj: Int) {
+    self.init(
+      encodedOffset: other.encodedOffset &+ adj,
+      transcodedOffset: other.transcodedOffset)
+    self._utf8Buffer = other._utf8Buffer
+    self._graphemeStrideCache = other._graphemeStrideCache
+  }
+
   /// Creates a new index at the specified UTF-16 offset.
   ///
   /// - Parameter offset: An offset in UTF-16 code units.
   @inlinable // FIXME(sil-serialize-all)
   public init(encodedOffset offset: Int) {
-    _compoundOffset = UInt64(offset << _Self._strideBits)
-    _cache = .utf16
+    self.init(encodedOffset: offset, transcodedOffset: 0)
   }
 
   @inlinable // FIXME(sil-serialize-all)
-  internal init(encodedOffset o: Int, transcodedOffset: Int = 0, _ c: _Cache) {
-    _compoundOffset = UInt64(o << _Self._strideBits | transcodedOffset)
-    _cache = c
+  internal init(
+    encodedOffset offset: Int, transcodedOffset: Int, buffer: _UTF8Buffer
+  ) {
+    self.init(encodedOffset: offset, transcodedOffset: transcodedOffset)
+    self._utf8Buffer = buffer
   }
-  
-  @inlinable // FIXME(sil-serialize-all)
-  internal static var _strideBits : Int { return 2 }
-  @inlinable // FIXME(sil-serialize-all)
-  internal static var _mask : UInt64 { return (1 &<< _Self._strideBits) &- 1 }
-  
-  @inlinable // FIXME(sil-serialize-all)
-  internal mutating func _setEncodedOffset(_ x: Int) {
-    _compoundOffset = UInt64(x << _Self._strideBits)
+
+  @inlinable
+  internal init(encodedOffset: Int, characterStride: Int) {
+    self.init(encodedOffset: encodedOffset, transcodedOffset: 0)
+    if characterStride < UInt16.max {
+      self._graphemeStrideCache = UInt16(truncatingIfNeeded: characterStride)
+    }
   }
-  
+
   /// The offset into a string's UTF-16 encoding for this index.
   @inlinable // FIXME(sil-serialize-all)
   public var encodedOffset : Int {
-    return Int(_compoundOffset >> _Self._strideBits)
+    return Int(truncatingIfNeeded: _compoundOffset &>> 2)
   }
 
   /// The offset of this index within whatever encoding this is being viewed as
   @inlinable // FIXME(sil-serialize-all)
-  internal var _transcodedOffset : Int {
-    get {
-      return Int(_compoundOffset & _Self._mask)
-    }
-    set {
-      let extended = UInt64(newValue)
-      _sanityCheck(extended <= _Self._mask)
-      _compoundOffset &= ~_Self._mask
-      _compoundOffset |= extended
-    }
+  internal var transcodedOffset: Int {
+    return Int(truncatingIfNeeded: _compoundOffset & 0x3)
   }
 }
 
@@ -130,27 +144,27 @@ extension String.Index {
   @inlinable // FIXME(sil-serialize-all)
   @available(swift, deprecated: 3.2)
   @available(swift, obsoleted: 4.0)
-  public // SPI(Foundation)    
+  public // SPI(Foundation)
   init(_position: Int) {
     self.init(encodedOffset: _position)
   }
-  
+
   @inlinable // FIXME(sil-serialize-all)
   @available(swift, deprecated: 3.2)
   @available(swift, obsoleted: 4.0)
-  public // SPI(Foundation)    
-  init(_offset: Int) {
-    self.init(encodedOffset: _offset)
+  public // SPI(Foundation)
+  init(_codeUnitOffset: Int) {
+    self.init(encodedOffset: _codeUnitOffset)
   }
-  
+
   @inlinable // FIXME(sil-serialize-all)
   @available(swift, deprecated: 3.2)
   @available(swift, obsoleted: 4.0)
-  public // SPI(Foundation)    
+  public // SPI(Foundation)
   init(_base: String.Index, in c: String.CharacterView) {
     self = _base
   }
-  
+
   /// The integer offset of this index in UTF-16 code units.
   @inlinable // FIXME(sil-serialize-all)
   @available(swift, deprecated: 3.2)
@@ -171,7 +185,7 @@ extension String.Index {
 }
 
 
-// backward compatibility for index interchange.  
+// backward compatibility for index interchange.
 extension Optional where Wrapped == String.Index {
   @inlinable // FIXME(sil-serialize-all)
   @available(

stdlib/public/core/StringRangeReplaceableCollection.swift

@@ -455,11 +455,8 @@ extension String {
 
   @inlinable // FIXME(sil-serialize-all)
   internal func _stride(of i: Index) -> Int {
-    if case .character(let stride) = i._cache {
-      // TODO: should _fastPath the case somehow
-      _sanityCheck(stride > 0)
-      return Int(stride)
-    }
+    if let stride = i.characterStride { return stride }
+
     let offset = i.encodedOffset
     return _visitGuts(_guts, args: offset,
       ascii: { ascii, offset in