From 53d4209fdaa204747f5bb894b616d9770cb767a9 Mon Sep 17 00:00:00 2001
From: Tudor Bosman <tudorb@fb.com>
Date: Fri, 6 Jul 2012 17:12:07 -0700
Subject: [PATCH] Add IOBufQueue::prepend, fix bug in IOBuf::prepend.

Summary:
IOBuf::prepend needs to increment length_.

Added IOBufQueue::prepend, which uses the headroom in the first buffer
instead of growing the queue at the head.

Test Plan: tests added

Reviewed By: simpkins@fb.com

FB internal diff: D513676
---
 folly/experimental/io/IOBuf.h                 |  1 +
 folly/experimental/io/IOBufQueue.cpp          | 31 +++++++++++++++++++
 folly/experimental/io/IOBufQueue.h            | 19 ++++++++++++
 folly/experimental/io/test/IOBufQueueTest.cpp | 19 ++++++++++++
 folly/experimental/io/test/IOBufTest.cpp      | 16 ++++++++--
 5 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/folly/experimental/io/IOBuf.h b/folly/experimental/io/IOBuf.h
index 91dbc395..a06f7a44 100644
--- a/folly/experimental/io/IOBuf.h
+++ b/folly/experimental/io/IOBuf.h
@@ -497,6 +497,7 @@ class IOBuf {
   void prepend(uint32_t amount) {
     CHECK(amount <= headroom());
     data_ -= amount;
+    length_ += amount;
   }
 
   /**
diff --git a/folly/experimental/io/IOBufQueue.cpp b/folly/experimental/io/IOBufQueue.cpp
index 35295bf8..bca8a1ce 100644
--- a/folly/experimental/io/IOBufQueue.cpp
+++ b/folly/experimental/io/IOBufQueue.cpp
@@ -69,6 +69,37 @@ IOBufQueue& IOBufQueue::operator=(IOBufQueue&& other) {
   return *this;
 }
 
+std::pair<void*, uint32_t>
+IOBufQueue::headroom() {
+  if (head_) {
+    return std::make_pair(head_->writableBuffer(), head_->headroom());
+  } else {
+    return std::make_pair(nullptr, 0);
+  }
+}
+
+void
+IOBufQueue::markPrepended(uint32_t n) {
+  if (n == 0) {
+    return;
+  }
+  assert(head_);
+  head_->prepend(n);
+  if (options_.cacheChainLength) {
+    chainLength_ += n;
+  }
+}
+
+void
+IOBufQueue::prepend(const void* buf, uint32_t n) {
+  auto p = headroom();
+  if (n > p.second) {
+    throw std::overflow_error("Not enough room to prepend");
+  }
+  memcpy(static_cast<char*>(p.first) + p.second - n, buf, n);
+  markPrepended(n);
+}
+
 void
 IOBufQueue::append(unique_ptr<IOBuf>&& buf) {
   if (!buf) {
diff --git a/folly/experimental/io/IOBufQueue.h b/folly/experimental/io/IOBufQueue.h
index 3cb60b49..8883417f 100644
--- a/folly/experimental/io/IOBufQueue.h
+++ b/folly/experimental/io/IOBufQueue.h
@@ -28,6 +28,9 @@ namespace folly {
  * An IOBufQueue encapsulates a chain of IOBufs and provides
  * convenience functions to append data to the back of the chain
  * and remove data from the front.
+ *
+ * You may also prepend data into the headroom of the first buffer in the
+ * chain, if any.
  */
 class IOBufQueue {
  public:
@@ -48,6 +51,22 @@ class IOBufQueue {
 
   explicit IOBufQueue(const Options& options = Options());
 
+  /**
+   * Return a space to prepend bytes and the amount of headroom available.
+   */
+  std::pair<void*, uint32_t> headroom();
+
+  /**
+   * Indicate that n bytes from the headroom have been used.
+   */
+  void markPrepended(uint32_t n);
+
+  /**
+   * Prepend an existing range; throws std::overflow_error if not enough
+   * room.
+   */
+  void prepend(const void* buf, uint32_t n);
+
   /**
    * Add a buffer or buffer chain to the end of this queue. The
    * queue takes ownership of buf.
diff --git a/folly/experimental/io/test/IOBufQueueTest.cpp b/folly/experimental/io/test/IOBufQueueTest.cpp
index 9562c204..221beaac 100644
--- a/folly/experimental/io/test/IOBufQueueTest.cpp
+++ b/folly/experimental/io/test/IOBufQueueTest.cpp
@@ -239,6 +239,25 @@ TEST(IOBufQueue, trim) {
   checkConsistency(queue);
 }
 
+TEST(IOBufQueue, Prepend) {
+  folly::IOBufQueue queue;
+
+  auto buf = folly::IOBuf::create(10);
+  buf->advance(5);
+  queue.append(std::move(buf));
+
+  queue.append(SCL(" World"));
+  queue.prepend(SCL("Hello"));
+
+  EXPECT_THROW(queue.prepend(SCL("x")), std::overflow_error);
+
+  auto out = queue.move();
+  out->coalesce();
+  EXPECT_EQ("Hello World",
+            StringPiece(reinterpret_cast<const char*>(out->data()),
+                        out->length()));
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);
diff --git a/folly/experimental/io/test/IOBufTest.cpp b/folly/experimental/io/test/IOBufTest.cpp
index 3b88acb2..46dfdcb5 100644
--- a/folly/experimental/io/test/IOBufTest.cpp
+++ b/folly/experimental/io/test/IOBufTest.cpp
@@ -35,6 +35,12 @@ void append(std::unique_ptr<IOBuf>& buf, StringPiece str) {
   buf->append(str.size());
 }
 
+void prepend(std::unique_ptr<IOBuf>& buf, StringPiece str) {
+  EXPECT_LE(str.size(), buf->headroom());
+  memcpy(buf->writableData() - str.size(), str.data(), str.size());
+  buf->prepend(str.size());
+}
+
 TEST(IOBuf, Simple) {
   unique_ptr<IOBuf> buf(IOBuf::create(100));
   uint32_t cap = buf->capacity();
@@ -43,13 +49,19 @@ TEST(IOBuf, Simple) {
   EXPECT_EQ(0, buf->length());
   EXPECT_EQ(cap, buf->tailroom());
 
-  append(buf, "hello");
+  append(buf, "world");
   buf->advance(10);
   EXPECT_EQ(10, buf->headroom());
   EXPECT_EQ(5, buf->length());
   EXPECT_EQ(cap - 15, buf->tailroom());
+
+  prepend(buf, "hello ");
+  EXPECT_EQ(4, buf->headroom());
+  EXPECT_EQ(11, buf->length());
+  EXPECT_EQ(cap - 15, buf->tailroom());
+
   const char* p = reinterpret_cast<const char*>(buf->data());
-  EXPECT_EQ("hello", std::string(p, buf->length()));
+  EXPECT_EQ("hello world", std::string(p, buf->length()));
 
   buf->clear();
   EXPECT_EQ(0, buf->headroom());
-- 
2.34.1