Implement faster pointer rebasing. (#1696)

Motivation:

I recently discovered that UnsafeRawBufferPointer.init(rebasing:) is
surprisingly expensive, with 7 traps and 11 branches. A simple
replacement can make it a lot cheaper, down to two traps and four
branches. This ends up having pretty drastic effects on
ByteBuffer-heavy NIO code, which often outlines the call to that
initializer and loses the ability to make a bunch of site-local
optimisations.

While this has been potentially fixed upstream with
swiftlang/swift#34879, there is no good reason to
wait until Swift 5.4 for this improvement.

Due to the niche use-case, I didn't bother doing this for _every_
rebasing in the program. In particular, there is at least one
UnsafeBufferPointer(rebasing:) that I didn't do this with, and there are
uses in both NIOTLS and NIOHTTP1 that I didn't change. While we can fix
those if we really need to, it would be nice to avoid this helper
proliferating too far through our codebase.

Modifications:

- Replaced the use of URBP.init(rebasing:) with a custom hand-rolled
  version that avoids Slice.count.

Result:

Cheaper code. One NIOHTTP2 benchmark sees a 2.9% speedup from this
change alone.

Author: Lukasa

Sources/NIO/ByteBuffer-aux.swift

@@ -34,7 +34,7 @@ extension ByteBuffer {
             // this is not technically correct because we shouldn't just bind
             // the memory to `UInt8` but it's not a real issue either and we
             // need to work around https://bugs.swift.org/browse/SR-9604
-            Array<UInt8>(UnsafeRawBufferPointer(rebasing: ptr[range]).bindMemory(to: UInt8.self))
+            Array<UInt8>(UnsafeRawBufferPointer(fastRebase: ptr[range]).bindMemory(to: UInt8.self))
         }
     }
 
@@ -139,7 +139,7 @@ extension ByteBuffer {
         }
         return self.withUnsafeReadableBytes { pointer in
             assert(range.lowerBound >= 0 && (range.upperBound - range.lowerBound) <= pointer.count)
-            return String(decoding: UnsafeRawBufferPointer(rebasing: pointer[range]), as: Unicode.UTF8.self)
+            return String(decoding: UnsafeRawBufferPointer(fastRebase: pointer[range]), as: Unicode.UTF8.self)
         }
     }
 
@@ -210,7 +210,7 @@ extension ByteBuffer {
         self.withVeryUnsafeMutableBytes { destCompleteStorage in
             assert(destCompleteStorage.count >= index + allBytesCount)
             let dest = destCompleteStorage[index ..< index + allBytesCount]
-            dispatchData.copyBytes(to: .init(rebasing: dest), count: dest.count)
+            dispatchData.copyBytes(to: .init(fastRebase: dest), count: dest.count)
         }
         return allBytesCount
     }
@@ -228,7 +228,7 @@ extension ByteBuffer {
             return nil
         }
         return self.withUnsafeReadableBytes { pointer in
-            return DispatchData(bytes: UnsafeRawBufferPointer(rebasing: pointer[range]))
+            return DispatchData(bytes: UnsafeRawBufferPointer(fastRebase: pointer[range]))
         }
     }
 
