add clone / insert methods

Summary:
cursor.clone() will clone a length of the chain.  insert(std::move(buf)) will splice in a length of chain
I want this to change thrift2 binary type to IOBuf: we will clone() the network data for zero-copy userspace data, and insert() it if the return value is a IOBuf.

Test Plan: added unittest

Reviewed By: afrind@fb.com

FB internal diff: D632073

diff --git a/folly/experimental/io/Cursor.h b/folly/experimental/io/Cursor.h
index fbd77e5aae636cb4696d3e178e3c0f212fd48075..5b6e8b2631f3dbec629bc1922eb4ec34eb4fe128 100644 (file)
--- a/folly/experimental/io/Cursor.h
+++ b/folly/experimental/io/Cursor.h
@@ -127,6 +127,12 @@ class CursorBase {
     }
   }
 
+  void clone(std::unique_ptr<folly::IOBuf>& buf, size_t length) {
+    if (UNLIKELY(cloneAtMost(buf, length) != length)) {
+      throw std::out_of_range("underflow");
+    }
+  }
+
   void skip(size_t length) {
     if (UNLIKELY(skipAtMost(length) != length)) {
       throw std::out_of_range("underflow");
@@ -155,6 +161,43 @@ class CursorBase {
     }
   }
 
+  size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
+    buf.reset(nullptr);
+
+    std::unique_ptr<folly::IOBuf> tmp;
+    size_t copied = 0;
+    for (;;) {
+      // Fast path: it all fits in one buffer.
+      size_t available = length();
+      if (LIKELY(available >= len)) {
+        tmp = crtBuf_->cloneOne();
+        tmp->trimStart(offset_);
+        tmp->trimEnd(tmp->length() - len);
+        offset_ += len;
+        if (!buf) {
+          buf = std::move(tmp);
+        } else {
+          buf->prependChain(std::move(tmp));
+        }
+        return copied + len;
+      }
+
+      tmp = crtBuf_->cloneOne();
+      tmp->trimStart(offset_);
+      if (!buf) {
+        buf = std::move(tmp);
+      } else {
+        buf->prependChain(std::move(tmp));
+      }
+
+      copied += available;
+      if (UNLIKELY(!tryAdvanceBuffer())) {
+        return copied;
+      }
+      len -= available;
+    }
+  }
+
   size_t skipAtMost(size_t len) {
     size_t skipped = 0;
     for (;;) {
@@ -293,6 +336,32 @@ class RWCursor
     }
   }
 
+  void insert(std::unique_ptr<folly::IOBuf> buf) {
+    folly::IOBuf* nextBuf;
+    if (this->offset_ == 0) {
+      // Can just prepend
+      nextBuf = buf.get();
+      this->crtBuf_->prependChain(std::move(buf));
+    } else {
+      std::unique_ptr<folly::IOBuf> remaining;
+      if (this->crtBuf_->length() - this->offset_ > 0) {
+        // Need to split current IOBuf in two.
+        remaining = this->crtBuf_->cloneOne();
+        remaining->trimStart(this->offset_);
+        nextBuf = remaining.get();
+        buf->prependChain(std::move(remaining));
+      } else {
+        // Can just append
+        nextBuf = this->crtBuf_->next();
+      }
+      this->crtBuf_->trimEnd(this->length());
+      this->crtBuf_->appendChain(std::move(buf));
+    }
+    // Jump past the new links
+    this->offset_ = 0;
+    this->crtBuf_ = nextBuf;
+  }
+
   uint8_t* writableData() {
     return this->crtBuf_->writableData() + this->offset_;
   }

diff --git a/folly/experimental/io/test/IOBufCursorTest.cpp b/folly/experimental/io/test/IOBufCursorTest.cpp
index 3ed91f04ee068a67ffdb76af353d611fac85bb66..6ef4d0a5ac70325759460e3363e574676ceb528f 100644 (file)
--- a/folly/experimental/io/test/IOBufCursorTest.cpp
+++ b/folly/experimental/io/test/IOBufCursorTest.cpp
@@ -236,6 +236,79 @@ TEST(IOBuf, PullAndPeek) {
   }
 }
 
+TEST(IOBuf, cloneAndInsert) {
+  std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
+  append(iobuf1, "he");
+  std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
+  append(iobuf2, "llo ");
+  std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
+  append(iobuf3, "world");
+  iobuf1->prependChain(std::move(iobuf2));
+  iobuf1->prependChain(std::move(iobuf3));
+  EXPECT_EQ(3, iobuf1->countChainElements());
+  EXPECT_EQ(11, iobuf1->computeChainDataLength());
+
+  std::unique_ptr<IOBuf> cloned;
+
+  Cursor(iobuf1.get()).clone(cloned, 3);
+  EXPECT_EQ(2, cloned->countChainElements());
+  EXPECT_EQ(3, cloned->computeChainDataLength());
+
+
+  EXPECT_EQ(11, Cursor(iobuf1.get()).cloneAtMost(cloned, 20));
+  EXPECT_EQ(3, cloned->countChainElements());
+  EXPECT_EQ(11, cloned->computeChainDataLength());
+
+
+  EXPECT_THROW({Cursor(iobuf1.get()).clone(cloned, 20);},
+               std::out_of_range);
+
+  {
+    // Check that inserting in the middle of an iobuf splits
+    RWPrivateCursor cursor(iobuf1.get());
+    Cursor(iobuf1.get()).clone(cloned, 3);
+    EXPECT_EQ(2, cloned->countChainElements());
+    EXPECT_EQ(3, cloned->computeChainDataLength());
+
+    cursor.skip(1);
+
+    cursor.insert(std::move(cloned));
+    EXPECT_EQ(6, iobuf1->countChainElements());
+    EXPECT_EQ(14, iobuf1->computeChainDataLength());
+    // Check that nextBuf got set correctly
+    cursor.read<uint8_t>();
+  }
+
+  {
+    // Check that inserting at the end doesn't create empty buf
+    RWPrivateCursor cursor(iobuf1.get());
+    Cursor(iobuf1.get()).clone(cloned, 1);
+    EXPECT_EQ(1, cloned->countChainElements());
+    EXPECT_EQ(1, cloned->computeChainDataLength());
+
+    cursor.skip(1);
+
+    cursor.insert(std::move(cloned));
+    EXPECT_EQ(7, iobuf1->countChainElements());
+    EXPECT_EQ(15, iobuf1->computeChainDataLength());
+    // Check that nextBuf got set correctly
+    cursor.read<uint8_t>();
+  }
+  {
+    // Check that inserting at the beginning doesn't create empty buf
+    RWPrivateCursor cursor(iobuf1.get());
+    Cursor(iobuf1.get()).clone(cloned, 1);
+    EXPECT_EQ(1, cloned->countChainElements());
+    EXPECT_EQ(1, cloned->computeChainDataLength());
+
+    cursor.insert(std::move(cloned));
+    EXPECT_EQ(8, iobuf1->countChainElements());
+    EXPECT_EQ(16, iobuf1->computeChainDataLength());
+    // Check that nextBuf got set correctly
+    cursor.read<uint8_t>();
+  }
+}
+
 TEST(IOBuf, Appender) {
   std::unique_ptr<IOBuf> head(IOBuf::create(10));
   append(head, "hello");