Add "unchecked" subscripts for buffers and arrays

Introduce `subscript(unchecked:)` and `subscript(uncheckedBounds:)`
operations for the various buffer and array types, allowing one to
explicitly opt out of bounds checking for release builds.

Author: DougGregor

stdlib/public/Differentiation/ArrayDifferentiation.swift

@@ -193,7 +193,7 @@ extension Array: Differentiable where Element: Differentiable {
 
 extension Array where Element: Differentiable {
   @usableFromInline
-  @derivative(of: subscript)
+  @derivative(of: subscript(_:))
   func _vjpSubscript(index: Int) -> (
     value: Element, pullback: (Element.TangentVector) -> TangentVector
   ) {
@@ -208,7 +208,7 @@ extension Array where Element: Differentiable {
   }
 
   @usableFromInline
-  @derivative(of: subscript)
+  @derivative(of: subscript(_:))
   func _jvpSubscript(index: Int) -> (
     value: Element, differential: (TangentVector) -> Element.TangentVector
   ) {

stdlib/public/core/Array.swift

@@ -400,6 +400,28 @@ extension Array {
     return _DependenceToken()
   }
 
+  /// Check that the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug buikds.
+  @inlinable
+  @_semantics("array.check_subscript")
+  @_effects(notEscaping self.**)
+  public // @testable
+  func _debugCheckSubscript(
+    _ index: Int, wasNativeTypeChecked: Bool
+  ) -> _DependenceToken {
+#if _runtime(_ObjC)
+    // There is no need to do bounds checking for the non-native case because
+    // ObjectiveC arrays do bounds checking by their own.
+    // And in the native-non-type-checked case, it's also not needed to do bounds
+    // checking here, because it's done in ArrayBuffer._getElementSlowPath.
+    if _fastPath(wasNativeTypeChecked) {
+      _buffer._native._debugCheckValidSubscript(index)
+    }
+#else
+    _buffer._debugCheckValidSubscript(index)
+#endif
+    return _DependenceToken()
+  }
   /// Check that the given `index` is valid for subscripting, i.e.
   /// `0 ≤ index < count`.
   ///
@@ -411,6 +433,17 @@ extension Array {
     _buffer._checkValidSubscriptMutating(index)
   }
 
+  /// Check that the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug configurations.
+  ///
+  /// - Precondition: The buffer must be uniquely referenced and native.
+  @_alwaysEmitIntoClient
+  @_semantics("array.check_subscript")
+  @_effects(notEscaping self.**)
+  internal func _debugCheckSubscript_mutating(_ index: Int) {
+    _buffer._debugCheckValidSubscriptMutating(index)
+  }
+  
   /// Check that the specified `index` is valid, i.e. `0 ≤ index ≤ count`.
   @inlinable
   @_semantics("array.check_index")
@@ -420,6 +453,16 @@ extension Array {
     _precondition(index >= startIndex, "Negative Array index is out of range")
   }
 
+  /// Check that the specified `index` is valid, i.e. `0 ≤ index ≤ count`,
+  /// but only in debug configurations.
+  @inlinable
+  @_semantics("array.check_index")
+  @_effects(notEscaping self.**)
+  internal func _debugCheckIndex(_ index: Int) {
+    _debugPrecondition(index <= endIndex, "Array index is out of range")
+    _debugPrecondition(index >= startIndex, "Negative Array index is out of range")
+  }
+  
   @_semantics("array.get_element")
   @_effects(notEscaping self.value**)
   @_effects(escaping self.value**.class*.value** -> return.value**)
@@ -761,6 +804,38 @@ extension Array: RandomAccessCollection, MutableCollection {
     }
   }
 
+  /// Accesses the element at the specified position without bounds
+  /// checking.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(unchecked index: Int) -> Element {
+    get {
+      // This call may be hoisted or eliminated by the optimizer.  If
+      // there is an inout violation, this value may be stale so needs to be
+      // checked again below.
+      let wasNativeTypeChecked = _hoistableIsNativeTypeChecked()
+
+      // Make sure the index is in range and wasNativeTypeChecked is
+      // still valid.
+      let token = _debugCheckSubscript(
+        index, wasNativeTypeChecked: wasNativeTypeChecked)
+
+      return _getElement(
+        index, wasNativeTypeChecked: wasNativeTypeChecked,
+        matchingSubscriptCheck: token)
+    }
+    _modify {
+      _makeMutableAndUnique() // makes the array native, too
+      _debugCheckSubscript_mutating(index)
+      let address = _buffer.mutableFirstElementAddress + index
+      defer { _endMutation() }
+      yield &address.pointee
+    }
+  }
+  
   /// Accesses a contiguous subrange of the array's elements.
   ///
   /// The returned `ArraySlice` instance uses the same indices for the same
@@ -804,6 +879,31 @@ extension Array: RandomAccessCollection, MutableCollection {
     }
   }
 
+  /// Accesses a contiguous subrange of the array's elements without
+  /// bounds checks on the range.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(uncheckedBounds bounds: Range<Int>) -> ArraySlice<Element> {
+    get {
+      _debugCheckIndex(bounds.lowerBound)
+      _debugCheckIndex(bounds.upperBound)
+      return ArraySlice(_buffer: _buffer[bounds])
+    }
+    set(rhs) {
+      _debugCheckIndex(bounds.lowerBound)
+      _debugCheckIndex(bounds.upperBound)
+      // If the replacement buffer has same identity, and the ranges match,
+      // then this was a pinned in-place modification, nothing further needed.
+      if self[bounds]._buffer.identity != rhs._buffer.identity
+      || bounds != rhs.startIndex..<rhs.endIndex {
+        self.replaceSubrange(bounds, with: rhs)
+      }
+    }
+  }
+  
   /// The number of elements in the array.
   @inlinable
   @_semantics("array.get_count")

