[Swift] Moves capacity outside of Storage (#8650)

- Cleans up capacity usage within the lib and moves it outside of the Storage

-  Use overflow operators

Author: mustiikhalil

swift/Sources/FlatBuffers/ByteBuffer.swift

@@ -40,11 +40,11 @@ public struct ByteBuffer {
 
     /// This storage doesn't own the memory, therefore, we won't deallocate on deinit.
     private let isOwned: Bool
+    /// Capacity of UInt8 the buffer can hold
+    private let capacity: Int
     /// Retained blob of data that requires the storage to retain a pointer to.
     @usableFromInline
     var retainedBlob: Blob
-    /// Capacity of UInt8 the buffer can hold
-    var capacity: Int
 
     @usableFromInline
     init(count: Int) {
@@ -179,11 +179,11 @@ public struct ByteBuffer {
   /// The size of the elements written to the buffer + their paddings
   private var _readerIndex: Int = 0
   /// Reader is the position of the current Writer Index (capacity - size)
-  public var reader: Int { _storage.capacity &- _readerIndex }
+  public var reader: Int { capacity &- _readerIndex }
   /// Current size of the buffer
   public var size: UOffset { UOffset(_readerIndex) }
   /// Current capacity for the buffer
-  public var capacity: Int { _storage.capacity }
+  public let capacity: Int
 
   /// Constructor that creates a Flatbuffer object from an InternalByteBuffer
   /// - Parameter
@@ -194,6 +194,7 @@ public struct ByteBuffer {
       blob: .byteBuffer(byteBuffer),
       capacity: byteBuffer.capacity)
     _readerIndex = Int(byteBuffer.size)
+    self.capacity = byteBuffer.capacity
   }
 
   /// Constructor that creates a Flatbuffer from unsafe memory region by copying
@@ -209,7 +210,8 @@ public struct ByteBuffer {
   {
     _storage = Storage(count: capacity)
     _storage.copy(from: memory, count: capacity)
-    _readerIndex = _storage.capacity
+    _readerIndex = capacity
+    self.capacity = capacity
   }
 
   /// Constructor that creates a Flatbuffer object from a UInt8
@@ -218,7 +220,8 @@ public struct ByteBuffer {
   @inline(__always)
   public init(bytes: [UInt8]) {
     _storage = Storage(blob: .array(bytes), capacity: bytes.count)
-    _readerIndex = _storage.capacity
+    _readerIndex = bytes.count
+    capacity = bytes.count
   }
 
   #if !os(WASI)
@@ -228,7 +231,8 @@ public struct ByteBuffer {
   @inline(__always)
   public init(data: Data) {
     _storage = Storage(blob: .data(data), capacity: data.count)
-    _readerIndex = _storage.capacity
+    _readerIndex = data.count
+    capacity = data.count
   }
 
   /// Constructor that creates a Flatbuffer object from a ContiguousBytes
@@ -241,7 +245,8 @@ public struct ByteBuffer {
     count: Int)
   {
     _storage = Storage(blob: .bytes(contiguousBytes), capacity: count)
-    _readerIndex = _storage.capacity
+    _readerIndex = count
+    self.capacity = count
   }
   #endif
 
@@ -259,7 +264,8 @@ public struct ByteBuffer {
     _storage = Storage(
       blob: .pointer(memory),
       capacity: capacity)
-    _readerIndex = _storage.capacity
+    _readerIndex = capacity
+    self.capacity = capacity
   }
 
   /// Creates a copy of the existing flatbuffer, by copying it to a different memory.
@@ -275,6 +281,7 @@ public struct ByteBuffer {
   {
     _storage = Storage(blob: blob, capacity: count)
     _readerIndex = removeBytes
+    capacity = count
   }
 
   /// Write stores an object into the buffer directly or indirectly.
@@ -289,11 +296,11 @@ public struct ByteBuffer {
   func write<T>(value: T, index: Int, direct: Bool = false) {
     var index = index
     if !direct {
-      index = _storage.capacity &- index
+      index = capacity &- index
     }
-    assert(index < _storage.capacity, "Write index is out of writing bound")
+    assert(index < capacity, "Write index is out of writing bound")
     assert(index >= 0, "Writer index should be above zero")
-    withUnsafePointer(to: value) { ptr in
+    _ = withUnsafePointer(to: value) { ptr in
       _storage.withUnsafeRawPointer {
         memcpy(
           $0.advanced(by: index),
@@ -325,7 +332,7 @@ public struct ByteBuffer {
     count: Int) -> [T]
   {
     assert(
-      index + count <= _storage.capacity,
+      index + count <= capacity,
       "Reading out of bounds is illegal")
 
     return _storage.readWithUnsafeRawPointer(position: index) {
@@ -348,7 +355,7 @@ public struct ByteBuffer {
     body: (UnsafeRawBufferPointer) throws -> T) rethrows -> T
   {
     assert(
-      index + count <= _storage.capacity,
+      index + count <= capacity,
       "Reading out of bounds is illegal")
     return try _storage.readWithUnsafeRawPointer(position: index) {
       try body(UnsafeRawBufferPointer(start: $0, count: count))
@@ -368,7 +375,7 @@ public struct ByteBuffer {
     type: String.Encoding = .utf8) -> String?
   {
     assert(
-      index + count <= _storage.capacity,
+      index + count <= capacity,
       "Reading out of bounds is illegal")
     return _storage.readWithUnsafeRawPointer(position: index) {
       let buf = UnsafeBufferPointer(
@@ -390,7 +397,7 @@ public struct ByteBuffer {
     count: Int) -> String?
   {
     assert(
-      index + count <= _storage.capacity,
+      index + count <= capacity,
       "Reading out of bounds is illegal")
     return _storage.readWithUnsafeRawPointer(position: index) {
       String(cString: $0.bindMemory(to: UInt8.self, capacity: count))
@@ -404,11 +411,11 @@ public struct ByteBuffer {
   public func duplicate(removing removeBytes: Int = 0) -> ByteBuffer {
     assert(removeBytes > 0, "Can NOT remove negative bytes")
     assert(
-      removeBytes < _storage.capacity,
+      removeBytes < capacity,
       "Can NOT remove more bytes than the ones allocated")
     return ByteBuffer(
       blob: _storage.retainedBlob,
-      count: _storage.capacity,
+      count: capacity,
       removing: _readerIndex &- removeBytes)
   }
 
@@ -456,7 +463,7 @@ extension ByteBuffer: CustomDebugStringConvertible {
   public var debugDescription: String {
     """
     buffer located at: \(_storage.retainedBlob), 
-    with capacity of \(_storage.capacity),
+    with capacity of \(capacity),
     { writtenSize: \(_readerIndex), readerSize: \(reader), 
     size: \(size) }
     """

swift/Sources/FlatBuffers/Mutable.swift

@@ -49,7 +49,7 @@ extension Mutable where Self == Table {
   ///   - index: index of the Element
   public func mutate<T: Scalar>(_ value: T, index: Int32) -> Bool {
     guard index != 0 else { return false }
-    return mutate(value: value, o: index + position)
+    return mutate(value: value, o: index &+ position)
   }
 
   /// Directly mutates the element by calling mutate
@@ -70,7 +70,7 @@ extension Mutable where Self == Struct {
   ///   - value: New value to be inserted to the buffer
   ///   - index: index of the Element
   public func mutate<T: Scalar>(_ value: T, index: Int32) -> Bool {
-    mutate(value: value, o: index + position)
+    mutate(value: value, o: index &+ position)
   }
 
   /// Directly mutates the element by calling mutate

swift/Sources/FlatBuffers/Struct.swift

@@ -45,7 +45,7 @@ public struct Struct {
   ///   - o: Current offset of the data
   /// - Returns: Data of Type T that conforms to type Scalar
   public func readBuffer<T: Scalar>(of type: T.Type, at o: Int32) -> T {
-    let r = bb.read(def: T.self, position: Int(o + position))
+    let r = bb.read(def: T.self, position: Int(o &+ position))
     return r
   }
 }

swift/Sources/FlatBuffers/Table.swift

@@ -167,7 +167,7 @@ public struct Table {
   public func vector(at o: Int32) -> Int32 {
     var o = o
     o &+= position
-    return o &+ bb.read(def: Int32.self, position: Int(o)) + 4
+    return o &+ bb.read(def: Int32.self, position: Int(o)) &+ 4
   }
 
   /// Reading an indirect offset of a table.
@@ -176,7 +176,7 @@ public struct Table {
   ///   - fbb: ByteBuffer
   /// - Returns: table offset
   static public func indirect(_ o: Int32, _ fbb: ByteBuffer) -> Int32 {
-    o + fbb.read(def: Int32.self, position: Int(o))
+    o &+ fbb.read(def: Int32.self, position: Int(o))
   }
 
   /// Gets a vtable value according to an table Offset and a field offset

swift/Sources/FlatBuffers/_InternalByteBuffer.swift

@@ -31,15 +31,12 @@ struct _InternalByteBuffer {
   final class Storage {
     /// pointer to the start of the buffer object in memory
     private(set) var memory: UnsafeMutableRawPointer
-    /// Capacity of UInt8 the buffer can hold
-    private(set) var capacity: Int
 
     @usableFromInline
     init(count: Int, alignment: Int) {
       memory = UnsafeMutableRawPointer.allocate(
         byteCount: count,
         alignment: alignment)
-      capacity = count
     }
 
     deinit {
@@ -54,15 +51,12 @@ struct _InternalByteBuffer {
     /// Reallocates the buffer incase the object to be written doesnt fit in the current buffer
     /// - Parameter size: Size of the current object
     @usableFromInline
-    func reallocate(_ size: Int, writerSize: Int, alignment: Int) {
-      let currentWritingIndex = capacity &- writerSize
-      while capacity <= writerSize &+ size {
-        capacity = capacity << 1
-      }
-
-      /// solution take from Apple-NIO
-      capacity = capacity.convertToPowerofTwo
-
+    func reallocate(
+      capacity: Int,
+      writerSize: Int,
+      currentWritingIndex: Int,
+      alignment: Int
+    ) {
       let newData = UnsafeMutableRawPointer.allocate(
         byteCount: capacity,
         alignment: alignment)
@@ -84,7 +78,7 @@ struct _InternalByteBuffer {
   /// Alignment of the current  memory being written to the buffer
   var alignment = 1
   /// Current Index which is being used to write to the buffer, it is written from the end to the start of the buffer
-  var writerIndex: Int { _storage.capacity &- _writerSize }
+  var writerIndex: Int { capacity &- _writerSize }
 
   /// Reader is the position of the current Writer Index (capacity - size)
   public var reader: Int { writerIndex }
@@ -94,14 +88,15 @@ struct _InternalByteBuffer {
   @usableFromInline
   var memory: UnsafeMutableRawPointer { _storage.memory }
   /// Current capacity for the buffer
-  public var capacity: Int { _storage.capacity }
+  public private(set) var capacity: Int
 
   /// Constructor that creates a Flatbuffer instance with a size
   /// - Parameter:
   ///   - size: Length of the buffer
   ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
   init(initialSize size: Int) {
     initialSize = size.convertToPowerofTwo
+    capacity = initialSize
     _storage = Storage(count: initialSize, alignment: alignment)
     _storage.initialize(for: initialSize)
   }
@@ -246,11 +241,11 @@ struct _InternalByteBuffer {
   func write<T>(value: T, index: Int, direct: Bool = false) {
     var index = index
     if !direct {
-      index = _storage.capacity &- index
+      index = capacity &- index
     }
-    assert(index < _storage.capacity, "Write index is out of writing bound")
+    assert(index < capacity, "Write index is out of writing bound")
     assert(index >= 0, "Writer index should be above zero")
-    withUnsafePointer(to: value) {
+    _ = withUnsafePointer(to: value) {
       memcpy(
         _storage.memory.advanced(by: index),
         $0,
@@ -262,10 +257,24 @@ struct _InternalByteBuffer {
   /// - Parameter size: size of object
   @discardableResult
   @usableFromInline
-  @inline(__always)
   mutating func ensureSpace(size: Int) -> Int {
-    if size &+ _writerSize > _storage.capacity {
-      _storage.reallocate(size, writerSize: _writerSize, alignment: alignment)
+    let expectedWriterIndex = size &+ _writerSize
+    if expectedWriterIndex > capacity {
+      
+      let currentWritingIndex = capacity &- _writerSize
+      while capacity <= expectedWriterIndex {
+        capacity = capacity << 1
+      }
+
+      /// solution take from Apple-NIO
+      capacity = capacity.convertToPowerofTwo
+
+      
+      _storage.reallocate(
+        capacity: capacity,
+        writerSize: _writerSize,
+        currentWritingIndex: currentWritingIndex,
+        alignment: alignment)
     }
     assert(size < FlatBufferMaxSize, "Buffer can't grow beyond 2 Gigabytes")
     return size
@@ -295,8 +304,9 @@ struct _InternalByteBuffer {
     _writerSize = 0
     alignment = 1
     if keepingCapacity {
-      _storage.initialize(for: _storage.capacity)
+      _storage.initialize(for: capacity)
     } else {
+      capacity = initialSize
       _storage = Storage(count: initialSize, alignment: alignment)
     }
   }
@@ -358,7 +368,7 @@ extension _InternalByteBuffer: CustomDebugStringConvertible {
 
   public var debugDescription: String {
     """
-    buffer located at: \(_storage.memory), with capacity of \(_storage.capacity)
+    buffer located at: \(_storage.memory), with capacity of \(capacity)
     { writerSize: \(_writerSize), readerSize: \(reader), writerIndex: \(
       writerIndex) }
     """