[Swift] Memory usage fix (#8643)

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

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)