stdlib/public/core/ArrayBuffer.swift

@@ -467,6 +467,14 @@ extension _ArrayBuffer {
     _native._checkValidSubscriptMutating(index)
   }
 
+  /// Traps unless the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`.
+  ///
+  /// - Precondition: The buffer must be mutable.
+  @_alwaysEmitIntoClient
+  internal func _debugCheckValidSubscriptMutating(_ index: Int) {
+    _native._debugCheckValidSubscriptMutating(index)
+  }
   /// The number of elements the buffer can store without reallocation.
   ///
   /// This property is obsolete. It's only used for the ArrayBufferProtocol and

stdlib/public/core/ContiguousArray.swift

@@ -91,6 +91,14 @@ extension ContiguousArray {
     _buffer._checkValidSubscript(index)
   }
 
+  /// Check that the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug builds.
+  @inlinable
+  @inline(__always)
+  internal func _debugCheckSubscript_native(_ index: Int) {
+    _buffer._debugCheckValidSubscript(index)
+  }
+  
   /// Check that the given `index` is valid for subscripting, i.e.
   /// `0 ≤ index < count`.
   ///
@@ -101,6 +109,16 @@ extension ContiguousArray {
     _buffer._checkValidSubscriptMutating(index)
   }
 
+  /// Check that the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug builds.
+  ///
+  /// - Precondition: The buffer must be uniquely referenced and native.
+  @_alwaysEmitIntoClient
+  @_semantics("array.check_subscript")
+  internal func _debugCheckSubscript_mutating(_ index: Int) {
+    _buffer._debugCheckValidSubscriptMutating(index)
+  }
+  
   /// Check that the specified `index` is valid, i.e. `0 ≤ index ≤ count`.
   @inlinable
   @_semantics("array.check_index")
@@ -109,6 +127,14 @@ extension ContiguousArray {
     _precondition(index >= startIndex, "Negative ContiguousArray index is out of range")
   }
 
+  /// Check that the specified `index` is valid, i.e. `0 ≤ index ≤ count`, but only in debug builds.
+  @inlinable
+  @_semantics("array.check_index")
+  internal func _debugCheckIndex(_ index: Int) {
+    _debugPrecondition(index <= endIndex, "ContiguousArray index is out of range")
+    _debugPrecondition(index >= startIndex, "Negative ContiguousArray index is out of range")
+  }
+  
   @inlinable
   @_semantics("array.get_element_address")
   internal func _getElementAddress(_ index: Int) -> UnsafeMutablePointer<Element> {
@@ -421,6 +447,27 @@ extension ContiguousArray: RandomAccessCollection, MutableCollection {
     }
   }
 
