add Cursor::readWhile() and skipWhile()

[folly.git] / folly / io / test / IOBufCursorTest.cpp

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index a5bc4eb982983c55b803d2406d8ed123ae09d165..08cb23d44eba9fabe18a31ec90dcdccc09f486a9 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -770,3 +770,96 @@ TEST(IOBuf, StringOperations) {
     EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
   }
 }
+
+TEST(IOBuf, ReadWhileTrue) {
+  auto isAlpha = [](uint8_t ch) {
+    return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z');
+  };
+  auto isDigit = [](uint8_t ch) { return (ch >= '0' && ch <= '9'); };
+
+  // Test reading alternating alphabetic and numeric strings
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(32));
+    Appender app(chain.get(), 0);
+    app.push(StringPiece("hello123world456"));
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
+    EXPECT_STREQ("123", curs.readWhile(isDigit).c_str());
+    EXPECT_STREQ("world", curs.readWhile(isAlpha).c_str());
+    EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
+    EXPECT_TRUE(curs.isAtEnd());
+  }
+
+  // The same, but also use skipWhile()
+  {
+    std::unique_ptr<IOBuf> chain(IOBuf::create(16));
+    Appender app(chain.get(), 0);
+    app.push(StringPiece("hello123world456"));
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
+    curs.skipWhile(isDigit);
+    curs.skipWhile(isAlpha);
+    EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
+    EXPECT_TRUE(curs.isAtEnd());
+  }
+
+  // Test readWhile() using data split across multiple buffers,
+  // including some empty buffers in the middle of the chain.
+  {
+    std::unique_ptr<IOBuf> chain;
+
+    // First element in the chain has "he"
+    auto buf = IOBuf::create(40);
+    Appender app(buf.get(), 0);
+    app.push(StringPiece("he"));
+    chain = std::move(buf);
+
+    // The second element has "ll", after 10 bytes of headroom
+    buf = IOBuf::create(40);
+    buf->advance(10);
+    app = Appender{buf.get(), 0};
+    app.push(StringPiece("ll"));
+    chain->prependChain(std::move(buf));
+
+    // The third element is empty
+    buf = IOBuf::create(40);
+    buf->advance(15);
+    chain->prependChain(std::move(buf));
+
+    // The fourth element has "o12"
+    buf = IOBuf::create(40);
+    buf->advance(37);
+    app = Appender{buf.get(), 0};
+    app.push(StringPiece("o12"));
+    chain->prependChain(std::move(buf));
+
+    // The fifth element has "3"
+    buf = IOBuf::create(40);
+    app = Appender{buf.get(), 0};
+    app.push(StringPiece("3"));
+    chain->prependChain(std::move(buf));
+
+    // The sixth element is empty
+    buf = IOBuf::create(40);
+    chain->prependChain(std::move(buf));
+
+    // The seventh element has "world456"
+    buf = IOBuf::create(40);
+    app = Appender{buf.get(), 0};
+    app.push(StringPiece("world456"));
+    chain->prependChain(std::move(buf));
+
+    // The eighth element is empty
+    buf = IOBuf::create(40);
+    chain->prependChain(std::move(buf));
+
+    Cursor curs(chain.get());
+    EXPECT_STREQ("hello", curs.readWhile(isAlpha).c_str());
+    EXPECT_STREQ("123", curs.readWhile(isDigit).c_str());
+    EXPECT_STREQ("world", curs.readWhile(isAlpha).c_str());
+    EXPECT_STREQ("456", curs.readWhile(isDigit).c_str());
+    EXPECT_TRUE(curs.isAtEnd());
+  }
+}