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 Mar 29, 2023
1 parent 7878118 commit 238c348
Show file tree
Hide file tree
Showing 7 changed files with 119 additions and 4 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,51 @@ 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 {
values = [try T(from: .empty)]
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
}

}
31 changes: 29 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,40 @@ 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)

let fullBuffer = WriteBuffer(capacity: structSize * 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 @@ -36,7 +36,7 @@ final class WriteBuffer {
init(capacity: Int = 0) {
self.capacity = 0

if capacity > 0 {
if capacity >= 0 {
expand(to: capacity)
}
}
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, [SimpleOptional2()])
}

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: "[email protected]")
]

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 238c348

Please sign in to comment.