Merge pull request #2424 from square/jszumski/swift-size-delimited-me…

…ssages

Swift: encode & decode size-delimited messages

wire-runtime-swift/src/main/swift/ProtoCodable/ProtoDecoder.swift

@@ -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
+    }
+
 }

wire-runtime-swift/src/main/swift/ProtoCodable/ProtoEncoder.swift

@@ -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)
+    }
+
 }

wire-runtime-swift/src/main/swift/ProtoCodable/WriteBuffer.swift

@@ -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)
@@ -63,6 +65,8 @@ final class WriteBuffer {
     }
 
     func append(_ value: [UInt8]) {
+        guard !value.isEmpty else { return }
+
         expandIfNeeded(adding: value.count)
 
         for byte in value {
@@ -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)

wire-runtime-swift/src/test/swift/ProtoDecoderTests.swift

@@ -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

wire-runtime-swift/src/test/swift/ProtoEncoderTests.swift

@@ -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())
+    }
 }

wire-runtime-swift/src/test/swift/RoundTripTests.swift

@@ -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)
+    }
 }

wire-runtime-swift/src/test/swift/WriteBufferTests.swift

@@ -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())
+    }
+
 }