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 #2424

Merged
merged 1 commit into from
Nov 12, 2024
Merged
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
48 changes: 48 additions & 0 deletions wire-runtime-swift/src/main/swift/ProtoCodable/ProtoDecoder.swift
Original file line number Diff line number Diff line change
Expand Up @@ -279,4 +279,52 @@ public final class ProtoDecoder {
}
return value
}

/** 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.readVarint() {
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
}

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

let writer = ProtoWriter(
data: .init(capacity: structSize),
outputFormatting: [],
outputFormatting: outputFormatting,
rootMessageProtoSyntax: syntax
)
writer.outputFormatting = outputFormatting

try encoder(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)
Copy link
Collaborator

Choose a reason for hiding this comment

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

FYI: We have some tricks inside the writer to avoid this write-then-copy inside of the ProtoWriter (see beginLengthDelimitedEncode). They're probably not worth doing here, but for future archaeologists who might be trying to optimize this I'll mention.

Basically, most messages will be in the size range that requires (IIRC) two bytes of varint, so we reserve two bytes, then write out value into the full buffer. If it happens to not take two bytes for its size, only then do we go back and move it to add more or less room for the prefix size.

In the future, we could extract this into an extension on WriteBuffer, like:

// This will reserve the length bytes, then write the buffer, then adjust and fill in the length
// bytes based on what was written.
buffer.writeLengthEncoded { tempBuffer in
   let writer = ProtoWriter(data: tempBuffer, ...)
   try value.encode(to: writer)
}

Again, all a good future enhancement. Not required now.

fullBuffer.append(writer.buffer)
}

return Data(fullBuffer, copyBytes: false)
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,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 @@ -63,6 +65,8 @@ final class WriteBuffer {
}

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

expandIfNeeded(adding: value.count)

for byte in value {
Expand All @@ -73,13 +77,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 @@ -25,6 +25,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 @@ -43,4 +43,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())
}
}
16 changes: 16 additions & 0 deletions wire-runtime-swift/src/test/swift/RoundTripTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -61,4 +61,20 @@ final class RoundTripTests: XCTestCase {
XCTAssertEqual(decodedEmpty, empty)
}

func testSizeDelimited() throws {
let values = [
Person3(name: "John Doe", id: 123),
Person3(name: "Jane Doe", id: 456) {
$0.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 @@ -59,4 +59,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())
}

}