From cb771e84cd1bfd21bfdfb3dd0bae03ecb3ec334f Mon Sep 17 00:00:00 2001
From: Adam Simpkins <simpkins@fb.com>
Date: Tue, 13 Nov 2012 20:58:10 -0800
Subject: [PATCH] add IOBuf::maybeCopyBuffer()

Summary:
Add a version of IOBuf::copyBuffer() which accepts a std::string, as
well as a maybeCopyBuffer() function which returns a null pointer if the
input string is empty.

In proxygen we have a few places where we construct an IOBuf from a
string configuration parameter, and we almost always want a null pointer
rather than an empty IOBuf chain if the string is empty.

Test Plan: Included unit tests for these new functions.

Reviewed By: tudorb@fb.com

FB internal diff: D630547
---
 folly/experimental/io/IOBuf.h            | 38 ++++++++++++++++++++++++
 folly/experimental/io/test/IOBufTest.cpp | 28 +++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/folly/experimental/io/IOBuf.h b/folly/experimental/io/IOBuf.h
index e48650ac..3a0abb8a 100644
--- a/folly/experimental/io/IOBuf.h
+++ b/folly/experimental/io/IOBuf.h
@@ -299,6 +299,29 @@ class IOBuf {
                                            uint32_t headroom=0,
                                            uint32_t minTailroom=0);
 
+  /**
+   * Convenience function to create a new IOBuf object that copies data from a
+   * user-supplied string, optionally allocating a given amount of
+   * headroom and tailroom.
+   *
+   * Beware when attempting to invoke this function with a constant string
+   * literal and a headroom argument: you will likely end up invoking the
+   * version of copyBuffer() above.  IOBuf::copyBuffer("hello", 3) will treat
+   * the first argument as a const void*, and will invoke the version of
+   * copyBuffer() above, with the size argument of 3.
+   */
+  static std::unique_ptr<IOBuf> copyBuffer(const std::string& buf,
+                                           uint32_t headroom=0,
+                                           uint32_t minTailroom=0);
+
+  /**
+   * A version of copyBuffer() that returns a null pointer if the input string
+   * is empty.
+   */
+  static std::unique_ptr<IOBuf> maybeCopyBuffer(const std::string& buf,
+                                                uint32_t headroom=0,
+                                                uint32_t minTailroom=0);
+
   /**
    * Convenience function to free a chain of IOBufs held by a unique_ptr.
    */
@@ -1076,6 +1099,21 @@ inline std::unique_ptr<IOBuf> IOBuf::copyBuffer(
   return buf;
 }
 
+inline std::unique_ptr<IOBuf> IOBuf::copyBuffer(const std::string& buf,
+                                                uint32_t headroom,
+                                                uint32_t minTailroom) {
+  return copyBuffer(buf.data(), buf.size(), headroom, minTailroom);
+}
+
+inline std::unique_ptr<IOBuf> IOBuf::maybeCopyBuffer(const std::string& buf,
+                                                     uint32_t headroom,
+                                                     uint32_t minTailroom) {
+  if (buf.empty()) {
+    return nullptr;
+  }
+  return copyBuffer(buf.data(), buf.size(), headroom, minTailroom);
+}
+
 } // folly
 
 #endif // FOLLY_IO_IOBUF_H_
diff --git a/folly/experimental/io/test/IOBufTest.cpp b/folly/experimental/io/test/IOBufTest.cpp
index b16511a4..f6f6fe28 100644
--- a/folly/experimental/io/test/IOBufTest.cpp
+++ b/folly/experimental/io/test/IOBufTest.cpp
@@ -550,6 +550,34 @@ TEST(IOBuf, copyBuffer) {
   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
                            buf->length()));
   EXPECT_LE(2, buf->tailroom());
+
+  buf = IOBuf::copyBuffer(s, 5, 7);
+  EXPECT_EQ(5, buf->headroom());
+  EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
+                           buf->length()));
+  EXPECT_LE(7, buf->tailroom());
+
+  std::string empty;
+  buf = IOBuf::copyBuffer(empty, 3, 6);
+  EXPECT_EQ(3, buf->headroom());
+  EXPECT_EQ(0, buf->length());
+  EXPECT_LE(6, buf->tailroom());
+}
+
+TEST(IOBuf, maybeCopyBuffer) {
+  std::string s("this is a test");
+  auto buf = IOBuf::maybeCopyBuffer(s, 1, 2);
+  EXPECT_EQ(1, buf->headroom());
+  EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
+                           buf->length()));
+  EXPECT_LE(2, buf->tailroom());
+
+  std::string empty;
+  buf = IOBuf::maybeCopyBuffer("", 5, 7);
+  EXPECT_EQ(nullptr, buf.get());
+
+  buf = IOBuf::maybeCopyBuffer("");
+  EXPECT_EQ(nullptr, buf.get());
 }
 
 namespace {
-- 
2.34.1