Add IOBuf::cloneCoalesced()

[folly.git] / folly / io / Compression.cpp

diff --git a/folly/io/Compression.cpp b/folly/io/Compression.cpp
index 1beea9cc9500959c63d27acf00d827a446119a5a..c626be74b3949e21adfe39fe9ce7eb0aafd112c9 100644 (file)
--- a/folly/io/Compression.cpp
+++ b/folly/io/Compression.cpp
@@ -56,13 +56,26 @@ std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
   uint64_t len = data->computeChainDataLength();
   if (len == 0) {
     return IOBuf::create(0);
-  } else if (len > maxUncompressedLength()) {
+  }
+  if (len > maxUncompressedLength()) {
     throw std::runtime_error("Codec: uncompressed length too large");
   }
 
   return doCompress(data);
 }
 
+std::string Codec::compress(const StringPiece data) {
+  const uint64_t len = data.size();
+  if (len == 0) {
+    return "";
+  }
+  if (len > maxUncompressedLength()) {
+    throw std::runtime_error("Codec: uncompressed length too large");
+  }
+
+  return doCompressString(data);
+}
+
 std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
                                          uint64_t uncompressedLength) {
   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
@@ -84,6 +97,28 @@ std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
   return doUncompress(data, uncompressedLength);
 }
 
+std::string Codec::uncompress(
+    const StringPiece data,
+    uint64_t uncompressedLength) {
+  if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
+    if (needsUncompressedLength()) {
+      throw std::invalid_argument("Codec: uncompressed length required");
+    }
+  } else if (uncompressedLength > maxUncompressedLength()) {
+    throw std::runtime_error("Codec: uncompressed length too large");
+  }
+
+  if (data.empty()) {
+    if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
+        uncompressedLength != 0) {
+      throw std::runtime_error("Codec: invalid uncompressed length");
+    }
+    return "";
+  }
+
+  return doUncompressString(data, uncompressedLength);
+}
+
 bool Codec::needsUncompressedLength() const {
   return doNeedsUncompressedLength();
 }
@@ -100,6 +135,30 @@ uint64_t Codec::doMaxUncompressedLength() const {
   return UNLIMITED_UNCOMPRESSED_LENGTH;
 }
 
+std::string Codec::doCompressString(const StringPiece data) {
+  const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
+  auto outputBuffer = doCompress(&inputBuffer);
+  std::string output;
+  output.reserve(outputBuffer->computeChainDataLength());
+  for (auto range : *outputBuffer) {
+    output.append(reinterpret_cast<const char*>(range.data()), range.size());
+  }
+  return output;
+}
+
+std::string Codec::doUncompressString(
+    const StringPiece data,
+    uint64_t uncompressedLength) {
+  const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
+  auto outputBuffer = doUncompress(&inputBuffer, uncompressedLength);
+  std::string output;
+  output.reserve(outputBuffer->computeChainDataLength());
+  for (auto range : *outputBuffer) {
+    output.append(reinterpret_cast<const char*>(range.data()), range.size());
+  }
+  return output;
+}
+
 namespace {
 
 /**
@@ -244,12 +303,11 @@ uint64_t LZ4Codec::doMaxUncompressedLength() const {
 }
 
 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
-  std::unique_ptr<IOBuf> clone;
+  IOBuf clone;
   if (data->isChained()) {
     // LZ4 doesn't support streaming, so we have to coalesce
-    clone = data->clone();
-    clone->coalesce();
-    data = clone.get();
+    clone = data->cloneCoalescedAsValue();
+    data = &clone;
   }
 
   uint32_t extraSize = encodeSize() ? kMaxVarintLength64 : 0;
@@ -286,12 +344,11 @@ std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
     const IOBuf* data,
     uint64_t uncompressedLength) {
-  std::unique_ptr<IOBuf> clone;
+  IOBuf clone;
   if (data->isChained()) {
     // LZ4 doesn't support streaming, so we have to coalesce
-    clone = data->clone();
-    clone->coalesce();
-    data = clone.get();
+    clone = data->cloneCoalescedAsValue();
+    data = &clone;
   }
 
   folly::io::Cursor cursor(data);