From d0d783867d295e3519545c91e6e24625a5cb341d Mon Sep 17 00:00:00 2001
From: TJ Yin <ytj@meta.com>
Date: Fri, 11 Oct 2024 12:53:08 -0700
Subject: [PATCH] add helper function to convert std::string to folly::IOBuf

Reviewed By: ot

Differential Revision: D63949628

fbshipit-source-id: ae3897246405d8605f803da69d9ec54ebcc5c17e
---
 folly/io/IOBuf.cpp          | 10 ++++++++++
 folly/io/IOBuf.h            |  9 +++++++++
 folly/io/test/IOBufTest.cpp |  7 +++++++
 3 files changed, 26 insertions(+)

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 4eba7aca154..b3667a602d3 100644
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -1328,6 +1328,16 @@ IOBuf::Iterator IOBuf::cend() const {
   return Iterator(nullptr, nullptr);
 }
 
+std::unique_ptr<IOBuf> IOBuf::fromString(std::unique_ptr<std::string> ptr) {
+  auto ret = takeOwnership(
+      ptr->data(),
+      ptr->size(),
+      [](void*, void* userData) { delete static_cast<std::string*>(userData); },
+      static_cast<void*>(ptr.get()));
+  std::ignore = ptr.release();
+  return ret;
+}
+
 folly::fbvector<struct iovec> IOBuf::getIov() const {
   folly::fbvector<struct iovec> iov;
   iov.reserve(countChainElements());
diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index 8bcc253d144..67394c09c27 100644
--- a/folly/io/IOBuf.h
+++ b/folly/io/IOBuf.h
@@ -1810,6 +1810,15 @@ class IOBuf {
    */
   std::string toString() const { return to<std::string>(); }
 
+  /**
+   * Create an IOBuf from a std::string. Avoids copying the contents of the
+   * string, at the cost of an extra allocation.
+   */
+  static std::unique_ptr<IOBuf> fromString(std::unique_ptr<std::string>);
+  static std::unique_ptr<IOBuf> fromString(std::string s) {
+    return fromString(std::make_unique<std::string>(std::move(s)));
+  }
+
   /**
    * Get an iovector suitable for e.g. writev()
    *
diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index a4b3a0290b0..230a1d1eecc 100644
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -1903,3 +1903,10 @@ TEST(IOBuf, MaybeSplitTail) {
   auto buf3 = buf2->maybeSplitTail();
   append(buf3, "!!!1");
 }
+
+TEST(IOBuf, FromString) {
+  EXPECT_EQ(folly::IOBuf::fromString("")->toString(), "");
+  EXPECT_EQ(folly::IOBuf::fromString("test")->toString(), "test");
+  auto longStr = std::string(1000, '0');
+  EXPECT_EQ(folly::IOBuf::fromString(longStr)->toString(), longStr);
+}