From b08abbbbf632efde42416a6aaff0709fe1b3546f Mon Sep 17 00:00:00 2001
From: Taylor Holliday <wtholliday@gmail.com>
Date: Mon, 20 Nov 2023 08:47:11 -0800
Subject: [PATCH] [Swift] Push contiguous bytes (#8157)

* Add version of push which takes ContiguousBytes

* Ensure overloads aren't ambiguous

* Add version of createVector

* Add version of push which takes ContiguousBytes

* Ensure overloads aren't ambiguous

* Add version of createVector

* Add similar conditional to other use of ContiguousBytes

* Attempt CI fix

* Use memcpy instead of copyMemory

memcpy is faster in tests

* Add testContiguousBytes

* Add benchmarks

* Add version of createVector

* Add benchmarks

* Update push to copy memory

Since we don't care about endianness, we can simply memcpy the array of scalars

* Remove function and benchmarks

Since we don't care about endianness, a FixedWidthInteger version of createVector isn't needed

* Improve naming

* Add doc comment
---
 swift/Sources/FlatBuffers/ByteBuffer.swift    | 28 ++++++++++++++++---
 .../FlatBuffers/FlatBufferBuilder.swift       | 16 +++++++++++
 .../benchmarks/Sources/benchmarks/main.swift  | 15 ++++++++++
 .../FlatBuffersMonsterWriterTests.swift       | 17 +++++++++++
 4 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/swift/Sources/FlatBuffers/ByteBuffer.swift b/swift/Sources/FlatBuffers/ByteBuffer.swift
index 7f5c37999f7..60d632d1b97 100644
--- a/swift/Sources/FlatBuffers/ByteBuffer.swift
+++ b/swift/Sources/FlatBuffers/ByteBuffer.swift
@@ -228,12 +228,32 @@ public struct ByteBuffer {
   @inline(__always)
   @usableFromInline
   mutating func push<T: Scalar>(elements: [T]) {
-    let size = elements.count &* MemoryLayout<T>.size
-    ensureSpace(size: size)
-    elements.reversed().forEach { s in
-      push(value: s, len: MemoryLayout.size(ofValue: s))
+    elements.withUnsafeBytes { ptr in
+      ensureSpace(size: ptr.count)
+      memcpy(
+        _storage.memory.advanced(by: writerIndex &- ptr.count),
+        UnsafeRawPointer(ptr.baseAddress!),
+        ptr.count)
+      self._writerSize = self._writerSize &+ ptr.count
+    }
+  }
+
+  /// Adds a `ContiguousBytes` to buffer memory
+  /// - Parameter value: bytes to copy
+  #if swift(>=5.0) && !os(WASI)
+  @inline(__always)
+  @usableFromInline
+  mutating func push(bytes: ContiguousBytes) {
+    bytes.withUnsafeBytes { ptr in
+      ensureSpace(size: ptr.count)
+      memcpy(
+        _storage.memory.advanced(by: writerIndex &- ptr.count),
+            UnsafeRawPointer(ptr.baseAddress!),
+            ptr.count)
+      self._writerSize = self._writerSize &+ ptr.count
     }
   }
+  #endif
 
   /// Adds an object of type NativeStruct into the buffer
   /// - Parameters:
diff --git a/swift/Sources/FlatBuffers/FlatBufferBuilder.swift b/swift/Sources/FlatBuffers/FlatBufferBuilder.swift
index f96ad611411..1fb66d55a0c 100644
--- a/swift/Sources/FlatBuffers/FlatBufferBuilder.swift
+++ b/swift/Sources/FlatBuffers/FlatBufferBuilder.swift
@@ -473,6 +473,22 @@ public struct FlatBufferBuilder {
     return endVector(len: size)
   }
 
+  #if swift(>=5.0) && !os(WASI)
+  @inline(__always)
+  /// Creates a vector of bytes in the buffer.
+  ///
+  /// Allows creating a vector from `Data` without copying to a `[UInt8]`
+  ///
+  /// - Parameter bytes: bytes to be written into the buffer
+  /// - Returns: ``Offset`` of the vector
+  mutating public func createVector(bytes: ContiguousBytes) -> Offset {
+    let size = bytes.withUnsafeBytes { ptr in ptr.count }
+    startVector(size, elementSize: MemoryLayout<UInt8>.size)
+    _bb.push(bytes: bytes)
+    return endVector(len: size)
+  }
+  #endif
+
   /// Creates a vector of type ``Enum`` into the ``ByteBuffer``
   ///
   /// ``createVector(_:)-9h189`` writes a vector of type ``Enum`` into
diff --git a/tests/swift/benchmarks/Sources/benchmarks/main.swift b/tests/swift/benchmarks/Sources/benchmarks/main.swift
index e41d5ce5757..b368add5893 100644
--- a/tests/swift/benchmarks/Sources/benchmarks/main.swift
+++ b/tests/swift/benchmarks/Sources/benchmarks/main.swift
@@ -32,6 +32,20 @@ benchmark("100Strings") {
   }
 }
 
+benchmark("100Bytes") {
+  var fb = FlatBufferBuilder(initialSize: 1<<20)
+  for _ in 0..<1_000_000 {
+    _ = fb.createVector(bytes)
+  }
+}
+
+benchmark("100Bytes-ContiguousBytes") {
+  var fb = FlatBufferBuilder(initialSize: 1<<20)
+  for _ in 0..<1_000_000 {
+    _ = fb.createVector(bytes: bytes)
+  }
+}
+
 benchmark("FlatBufferBuilder.add") {
   var fb = FlatBufferBuilder(initialSize: 1024 * 1024 * 32)
   for _ in 0..<1_000_000 {
@@ -73,6 +87,7 @@ benchmark("structs") {
 }
 
 let str = (0...99).map { _ -> String in "x" }.joined()
+let bytes: [UInt8] = Array(repeating: 42, count: 100)
 
 @usableFromInline
 struct AA: NativeStruct {
diff --git a/tests/swift/tests/Tests/FlatBuffers.Test.SwiftTests/FlatBuffersMonsterWriterTests.swift b/tests/swift/tests/Tests/FlatBuffers.Test.SwiftTests/FlatBuffersMonsterWriterTests.swift
index d9da317db77..b72af55a521 100644
--- a/tests/swift/tests/Tests/FlatBuffers.Test.SwiftTests/FlatBuffersMonsterWriterTests.swift
+++ b/tests/swift/tests/Tests/FlatBuffers.Test.SwiftTests/FlatBuffersMonsterWriterTests.swift
@@ -484,4 +484,21 @@ class FlatBuffersMonsterWriterTests: XCTestCase {
     #endif
   }
 
+  func testContiguousBytes() {
+    let byteArray: [UInt8] = [3, 1, 4, 1, 5, 9]
+    var fbb = FlatBufferBuilder(initialSize: 1)
+    let name = fbb.create(string: "Frodo")
+    let bytes = fbb.createVector(bytes: byteArray)
+    let root = Monster.createMonster(
+      &fbb,
+      nameOffset: name,
+      inventoryVectorOffset: bytes)
+    fbb.finish(offset: root)
+    var buffer = fbb.sizedBuffer
+    let monster: Monster = getRoot(byteBuffer: &buffer)
+    let values = monster.inventory
+
+    XCTAssertEqual(byteArray, values)
+  }
+
 }