@@ -396,7 +396,7 @@ extension ByteBuffer {
         precondition(count >= 0, "Can't write fewer than 0 bytes")
         self.reserveCapacity(index + count)
         self.withVeryUnsafeMutableBytes { pointer in
-            let dest = UnsafeMutableRawBufferPointer(rebasing: pointer[index ..< index+count])
+            let dest = UnsafeMutableRawBufferPointer(fastRebase: pointer[index ..< index+count])
             _ = dest.initializeMemory(as: UInt8.self, repeating: byte)
         }
         return count

Sources/NIO/ByteBuffer-core.swift

@@ -379,7 +379,7 @@ public struct ByteBuffer {
 
     @inlinable
     mutating func _setBytesAssumingUniqueBufferAccess(_ bytes: UnsafeRawBufferPointer, at index: _Index) {
-        let targetPtr = UnsafeMutableRawBufferPointer(rebasing: self._slicedStorageBuffer.dropFirst(Int(index)))
+        let targetPtr = UnsafeMutableRawBufferPointer(fastRebase: self._slicedStorageBuffer.dropFirst(Int(index)))
         targetPtr.copyMemory(from: bytes)
     }
 
@@ -389,7 +389,7 @@ public struct ByteBuffer {
     mutating func _setSlowPath<Bytes: Sequence>(bytes: Bytes, at index: _Index) -> _Capacity where Bytes.Element == UInt8 {
         func ensureCapacityAndReturnStorageBase(capacity: Int) -> UnsafeMutablePointer<UInt8> {
             self._ensureAvailableCapacity(_Capacity(capacity), at: index)
-            let newBytesPtr = UnsafeMutableRawBufferPointer(rebasing: self._slicedStorageBuffer[Int(index) ..< Int(index) + Int(capacity)])
+            let newBytesPtr = UnsafeMutableRawBufferPointer(fastRebase: self._slicedStorageBuffer[Int(index) ..< Int(index) + Int(capacity)])
             return newBytesPtr.bindMemory(to: UInt8.self).baseAddress!
         }
         let underestimatedByteCount = bytes.underestimatedCount
@@ -513,7 +513,7 @@ public struct ByteBuffer {
     public mutating func withUnsafeMutableReadableBytes<T>(_ body: (UnsafeMutableRawBufferPointer) throws -> T) rethrows -> T {
         self._copyStorageAndRebaseIfNeeded()
         let readerIndex = self.readerIndex
-        return try body(.init(rebasing: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]))
+        return try body(.init(fastRebase: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]))
     }
 
     /// Yields the bytes currently writable (`bytesWritable` = `capacity` - `writerIndex`). Before reading those bytes you must first
@@ -529,7 +529,7 @@ public struct ByteBuffer {
     @inlinable
     public mutating func withUnsafeMutableWritableBytes<T>(_ body: (UnsafeMutableRawBufferPointer) throws -> T) rethrows -> T {
         self._copyStorageAndRebaseIfNeeded()
-        return try body(.init(rebasing: self._slicedStorageBuffer.dropFirst(self.writerIndex)))
+        return try body(.init(fastRebase: self._slicedStorageBuffer.dropFirst(self.writerIndex)))
     }
 
     /// This vends a pointer of the `ByteBuffer` at the `writerIndex` after ensuring that the buffer has at least `minimumWritableBytes` of writable bytes available.
@@ -587,7 +587,7 @@ public struct ByteBuffer {
     @inlinable
     public func withUnsafeReadableBytes<T>(_ body: (UnsafeRawBufferPointer) throws -> T) rethrows -> T {
         let readerIndex = self.readerIndex
-        return try body(.init(rebasing: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]))
+        return try body(.init(fastRebase: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]))
     }
 
     /// Yields a buffer pointer containing this `ByteBuffer`'s readable bytes. You may hold a pointer to those bytes
@@ -605,7 +605,7 @@ public struct ByteBuffer {
     public func withUnsafeReadableBytesWithStorageManagement<T>(_ body: (UnsafeRawBufferPointer, Unmanaged<AnyObject>) throws -> T) rethrows -> T {
         let storageReference: Unmanaged<AnyObject> = Unmanaged.passUnretained(self._storage)
         let readerIndex = self.readerIndex
-        return try body(.init(rebasing: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]),
+        return try body(.init(fastRebase: self._slicedStorageBuffer[readerIndex ..< readerIndex + self.readableBytes]),
                         storageReference)
     }
 

Sources/NIO/ByteBuffer-int.swift

@@ -62,7 +62,7 @@ extension ByteBuffer {
         return self.withUnsafeReadableBytes { ptr in
             var value: T = 0
             withUnsafeMutableBytes(of: &value) { valuePtr in
-                valuePtr.copyMemory(from: UnsafeRawBufferPointer(rebasing: ptr[range]))
+                valuePtr.copyMemory(from: UnsafeRawBufferPointer(fastRebase: ptr[range]))
             }
             return _toEndianness(value: value, endianness: endianness)
         }

Sources/NIO/ControlMessage.swift

@@ -52,7 +52,7 @@ struct UnsafeControlMessageStorage: Collection {
     /// Get the part of the buffer for use with a message.
     public subscript(position: Int) -> UnsafeMutableRawBufferPointer {
         return UnsafeMutableRawBufferPointer(
-            rebasing: self.buffer[(position * self.bytesPerMessage)..<((position+1) * self.bytesPerMessage)])
+            fastRebase: self.buffer[(position * self.bytesPerMessage)..<((position+1) * self.bytesPerMessage)])
     }
 
     var startIndex: Int { return 0 }
@@ -239,7 +239,7 @@ struct UnsafeOutboundControlBytes {
     private mutating func appendGenericControlMessage<PayloadType>(level: CInt,
                                                                    type: CInt,
                                                                    payload: PayloadType) {
-        let writableBuffer = UnsafeMutableRawBufferPointer(rebasing: self.controlBytes[writePosition...])
+        let writableBuffer = UnsafeMutableRawBufferPointer(fastRebase: self.controlBytes[writePosition...])
 
         let requiredSize = NIOBSDSocketControlMessage.space(payloadSize: MemoryLayout.stride(ofValue: payload))
         precondition(writableBuffer.count >= requiredSize, "Insufficient size for cmsghdr and data")
@@ -263,7 +263,7 @@ struct UnsafeOutboundControlBytes {
         if writePosition == 0 {
             return UnsafeMutableRawBufferPointer(start: nil, count: 0)
         }
-        return UnsafeMutableRawBufferPointer(rebasing: self.controlBytes[0 ..< self.writePosition])
+        return UnsafeMutableRawBufferPointer(fastRebase: self.controlBytes[0 ..< self.writePosition])
     }
 
 }

Sources/NIO/PointerHelpers.swift

@@ -0,0 +1,41 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+// MARK: Rebasing shims
+
+// These methods are shimmed in to NIO until https://github.com/apple/swift/pull/34879 is resolved.
+// They address the fact that the current rebasing initializers are surprisingly expensive and do excessive
+// checked arithmetic. This expense forces them to often be outlined, reducing the ability to optimise out
+// further preconditions and branches.
+extension UnsafeRawBufferPointer {
+    @inlinable
+    init(fastRebase slice: Slice<UnsafeRawBufferPointer>) {
+        let base = slice.base.baseAddress?.advanced(by: slice.startIndex)
+        self.init(start: base, count: slice.endIndex &- slice.startIndex)
+    }
+
+    @inlinable
+    init(fastRebase slice: Slice<UnsafeMutableRawBufferPointer>) {
+        let base = slice.base.baseAddress?.advanced(by: slice.startIndex)
+        self.init(start: base, count: slice.endIndex &- slice.startIndex)
+    }
+}
+
+extension UnsafeMutableRawBufferPointer {
+    @inlinable
+    init(fastRebase slice: Slice<UnsafeMutableRawBufferPointer>) {
+        let base = slice.base.baseAddress?.advanced(by: slice.startIndex)
+        self.init(start: base, count: slice.endIndex &- slice.startIndex)
+    }
+}