[Swift] Memory usage fix (#8643)

Allows a complete reset for the underlying memory of the
_InternalByteBuffers within FlatBuffers and FlexBuffers.

Package.swift

@@ -47,13 +47,11 @@ let package = Package(
     .testTarget(
       name: "FlatbuffersTests",
       dependencies: .dependencies,
-      path: "tests/swift/Tests/Flatbuffers"
-    ),
+      path: "tests/swift/Tests/Flatbuffers"),
     .testTarget(
       name: "FlexbuffersTests",
       dependencies: ["FlexBuffers"],
-      path: "tests/swift/Tests/Flexbuffers"
-    )
+      path: "tests/swift/Tests/Flexbuffers"),
   ])
 
 extension Array where Element == Package.Dependency {
@@ -75,7 +73,7 @@ extension Array where Element == PackageDescription.Target.Dependency {
     // Test only Dependency
     [
       .product(name: "GRPC", package: "grpc-swift"),
-      "FlatBuffers"
+      "FlatBuffers",
     ]
     #endif
   }

swift/Sources/FlatBuffers/Enum.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlatBuffers/FlatBufferBuilder.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif
@@ -72,8 +73,8 @@ public struct FlatBufferBuilder {
       data.append(
         ptr.baseAddress!.bindMemory(
           to: UInt8.self,
-          capacity: _bb.capacity),
-        count: _bb.capacity)
+          capacity: ptr.count),
+        count: ptr.count)
       return data
     }
   }
@@ -143,13 +144,13 @@ public struct FlatBufferBuilder {
   }
 
   /// Clears the builder and the buffer from the written data.
