From: Subodh Iyengar <subodh@fb.com>
Date: Wed, 5 Apr 2017 22:15:49 +0000 (-0700)
Subject: Add trimAtMost functions
X-Git-Tag: v2017.04.10.00~14
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=120cc11d827f07e1bcbcc12132d871f167d5342b;p=folly.git

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
---

diff --git a/folly/io/IOBufQueue.cpp b/folly/io/IOBufQueue.cpp
index 6e4715d7..ce3cb8a5 100644
--- 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
index be334d38..c90deaf8 100644
--- 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
index bc2a8901..8050812e 100644
--- 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);