Add trimAtMost functions

Summary: Add trimAtMost start and trimAtMost end. These are useful functions to have if you don't care about trimming an exact amount. Reviewed By: yfeldblum Differential Revision: D4658087 fbshipit-source-id: f7dd62c0f4837503e59cdec75fdad433be20306c
mtabebe · Apr 5, 2017 · 120cc11 · 120cc11
1 parent 2d2d3d3
commit 120cc11
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 4 deletions.
diff --git a/folly/io/IOBufQueue.cpp b/folly/io/IOBufQueue.cpp
@@ -216,31 +216,50 @@ unique_ptr<IOBuf> IOBufQueue::split(size_t n, bool throwOnUnderflow) {
 }
 
 void IOBufQueue::trimStart(size_t amount) {
+  auto trimmed = trimStartAtMost(amount);
+  if (trimmed != amount) {
+    throw std::underflow_error(
+        "Attempt to trim more bytes than are present in IOBufQueue");
+  }
+}
+
+size_t IOBufQueue::trimStartAtMost(size_t amount) {
+  auto original = amount;
   while (amount > 0) {
     if (!head_) {
-      throw std::underflow_error(
-        "Attempt to trim more bytes than are present in IOBufQueue");
+      break;
     }
     if (head_->length() > amount) {
       head_->trimStart(amount);
       chainLength_ -= amount;
+      amount = 0;
       break;
     }
     amount -= head_->length();
     chainLength_ -= head_->length();
     head_ = head_->pop();
   }
+  return original - amount;
 }
 
 void IOBufQueue::trimEnd(size_t amount) {
+  auto trimmed = trimEndAtMost(amount);
+  if (trimmed != amount) {
+    throw std::underflow_error(
+        "Attempt to trim more bytes than are present in IOBufQueue");
+  }
+}
+
+size_t IOBufQueue::trimEndAtMost(size_t amount) {
+  auto original = amount;
   while (amount > 0) {
     if (!head_) {
-      throw std::underflow_error(
-        "Attempt to trim more bytes than are present in IOBufQueue");
+      break;
     }
     if (head_->prev()->length() > amount) {
       head_->prev()->trimEnd(amount);
       chainLength_ -= amount;
+      amount = 0;
       break;
     }
     amount -= head_->prev()->length();
@@ -252,6 +271,7 @@ void IOBufQueue::trimEnd(size_t amount) {
       head_.reset();
     }
   }
+  return original - amount;
 }
 
 std::unique_ptr<folly::IOBuf> IOBufQueue::pop_front() {

diff --git a/folly/io/IOBufQueue.h b/folly/io/IOBufQueue.h
@@ -213,12 +213,24 @@ class IOBufQueue {
    */
   void trimStart(size_t amount);
 
+  /**
+   * Similar to trimStart, but will trim at most amount bytes and returns
+   * the number of bytes trimmed.
+   */
+  size_t trimStartAtMost(size_t amount);
+
   /**
    * Similar to IOBuf::trimEnd, but works on the whole queue.  Will
    * pop off buffers that have been completely trimmed.
    */
   void trimEnd(size_t amount);
 
+  /**
+   * Similar to trimEnd, but will trim at most amount bytes and returns
+   * the number of bytes trimmed.
+   */
+  size_t trimEndAtMost(size_t amount);
+
   /**
    * Transfer ownership of the queue's entire IOBuf chain to the caller.
    */

diff --git a/folly/io/test/IOBufQueueTest.cpp b/folly/io/test/IOBufQueueTest.cpp
@@ -268,6 +268,76 @@ TEST(IOBufQueue, Trim) {
   checkConsistency(queue);
 }
 
+TEST(IOBufQueue, TrimStartAtMost) {
+  IOBufQueue queue(clOptions);
+  unique_ptr<IOBuf> a = IOBuf::create(4);
+  a->append(4);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(6);
+  a->append(6);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(8);
+  a->append(8);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(10);
+  a->append(10);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(28, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimStartAtMost(1);
+  checkConsistency(queue);
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(27, queue.front()->computeChainDataLength());
+  EXPECT_EQ(3, queue.front()->length());
+
+  queue.trimStartAtMost(50);
+  checkConsistency(queue);
+  EXPECT_EQ(nullptr, queue.front());
+  EXPECT_EQ(0, queue.chainLength());
+}
+
+TEST(IOBufQueue, TrimEndAtMost) {
+  IOBufQueue queue(clOptions);
+  unique_ptr<IOBuf> a = IOBuf::create(4);
+  a->append(4);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(6);
+  a->append(6);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(8);
+  a->append(8);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(10);
+  a->append(10);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(28, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimEndAtMost(1);
+  checkConsistency(queue);
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(27, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimEndAtMost(50);
+  checkConsistency(queue);
+  EXPECT_EQ(nullptr, queue.front());
+  EXPECT_EQ(0, queue.chainLength());
+}
+
 TEST(IOBufQueue, TrimPack) {
   IOBufQueue queue(clOptions);
   unique_ptr<IOBuf> a = IOBuf::create(64);