Moves the internal stack to use a pointer stack instead of the native array for improved performance

Author: mustiikhalil

swift/Sources/FlexBuffers/Reader/Reference.swift

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import Common
 import Foundation
 
 enum FlexBuffersErrors: Error {
@@ -23,6 +24,10 @@ enum FlexBuffersErrors: Error {
 
 @inline(__always)
 public func getRoot(buffer: ByteBuffer) throws -> Reference? {
+  assert(
+    isLitteEndian,
+    "Swift FlexBuffers currently only supports little-endian systems")
+
   let end = buffer.count
   if buffer.count < 3 {
     throw FlexBuffersErrors.sizeOfBufferIsTooSmall

swift/Sources/FlexBuffers/Writer/FlexBuffersWriter.swift

@@ -37,12 +37,19 @@ public struct FlexBuffersWriter {
   private var hasDuplicatedKeys = false
   private var minBitWidth: BitWidth = .w8
   private var _bb: _InternalByteBuffer
+  #if os(Wasm)
   private var stack: [Value] = []
+  #else
+  private var stack: Stack = Stack()
+  #endif
   private var keyPool: [Int: UInt] = [:]
   private var stringPool: [Int: UInt] = [:]
   private var flags: BuilderFlag
 
   public init(initialSize: Int = 1024, flags: BuilderFlag = .shareKeys) {
+    assert(
+      isLitteEndian,
+      "Swift FlexBuffers currently only supports little-endian systems")
     _bb = _InternalByteBuffer(initialSize: initialSize)
     self.flags = flags
   }
@@ -148,7 +155,7 @@ public struct FlexBuffersWriter {
       typed: typed,
       fixed: fixed,
       keys: nil)
-    stack = Array(stack[..<start])
+    stack.popLast(start)
     stack.append(vec)
     return vec.u
   }
@@ -217,7 +224,7 @@ public struct FlexBuffersWriter {
       typed: false,
       fixed: false,
       keys: keys)
-    stack = Array(stack[..<start])
+    stack.popLast(start)
     stack.append(vec)
     return numericCast(vec.u)
   }
@@ -721,7 +728,7 @@ public struct FlexBuffersWriter {
         assert(
           vectorType == stack[i].type,
           """
-          If you get this assert you are writing a typed vector 
+          If you get this assert you are writing a typed vector
           with elements that are not all the same type
           """)
       }
@@ -830,13 +837,14 @@ public struct FlexBuffersWriter {
       let key, value: Value
     }
 
-    stack[start...].withUnsafeMutableBytes { buffer in
+    stack.withUnsafeMutableBytes(start: start) { buffer in
       var ptr = buffer.assumingMemoryBound(to: TwoValue.self)
       ptr.sort { a, b in
         let aMem = _bb.memory.advanced(by: numericCast(a.key.u))
           .assumingMemoryBound(to: CChar.self)
         let bMem = _bb.memory.advanced(by: numericCast(b.key.u))
           .assumingMemoryBound(to: CChar.self)
+
         let comp = strcmp(aMem, bMem)
         if (comp == 0) && a != b { hasDuplicatedKeys = true }
         return comp < 0
@@ -897,3 +905,145 @@ extension FlexBuffersWriter {
     return endMap(start: start)
   }
 }
+
+private extension Array where Element == Value {
+  mutating func popLast(_ val: Int) {
+    self = Array(self[..<val])
+  }
+
+  @discardableResult
+  mutating func withUnsafeMutableBytes<R>(
+    start: Int,
+    _ body: (UnsafeMutableRawBufferPointer) throws -> R) rethrows -> R
+  {
+    try self[start...].withUnsafeMutableBytes { buffer in
+      try body(buffer)
+    }
+  }
+}
+
+fileprivate struct Stack: RandomAccessCollection {
+  typealias Element = Value
+  typealias Index = Int
+
+  private final class Storage {
+    var memory: UnsafeMutableRawPointer
+
+    init(capacity: Int, alignment: Int) {
+      memory = .allocate(byteCount: capacity, alignment: alignment)
+      memset(memory, 0, capacity)
+    }
+
+    deinit {
+      memory.deallocate()
+    }
+  }
+
+  private static let initialCapacity = 10 &* MemoryLayout<Value>.stride
+  private let storage: Storage
+  private var capacity: Int
+  private(set) var count: Int
+
+  var startIndex: Int {
+    0
+  }
+
+  var endIndex: Int {
+    count
+  }
+
+  init() {
+    count = 0
+    capacity = Self.initialCapacity
+
+    storage = Storage(
+      capacity: capacity,
+      alignment: MemoryLayout<Value>.alignment)
+  }
+
+  @inline(__always)
+  subscript(position: Int) -> Value {
+    get {
+      storage.memory.advanced(by: position &* MemoryLayout<Value>.stride)
+        .assumingMemoryBound(to: Value.self).pointee
+    }
+    set {
+      storage.memory.advanced(by: position &* MemoryLayout<Value>.stride)
+        .assumingMemoryBound(to: Value.self).pointee = newValue
+    }
+  }
+
+  @inline(__always)
+  mutating func popLast(_ val: Int) {
+    count = if val < 0 {
+      0
+    } else {
+      val
+    }
+  }
+
+  mutating func append(_ value: Value) {
+    let writePosition = count &* MemoryLayout<Value>.stride
+    if writePosition >= capacity {
+      reallocate(writePosition: writePosition)
+    }
+
+    storage.memory.advanced(by: writePosition).storeBytes(
+      of: value,
+      as: Value.self)
+    count += 1
+  }
+
+  mutating func removeAll(keepingCapacity keepCapacity: Bool = false) {
+    count = 0
+    if !keepCapacity {
+      capacity = Self.initialCapacity
+      storage.memory = UnsafeMutableRawPointer.allocate(
+        byteCount: capacity,
+        alignment: MemoryLayout<Value>.alignment)
+    }
+    memset(storage.memory, 0, capacity)
+  }
+
+  @discardableResult
+  mutating func withUnsafeMutableBytes<R>(
+    start: Int,
+    _ body: (UnsafeMutableRawBufferPointer) throws -> R) rethrows -> R
+  {
+    let startingPosition = start &* MemoryLayout<Value>.stride
+    let pointer = storage.memory.advanced(by: startingPosition)
+    return try body(UnsafeMutableRawBufferPointer(
+      start: pointer,
+      count: (count &* MemoryLayout<Value>.stride) &- startingPosition))
+  }
+
+  @discardableResult
+  mutating func withUnsafeMutableBytes<R>(
+    _ body: (UnsafeMutableRawBufferPointer) throws
+      -> R) rethrows -> R
+  {
+    return try body(UnsafeMutableRawBufferPointer(
+      start: storage.memory,
+      count: count &* MemoryLayout<Value>.stride))
+  }
+
+  mutating private func reallocate(writePosition: Int) {
+    while capacity <= writePosition {
+      capacity = capacity << 1
+    }
+
+    /// solution take from Apple-NIO
+    capacity = capacity.convertToPowerofTwo
+
+    let newData = UnsafeMutableRawPointer.allocate(
+      byteCount: capacity,
+      alignment: MemoryLayout<Value>.alignment)
+    memset(newData, 0, capacity)
+    memcpy(
+      newData,
+      storage.memory,
+      writePosition)
+    storage.memory.deallocate()
+    storage.memory = newData
+  }
+}

swift/Sources/FlexBuffers/_InternalByteBuffer.swift

@@ -60,7 +60,7 @@ struct _InternalByteBuffer {
       let newData = UnsafeMutableRawPointer.allocate(
         byteCount: capacity,
         alignment: alignment)
-      memset(newData, 0, capacity &- writerSize)
+      memset(newData, 0, capacity)
       memcpy(
         newData,
         memory,

tests/swift/Wasm.tests/Package.swift

@@ -24,7 +24,7 @@ let package = Package(
   ],
   dependencies: [
     .package(path: "../../.."),
-    .package(url: "https://github.com/swiftwasm/WasmKit", exact: "0.1.6")
+    .package(url: "https://github.com/swiftwasm/WasmKit", exact: "0.1.6"),
   ],
   targets: [
     .target(name: "Wasm"),
@@ -37,5 +37,5 @@ let package = Package(
       name: "FlexBuffers.Test.Swift.WasmTests",
       dependencies: [
         .product(name: "FlexBuffers", package: "flatbuffers"),
-      ])
+      ]),
   ])