From 6a3f8e8c3eb282fd81db6037c8e3db48149b6a0a Mon Sep 17 00:00:00 2001
From: Borys <borys@dragonflydb.io>
Date: Mon, 14 Oct 2024 14:20:07 +0300
Subject: [PATCH] chore: add prefetch read for DenseSet::Grow()

---
 src/core/dense_set.cc | 85 +++++++++++++++++++++++++++++--------------
 src/core/dense_set.h  |  2 +-
 2 files changed, 58 insertions(+), 29 deletions(-)

diff --git a/src/core/dense_set.cc b/src/core/dense_set.cc
index 70dee1e060e1..27733a9c8c05 100644
--- a/src/core/dense_set.cc
+++ b/src/core/dense_set.cc
@@ -312,46 +312,65 @@ void DenseSet::GrowBatch(uint32_t len, GrowItem* items,
   while (len) {
     unsigned dest_id = 0;
     for (uint32_t i = 0; i < len; ++i) {
-      DensePtr* curr = &items[i].ptr;
+      auto& item = items[i];
+      if (!item.obj.IsEmpty()) {
+        DensePtr* curr = &item.obj;
 
-      if (ExpireIfNeeded(nullptr, curr)) {
-        // if curr has disappeared due to expiry and prev was converted from Link to a
-        // regular DensePtr
-      }
+        if (ExpireIfNeeded(nullptr, curr)) {
+          // if curr has disappeared due to expiry and prev was converted from Link to a
+          // regular DensePtr
+        }
 
-      if (curr->IsEmpty())
-        continue;
-      void* ptr = curr->GetObject();
+        if (curr->IsEmpty())
+          continue;
+        void* ptr = curr->GetObject();
 
-      DCHECK(ptr != nullptr && ObjectAllocSize(ptr));
+        DCHECK(ptr != nullptr && ObjectAllocSize(ptr));
 
-      uint32_t bid = BucketId(ptr, 0);
+        uint32_t bid = BucketId(ptr, 0);
 
-      // if the item does not move from the current chain, ensure
-      // it is not marked as displaced and move to the next item in the chain
+        // if the item does not move from the current chain, ensure
+        // it is not marked as displaced and move to the next item in the chain
 
-      auto dest = new_entries->begin() + bid;
-      DensePtr dptr = *curr;
+        auto dest = new_entries->begin() + bid;
+        DensePtr dptr = *curr;
 
-      if (curr->IsObject()) {
-        curr->Reset();  // reset the original placeholder (.next or root)
+        if (curr->IsObject()) {
+          curr->Reset();  // reset the original placeholder (.next or root)
 
-        DVLOG(2) << " Pushing to " << bid << " " << dptr.GetObject();
-        DCHECK_EQ(BucketId(dptr.GetObject(), 0), bid);
-        PushFront(dest, dptr);
+          DVLOG(2) << " Pushing to " << bid << " " << dptr.GetObject();
+          DCHECK_EQ(BucketId(dptr.GetObject(), 0), bid);
+          PushFront(dest, dptr);
 
-        dest->ClearDisplaced();
+          dest->ClearDisplaced();
 
-        continue;
-      }  // if IsObject
+          continue;
+        }  // if IsObject
 
-      *curr = *dptr.Next();
-      DCHECK(!curr->IsEmpty());
+        *curr = *dptr.Next();
+        if (curr->IsLink()) {
+          PREFETCH_READ(curr->AsLink());
+        }
 
-      PushFront(dest, dptr);
-      dest->ClearDisplaced();
+        PREFETCH_READ(curr->Raw());
+        DCHECK(!curr->IsEmpty());
 
-      items[dest_id++] = {*curr, nullptr};
+        PushFront(dest, dptr);
+        dest->ClearDisplaced();
+      } else {
+        if (!item.ptr.IsEmpty()) {
+          auto link = item.ptr.AsLink();
+          PREFETCH_READ(link->next.Raw());
+          PREFETCH_READ(link->Raw());
+
+          item.obj = item.ptr;
+          item.ptr.Reset();
+        } else {
+          continue;
+        }
+      }
+
+      items[dest_id++] = {item.ptr, item.obj};
     }
     // update the length of the batch for the next iteration.
     len = dest_id;
@@ -460,12 +479,22 @@ void DenseSet::Fill(DenseSet* other) const {
 
 void DenseSet::Grow(size_t new_size) {
   decltype(entries_) ne(new_size, entries_.get_allocator());
+  const auto kMaxBatchLen = 32;
   GrowItem items[kMaxBatchLen];
   uint32_t len = 0;
   // perform rehashing of items in the set
   for (auto& entry : entries_) {
+    PREFETCH_READ(entry.Raw());
     if (!entry.IsEmpty()) {
-      items[len++] = {entry, entry.IsLink() ? nullptr : entry.Raw()};
+      auto& item = items[len++];
+      if (entry.IsLink()) {
+        item.ptr = entry;
+        item.obj.Reset();
+      } else {
+        item.ptr.Reset();
+        item.obj = entry;
+      }
+
       if (len == kMaxBatchLen) {
         GrowBatch(len, items, &ne);
         len = 0;
diff --git a/src/core/dense_set.h b/src/core/dense_set.h
index 108df0b07c49..503044f943c7 100644
--- a/src/core/dense_set.h
+++ b/src/core/dense_set.h
@@ -343,7 +343,7 @@ class DenseSet {
 
   struct GrowItem {
     DensePtr ptr;
-    void* obj = nullptr;
+    DensePtr obj;
   };
   void GrowBatch(uint32_t len, GrowItem* items,
                  std::vector<DensePtr, DensePtrAllocator>* new_entries);