From ce4efaadff37f90dc707e84e0b64c3034739c81a Mon Sep 17 00:00:00 2001
From: Bo Liu <bol@fb.com>
Date: Tue, 18 Feb 2014 15:25:13 -0800
Subject: [PATCH] add clone() of stack allocated memory for folly::Cursor and
 folly::IOBuf

Summary: as title

Test Plan: fbconfig folly/io/test/ && fbmake runtests

Reviewed By: simpkins@fb.com

FB internal diff: D1176922
---
 folly/io/Cursor.h  | 50 ++++++++++++++++++++++++++++++++--------------
 folly/io/IOBuf.cpp | 30 +++++++++++++++++++---------
 folly/io/IOBuf.h   | 12 +++++++++++
 3 files changed, 68 insertions(+), 24 deletions(-)

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 5f65fb0e..67a86243 100644
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -27,6 +27,7 @@
 #include "folly/io/IOBuf.h"
 #include "folly/io/IOBufQueue.h"
 #include "folly/Likely.h"
+#include "folly/Memory.h"
 
 /**
  * Cursor class for fast iteration over IOBuf chains.
@@ -221,6 +222,12 @@ class CursorBase {
     }
   }
 
+  void clone(folly::IOBuf& buf, size_t len) {
+    if (UNLIKELY(cloneAtMost(buf, len) != len)) {
+      throw std::out_of_range("underflow");
+    }
+  }
+
   void skip(size_t len) {
     if (UNLIKELY(skipAtMost(len) != len)) {
       throw std::out_of_range("underflow");
@@ -249,33 +256,38 @@ class CursorBase {
     }
   }
 
-  size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
-    buf.reset(nullptr);
+  size_t cloneAtMost(folly::IOBuf& buf, size_t len) {
+    buf = folly::IOBuf();
 
     std::unique_ptr<folly::IOBuf> tmp;
     size_t copied = 0;
-    for (;;) {
+    for (int loopCount = 0; true; ++loopCount) {
       // 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);
+        if (loopCount == 0) {
+          crtBuf_->cloneOneInto(buf);
+          buf.trimStart(offset_);
+          buf.trimEnd(buf.length() - len);
         } else {
-          buf->prependChain(std::move(tmp));
+          tmp = crtBuf_->cloneOne();
+          tmp->trimStart(offset_);
+          tmp->trimEnd(tmp->length() - len);
+          buf.prependChain(std::move(tmp));
         }
+
+        offset_ += len;
         return copied + len;
       }
 
-      tmp = crtBuf_->cloneOne();
-      tmp->trimStart(offset_);
-      if (!buf) {
-        buf = std::move(tmp);
+
+      if (loopCount == 0) {
+        crtBuf_->cloneOneInto(buf);
+        buf.trimStart(offset_);
       } else {
-        buf->prependChain(std::move(tmp));
+        tmp = crtBuf_->cloneOne();
+        tmp->trimStart(offset_);
+        buf.prependChain(std::move(tmp));
       }
 
       copied += available;
@@ -286,6 +298,14 @@ class CursorBase {
     }
   }
 
+  size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
+    if (!buf) {
+      buf = make_unique<folly::IOBuf>();
+    }
+
+    return cloneAtMost(*buf, len);
+  }
+
   size_t skipAtMost(size_t len) {
     size_t skipped = 0;
     for (;;) {
diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index b3cd4125..3bcef5bd 100644
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -459,27 +459,39 @@ void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) {
 }
 
 unique_ptr<IOBuf> IOBuf::clone() const {
-  unique_ptr<IOBuf> newHead(cloneOne());
+  unique_ptr<IOBuf> ret = make_unique<IOBuf>();
+  cloneInto(*ret);
+  return ret;
+}
+
+unique_ptr<IOBuf> IOBuf::cloneOne() const {
+  unique_ptr<IOBuf> ret = make_unique<IOBuf>();
+  cloneOneInto(*ret);
+  return ret;
+}
+
+void IOBuf::cloneInto(IOBuf& other) const {
+  IOBuf tmp;
+  cloneOneInto(tmp);
 
   for (IOBuf* current = next_; current != this; current = current->next_) {
-    newHead->prependChain(current->cloneOne());
+    tmp.prependChain(current->cloneOne());
   }
 
-  return newHead;
+  other = std::move(tmp);
 }
 
-unique_ptr<IOBuf> IOBuf::cloneOne() const {
+void IOBuf::cloneOneInto(IOBuf& other) const {
   if (sharedInfo_) {
     flags_ |= kFlagMaybeShared;
   }
-  unique_ptr<IOBuf> iobuf(new IOBuf(static_cast<ExtBufTypeEnum>(type_),
-                                    flags_, buf_, capacity_,
-                                    data_, length_,
-                                    sharedInfo_));
+  other = IOBuf(static_cast<ExtBufTypeEnum>(type_),
+                flags_, buf_, capacity_,
+                data_, length_,
+                sharedInfo_);
   if (sharedInfo_) {
     sharedInfo_->refcount.fetch_add(1, std::memory_order_acq_rel);
   }
-  return iobuf;
 }
 
 void IOBuf::unshareOneSlow() {
diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index 8e063822..ba216061 100644
--- a/folly/io/IOBuf.h
+++ b/folly/io/IOBuf.h
@@ -1026,6 +1026,18 @@ class IOBuf {
    */
   std::unique_ptr<IOBuf> cloneOne() const;
 
+  /**
+   * Similar to Clone(). But use other as the head node. Other nodes in the
+   * chain (if any) will be allocted on heap.
+   */
+  void cloneInto(IOBuf& other) const;
+
+  /**
+   * Similar to CloneOne(). But to fill an existing IOBuf instead of a new
+   * IOBuf.
+   */
+  void cloneOneInto(IOBuf& other) const;
+
   /**
    * Return an iovector suitable for e.g. writev()
    *
-- 
2.34.1