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

diff --git a/folly/experimental/io/IOBuf.h b/folly/experimental/io/IOBuf.h
index 91dbc395de9cb08cee273d428c3b4a09c36b5642..a06f7a44f215c521d50168224df323534295b191 100644 (file)
--- 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 35295bf87fcba5ff1426d7580afe6192c272a7fe..bca8a1ce675acf606378001174bd0beec5dbd99c 100644 (file)
--- 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 3cb60b4998b4f89b02aa5ad0948e1064c66ab7f1..8883417f0373d22dbbff2608368e7eb11333c214 100644 (file)
--- 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 9562c2040bc62c966dbbc4e56e5ed7cee981ce73..221beaaccae37d8df26ded81d7fa12991f20fe56 100644 (file)
--- 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 3b88acb23267465484836b9fdef7274f252d3c97..46dfdcb51e8284f44141551a6ae11ebcb6d82bae 100644 (file)
--- 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());