| /* |
| * Copyright 2024 Google Inc. All rights reserved. |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); |
| * you may not use this file except in compliance with the License. |
| * You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| import Foundation |
| |
| /// `ByteBuffer` is the interface that stores the data for a `Flatbuffers` object |
| /// it allows users to write and read data directly from memory thus the use of its |
| /// functions should be used |
| @frozen |
| public struct ByteBuffer { |
| |
| /// Storage is a container that would hold the memory pointer to solve the issue of |
| /// deallocating the memory that was held by (memory: UnsafeMutableRawPointer) |
| @usableFromInline |
| final class Storage { |
| @usableFromInline |
| enum Blob { |
| #if !os(WASI) |
| case data(Data) |
| case bytes(ContiguousBytes) |
| #endif |
| |
| case byteBuffer(_InternalByteBuffer) |
| case array([UInt8]) |
| case pointer(UnsafeMutableRawPointer) |
| } |
| |
| /// 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 |
| |
| @usableFromInline |
| init(count: Int) { |
| let memory = UnsafeMutableRawPointer.allocate( |
| byteCount: count, |
| alignment: MemoryLayout<UInt8>.alignment) |
| capacity = count |
| retainedBlob = .pointer(memory) |
| isOwned = true |
| } |
| |
| @usableFromInline |
| init(blob: Blob, capacity count: Int) { |
| capacity = count |
| retainedBlob = blob |
| isOwned = false |
| } |
| |
| deinit { |
| guard isOwned else { return } |
| switch retainedBlob { |
| case .pointer(let unsafeMutableRawPointer): |
| unsafeMutableRawPointer.deallocate() |
| default: break |
| } |
| } |
| |
| @usableFromInline |
| func copy(from ptr: UnsafeRawPointer, count: Int) { |
| assert( |
| isOwned, |
| "copy should NOT be called on a buffer that is built by assumingMemoryBound") |
| withUnsafeRawPointer { |
| $0.copyMemory(from: ptr, byteCount: count) |
| } |
| } |
| |
| @usableFromInline |
| func initialize(for size: Int) { |
| assert( |
| isOwned, |
| "initalize should NOT be called on a buffer that is built by assumingMemoryBound") |
| withUnsafeRawPointer { |
| memset($0, 0, size) |
| } |
| } |
| |
| @discardableResult |
| @inline(__always) |
| func withUnsafeBytes<T>( |
| _ body: (UnsafeRawBufferPointer) throws |
| -> T |
| ) rethrows -> T { |
| switch retainedBlob { |
| case .byteBuffer(let byteBuffer): |
| return try byteBuffer.withUnsafeBytes(body) |
| #if !os(WASI) |
| case .data(let data): |
| return try data.withUnsafeBytes(body) |
| case .bytes(let contiguousBytes): |
| return try contiguousBytes.withUnsafeBytes(body) |
| #endif |
| case .array(let array): |
| return try array.withUnsafeBytes(body) |
| case .pointer(let ptr): |
| return try body(UnsafeRawBufferPointer(start: ptr, count: capacity)) |
| } |
| } |
| |
| @discardableResult |
| @inline(__always) |
| func withUnsafeRawPointer<T>( |
| _ body: (UnsafeMutableRawPointer) throws |
| -> T |
| ) rethrows -> T { |
| switch retainedBlob { |
| case .byteBuffer(let byteBuffer): |
| return try byteBuffer.withUnsafeRawPointer(body) |
| #if !os(WASI) |
| case .data(let data): |
| return |
| try data |
| .withUnsafeBytes { |
| try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!)) |
| } |
| case .bytes(let contiguousBytes): |
| return |
| try contiguousBytes |
| .withUnsafeBytes { |
| try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!)) |
| } |
| #endif |
| case .array(let array): |
| return |
| try array |
| .withUnsafeBytes { |
| try body(UnsafeMutableRawPointer(mutating: $0.baseAddress!)) |
| } |
| case .pointer(let ptr): |
| return try body(ptr) |
| } |
| } |
| |
| @discardableResult |
| @inline(__always) |
| func readWithUnsafeRawPointer<T>( |
| position: Int, |
| _ body: (UnsafeRawPointer) throws -> T |
| ) rethrows -> T { |
| switch retainedBlob { |
| case .byteBuffer(let byteBuffer): |
| return try byteBuffer.readWithUnsafeRawPointer(position: position, body) |
| #if !os(WASI) |
| case .data(let data): |
| return try data.withUnsafeBytes { |
| try body($0.baseAddress!.advanced(by: position)) |
| } |
| case .bytes(let contiguousBytes): |
| return try contiguousBytes.withUnsafeBytes { |
| try body($0.baseAddress!.advanced(by: position)) |
| } |
| #endif |
| case .array(let array): |
| return try array.withUnsafeBytes { |
| try body($0.baseAddress!.advanced(by: position)) |
| } |
| case .pointer(let ptr): |
| return try body(ptr.advanced(by: position)) |
| } |
| } |
| } |
| |
| @usableFromInline var _storage: Storage |
| |
| /// 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 { capacity &- _readerIndex } |
| /// Current size of the buffer |
| public var size: UOffset { UOffset(_readerIndex) } |
| /// Current capacity for the buffer |
| public let capacity: Int |
| |
| /// Constructor that creates a Flatbuffer object from an InternalByteBuffer |
| /// - Parameter |
| /// - bytes: Array of UInt8 |
| @inline(__always) |
| init(byteBuffer: _InternalByteBuffer) { |
| _storage = Storage( |
| 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 |
| /// the underlying data to a new pointer |
| /// |
| /// - Parameters: |
| /// - copyingMemoryBound: The unsafe memory region |
| /// - capacity: The size of the given memory region |
| @inline(__always) |
| public init( |
| copyingMemoryBound memory: UnsafeRawPointer, |
| capacity: Int |
| ) { |
| _storage = Storage(count: capacity) |
| _storage.copy(from: memory, count: capacity) |
| _readerIndex = capacity |
| self.capacity = capacity |
| } |
| |
| /// Constructor that creates a Flatbuffer object from a UInt8 |
| /// - Parameter |
| /// - bytes: Array of UInt8 |
| @inline(__always) |
| public init(bytes: [UInt8]) { |
| _storage = Storage(blob: .array(bytes), capacity: bytes.count) |
| _readerIndex = bytes.count |
| capacity = bytes.count |
| } |
| |
| #if !os(WASI) |
| /// Constructor that creates a Flatbuffer from the Swift Data type object |
| /// - Parameter |
| /// - data: Swift data Object |
| @inline(__always) |
| public init(data: Data) { |
| _storage = Storage(blob: .data(data), capacity: data.count) |
| _readerIndex = data.count |
| capacity = data.count |
| } |
| |
| /// Constructor that creates a Flatbuffer object from a ContiguousBytes |
| /// - Parameters: |
| /// - contiguousBytes: Binary stripe to use as the buffer |
| /// - count: amount of readable bytes |
| @inline(__always) |
| public init<Bytes: ContiguousBytes>( |
| contiguousBytes: Bytes, |
| count: Int |
| ) { |
| _storage = Storage(blob: .bytes(contiguousBytes), capacity: count) |
| _readerIndex = count |
| self.capacity = count |
| } |
| #endif |
| |
| /// Constructor that creates a Flatbuffer from unsafe memory region without copying |
| /// **NOTE** Needs a call to `memory.deallocate()` later on to free the memory |
| /// |
| /// - Parameters: |
| /// - assumingMemoryBound: The unsafe memory region |
| /// - capacity: The size of the given memory region |
| @inline(__always) |
| public init( |
| assumingMemoryBound memory: UnsafeMutableRawPointer, |
| capacity: Int |
| ) { |
| _storage = Storage( |
| blob: .pointer(memory), |
| capacity: capacity) |
| _readerIndex = capacity |
| self.capacity = capacity |
| } |
| |
| /// Creates a copy of the existing flatbuffer, by copying it to a different memory. |
| /// - Parameters: |
| /// - memory: Current memory of the buffer |
| /// - count: count of bytes |
| /// - removeBytes: Removes a number of bytes from the current size |
| @inline(__always) |
| init( |
| blob: Storage.Blob, |
| count: Int, |
| removing removeBytes: Int |
| ) { |
| _storage = Storage(blob: blob, capacity: count) |
| _readerIndex = removeBytes |
| capacity = count |
| } |
| |
| /// Write stores an object into the buffer directly or indirectly. |
| /// |
| /// Direct: ignores the capacity of buffer which would mean we are referring to the direct point in memory |
| /// indirect: takes into respect the current capacity of the buffer (capacity - index), writing to the buffer from the end |
| /// - Parameters: |
| /// - value: Value that needs to be written to the buffer |
| /// - index: index to write to |
| /// - direct: Should take into consideration the capacity of the buffer |
| @inline(__always) |
| func write<T>(value: T, index: Int, direct: Bool = false) { |
| var index = index |
| if !direct { |
| index = capacity &- index |
| } |
| assert(index < capacity, "Write index is out of writing bound") |
| assert(index >= 0, "Writer index should be above zero") |
| _ = withUnsafePointer(to: value) { ptr in |
| _storage.withUnsafeRawPointer { |
| memcpy( |
| $0.advanced(by: index), |
| ptr, |
| MemoryLayout<T>.size) |
| } |
| } |
| } |
| |
| /// Reads an object from the buffer |
| /// - Parameters: |
| /// - def: Type of the object |
| /// - position: the index of the object in the buffer |
| @inline(__always) |
| public func read<T>(def: T.Type, position: Int) -> T { |
| _storage.readWithUnsafeRawPointer(position: position) { |
| $0.bindMemory(to: T.self, capacity: 1) |
| .pointee |
| } |
| } |
| |
| /// Reads a slice from the memory assuming a type of T |
| /// - Parameters: |
| /// - index: index of the object to be read from the buffer |
| /// - count: count of bytes in memory |
| @inline(__always) |
| public func readSlice<T>( |
| index: Int, |
| count: Int |
| ) -> [T] { |
| assert( |
| index + count <= capacity, |
| "Reading out of bounds is illegal") |
| |
| return _storage.readWithUnsafeRawPointer(position: index) { |
| let buf = UnsafeBufferPointer( |
| start: $0.bindMemory(to: T.self, capacity: count), |
| count: count) |
| return Array(buf) |
| } |
| } |
| |
| /// Provides a pointer towards the underlying primitive types |
| /// - Parameters: |
| /// - index: index of the object to be read from the buffer |
| /// - count: count of bytes in memory |
| @discardableResult |
| @inline(__always) |
| public func withUnsafePointerToSlice<T>( |
| index: Int, |
| count: Int, |
| body: (UnsafeRawBufferPointer) throws -> T |
| ) rethrows -> T { |
| assert( |
| index + count <= capacity, |
| "Reading out of bounds is illegal") |
| return try _storage.readWithUnsafeRawPointer(position: index) { |
| try body(UnsafeRawBufferPointer(start: $0, count: count)) |
| } |
| } |
| |
| #if !os(WASI) |
| /// Reads a string from the buffer and encodes it to a swift string |
| /// - Parameters: |
| /// - index: index of the string in the buffer |
| /// - count: length of the string |
| /// - type: Encoding of the string |
| @inline(__always) |
| public func readString( |
| at index: Int, |
| count: Int, |
| type: String.Encoding = .utf8 |
| ) -> String? { |
| assert( |
| index + count <= capacity, |
| "Reading out of bounds is illegal") |
| return _storage.readWithUnsafeRawPointer(position: index) { |
| let buf = UnsafeBufferPointer( |
| start: $0.bindMemory(to: UInt8.self, capacity: count), |
| count: count) |
| return String( |
| bytes: buf, |
| encoding: type) |
| } |
| } |
| #else |
| /// Reads a string from the buffer and encodes it to a swift string |
| /// - Parameters: |
| /// - index: index of the string in the buffer |
| /// - count: length of the string |
| @inline(__always) |
| public func readString( |
| at index: Int, |
| count: Int |
| ) -> String? { |
| assert( |
| index + count <= capacity, |
| "Reading out of bounds is illegal") |
| return _storage.readWithUnsafeRawPointer(position: index) { |
| String(cString: $0.bindMemory(to: UInt8.self, capacity: count)) |
| } |
| } |
| #endif |
| |
| /// Creates a new Flatbuffer object that's duplicated from the current one |
| /// - Parameter removeBytes: the amount of bytes to remove from the current Size |
| @inline(__always) |
| public func duplicate(removing removeBytes: Int = 0) -> ByteBuffer { |
| assert(removeBytes > 0, "Can NOT remove negative bytes") |
| assert( |
| removeBytes < capacity, |
| "Can NOT remove more bytes than the ones allocated") |
| return ByteBuffer( |
| blob: _storage.retainedBlob, |
| count: capacity, |
| removing: _readerIndex &- removeBytes) |
| } |
| |
| /// SkipPrefix Skips the first 4 bytes in case one of the following |
| /// functions are called `getPrefixedSizeCheckedRoot` & `getPrefixedSizeRoot` |
| /// which allows us to skip the first 4 bytes instead of recreating the buffer |
| @discardableResult |
| @usableFromInline |
| @inline(__always) |
| mutating func skipPrefix() -> Int32 { |
| _readerIndex = _readerIndex &- MemoryLayout<Int32>.size |
| return read(def: Int32.self, position: 0) |
| } |
| |
| @discardableResult |
| @inline(__always) |
| public func withUnsafeBytes<T>( |
| body: (UnsafeRawBufferPointer) throws |
| -> T |
| ) rethrows -> T { |
| try _storage.withUnsafeBytes(body) |
| } |
| |
| @discardableResult |
| @inline(__always) |
| func withUnsafeMutableRawPointer<T>( |
| body: (UnsafeMutableRawPointer) throws |
| -> T |
| ) rethrows -> T { |
| try _storage.withUnsafeRawPointer(body) |
| } |
| |
| @discardableResult |
| @inline(__always) |
| func readWithUnsafeRawPointer<T>( |
| position: Int, |
| _ body: (UnsafeRawPointer) throws -> T |
| ) rethrows -> T { |
| try _storage.readWithUnsafeRawPointer(position: position, body) |
| } |
| } |
| |
| extension ByteBuffer: CustomDebugStringConvertible { |
| |
| public var debugDescription: String { |
| """ |
| buffer located at: \(_storage.retainedBlob), |
| with capacity of \(capacity), |
| { writtenSize: \(_readerIndex), readerSize: \(reader), |
| size: \(size) } |
| """ |
| } |
| } |