[Swift] Adds a way to store external data and use them directly within the ByteBuffer

benchmarks/swift/Benchmarks/FlatbuffersBenchmarks/FlatbuffersBenchmarks.swift

@@ -15,8 +15,8 @@
  */
 
 import Benchmark
-import CoreFoundation
 import FlatBuffers
+import Foundation
 
 @usableFromInline
 struct AA: NativeStruct {
@@ -29,6 +29,15 @@ struct AA: NativeStruct {
 }
 
 let benchmarks = {
+  let oneGB: Int32 = 1_024_000_000
+  let data = {
+    var array = [8888.88, 8888.88]
+    var data = Data()
+    array.withUnsafeBytes { ptr in
+      data.append(contentsOf: ptr)
+    }
+    return data
+  }()
   let ints: [Int] = Array(repeating: 42, count: 100)
   let bytes: [UInt8] = Array(repeating: 42, count: 100)
   let str10 = (0...9).map { _ -> String in "x" }.joined()
@@ -73,12 +82,25 @@ let benchmarks = {
 
   Benchmark("Allocating 1GB", configuration: singleConfiguration) { benchmark in
     for _ in benchmark.scaledIterations {
-      blackHole(FlatBufferBuilder(initialSize: 1_024_000_000))
+      blackHole(FlatBufferBuilder(initialSize: oneGB))
+    }
+  }
+
+  Benchmark(
+    "Allocating ByteBuffer 1GB",
+    configuration: singleConfiguration)
+  { benchmark in
+    let memory = UnsafeMutableRawPointer.allocate(
+      byteCount: 1_024_000_000,
+      alignment: 1)
+    benchmark.startMeasurement()
+    for _ in benchmark.scaledIterations {
+      blackHole(ByteBuffer(assumingMemoryBound: memory, capacity: Int(oneGB)))
     }
   }
 
   Benchmark("Clearing 1GB", configuration: singleConfiguration) { benchmark in
-    var fb = FlatBufferBuilder(initialSize: 1_024_000_000)
+    var fb = FlatBufferBuilder(initialSize: oneGB)
     benchmark.startMeasurement()
     for _ in benchmark.scaledIterations {
       blackHole(fb.clear())
@@ -198,4 +220,11 @@ let benchmarks = {
       blackHole(fb.endTable(at: s))
     }
   }
+
+  Benchmark("Reading Doubles") { benchmark in
+    let byteBuffer = ByteBuffer(data: data)
+    for _ in benchmark.scaledIterations {
+      blackHole(byteBuffer.read(def: Double.self, position: 0))
+    }
+  }
 }

swift.swiftformat

@@ -24,4 +24,4 @@
 --exclude **/swift_code_*.swift
 --exclude **/*.grpc.swift
 
---header "/*\n * Copyright {year} Google Inc. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */"
+--header "/*\n * Copyright 2024 Google Inc. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */"

swift/Sources/FlatBuffers/ByteBuffer.swift

@@ -26,8 +26,22 @@ public struct ByteBuffer {
   /// deallocating the memory that was held by (memory: UnsafeMutableRawPointer)
   @usableFromInline
   final class Storage {
-    // This storage doesn't own the memory, therefore, we won't deallocate on deinit.
+    @usableFromInline
+    enum Blob {
+      #if !os(WASI)
+      case data(Data)
+      case bytes(ContiguousBytes)
+      #endif
+
+      case array([UInt8])
+      case pointer(UnsafeMutableRawPointer)
+      case none
+    }
+
+    /// This storage doesn't own the memory, therefore, we won't deallocate on deinit.
     private let unowned: Bool
+    /// Retained blob of data that requires the storage to retain a pointer to.
+    let retainedBlob: Blob
     /// pointer to the start of the buffer object in memory
     var memory: UnsafeMutableRawPointer
     /// Capacity of UInt8 the buffer can hold
@@ -40,13 +54,40 @@ public struct ByteBuffer {
         alignment: alignment)
       capacity = count
       unowned = false
+      retainedBlob = .none
     }
 
     @usableFromInline
-    init(memory: UnsafeMutableRawPointer, capacity: Int, unowned: Bool) {
-      self.memory = memory
-      self.capacity = capacity
-      self.unowned = unowned
+    init(blob: Blob, capacity count: Int) {
+      capacity = count
+      retainedBlob = blob
+      unowned = true
+      memory = UnsafeMutableRawPointer.allocate(byteCount: 2, alignment: 1)
+
+      switch blob {
+      #if !os(WASI)
+      case .data(var data):
+        data.withUnsafeBytes { ptr in
+          self.memory = UnsafeMutableRawPointer(mutating: ptr.baseAddress!)
+          self.capacity = ptr.count
+        }
+      case .bytes(var bytes):
+        bytes.withUnsafeBytes { ptr in
+          self.memory = UnsafeMutableRawPointer(mutating: ptr.baseAddress!)
+          self.capacity = ptr.count
+        }
+      #endif
+      case .array(var array):
+        array.withUnsafeBytes { ptr in
+          self.memory = UnsafeMutableRawPointer(mutating: ptr.baseAddress!)
+          self.capacity = ptr.count
+        }
+      case .pointer(let unsafeMutableRawPointer):
+        memory = unsafeMutableRawPointer
+      case .none:
+        fatalError(
+          "This initializer should only work when there is a data blob")
+      }
     }
 
     deinit {
@@ -113,123 +154,98 @@ public struct ByteBuffer {
   public var memory: UnsafeMutableRawPointer { _storage.memory }
   /// Current capacity for the buffer
   public var capacity: Int { _storage.capacity }
-  /// Crash if the trying to read an unaligned buffer instead of allowing users to read them.
-  public let allowReadingUnalignedBuffers: Bool
 
   /// Constructor that creates a Flatbuffer object from a UInt8
   /// - Parameter
   ///   - bytes: Array of UInt8
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
-  public init(
-    bytes: [UInt8],
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
-  {
-    var b = bytes
-    _storage = Storage(count: bytes.count, alignment: alignment)
+  @inline(__always)
+  public init(bytes: [UInt8]) {
+    _storage = Storage(blob: .array(bytes), capacity: bytes.count)
     _writerSize = _storage.capacity
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
-    b.withUnsafeMutableBytes { bufferPointer in
-      _storage.copy(from: bufferPointer.baseAddress!, count: bytes.count)
-    }
   }
 
   #if !os(WASI)
   /// Constructor that creates a Flatbuffer from the Swift Data type object
   /// - Parameter
   ///   - data: Swift data Object
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
-  public init(
-    data: Data,
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
-  {
-    var b = data
-    _storage = Storage(count: data.count, alignment: alignment)
+  @inline(__always)
+  public init(data: Data) {
+    _storage = Storage(blob: .data(data), capacity: data.count)
     _writerSize = _storage.capacity
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
-    b.withUnsafeMutableBytes { bufferPointer in
-      _storage.copy(from: bufferPointer.baseAddress!, count: data.count)
-    }
   }
   #endif
 
   /// Constructor that creates a Flatbuffer instance with a size
   /// - Parameter:
   ///   - size: Length of the buffer
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
+  @inline(__always)
   init(initialSize size: Int) {
     let size = size.convertToPowerofTwo
     _storage = Storage(count: size, alignment: alignment)
     _storage.initialize(for: size)
-    allowReadingUnalignedBuffers = false
   }
 
-  #if swift(>=5.0) && !os(WASI)
+  #if !os(WASI)
   /// Constructor that creates a Flatbuffer object from a ContiguousBytes
   /// - Parameters:
   ///   - contiguousBytes: Binary stripe to use as the buffer
   ///   - count: amount of readable bytes
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
+  @inline(__always)
   public init<Bytes: ContiguousBytes>(
     contiguousBytes: Bytes,
-    count: Int,
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
+    count: Int)
   {
-    _storage = Storage(count: count, alignment: alignment)
+    _storage = Storage(blob: .bytes(contiguousBytes), capacity: count)
     _writerSize = _storage.capacity
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
-    contiguousBytes.withUnsafeBytes { buf in
-      _storage.copy(from: buf.baseAddress!, count: buf.count)
-    }
   }
   #endif
 
   /// Constructor that creates a Flatbuffer from unsafe memory region without copying
-  /// - Parameter:
+  /// **NOTE** Needs a call to `memory.deallocate()` later on to free the memory
+  ///
+  /// - Parameters:
   ///   - assumingMemoryBound: The unsafe memory region
   ///   - capacity: The size of the given memory region
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
+  @inline(__always)
   public init(
     assumingMemoryBound memory: UnsafeMutableRawPointer,
-    capacity: Int,
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
+    capacity: Int)
   {
-    _storage = Storage(memory: memory, capacity: capacity, unowned: true)
+    _storage = Storage(
+      blob: .pointer(memory),
+      capacity: capacity)
     _writerSize = capacity
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
   }
 
-  /// Creates a copy of the buffer that's being built by calling sizedBuffer
+  /// Constructor that creates a Flatbuffer from unsafe memory region by copying
+  /// the underlying data to a new pointer
+  ///
   /// - Parameters:
-  ///   - memory: Current memory of the buffer
-  ///   - count: count of bytes
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
-  init(
-    memory: UnsafeMutableRawPointer,
-    count: Int,
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
+  ///   - copyingMemoryBound: The unsafe memory region
+  ///   - capacity: The size of the given memory region
+  @inline(__always)
+  public init(
+    copyingMemoryBound memory: UnsafeMutableRawPointer,
+    capacity: Int)
   {
-    _storage = Storage(count: count, alignment: alignment)
-    _storage.copy(from: memory, count: count)
+    _storage = Storage(count: capacity, alignment: alignment)
+    _storage.copy(from: memory, count: capacity)
     _writerSize = _storage.capacity
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
   }
 
   /// Creates a copy of the existing flatbuffer, by copying it to a different memory.
   /// - Parameters:
   ///   - memory: Current memory of the buffer
   ///   - count: count of bytes
   ///   - removeBytes: Removes a number of bytes from the current size
-  ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
+  @inline(__always)
   init(
-    memory: UnsafeMutableRawPointer,
+    blob: Storage.Blob,
     count: Int,
-    removing removeBytes: Int,
-    allowReadingUnalignedBuffers allowUnalignedBuffers: Bool = false)
+    removing removeBytes: Int)
   {
-    _storage = Storage(count: count, alignment: alignment)
-    _storage.copy(from: memory, count: count)
+    _storage = Storage(blob: blob, capacity: count)
     _writerSize = removeBytes
-    allowReadingUnalignedBuffers = allowUnalignedBuffers
   }
 
   /// Fills the buffer with padding by adding to the writersize
@@ -429,10 +445,8 @@ public struct ByteBuffer {
   ///   - position: the index of the object in the buffer
   @inline(__always)
   public func read<T>(def: T.Type, position: Int) -> T {
-    if allowReadingUnalignedBuffers {
-      return _storage.memory.advanced(by: position).loadUnaligned(as: T.self)
-    }
-    return _storage.memory.advanced(by: position).load(as: T.self)
+    _storage.memory.advanced(by: position).assumingMemoryBound(to: T.self)
+      .pointee
   }
 
   /// Reads a slice from the memory assuming a type of T
@@ -502,7 +516,7 @@ public struct ByteBuffer {
       removeBytes < _storage.capacity,
       "Can NOT remove more bytes than the ones allocated")
     return ByteBuffer(
-      memory: _storage.memory,
+      blob: _storage.retainedBlob,
       count: _storage.capacity,
       removing: _writerSize &- removeBytes)
   }

swift/Sources/FlatBuffers/FlatBufferBuilder.swift

@@ -98,8 +98,8 @@ public struct FlatBufferBuilder {
   public var sizedBuffer: ByteBuffer {
     assert(finished, "Data shouldn't be called before finish()")
     return ByteBuffer(
-      memory: _bb.memory.advanced(by: _bb.reader),
-      count: Int(_bb.size))
+      copyingMemoryBound: _bb.memory.advanced(by: _bb.reader),
+      capacity: Int(_bb.size))
   }
 
   // MARK: - Init

tests/swift/tests/Tests/FlatBuffers.Test.SwiftTests/ByteBufferTests.swift

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2024 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import XCTest
+@testable import FlatBuffers
+
+final class ByteBufferTests: XCTestCase {
+  func testCopyingMemory() {
+    let count = 100
+    let ptr = UnsafeMutableRawPointer.allocate(byteCount: count, alignment: 1)
+    let byteBuffer = ByteBuffer(copyingMemoryBound: ptr, capacity: count)
+    XCTAssertNotEqual(byteBuffer.memory, ptr)
+  }
+
+  func testSamePointer() {
+    let count = 100
+    let ptr = UnsafeMutableRawPointer.allocate(byteCount: count, alignment: 1)
+    let byteBuffer = ByteBuffer(assumingMemoryBound: ptr, capacity: count)
+    XCTAssertEqual(byteBuffer.memory, ptr)
+  }
+
+  func testSameDataPtr() {
+    let count = 100
+    let ptr = Data(repeating: 0, count: count)
+    let byteBuffer = ByteBuffer(data: ptr)
+    ptr.withUnsafeBytes { ptr in
+      XCTAssertEqual(byteBuffer.memory, ptr.baseAddress)
+    }
+  }
+
+  func testSameArrayPtr() {
+    let count = 100
+    let ptr: [UInt8] = Array(repeating: 0, count: count)
+    let byteBuffer = ByteBuffer(bytes: ptr)
+    ptr.withUnsafeBytes { ptr in
+      XCTAssertEqual(byteBuffer.memory, ptr.baseAddress)
+    }
+  }
+}