-  mutating public func clear() {
+  mutating public func clear(keepingCapacity: Bool = false) {
     _minAlignment = 0
     isNested = false
-    stringOffsetMap.removeAll(keepingCapacity: true)
-    _vtables.removeAll(keepingCapacity: true)
-    _vtableStorage.clear()
-    _bb.clear()
+    stringOffsetMap.removeAll(keepingCapacity: keepingCapacity)
+    _vtables.removeAll(keepingCapacity: keepingCapacity)
+    _vtableStorage.reset(keepingCapacity: keepingCapacity)
+    _bb.clear(keepingCapacity: keepingCapacity)
   }
 
   // MARK: - Create Tables
@@ -852,10 +853,6 @@ extension FlatBufferBuilder: CustomDebugStringConvertible {
   /// VTableStorage is a class to contain the VTable buffer that would be serialized into buffer
   @usableFromInline
   internal class VTableStorage {
-    /// Memory check since deallocating each time we want to clear would be expensive
-    /// and memory leaks would happen if we dont deallocate the first allocated memory.
-    /// memory is promised to be available before adding `FieldLoc`
-    private var memoryInUse = false
     /// Size of FieldLoc in memory
     let size = MemoryLayout<FieldLoc>.stride
     /// Memeory buffer
@@ -905,6 +902,24 @@ extension FlatBufferBuilder: CustomDebugStringConvertible {
       maxOffset = max(loc.position, maxOffset)
     }
 
+    /// Clears the data stored related to the encoded buffer
+    @inline(__always)
+    func reset(keepingCapacity: Bool) {
+      maxOffset = 0
+      numOfFields = 0
+      writtenIndex = 0
+      if keepingCapacity {
+        memset(memory.baseAddress!, 0, memory.count)
+      } else {
+        capacity = 0
+        let memory = UnsafeMutableRawBufferPointer.allocate(
+          byteCount: 0,
+          alignment: 0)
+        self.memory.deallocate()
+        self.memory = memory
+      }
+    }
+
     /// Clears the data stored related to the encoded buffer
     @inline(__always)
     func clear() {

swift/Sources/FlatBuffers/Mutable.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlatBuffers/Struct.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlatBuffers/Table.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlatBuffers/_InternalByteBuffer.swift

@@ -29,12 +29,10 @@ struct _InternalByteBuffer {
   /// 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.
-    private let unowned: Bool
     /// pointer to the start of the buffer object in memory
-    var memory: UnsafeMutableRawPointer
+    private(set) var memory: UnsafeMutableRawPointer
     /// Capacity of UInt8 the buffer can hold
-    var capacity: Int
+    private(set) var capacity: Int
 
     @usableFromInline
     init(count: Int, alignment: Int) {
@@ -42,20 +40,14 @@ struct _InternalByteBuffer {
         byteCount: count,
         alignment: alignment)
       capacity = count
-      unowned = false
     }
 
     deinit {
-      if !unowned {
-        memory.deallocate()
-      }
+      memory.deallocate()
     }
 
     @usableFromInline
     func initialize(for size: Int) {
-      assert(
-        !unowned,
-        "initalize should NOT be called on a buffer that is built by assumingMemoryBound")
       memset(memory, 0, size)
     }
 
@@ -86,6 +78,7 @@ struct _InternalByteBuffer {
 
   @usableFromInline var _storage: Storage
 
+  private let initialSize: Int
   /// The size of the elements written to the buffer + their paddings
   private var _writerSize: Int = 0
   /// Alignment of the current  memory being written to the buffer
@@ -108,9 +101,9 @@ struct _InternalByteBuffer {
   ///   - size: Length of the buffer
   ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
   init(initialSize size: Int) {
-    let size = size.convertToPowerofTwo
-    _storage = Storage(count: size, alignment: alignment)
-    _storage.initialize(for: size)
+    initialSize = size.convertToPowerofTwo
+    _storage = Storage(count: initialSize, alignment: alignment)
+    _storage.initialize(for: initialSize)
   }
 
   /// Fills the buffer with padding by adding to the writersize
@@ -298,10 +291,14 @@ struct _InternalByteBuffer {
 
   /// Clears the current instance of the buffer, replacing it with new memory
   @inline(__always)
-  mutating public func clear() {
+  mutating public func clear(keepingCapacity: Bool = false) {
     _writerSize = 0
     alignment = 1
-    _storage.initialize(for: _storage.capacity)
+    if keepingCapacity {
+      _storage.initialize(for: _storage.capacity)
+    } else {
+      _storage = Storage(count: initialSize, alignment: alignment)
+    }
   }
 
   /// Reads an object from the buffer

swift/Sources/FlexBuffers/Utils/Value.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlexBuffers/Utils/functions.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif

swift/Sources/FlexBuffers/Writer/FlexBuffersWriter.swift

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #if canImport(Common)
 import Common
 #endif
@@ -72,14 +73,32 @@ public struct FlexBuffersWriter {
     return ByteBuffer(byteBuffer: _bb)
   }
 
+  #if !os(WASI)
+  /// Data representation of the buffer
+  ///
+  /// Should only be used after ``finish(offset:addPrefix:)`` is called
+  public var data: Data {
+    assert(finished, "Data shouldn't be called before finish()")
+    return _bb.withUnsafeSlicedBytes { ptr in
+      var data = Data()
+      data.append(
+        ptr.baseAddress!.bindMemory(
+          to: UInt8.self,
+          capacity: ptr.count),
+        count: ptr.count)
+      return data
+    }
+  }
+  #endif
+
   /// Resets the internal state. Automatically called before building a new flexbuffer.
-  public mutating func reset() {
-    _bb.clear()
-    stack.removeAll(keepingCapacity: true)
+  public mutating func reset(keepingCapacity: Bool = false) {
+    _bb.clear(keepingCapacity: keepingCapacity)
+    stack.removeAll(keepingCapacity: keepingCapacity)
     finished = false
     minBitWidth = .w8
-    keyPool.removeAll()
-    stringPool.removeAll()
+    keyPool.removeAll(keepingCapacity: keepingCapacity)
+    stringPool.removeAll(keepingCapacity: keepingCapacity)
   }
 
   // MARK: - Storing root

swift/Sources/FlexBuffers/_InternalByteBuffer.swift

@@ -30,8 +30,6 @@ struct _InternalByteBuffer {
   /// 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.
-    private let unowned: Bool
     /// pointer to the start of the buffer object in memory
     var memory: UnsafeMutableRawPointer
     /// Capacity of UInt8 the buffer can hold
@@ -43,35 +41,25 @@ struct _InternalByteBuffer {
         byteCount: count,
         alignment: alignment)
       capacity = count
-      unowned = false
     }
 
     @usableFromInline
     init(memory: UnsafeMutableRawPointer, capacity: Int, unowned: Bool) {
       self.memory = memory
       self.capacity = capacity
-      self.unowned = unowned
     }
 
     deinit {
-      if !unowned {
-        memory.deallocate()
-      }
+      memory.deallocate()
     }
 
     @usableFromInline
     func copy(from ptr: UnsafeRawPointer, count: Int) {
-      assert(
-        !unowned,
-        "copy should NOT be called on a buffer that is built by assumingMemoryBound")
       memory.copyMemory(from: ptr, byteCount: count)
     }
 
     @usableFromInline
     func initialize(for size: Int) {
-      assert(
-        !unowned,
-        "initalize should NOT be called on a buffer that is built by assumingMemoryBound")
       memset(memory, 0, size)
     }
 
@@ -100,6 +88,8 @@ struct _InternalByteBuffer {
   }
 
   @usableFromInline var _storage: Storage
+  // Initial size of the internal storage
+  private let initialSize: Int
   /// The size of the elements written to the buffer + their paddings
   var writerIndex: Int = 0
   /// Alignment of the current  memory being written to the buffer
@@ -122,17 +112,21 @@ struct _InternalByteBuffer {
   ///   - size: Length of the buffer
   ///   - allowReadingUnalignedBuffers: allow reading from unaligned buffer
   init(initialSize size: Int) {
-    let size = size.convertToPowerofTwo
-    _storage = Storage(count: size, alignment: alignment)
-    _storage.initialize(for: size)
+    initialSize = size.convertToPowerofTwo
+    _storage = Storage(count: initialSize, alignment: alignment)
+    _storage.initialize(for: initialSize)
   }
 
   /// Clears the current instance of the buffer, replacing it with new memory
   @inline(__always)
-  mutating public func clear() {
+  mutating public func clear(keepingCapacity: Bool = false) {
     writerIndex = 0
     alignment = 1
-    _storage.initialize(for: _storage.capacity)
+    if keepingCapacity {
+      _storage.initialize(for: _storage.capacity)
+    } else {
+      _storage = Storage(count: initialSize, alignment: alignment)
+    }
   }
 
   @inline(__always)

tests/swift/Tests/Flatbuffers/FlatBuffersMonsterWriterTests.swift

@@ -43,6 +43,34 @@ class FlatBuffersMonsterWriterTests: XCTestCase {
     readVerifiedMonster(fb: _data)
   }
 
+  func testCreateMonsterData() {
+    let bytes = createMonster(withPrefix: false)
+    var buffer = ByteBuffer(data: bytes.data)
+    let monster: MyGame_Example_Monster = getRoot(byteBuffer: &buffer)
+    readMonster(monster: monster)
+    mutateMonster(fb: bytes.buffer)
+    readMonster(monster: monster)
+  }
+
+  func testCreateMonsterResetTests() {
+    var builder = createMonster(withPrefix: false)
+    var buffer = ByteBuffer(data: builder.data)
+    let monster: MyGame_Example_Monster = getRoot(byteBuffer: &buffer)
+    readMonster(monster: monster)
+    builder.clear()
+    XCTAssertEqual(builder.capacity, 1)
+    XCTAssertEqual(builder.size, 0)
+
+    write(fbb: &builder, prefix: false)
+    var _buffer = ByteBuffer(data: builder.data)
+    XCTAssertEqual(_buffer.capacity, 304)
+    let _monster: MyGame_Example_Monster = getRoot(byteBuffer: &_buffer)
+    readMonster(monster: _monster)
+    builder.clear(keepingCapacity: true)
+    XCTAssertEqual(builder.capacity, 512)
+    XCTAssertEqual(builder.size, 0)
+  }
+
   func testCreateMonster() {
     let bytes = createMonster(withPrefix: false)
     // swiftformat:disable all
@@ -257,6 +285,11 @@ class FlatBuffersMonsterWriterTests: XCTestCase {
 
   func createMonster(withPrefix prefix: Bool) -> FlatBufferBuilder {
     var fbb = FlatBufferBuilder(initialSize: 1)
+    write(fbb: &fbb, prefix: prefix)
+    return fbb
+  }
+
+  func write(fbb: inout FlatBufferBuilder, prefix: Bool = false) {
     let names = [
       fbb.create(string: "Frodo"),
       fbb.create(string: "Barney"),
@@ -313,7 +346,6 @@ class FlatBuffersMonsterWriterTests: XCTestCase {
     Monster.addVectorOf(testarrayoftables: sortedArray, &fbb)
     let end = Monster.endMonster(&fbb, start: mStart)
     Monster.finish(&fbb, end: end, prefix: prefix)
-    return fbb
   }
 
   func mutateMonster(fb: ByteBuffer) {

tests/swift/Tests/Flexbuffers/FlexBuffersReaderTests.swift

@@ -15,9 +15,10 @@
  */
 
 import Common
-import FlexBuffers
 import XCTest
 
+@testable import FlexBuffers
+
 final class FlexBuffersReaderTests: XCTestCase {
 
   func testReadingProperBuffer() throws {
@@ -30,6 +31,29 @@ final class FlexBuffersReaderTests: XCTestCase {
     try validate(buffer: buf)
   }
 
+  func testReset() throws {
+    var fbx = FlexBuffersWriter(
+      initialSize: 8,
+      flags: .shareKeysAndStrings)
+    write(fbx: &fbx)
+
+    try validate(buffer: ByteBuffer(data: fbx.data))
+    XCTAssertEqual(fbx.capacity, 512)
+    fbx.reset()
+    XCTAssertEqual(fbx.writerIndex, 0)
+    XCTAssertEqual(fbx.capacity, 8)
+
+    write(fbx: &fbx)
+    try validate(buffer: ByteBuffer(data: fbx.data))
+    fbx.reset(keepingCapacity: true)
+    XCTAssertEqual(fbx.writerIndex, 0)
+    XCTAssertEqual(fbx.capacity, 512)
+
+    write(fbx: &fbx)
+    try validate(buffer: ByteBuffer(data: fbx.data))
+    XCTAssertEqual(fbx.capacity, 512)
+  }
+
   private func validate(buffer buf: ByteBuffer) throws {
     let reference = try getRoot(buffer: buf)!
     XCTAssertEqual(reference.type, .map)

tests/swift/Tests/Flexbuffers/Mocks.swift

@@ -32,7 +32,11 @@ func createProperBuffer() -> FlexBuffersWriter {
   var fbx = FlexBuffersWriter(
     initialSize: 8,
     flags: .shareKeysAndStrings)
+  write(fbx: &fbx)
+  return fbx
+}
 
+func write(fbx: inout FlexBuffersWriter) {
   fbx.map { map in
     map.vector(key: "vec") { v in
       v.add(int64: -100)
@@ -57,5 +61,4 @@ func createProperBuffer() -> FlexBuffersWriter {
   }
 
   fbx.finish()
-  return fbx
 }