From 424c65b4e5db453d43bb9b91187acfd99acf80ee Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Wed, 25 Oct 2023 23:07:57 +0300
Subject: [PATCH] fix(bptree): Reverse iteration (#2075)

Co-authored-by: Roman Gershman <roman@dragonflydb.io>
---
 src/core/bptree_set.h       |  2 +-
 src/core/bptree_set_test.cc | 68 +++++++++++++++++++++++--------------
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/src/core/bptree_set.h b/src/core/bptree_set.h
index 9c9919584..12f215fe3 100644
--- a/src/core/bptree_set.h
+++ b/src/core/bptree_set.h
@@ -420,7 +420,7 @@ bool BPTree<T, Policy>::IterateReverse(uint32_t rank_start, uint32_t rank_end,
   assert(rank_start <= rank_end && rank_end < count_);
 
   BPTreePath path;
-  ToRank(rank_end, &path);
+  ToRank(count_ - 1 - rank_start, &path);
   for (uint32_t i = rank_start; i <= rank_end; ++i) {
     if (!cb(path.Terminal()))
       return false;
diff --git a/src/core/bptree_set_test.cc b/src/core/bptree_set_test.cc
index f3ee5ef9f..9c330e47c 100644
--- a/src/core/bptree_set_test.cc
+++ b/src/core/bptree_set_test.cc
@@ -4,6 +4,7 @@
 #include "core/bptree_set.h"
 
 #include <absl/container/btree_set.h>
+#include <gmock/gmock.h>
 #include <mimalloc.h>
 
 #include <random>
@@ -248,31 +249,6 @@ TEST_F(BPTreeSetTest, Iterate) {
     ASSERT_EQ(to - from + 1, cnt);
     ASSERT_TRUE(res);
   }
-
-  // Reverse iteration
-  cnt = 0;
-  res = bptree_.IterateReverse(5845, 6849, [&](uint64_t val) {
-    if ((6849 - cnt) * 2 != val)
-      return false;
-    ++cnt;
-    return true;
-  });
-  ASSERT_EQ(6849 - 5845 + 1, cnt);
-  ASSERT_TRUE(res);
-
-  for (unsigned j = 0; j < 10; ++j) {
-    cnt = 0;
-    unsigned from = generator_() % kNumElems;
-    unsigned to = from + generator_() % (kNumElems - from);
-    res = bptree_.IterateReverse(from, to, [&](uint64_t val) {
-      if ((to - cnt) * 2 != val)
-        return false;
-      ++cnt;
-      return true;
-    });
-    ASSERT_EQ(to - from + 1, cnt);
-    ASSERT_TRUE(res);
-  }
 }
 
 TEST_F(BPTreeSetTest, Ranges) {
@@ -352,6 +328,48 @@ TEST_F(BPTreeSetTest, InsertSDS) {
   }
 }
 
+TEST_F(BPTreeSetTest, ReverseIterate) {
+  vector<ZsetPolicy::KeyT> vals;
+  for (int i = -1000; i < 1000; ++i) {
+    sds s = sdsempty();
+
+    s = sdscatfmt(s, "a%u", i);
+    vals.emplace_back(ZsetPolicy::KeyT{.d = (double)i, .s = s});
+  }
+
+  SDSTree tree(&mi_alloc_);
+  for (auto v : vals) {
+    ASSERT_TRUE(tree.Insert(v));
+    {
+      double score = 0;
+      tree.IterateReverse(0, 0, [&score](auto i) {
+        score = i.d;
+        return false;
+      });
+      EXPECT_EQ(score, v.d);
+    }
+    {
+      double score = 0;
+      tree.Iterate(0, 0, [&score](auto i) {
+        score = i.d;
+        return false;
+      });
+      EXPECT_EQ(score, vals[0].d);
+    }
+  }
+
+  vector<int> res;
+  tree.IterateReverse(0, 1, [&](auto i) {
+    res.push_back(i.d);
+    return true;
+  });
+  EXPECT_THAT(res, testing::ElementsAre(999, 998));
+
+  for (auto v : vals) {
+    sdsfree(v.s);
+  }
+}
+
 static string RandomString(mt19937& rand, unsigned len) {
   const string_view alpanum = "1234567890abcdefghijklmnopqrstuvwxyz";
   string ret;