From 1c77ef4fa650ac53ad9c0ce51b0b73e1f9d1ab87 Mon Sep 17 00:00:00 2001
From: Neel Goyal <ngoyal@fb.com>
Date: Tue, 5 Jan 2016 09:33:55 -0800
Subject: [PATCH] Add canAdvance to cursor and tests

Summary:
Determine if the cursor can advance N bytes. This is useful if
applications want to check before reading so an exception isn't thrown.
It tries to walk the minimal amount of links needed in the chain.

I had a task that could have used this, though caching totalLength and
macro magic ended up being the implementation chosen. I think this just
adds to the cursor API.

Reviewed By: djwatson

Differential Revision: D2728498

fb-gh-sync-id: 8657653b82a48828cccab143653dc169ef715702
---
 folly/io/Cursor.h                 | 21 +++++++++++++++++++++
 folly/io/test/IOBufCursorTest.cpp | 18 ++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index d38d4dfb..477f8ec7 100644
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -108,6 +108,27 @@ class CursorBase {
     return end - *this;
   }
 
+  /**
+   * Return true if the cursor could advance the specified number of bytes
+   * from its current position.
+   * This is useful for applications that want to do checked reads instead of
+   * catching exceptions and is more efficient than using totalLength as it
+   * walks the minimal set of buffers in the chain to determine the result.
+   */
+  bool canAdvance(size_t amount) const {
+    const IOBuf* nextBuf = crtBuf_;
+    size_t available = length();
+    do {
+      if (available >= amount) {
+        return true;
+      }
+      amount -= available;
+      nextBuf = nextBuf->next();
+      available = nextBuf->length();
+    } while (nextBuf != buffer_);
+    return false;
+  }
+
   /*
    * Return true if the cursor is at the end of the entire IOBuf chain.
    */
diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 31801b8a..459b222a 100644
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -618,6 +618,24 @@ TEST(IOBuf, CursorOperators) {
     c.skip(5);
     EXPECT_TRUE(c.isAtEnd());
   }
+
+  // Test canAdvance with a chain of items
+  {
+    auto chain = IOBuf::create(10);
+    chain->append(10);
+    chain->appendChain(chain->clone());
+    EXPECT_EQ(2, chain->countChainElements());
+    EXPECT_EQ(20, chain->computeChainDataLength());
+
+    Cursor c(chain.get());
+    for (size_t i = 0; i <= 20; ++i) {
+      EXPECT_TRUE(c.canAdvance(i));
+    }
+    EXPECT_FALSE(c.canAdvance(21));
+    c.skip(10);
+    EXPECT_TRUE(c.canAdvance(10));
+    EXPECT_FALSE(c.canAdvance(11));
+  }
 }
 
 TEST(IOBuf, StringOperations) {
-- 
2.34.1