chore: Add QList::Iterator (#4082)

Signed-off-by: Roman Gershman <[email protected]>

src/core/qlist.cc

@@ -60,6 +60,8 @@ namespace dfly {
 
 namespace {
 
+enum IterDir : uint8_t { FWD = 1, REV = 0 };
+
 /* This is for test suite development purposes only, 0 means disabled. */
 static size_t packed_threshold = 0;
 
@@ -183,8 +185,8 @@ bool CompressNode(quicklistNode* node) {
 
 /* Uncompress the listpack in 'node' and update encoding details.
  * Returns 1 on successful decode, 0 on failure to decode. */
-bool DecompressNode(quicklistNode* node) {
-  node->recompress = 0;
+bool DecompressNode(bool recompress, quicklistNode* node) {
+  node->recompress = int(recompress);
 
   void* decompressed = zmalloc(node->sz);
   quicklistLZF* lzf = (quicklistLZF*)node->entry;
@@ -200,9 +202,9 @@ bool DecompressNode(quicklistNode* node) {
 }
 
 /* Decompress only compressed nodes. */
-void DecompressNodeIfNeeded(quicklistNode* node) {
+void DecompressNodeIfNeeded(bool recompress, quicklistNode* node) {
   if ((node) && (node)->encoding == QUICKLIST_NODE_ENCODING_LZF) {
-    DecompressNode(node);
+    DecompressNode(recompress, node);
   }
 }
 
@@ -391,8 +393,8 @@ void QList::Compress(quicklistNode* node) {
   int depth = 0;
   int in_depth = 0;
   while (depth++ < compress_) {
-    DecompressNodeIfNeeded(forward);
-    DecompressNodeIfNeeded(reverse);
+    DecompressNodeIfNeeded(false, forward);
+    DecompressNodeIfNeeded(false, reverse);
 
     if (forward == node || reverse == node)
       in_depth = 1;
@@ -414,4 +416,87 @@ void QList::Compress(quicklistNode* node) {
   CompressNode(reverse);
 }
 
+auto QList::GetIterator(Where where) -> Iterator {
+  Iterator it;
+  it.owner_ = this;
+  it.zi_ = NULL;
+  if (where == HEAD) {
+    it.current_ = head_;
+    it.offset_ = 0;
+    it.direction_ = FWD;
+  } else {
+    it.current_ = tail_;
+    it.offset_ = -1;
+    it.direction_ = REV;
+  }
+
+  return it;
+}
+
+bool QList::Iterator::Next() {
+  DCHECK(current_);
+
+  unsigned char* (*nextFn)(unsigned char*, unsigned char*) = NULL;
+  int offset_update = 0;
+
+  int plain = QL_NODE_IS_PLAIN(current_);
+  if (!zi_) {
+    /* If !zi, use current index. */
+    DecompressNodeIfNeeded(true, current_);
+    if (ABSL_PREDICT_FALSE(plain))
+      zi_ = current_->entry;
+    else
+      zi_ = lpSeek(current_->entry, offset_);
+  } else if (ABSL_PREDICT_FALSE(plain)) {
+    zi_ = NULL;
+  } else {
+    /* else, use existing iterator offset and get prev/next as necessary. */
+    if (direction_ == FWD) {
+      nextFn = lpNext;
+      offset_update = 1;
+    } else {
+      DCHECK_EQ(REV, direction_);
+      nextFn = lpPrev;
+      offset_update = -1;
+    }
+    zi_ = nextFn(current_->entry, zi_);
+    offset_ += offset_update;
+  }
+
+  if (zi_)
+    return true;
+
+  // Retry again with the next node.
+  owner_->Compress(current_);
+
+  if (direction_ == FWD) {
+    /* Forward traversal, Jumping to start of next node */
+    current_ = current_->next;
+    offset_ = 0;
+  } else {
+    /* Reverse traversal, Jumping to end of previous node */
+    DCHECK_EQ(REV, direction_);
+
+    current_ = current_->prev;
+    offset_ = -1;
+  }
+
+  return current_ ? Next() : false;
+}
+
+auto QList::Iterator::Get() const -> Entry {
+  int plain = QL_NODE_IS_PLAIN(current_);
+  if (ABSL_PREDICT_FALSE(plain)) {
+    char* str = reinterpret_cast<char*>(current_->entry);
+    return Entry(str, current_->sz);
+  }
+
+  /* Populate value from existing listpack position */
+  unsigned int sz = 0;
+  long long val;
+  uint8_t* ptr = lpGetValue(zi_, &sz, &val);
+
+  return ptr ? Entry(reinterpret_cast<char*>(ptr), sz) : Entry(val);
+}
+
 }  // namespace dfly

src/core/qlist.h

@@ -18,12 +18,18 @@ class QList {
  public:
   enum Where { TAIL, HEAD };
 
+  // Provides wrapper around the references to the listpack entries.
   struct Entry {
     Entry(const char* value, size_t length) : value{value}, length{length} {
     }
     Entry(long long longval) : value{nullptr}, longval{longval} {
     }
 
+    // Assumes value is not null.
+    std::string_view view() const {
+      return {value, length};
+    }
+
     const char* value;
     union {
       size_t length;
@@ -56,6 +62,27 @@ class QList {
 
   void Iterate(IterateFunc cb, long start, long end) const;
 
+  class Iterator {
+   public:
+    Entry Get() const;
+
+    // Returns false if no more entries.
+    bool Next();
+
+   private:
+    QList* owner_;
+    quicklistNode* current_;
+    unsigned char* zi_; /* points to the current element */
+    long offset_;       /* offset in current listpack */
+    uint8_t direction_;
+
+    friend class QList;
+  };
+
+  // Returns an iterator to tail or the head of the list.
+  // To follow the quicklist interface, the iterator is not valid until Next() is called.
+  Iterator GetIterator(Where where);
+
  private:
   bool AllowCompression() const {
     return compress_ != 0;

src/core/qlist_test.cc

@@ -12,6 +12,7 @@ extern "C" {
 }
 
 namespace dfly {
+using namespace std;
 
 class QListTest : public ::testing::Test {
  protected:
@@ -33,8 +34,26 @@ TEST_F(QListTest, Basic) {
   ql_.Push("abc", QList::HEAD);
   EXPECT_EQ(1, ql_.Size());
 
+  auto it = ql_.GetIterator(QList::HEAD);
+  ASSERT_TRUE(it.Next());  // Needed to initialize the iterator.
+
+  QList::Entry entry = it.Get();
+  EXPECT_EQ("abc", entry.view());
+
+  ASSERT_FALSE(it.Next());
+
   ql_.Push("def", QList::TAIL);
   EXPECT_EQ(2, ql_.Size());
+
+  it = ql_.GetIterator(QList::TAIL);
+  ASSERT_TRUE(it.Next());
+  entry = it.Get();
+  EXPECT_EQ("def", entry.view());
+  ASSERT_TRUE(it.Next());
+
+  entry = it.Get();
+  EXPECT_EQ("abc", entry.view());
+  ASSERT_FALSE(it.Next());
 }
 
 };  // namespace dfly