+  /// Accesses the element at the specified position without
+  /// bounds checking.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(unchecked index: Int) -> Element {
+    get {
+      _debugCheckSubscript_native(index)
+      return _buffer.getElement(index)
+    }
+    _modify {
+      _makeMutableAndUnique()
+      _debugCheckSubscript_mutating(index)
+      let address = _buffer.mutableFirstElementAddress + index
+      defer { _endMutation() }
+      yield &address.pointee
+    }
+  }
+
   /// Accesses a contiguous subrange of the array's elements.
   ///
   /// The returned `ArraySlice` instance uses the same indices for the same
@@ -464,6 +511,30 @@ extension ContiguousArray: RandomAccessCollection, MutableCollection {
     }
   }
 
+  /// Accesses a contiguous subrange of the array's elements.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(uncheckedBounds bounds: Range<Int>) -> ArraySlice<Element> {
+    get {
+      _debugCheckIndex(bounds.lowerBound)
+      _debugCheckIndex(bounds.upperBound)
+      return ArraySlice(_buffer: _buffer[bounds])
+    }
+    set(rhs) {
+      _debugCheckIndex(bounds.lowerBound)
+      _debugCheckIndex(bounds.upperBound)
+      // If the replacement buffer has same identity, and the ranges match,
+      // then this was a pinned in-place modification, nothing further needed.
+      if self[bounds]._buffer.identity != rhs._buffer.identity
+      || bounds != rhs.startIndex..<rhs.endIndex {
+        self.replaceSubrange(bounds, with: rhs)
+      }
+    }
+  }
+
   /// The number of elements in the array.
   @inlinable
   public var count: Int {

stdlib/public/core/ContiguousArrayBuffer.swift

@@ -669,6 +669,19 @@ internal struct _ContiguousArrayBuffer<Element>: _ArrayBufferProtocol {
     )
   }
 
+  /// Traps unless the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug mode.
+  ///
+  /// - Precondition: The buffer must be immutable.
+  @inlinable
+  @inline(__always)
+  internal func _debugCheckValidSubscript(_ index: Int) {
+    _debugPrecondition(
+      (index >= 0) && (index < immutableCount),
+      "Index out of range"
+    )
+  }
+
   /// Traps unless the given `index` is valid for subscripting, i.e.
   /// `0 ≤ index < count`.
   ///
@@ -682,6 +695,19 @@ internal struct _ContiguousArrayBuffer<Element>: _ArrayBufferProtocol {
     )
   }
 
+  /// Traps unless the given `index` is valid for subscripting, i.e.
+  /// `0 ≤ index < count`, but only in debug builds.
+  ///
+  /// - Precondition: The buffer must be mutable.
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  internal func _debugCheckValidSubscriptMutating(_ index: Int) {
+    _debugPrecondition(
+      (index >= 0) && (index < mutableCount),
+      "Index out of range"
+    )
+  }
+
   /// The number of elements the buffer can store without reallocation.
   ///
   /// This property is obsolete. It's only used for the ArrayBufferProtocol and

stdlib/public/core/SliceBuffer.swift

@@ -307,6 +307,14 @@ internal struct _SliceBuffer<Element>
       index >= startIndex && index < endIndex, "Index out of bounds")
   }
 
+  /// Traps unless the given `index` is valid for subscripting, i.e.
+  /// `startIndex ≤ index < endIndex`, but only in debug builkds.
+  @inlinable
+  internal func _debugCheckValidSubscript(_ index: Int) {
+    _debugPrecondition(
+      index >= startIndex && index < endIndex, "Index out of bounds")
+  }
+
   @inlinable
   internal var capacity: Int {
     let count = self.count

stdlib/public/core/UnsafeBufferPointer.swift.gyb

@@ -353,6 +353,28 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
 %end
   }
 
