Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Swift: encode & decode size-delimited messages #10

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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: Foundation.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 }
Comment on lines +180 to +181
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

logic: Breaking on size 0 might skip valid empty messages. Consider handling this case explicitly


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 }
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

style: Consider using isEmpty instead of count > 0 for consistency.


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: Foundation.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, Foundation.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 @@ -62,4 +62,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(Foundation.Data(buffer, copyBytes: true), Foundation.Data(hexEncoded: "0011"))
}

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

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

}