ByteBuffer: allow chaining of AppendOrDie() am: 937164df3f am: d031d2bc48
am: ece435d5d4

Change-Id: I371d9610468c2a830b636da29d12556155dc5b80
diff --git a/byte_buffer.h b/byte_buffer.h
index 18c74b7..2fdff62 100644
--- a/byte_buffer.h
+++ b/byte_buffer.h
@@ -31,10 +31,10 @@
 // The buffer tracks its (populated) size, and does not require dynamic
 // memory allocation.
 //
-// Typical usage would be as follows:
-//     ByteBuffer buffer;
-//     buffer.AppendOrDie(header.data(), header.size());
-//     buffer.AppendOrDie(body.data(), body.size());
+// Usage could be as follows:
+//     const auto& buffer = ByteBuffer<1024>()
+//         .AppendOrDie(header.data(), header.size())
+//         .AppendOrDie(body.data(), body.size());
 //     write(fd, buffer.data(), buffer.size());
 template <size_t SizeBytes>
 class ByteBuffer {
@@ -42,11 +42,14 @@
   ByteBuffer() : write_pos_(0) {}
 
   // Appends data to the end of this buffer. Aborts if the available
-  // space in the buffer is less than |data_len|.
-  void AppendOrDie(NONNULL const void* data, size_t data_len) {
+  // space in the buffer is less than |data_len|. Returns a reference to
+  // the ByteBuffer, to support chaining.
+  ByteBuffer<SizeBytes>& AppendOrDie(NONNULL const void* data,
+                                     size_t data_len) {
     CHECK(data_len <= raw_buffer_.size() - write_pos_);
     std::memcpy(raw_buffer_.data() + write_pos_, data, data_len);
     write_pos_ += data_len;
+    return *this;
   }
 
   // Returns a pointer to the head of this buffer.
diff --git a/command_processor.cpp b/command_processor.cpp
index 2e9c1fb..67484de 100644
--- a/command_processor.cpp
+++ b/command_processor.cpp
@@ -162,10 +162,10 @@
   tstamp_header.since_boot_with_sleep = os_->GetTimestamp(CLOCK_BOOTTIME);
   tstamp_header.since_epoch = os_->GetTimestamp(CLOCK_REALTIME);
 
-  ByteBuffer<sizeof(TimestampHeader) + protocol::kMaxMessageSize> message_buf;
-  message_buf.AppendOrDie(&tstamp_header, sizeof(tstamp_header));
-  message_buf.AppendOrDie(command_buffer, command_len);
-
+  const auto& message_buf =
+      ByteBuffer<sizeof(TimestampHeader) + protocol::kMaxMessageSize>()
+          .AppendOrDie(&tstamp_header, sizeof(tstamp_header))
+          .AppendOrDie(command_buffer, command_len);
   bool did_write = current_log_buffer_.Append(
       message_buf.data(),
       SAFELY_CLAMP(message_buf.size(), uint16_t, 0, GetMaxVal<uint16_t>()));
diff --git a/tests/byte_buffer_unittest.cpp b/tests/byte_buffer_unittest.cpp
index 71830f6..9a1233b 100644
--- a/tests/byte_buffer_unittest.cpp
+++ b/tests/byte_buffer_unittest.cpp
@@ -70,6 +70,16 @@
   EXPECT_EQ(0, std::memcmp(copy.data(), buffer_.data(), buffer_.size()));
 }
 
+TEST_F(ByteBufferTest, ChainingWorks) {
+  const std::string message1{"hello"};
+  const std::string message2{"world"};
+  buffer_.AppendOrDie(message1.data(), message1.size())
+      .AppendOrDie(message2.data(), message2.size());
+
+  const std::string expected{"helloworld"};
+  EXPECT_EQ(0, std::memcmp(buffer_.data(), expected.data(), expected.size()));
+}
+
 TEST_F(ByteBufferTest, CopyingWorks) {
   const std::string message1{"hello"};
   buffer_.AppendOrDie(message1.data(), message1.size());
@@ -103,9 +113,9 @@
 }
 
 TEST_F(ByteBufferTest, SizeIsCorrectAfterMultipleWrites) {
-  buffer_.AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size());
-  buffer_.AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size());
-  buffer_.AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size());
+  buffer_.AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size())
+      .AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size())
+      .AppendOrDie(kSmallestMessage.data(), kSmallestMessage.size());
   EXPECT_EQ(3 * kSmallestMessage.size(), buffer_.size());
 }
 
diff --git a/tests/command_processor_unittest.cpp b/tests/command_processor_unittest.cpp
index c221c46..fcf9fbc 100644
--- a/tests/command_processor_unittest.cpp
+++ b/tests/command_processor_unittest.cpp
@@ -106,12 +106,11 @@
     command.payload_len =
         SAFELY_CLAMP(payload_length, uint16_t, 0, kMaxPayloadLength);
 
-    CommandBuffer buf;
-    buf.AppendOrDie(&command, sizeof(command));
-    buf.AppendOrDie(&ascii_message_header, sizeof(ascii_message_header));
-    buf.AppendOrDie(tag.data(), tag.length());
-    buf.AppendOrDie(message.data(), message.length());
-    return buf;
+    return CommandBuffer()
+        .AppendOrDie(&command, sizeof(command))
+        .AppendOrDie(&ascii_message_header, sizeof(ascii_message_header))
+        .AppendOrDie(tag.data(), tag.length())
+        .AppendOrDie(message.data(), message.length());
   }
 
   CommandBuffer BuildAsciiMessageCommand(const std::string& tag,
@@ -144,9 +143,7 @@
     command.opcode = protocol::Opcode::kDumpBuffers;
     command.payload_len = 0;
 
-    CommandBuffer buf;
-    buf.AppendOrDie(&command, sizeof(command));
-
+    const auto& buf = CommandBuffer().AppendOrDie(&command, sizeof(command));
     constexpr int kFakeFd = 100;
     return command_processor_->ProcessCommand(buf.data(), buf.size(), kFakeFd);
   }
@@ -225,9 +222,7 @@
       &invalid_opcode, sizeof(invalid_opcode));
   command.payload_len = 0;
 
-  CommandBuffer buf;
-  buf.AppendOrDie(&command, sizeof(command));
-
+  const auto& buf = CommandBuffer().AppendOrDie(&command, sizeof(command));
   constexpr int kFakeFd = 100;
   EXPECT_FALSE(
       command_processor_->ProcessCommand(buf.data(), buf.size(), kFakeFd));
diff --git a/tests/memory_reader_unittest.cpp b/tests/memory_reader_unittest.cpp
index 014d0e6..d3d1c29 100644
--- a/tests/memory_reader_unittest.cpp
+++ b/tests/memory_reader_unittest.cpp
@@ -93,9 +93,9 @@
     uint64_t a;
     uint64_t b;
   } second_message{0, GetMaxVal(second_message.b)};
-  ByteBuffer<sizeof(first_message) + sizeof(second_message)> buf;
-  buf.AppendOrDie(&first_message, sizeof(first_message));
-  buf.AppendOrDie(&second_message, sizeof(second_message));
+  const auto& buf = ByteBuffer<sizeof(first_message) + sizeof(second_message)>()
+                        .AppendOrDie(&first_message, sizeof(first_message))
+                        .AppendOrDie(&second_message, sizeof(second_message));
 
   MemoryReader memory_reader(buf.data(), buf.size());
   const auto& copy_of_first_message = memory_reader.CopyOutOrDie<M1>();