Skip to content

Commit

Permalink
Merge pull request #2424 from square/jszumski/swift-size-delimited-me…
Browse files Browse the repository at this point in the history
…ssages

Swift: encode & decode size-delimited messages
  • Loading branch information
jszumski authored Nov 12, 2024
2 parents 4cfd464 + 986c030 commit a262267
Show file tree
Hide file tree
Showing 7 changed files with 127 additions and 2 deletions.
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)
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 = "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 @@ -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())
}

}

0 comments on commit a262267

Please sign in to comment.