From 564c79afcb0f462dcf8eb73f2627fe4a96e774f5 Mon Sep 17 00:00:00 2001 From: Kyle Nekritz Date: Thu, 3 Nov 2016 15:53:50 -0700 Subject: [PATCH] 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 --- folly/io/Cursor.h | 68 ++++++++++++++++++ folly/io/test/IOBufCursorTest.cpp | 115 ++++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+) diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h index b32670a2..bd5fc235 100644 --- 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(this)->advanceDone(); + } + } + Derived& operator+=(size_t offset) { Derived* p = static_cast(this); p->skip(offset); @@ -164,6 +175,17 @@ class CursorBase { return other; } + Derived& operator-=(size_t offset) { + Derived* p = static_cast(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(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 5c607ad2..9778edf5 100644 --- 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 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 chain(IOBuf::create(10)); + chain->append(10); + + std::unique_ptr 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 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 chain(IOBuf::create(10)); + chain->append(10); + + std::unique_ptr 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 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); +} -- 2.34.1