Merge pull request #37424 from Lukasa/cb-fewer-checks-in-ubp

Remove as much checked math as possible from buffer pointers.

Author: lorentey

stdlib/public/core/UnsafeBufferPointer.swift.gyb

@@ -205,7 +205,14 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    return i + 1
+    // NOTE: Wrapping math because we allow unsafe buffer pointers not to verify
+    // index preconditions in release builds. Our (optimistic) assumption is
+    // that the caller is already ensuring that indices are valid, so we can
+    // elide the usual checks to help the optimizer generate better code.
+    // However, we still check for overflow in debug mode.
+    let result = i.addingReportingOverflow(1)
+    _debugPrecondition(!result.overflow)
+    return result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -215,7 +222,10 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    i += 1
+    // See note on wrapping arithmetic in `index(after:)` above.
+    let result = i.addingReportingOverflow(1)
+    _debugPrecondition(!result.overflow)
+    i = result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -225,7 +235,10 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    return i - 1
+    // See note on wrapping arithmetic in `index(after:)` above.
+    let result = i.subtractingReportingOverflow(1)
+    _debugPrecondition(!result.overflow)
+    return result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -235,7 +248,10 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    i -= 1
+    // See note on wrapping arithmetic in `index(after:)` above.
+    let result = i.subtractingReportingOverflow(1)
+    _debugPrecondition(!result.overflow)
+    i = result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -245,7 +261,10 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    return i + n
+    // See note on wrapping arithmetic in `index(after:)` above.
+    let result = i.addingReportingOverflow(n)
+    _debugPrecondition(!result.overflow)
+    return result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -255,11 +274,18 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
     // optimizer is not capable of creating partial specializations yet.
     // NOTE: Range checks are not performed here, because it is done later by
     // the subscript function.
-    let l = limit - i
+    // See note on wrapping arithmetic in `index(after:)` above.
+    let maxOffset = limit.subtractingReportingOverflow(i)
+    _debugPrecondition(!maxOffset.overflow)
+    let l = maxOffset.partialValue
+
     if n > 0 ? l >= 0 && l < n : l <= 0 && n < l {
       return nil
     }
-    return i + n
+
+    let result = i.addingReportingOverflow(n)
+    _debugPrecondition(!result.overflow)
+    return result.partialValue
   }
 
   @inlinable // unsafe-performance
@@ -295,7 +321,8 @@ extension Unsafe${Mutable}BufferPointer: ${Mutable}Collection, RandomAccessColle
 
   @inlinable // unsafe-performance
   public var indices: Indices {
-    return startIndex..<endIndex
+    // Not checked because init forbids negative count.
+    return Indices(uncheckedBounds: (startIndex, endIndex))
   }
 
   /// Accesses the element at the specified position.

stdlib/public/core/UnsafeRawBufferPointer.swift.gyb

@@ -212,7 +212,8 @@ extension Unsafe${Mutable}RawBufferPointer: ${Mutable}Collection {
 
   @inlinable
   public var indices: Indices {
-    return startIndex..<endIndex
+    // Not checked because init forbids negative count.
+    return Indices(uncheckedBounds: (startIndex, endIndex))
   }
 
   /// Accesses the byte at the given offset in the memory region as a `UInt8`
@@ -293,7 +294,10 @@ extension Unsafe${Mutable}RawBufferPointer: ${Mutable}Collection {
   @inlinable
   public var count: Int {
     if let pos = _position {
-      return _assumeNonNegative(_end! - pos)
+      // Unsafely unwrapped because init forbids end being nil if _position
+      // isn't.
+      _internalInvariant(_end != nil)
+      return _assumeNonNegative(_end._unsafelyUnwrappedUnchecked - pos)
     }
     return 0
   }
@@ -631,7 +635,7 @@ extension Unsafe${Mutable}RawBufferPointer {
       return UnsafeMutableBufferPointer<T>(start: nil, count: 0)
     }
 
-    let count = (_end! - base) / MemoryLayout<T>.stride
+    let count = (_end._unsafelyUnwrappedUnchecked - base) / MemoryLayout<T>.stride
     let typed = base.initializeMemory(
       as: type, repeating: repeatedValue, count: count)
     return UnsafeMutableBufferPointer<T>(start: typed, count: count)
@@ -680,9 +684,10 @@ extension Unsafe${Mutable}RawBufferPointer {
       return (it, UnsafeMutableBufferPointer(start: nil, count: 0))
     }  
 
+    _internalInvariant(_end != nil)
     for p in stride(from: base, 
       // only advance to as far as the last element that will fit
-      to: base + count - elementStride + 1, 
+      to: _end._unsafelyUnwrappedUnchecked - elementStride + 1, 
       by: elementStride
     ) {
       // underflow is permitted -- e.g. a sequence into

test/SILOptimizer/unsafebufferpointer.swift

@@ -80,7 +80,7 @@ public func testSubscript(_ ubp: UnsafeBufferPointer<Float>) -> Float {
 // CHECK: [[CMP:%[0-9]+]] = icmp eq
 // CHECK: br i1 [[CMP]], label %[[RET:.*]], label %[[LOOP]]
 //
-// CHECK: [[RET]]: ; preds = %[[LOOP]], %entry
+// CHECK: [[RET]]: ; preds = %[[LOOP]], %{{.*}}
 // CHECK: ret i64
 public func testSubscript(_ ubp: UnsafeRawBufferPointer) -> Int64 {
   var sum: Int64 = 0