+  /// Accesses the element at the specified position without bounds
+  /// checking.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(unchecked i: Int) -> Element {
+    get {
+      _debugPrecondition(i >= 0)
+      _debugPrecondition(i < endIndex)
+      return _position._unsafelyUnwrappedUnchecked[i]
+    }
+%if Mutable:
+    nonmutating _modify {
+      _debugPrecondition(i >= 0)
+      _debugPrecondition(i < endIndex)
+      yield &_position._unsafelyUnwrappedUnchecked[i]
+    }
+%end
+  }
+
   // Skip all debug and runtime checks
 
   @inlinable // unsafe-performance
@@ -433,6 +455,38 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     }
 %  end
   }
+
+  /// Accesses a contiguous subrange of the buffer's elements without
+  /// bounds checking.
+  ///
+  /// This unsafe operation should only be used when performance analysis
+  /// has determined that the bounds checks are not being eliminated
+  /// by the optimizer despite being ensured by a higher-level invariant.
+  @inlinable @_alwaysEmitIntoClient
+  public subscript(uncheckedBounds bounds: Range<Int>)
+    -> Slice<Unsafe${Mutable}BufferPointer<Element>>
+  {
+    get {
+      _debugPrecondition(bounds.lowerBound >= startIndex)
+      _debugPrecondition(bounds.upperBound <= endIndex)
+      return Slice(
+        base: self, bounds: bounds)
+    }
+%  if Mutable:
+    nonmutating set {
+      _debugPrecondition(bounds.lowerBound >= startIndex)
+      _debugPrecondition(bounds.upperBound <= endIndex)
+      _debugPrecondition(bounds.count == newValue.count)
+
+      if !newValue.isEmpty {
+        (_position! + bounds.lowerBound).update(
+          from: newValue.base._position! + newValue.startIndex,
+          count: newValue.count)
+      }
+    }
+%  end
+  }
+
 % if mutable:
 
   /// Exchanges the values at the specified indices of the buffer.

test/stdlib/BoundsCheckTraps.swift

@@ -88,4 +88,52 @@ BoundsCheckTraps.test("Character")
   _blackHole(char)
 }
 
+BoundsCheckTraps.test("ArrayUnchecked")
+  .skip(.custom(
+    { _isFastAssertConfiguration() || _isReleaseAssertConfiguration() || _isReleaseAssertWithBoundsSafetyConfiguration() },
+    reason: "this trap is not guaranteed to happen"))
+  .code {
+  expectCrashLater()
+  var array = [1, 2]
+  array.append(3)
+  let value = array[unchecked: 3]
+  _blackHole(value)
+}
+
+BoundsCheckTraps.test("ArrayUncheckedBounds")
+  .skip(.custom(
+    { _isFastAssertConfiguration() || _isReleaseAssertConfiguration() || _isReleaseAssertWithBoundsSafetyConfiguration() },
+    reason: "this trap is not guaranteed to happen"))
+  .code {
+  expectCrashLater()
+  var array = [1, 2]
+  array.append(3)
+  let value = array[uncheckedBounds: 3..<4]
+  _blackHole(value)
+}
+
+BoundsCheckTraps.test("ContiguousArrayUnchecked")
+  .skip(.custom(
+    { _isFastAssertConfiguration() || _isReleaseAssertConfiguration() || _isReleaseAssertWithBoundsSafetyConfiguration() },
+    reason: "this trap is not guaranteed to happen"))
+  .code {
+  expectCrashLater()
+  var array: ContiguousArray = [1, 2]
+  array.append(3)
+  let value = array[unchecked: 3]
+  _blackHole(value)
+}
+
+BoundsCheckTraps.test("ContiguousArrayUncheckedBounds")
+  .skip(.custom(
+    { _isFastAssertConfiguration() || _isReleaseAssertConfiguration() || _isReleaseAssertWithBoundsSafetyConfiguration() },
+    reason: "this trap is not guaranteed to happen"))
+  .code {
+  expectCrashLater()
+  var array: ContiguousArray = [1, 2]
+  array.append(3)
+  let value = array[uncheckedBounds: 3..<4]
+  _blackHole(value)
+}
+
 runAllTests()