Allow folly::io::Cursor to move backwards.

Summary: This is helpful for parsing data IOBufs in reverse.

Reviewed By: siyengar

Differential Revision: D4082810

fbshipit-source-id: 6a201d25e7d22befe28f92e4a1d7aa41ad7d6817

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index b32670a27f5e62f0123c08c72b80bbbcc78275ea..bd5fc235dc3bb94c02d1c0bd41da3c6464a27793 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -153,6 +153,17 @@ class CursorBase {
     return true;
   }
 
+  /**
+   * Advances the cursor to the end of the entire IOBuf chain.
+   */
+  void advanceToEnd() {
+    offset_ = buffer_->prev()->length();
+    if (crtBuf_ != buffer_->prev()) {
+      crtBuf_ = buffer_->prev();
+      static_cast<Derived*>(this)->advanceDone();
+    }
+  }
+
   Derived& operator+=(size_t offset) {
     Derived* p = static_cast<Derived*>(this);
     p->skip(offset);
@@ -164,6 +175,17 @@ class CursorBase {
     return other;
   }
 
+  Derived& operator-=(size_t offset) {
+    Derived* p = static_cast<Derived*>(this);
+    p->retreat(offset);
+    return *p;
+  }
+  Derived operator-(size_t offset) const {
+    Derived other(*this);
+    other.retreat(offset);
+    return other;
+  }
+
   /**
    * Compare cursors for equality/inequality.
    *
@@ -279,6 +301,22 @@ class CursorBase {
     }
   }
 
+  size_t retreatAtMost(size_t len) {
+    if (len <= offset_) {
+      offset_ -= len;
+      return len;
+    }
+    return retreatAtMostSlow(len);
+  }
+
+  void retreat(size_t len) {
+    if (len <= offset_) {
+      offset_ -= len;
+    } else {
+      retreatSlow(len);
+    }
+  }
+
   size_t pullAtMost(void* buf, size_t len) {
     // Fast path: it all fits in one buffer.
     if (LIKELY(length() >= len)) {
@@ -455,6 +493,17 @@ class CursorBase {
     return true;
   }
 
+  bool tryRetreatBuffer() {
+    if (UNLIKELY(crtBuf_ == buffer_)) {
+      offset_ = 0;
+      return false;
+    }
+    crtBuf_ = crtBuf_->prev();
+    offset_ = crtBuf_->length();
+    static_cast<Derived*>(this)->advanceDone();
+    return true;
+  }
+
   void advanceBufferIfEmpty() {
     if (length() == 0) {
       tryAdvanceBuffer();
@@ -522,6 +571,25 @@ class CursorBase {
     }
   }
 
+  size_t retreatAtMostSlow(size_t len) {
+    size_t retreated = 0;
+    for (size_t available; (available = offset_) < len;) {
+      retreated += available;
+      if (UNLIKELY(!tryRetreatBuffer())) {
+        return retreated;
+      }
+      len -= available;
+    }
+    offset_ -= len;
+    return retreated + len;
+  }
+
+  void retreatSlow(size_t len) {
+    if (UNLIKELY(retreatAtMostSlow(len) != len)) {
+      std::__throw_out_of_range("underflow");
+    }
+  }
+
   void advanceDone() {
   }
 

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 5c607ad20a93497ba9f453b09ecdd91260c07e3d..9778edf536623e6620bb600d65268b9b5a66964a 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -862,3 +862,118 @@ TEST(IOBuf, ReadWhileTrue) {
     EXPECT_TRUE(curs.isAtEnd());
   }
 }
+
+TEST(IOBuf, TestAdvanceToEndSingle) {
+  std::unique_ptr<IOBuf> chain(IOBuf::create(10));
+  chain->append(10);
+
+  Cursor curs(chain.get());
+  curs.advanceToEnd();
+  EXPECT_TRUE(curs.isAtEnd());
+  EXPECT_EQ(curs - chain.get(), 10);
+}
+
+TEST(IOBuf, TestAdvanceToEndMulti) {
+  std::unique_ptr<IOBuf> chain(IOBuf::create(10));
+  chain->append(10);
+
+  std::unique_ptr<IOBuf> buf(IOBuf::create(5));
+  buf->append(5);
+  chain->prependChain(std::move(buf));
+
+  buf = IOBuf::create(20);
+  buf->append(20);
+  chain->prependChain(std::move(buf));
+
+  Cursor curs(chain.get());
+  curs.advanceToEnd();
+  EXPECT_TRUE(curs.isAtEnd());
+  EXPECT_EQ(curs - chain.get(), 35);
+
+  curs.reset(chain.get());
+  curs.skip(12);
+  curs.advanceToEnd();
+  EXPECT_TRUE(curs.isAtEnd());
+}
+
+TEST(IOBuf, TestRetreatSingle) {
+  std::unique_ptr<IOBuf> chain(IOBuf::create(20));
+  chain->append(20);
+
+  Cursor curs(chain.get());
+  EXPECT_EQ(curs.retreatAtMost(0), 0);
+  EXPECT_EQ(curs.totalLength(), 20);
+  EXPECT_EQ(curs.retreatAtMost(5), 0);
+  EXPECT_EQ(curs.totalLength(), 20);
+  EXPECT_EQ(curs.retreatAtMost(25), 0);
+  EXPECT_EQ(curs.totalLength(), 20);
+
+  curs.retreat(0);
+  EXPECT_THROW(curs.retreat(5), std::out_of_range);
+  curs.reset(chain.get());
+  EXPECT_THROW(curs.retreat(25), std::out_of_range);
+  curs.reset(chain.get());
+
+  curs.advanceToEnd();
+  curs.retreat(5);
+  EXPECT_EQ(curs.totalLength(), 5);
+  curs.retreat(10);
+  EXPECT_EQ(curs.totalLength(), 15);
+  EXPECT_THROW(curs.retreat(10), std::out_of_range);
+
+  curs.reset(chain.get());
+  curs.advanceToEnd();
+  EXPECT_EQ(curs.retreatAtMost(5), 5);
+  EXPECT_EQ(curs.totalLength(), 5);
+  EXPECT_EQ(curs.retreatAtMost(10), 10);
+  EXPECT_EQ(curs.totalLength(), 15);
+  EXPECT_EQ(curs.retreatAtMost(10), 5);
+  EXPECT_EQ(curs.totalLength(), 20);
+}
+
+TEST(IOBuf, TestRetreatMulti) {
+  std::unique_ptr<IOBuf> chain(IOBuf::create(10));
+  chain->append(10);
+
+  std::unique_ptr<IOBuf> buf(IOBuf::create(5));
+  buf->append(5);
+  chain->prependChain(std::move(buf));
+
+  buf = IOBuf::create(20);
+  buf->append(20);
+  chain->prependChain(std::move(buf));
+
+  Cursor curs(chain.get());
+  EXPECT_EQ(curs.retreatAtMost(10), 0);
+  EXPECT_THROW(curs.retreat(10), std::out_of_range);
+  curs.reset(chain.get());
+
+  curs.advanceToEnd();
+  curs.retreat(20);
+  EXPECT_EQ(curs.totalLength(), 20);
+  EXPECT_EQ(curs.length(), 20);
+  curs.retreat(1);
+  EXPECT_EQ(curs.totalLength(), 21);
+  EXPECT_EQ(curs.length(), 1);
+  EXPECT_EQ(curs.retreatAtMost(50), 14);
+  EXPECT_EQ(curs.totalLength(), 35);
+
+  curs.advanceToEnd();
+  curs.retreat(30);
+  EXPECT_EQ(curs.totalLength(), 30);
+}
+
+TEST(IOBuf, TestRetreatOperators) {
+  std::unique_ptr<IOBuf> chain(IOBuf::create(20));
+  chain->append(20);
+
+  Cursor curs(chain.get());
+  curs.advanceToEnd();
+  curs -= 5;
+  EXPECT_EQ(curs.totalLength(), 5);
+
+  curs.advanceToEnd();
+  auto retreated = curs - 5;
+  EXPECT_EQ(retreated.totalLength(), 5);
+  EXPECT_EQ(curs.totalLength(), 0);
+}