Skip to content

Commit

Permalink
Add the ability to encode and decode a size delimited message collect…
Browse files Browse the repository at this point in the history
…ion in Swift.
  • Loading branch information
jszumski committed Jul 12, 2023
1 parent 7878118 commit 41fdf6a
Show file tree
Hide file tree
Showing 7 changed files with 128 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,12 @@ public final class ProtoDecoder {

// MARK: - Public Methods

/// Decodes the provided data into an instance of the requested type.
///
/// - Parameters:
/// - type: the type to decode
/// - data: the serialized data for the message
/// - Returns: the decoded message
public func decode<T: ProtoDecodable>(_ type: T.Type, from data: Data) throws -> T {
var value: T?
try data.withUnsafeBytes { buffer in
Expand All @@ -148,5 +154,50 @@ public final class ProtoDecoder {
return unwrappedValue
}

}
/// Decodes the provided size-delimited data into instances of the requested type.
///
/// A size-delimited collection of messages is a sequence of varint + message pairs
/// where the varint indicates the size of the subsequent message.
///
/// - Parameters:
/// - type: the type to decode
/// - data: the serialized size-delimited data for the messages
/// - Returns: an array of the decoded messages
public func decodeSizeDelimited<T: ProtoDecodable>(_ type: T.Type, from data: Data) throws -> [T] {
var values: [T] = []

try data.withUnsafeBytes { buffer in
// Handle the empty-data case.
guard let baseAddress = buffer.baseAddress, buffer.count > 0 else {
return
}

let fullBuffer = ReadBuffer(
storage: baseAddress.bindMemory(to: UInt8.self, capacity: buffer.count),
count: buffer.count
)

while fullBuffer.isDataRemaining, let size = try? fullBuffer.readVarint64() {
if size == 0 { break }

let messageBuffer = ReadBuffer(
storage: fullBuffer.pointer,
count: Int(size)
)

let reader = ProtoReader(
buffer: messageBuffer,
enumDecodingStrategy: enumDecodingStrategy
)

values.append(try reader.decode(type))

// Advance the buffer before reading the next item in the stream
_ = try fullBuffer.readBuffer(count: Int(size))
}
}

return values
}

}
34 changes: 32 additions & 2 deletions wire-runtime-swift/src/main/swift/ProtoCodable/ProtoEncoder.swift
Original file line number Diff line number Diff line change
Expand Up @@ -82,13 +82,43 @@ public final class ProtoEncoder {

let writer = ProtoWriter(
data: .init(capacity: structSize),
outputFormatting: [],
outputFormatting: outputFormatting,
rootMessageProtoSyntax: T.self.protoSyntax ?? .proto2
)
writer.outputFormatting = outputFormatting

try value.encode(to: writer)

return Data(writer.buffer, copyBytes: false)
}

public func encodeSizeDelimited<T: ProtoEncodable>(_ values: [T]) throws -> Data {
// Use the size of the struct as an initial estimate for the space needed.
let structSize = MemoryLayout.size(ofValue: T.self)

// Reserve space for the largest varint size
let varintSize = 8

let fullBuffer = WriteBuffer(capacity: (structSize + varintSize) * values.count)

for value in values {
let writer = ProtoWriter(
data: .init(),
outputFormatting: outputFormatting,
rootMessageProtoSyntax: T.self.protoSyntax ?? .proto2
)

try value.encode(to: writer)

if writer.buffer.count == 0 {
continue
}

// write this value's size + contents to the main buffer
fullBuffer.writeVarint(UInt64(writer.buffer.count), at: fullBuffer.count)
fullBuffer.append(writer.buffer)
}

return Data(fullBuffer, copyBytes: false)
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ final class WriteBuffer {
// MARK: - Public Methods

func append(_ data: Data) {
guard !data.isEmpty else { return }

expandIfNeeded(adding: data.count)

data.copyBytes(to: storage.advanced(by: count), count: data.count)
Expand All @@ -64,6 +66,8 @@ final class WriteBuffer {
}

func append(_ value: [UInt8]) {
guard !value.isEmpty else { return }

expandIfNeeded(adding: value.count)

for byte in value {
Expand All @@ -74,13 +78,17 @@ final class WriteBuffer {

func append(_ value: WriteBuffer) {
precondition(value !== self)
guard value.count > 0 else { return }

expandIfNeeded(adding: value.count)

memcpy(storage.advanced(by: count), value.storage, value.count)
count += value.count
}

func append(_ value: UnsafeRawBufferPointer) {
guard value.count > 0 else { return }

expandIfNeeded(adding: value.count)

memcpy(storage.advanced(by: count), value.baseAddress, value.count)
Expand Down
7 changes: 7 additions & 0 deletions wire-runtime-swift/src/test/swift/ProtoDecoderTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,13 @@ final class ProtoDecoderTests: XCTestCase {
XCTAssertEqual(object, SimpleOptional2())
}

func testDecodeEmptySizeDelimitedData() throws {
let decoder = ProtoDecoder()
let object = try decoder.decodeSizeDelimited(SimpleOptional2.self, from: Data())

XCTAssertEqual(object, [])
}

func testDecodeEmptyDataTwice() throws {
let decoder = ProtoDecoder()
// The empty message case is optimized to reuse objects, so make sure
Expand Down
8 changes: 8 additions & 0 deletions wire-runtime-swift/src/test/swift/ProtoEncoderTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -44,4 +44,12 @@ final class ProtoEncoderTests: XCTestCase {

XCTAssertEqual(jsonString, "{}")
}

func testEncodeEmptySizeDelimitedMessage() throws {
let object = EmptyMessage()
let encoder = ProtoEncoder()
let data = try encoder.encodeSizeDelimited([object])

XCTAssertEqual(data, Data())
}
}
14 changes: 14 additions & 0 deletions wire-runtime-swift/src/test/swift/RoundTripTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -60,4 +60,18 @@ final class RoundTripTests: XCTestCase {
XCTAssertEqual(decodedEmpty, empty)
}

func testSizeDelimited() throws {
let values = [
Person3(name: "John Doe", id: 123),
Person3(name: "Jane Doe", id: 456, email: "jdoe@example.com")
]

let encoder = ProtoEncoder()
let data = try encoder.encodeSizeDelimited(values)

let decoder = ProtoDecoder()
let decodedValues = try decoder.decodeSizeDelimited(Person3.self, from: data)

XCTAssertEqual(decodedValues, values)
}
}
7 changes: 7 additions & 0 deletions wire-runtime-swift/src/test/swift/WriteBufferTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -60,4 +60,11 @@ final class WriteBufferTests: XCTestCase {
XCTAssertEqual(Data(buffer, copyBytes: true), Data(hexEncoded: "0011"))
}

func testAppendEmptyFirst() {
let buffer = WriteBuffer()
buffer.append(Data())

XCTAssertEqual(Data(buffer, copyBytes: true), Data())
}

}

0 comments on commit 41fdf6a

Please sign in to comment.