From 57efb5005bd0e769ecee9c2bc75125f8ea340730 Mon Sep 17 00:00:00 2001
From: Geoffrey Yu <geoffxy@gmail.com>
Date: Tue, 7 Sep 2021 13:25:01 -0400
Subject: [PATCH] Handle empty data nodes when finding last key (#19)

- Traverse backward in Alex<>::get_payload_last_no_greater_than() and
  Alex<>::find_last_no_greater_than().

- Add test that fails on `master`.
---
 src/core/alex.h      | 43 ++++++++++++++++++++++++-------------------
 test/unittest_alex.h | 42 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/src/core/alex.h b/src/core/alex.h
index 562394c..d6be807 100644
--- a/src/core/alex.h
+++ b/src/core/alex.h
@@ -994,18 +994,20 @@ class Alex {
   typename self_type::Iterator find_last_no_greater_than(const T& key) {
     stats_.num_lookups++;
     data_node_type* leaf = get_leaf(key);
-    int idx = leaf->upper_bound(key) - 1;
-    if (idx == -1) {
-      if (leaf->prev_leaf_) {
-        // Edge case: need to check previous data node
-        data_node_type* prev_leaf = leaf->prev_leaf_;
-        int last_pos = prev_leaf->last_pos();
-        return Iterator(prev_leaf, last_pos);
-      } else {
+    const int idx = leaf->upper_bound(key) - 1;
+    if (idx >= 0) {
+      return Iterator(leaf, idx);
+    }
+
+    // Edge case: need to check previous data node(s)
+    while (true) {
+      if (leaf->prev_leaf_ == nullptr) {
         return Iterator(leaf, 0);
       }
-    } else {
-      return Iterator(leaf, idx);
+      leaf = leaf->prev_leaf_;
+      if (leaf->num_keys_ > 0) {
+        return Iterator(leaf, leaf->last_pos());
+      }
     }
   }
 
@@ -1015,17 +1017,20 @@ class Alex {
   P* get_payload_last_no_greater_than(const T& key) {
     stats_.num_lookups++;
     data_node_type* leaf = get_leaf(key);
-    int idx = leaf->upper_bound(key) - 1;
-    if (idx == -1) {
-      if (leaf->prev_leaf_) {
-        // Edge case: need to check previous data node
-        data_node_type* prev_leaf = leaf->prev_leaf_;
-        return &(prev_leaf->get_payload(prev_leaf->last_pos()));
-      } else {
+    const int idx = leaf->upper_bound(key) - 1;
+    if (idx >= 0) {
+      return &(leaf->get_payload(idx));
+    }
+
+    // Edge case: Need to check previous data node(s)
+    while (true) {
+      if (leaf->prev_leaf_ == nullptr) {
         return &(leaf->get_payload(leaf->first_pos()));
       }
-    } else {
-      return &(leaf->get_payload(idx));
+      leaf = leaf->prev_leaf_;
+      if (leaf->num_keys_ > 0) {
+        return &(leaf->get_payload(leaf->last_pos()));
+      }
     }
   }
 
diff --git a/test/unittest_alex.h b/test/unittest_alex.h
index 0be6884..e83439d 100644
--- a/test/unittest_alex.h
+++ b/test/unittest_alex.h
@@ -297,6 +297,46 @@ TEST_CASE("TestFindLastNoGreaterThan") {
   CHECK_EQ(values[0].first, it.key());
 }
 
+TEST_CASE("TestLargeFindLastNoGreaterThan") {
+  Alex<uint64_t, uint64_t> index;
+  index.insert(std::make_pair(0ULL, 0ULL));
+
+  const uint64_t keys_per_segment = 80523;
+  const uint64_t step_size = 43206176;
+  const uint64_t num_segments = 16;
+  const uint64_t start_keys[] = {698631712, 658125922, 660826308, 663526694,
+                                 666227080, 668927466, 671627852, 674328238,
+                                 677028624, 679729010, 682429396, 685129782,
+                                 687830168, 690530554, 693230940, 695931326};
+
+  uint64_t max_key = 0;
+  uint64_t max_key_value = 0;
+  for (uint64_t segment = 0; segment < num_segments; ++segment) {
+    uint64_t curr_key = start_keys[segment];
+    for (uint64_t i = 0; i < keys_per_segment; ++i) {
+      if (curr_key > max_key) {
+        max_key = curr_key;
+        max_key_value = i + 1;
+      }
+
+      index.insert(curr_key, i + 1);
+      curr_key += step_size;
+    }
+  }
+
+  // This key is larger than all keys in the index.
+  const uint64_t test_key = 3650322694401;
+  CHECK_GT(test_key, max_key);
+
+  auto it = index.find_last_no_greater_than(test_key);
+  CHECK(!it.is_end());
+  CHECK_EQ(max_key, it.key());
+
+  const uint64_t *p = index.get_payload_last_no_greater_than(test_key);
+  CHECK(p);
+  CHECK_EQ(max_key_value, *p);
+}
+
 TEST_CASE("TestReadModifyWrite") {
   Alex<int, int> index;
 
@@ -482,4 +522,4 @@ TEST_CASE("TestRangeScan") {
   CHECK_EQ(results2.size(), 90);
   CHECK_EQ(sum2, 4905);
 }
-};
\ No newline at end of file
+};