Add index based iterator to QList interface (#4083)

chore: add index based iterator to the interface

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

src/core/qlist.cc

@@ -265,14 +265,6 @@ size_t QList::MallocUsed() const {
   return res + count_ * 16;  // we account for each member 16 bytes.
 }
 
-string QList::Peek(Where where) const {
-  return {};
-}
-
-optional<string> QList::Get(long index) const {
-  return nullopt;
-}
-
 void QList::Iterate(IterateFunc cb, long start, long end) const {
 }
 
@@ -433,11 +425,61 @@ auto QList::GetIterator(Where where) -> Iterator {
   return it;
 }
 
-bool QList::Iterator::Next() {
-  DCHECK(current_);
+auto QList::GetIterator(long idx) -> Iterator {
+  quicklistNode* n;
+  unsigned long long accum = 0;
+  unsigned long long index;
+  int forward = idx < 0 ? 0 : 1; /* < 0 -> reverse, 0+ -> forward */
+
+  index = forward ? idx : (-idx) - 1;
+  if (index >= count_)
+    return {};
+
+  /* Seek in the other direction if that way is shorter. */
+  int seek_forward = forward;
+  unsigned long long seek_index = index;
+  if (index > (count_ - 1) / 2) {
+    seek_forward = !forward;
+    seek_index = count_ - 1 - index;
+  }
+
+  n = seek_forward ? head_ : tail_;
+  while (ABSL_PREDICT_TRUE(n)) {
+    if ((accum + n->count) > seek_index) {
+      break;
+    } else {
+      accum += n->count;
+      n = seek_forward ? n->next : n->prev;
+    }
+  }
+
+  if (!n)
+    return {};
 
-  unsigned char* (*nextFn)(unsigned char*, unsigned char*) = NULL;
-  int offset_update = 0;
+  /* Fix accum so it looks like we seeked in the other direction. */
+  if (seek_forward != forward)
+    accum = count_ - n->count - accum;
+
+  Iterator iter;
+  iter.owner_ = this;
+  iter.direction_ = forward ? FWD : REV;
+  iter.current_ = n;
+
+  if (forward) {
+    /* forward = normal head-to-tail offset. */
+    iter.offset_ = index - accum;
+  } else {
+    /* reverse = need negative offset for tail-to-head, so undo
+     * the result of the original index = (-idx) - 1 above. */
+    iter.offset_ = (-index) - 1 + accum;
+  }
+
+  return iter;
+}
+
+bool QList::Iterator::Next() {
+  if (!current_)
+    return false;
 
   int plain = QL_NODE_IS_PLAIN(current_);
   if (!zi_) {
@@ -450,6 +492,9 @@ bool QList::Iterator::Next() {
   } else if (ABSL_PREDICT_FALSE(plain)) {
     zi_ = NULL;
   } else {
+    unsigned char* (*nextFn)(unsigned char*, unsigned char*) = NULL;
+    int offset_update = 0;
+
     /* else, use existing iterator offset and get prev/next as necessary. */
     if (direction_ == FWD) {
       nextFn = lpNext;
@@ -491,6 +536,8 @@ auto QList::Iterator::Get() const -> Entry {
     return Entry(str, current_->sz);
   }
 
+  DCHECK(zi_);
+
   /* Populate value from existing listpack position */
   unsigned int sz = 0;
   long long val;

src/core/qlist.h

@@ -55,11 +55,6 @@ class QList {
 
   size_t MallocUsed() const;
 
-  // Peeks at the head or tail of the list. Precondition: list is not empty.
-  std::string Peek(Where where) const;
-
-  std::optional<std::string> Get(long index) const;
-
   void Iterate(IterateFunc cb, long start, long end) const;
 
   class Iterator {
@@ -70,19 +65,26 @@ class QList {
     bool Next();
 
    private:
-    QList* owner_;
-    quicklistNode* current_;
-    unsigned char* zi_; /* points to the current element */
-    long offset_;       /* offset in current listpack */
-    uint8_t direction_;
+    QList* owner_ = nullptr;
+    quicklistNode* current_ = nullptr;
+    unsigned char* zi_ = nullptr; /* points to the current element */
+    long offset_ = 0;             /* offset in current listpack */
+    uint8_t direction_ = 1;
 
     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.
+  // To mirror the quicklist interface, the iterator is not valid until Next() is called.
+  // TODO: to fix this.
   Iterator GetIterator(Where where);
 
+  // Returns an iterator at a specific index 'idx',
+  // or Invalid iterator if index is out of range.
+  // negative index - means counting from the tail.
+  // Requires calling subsequent Next() to initialize the iterator.
+  Iterator GetIterator(long idx);
+
  private:
   bool AllowCompression() const {
     return compress_ != 0;

src/core/qlist_test.cc

@@ -37,8 +37,7 @@ TEST_F(QListTest, Basic) {
   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());
+  EXPECT_EQ("abc", it.Get().view());
 
   ASSERT_FALSE(it.Next());
 
@@ -47,13 +46,18 @@ TEST_F(QListTest, Basic) {
 
   it = ql_.GetIterator(QList::TAIL);
   ASSERT_TRUE(it.Next());
-  entry = it.Get();
-  EXPECT_EQ("def", entry.view());
-  ASSERT_TRUE(it.Next());
+  EXPECT_EQ("def", it.Get().view());
 
-  entry = it.Get();
-  EXPECT_EQ("abc", entry.view());
+  ASSERT_TRUE(it.Next());
+  EXPECT_EQ("abc", it.Get().view());
   ASSERT_FALSE(it.Next());
+
+  it = ql_.GetIterator(0);
+  ASSERT_TRUE(it.Next());
+  EXPECT_EQ("abc", it.Get().view());
+  it = ql_.GetIterator(-1);
+  ASSERT_TRUE(it.Next());
+  EXPECT_EQ("def", it.Get().view());
 }
 
 };  // namespace dfly