From 2f13e119758924146970929aa42bf34022d12327 Mon Sep 17 00:00:00 2001
From: Stepan Bagritsevich <stefan@dragonflydb.io>
Date: Wed, 4 Dec 2024 21:37:06 +0400
Subject: [PATCH] chore(rax_tree): Introduce raxFreeWithCallback call in
 RaxTreeMap destructor

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>
---
 src/core/search/rax_tree.h | 16 ++++++++++------
 src/redis/rax.c            | 27 +++++++++++++++++++++------
 src/redis/rax.h            |  1 +
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/core/search/rax_tree.h b/src/core/search/rax_tree.h
index 1f5662faad5c..78f545ce55c9 100644
--- a/src/core/search/rax_tree.h
+++ b/src/core/search/rax_tree.h
@@ -109,12 +109,16 @@ template <typename V> struct RaxTreeMap {
   }
 
   ~RaxTreeMap() {
-    for (auto it = begin(); it != end(); ++it) {
-      V* ptr = &(*it).second;
-      std::allocator_traits<decltype(alloc_)>::destroy(alloc_, ptr);
-      alloc_.deallocate(ptr, 1);
-    }
-    raxFree(tree_);
+    using Allocator = decltype(alloc_);
+
+    auto free_callback = [](void* data, void* context) {
+      Allocator* allocator = static_cast<Allocator*>(context);
+      V* ptr = static_cast<V*>(data);
+      std::allocator_traits<Allocator>::destroy(*allocator, ptr);
+      allocator->deallocate(ptr, 1);
+    };
+
+    raxFreeWithCallbackAndArgument(tree_, free_callback, &alloc_);
   }
 
   size_t size() const {
diff --git a/src/redis/rax.c b/src/redis/rax.c
index 5caba09ee0e0..0e7b1632a8b7 100644
--- a/src/redis/rax.c
+++ b/src/redis/rax.c
@@ -1221,29 +1221,44 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
 
 /* This is the core of raxFree(): performs a depth-first scan of the
  * tree and releases all the nodes found. */
-void raxRecursiveFree(rax *rax, raxNode *n, void (*free_callback)(void*)) {
+void raxRecursiveFree(rax *rax, raxNode *n, void (*free_callback)(void*, void*), void* argument) {
     debugnode("free traversing",n);
     int numchildren = n->iscompr ? 1 : n->size;
     raxNode **cp = raxNodeLastChildPtr(n);
     while(numchildren--) {
         raxNode *child;
         memcpy(&child,cp,sizeof(child));
-        raxRecursiveFree(rax,child,free_callback);
+        raxRecursiveFree(rax,child,free_callback,argument);
         cp--;
     }
     debugnode("free depth-first",n);
     if (free_callback && n->iskey && !n->isnull)
-        free_callback(raxGetData(n));
+        free_callback(raxGetData(n),argument);
     rax_free(n);
     rax->numnodes--;
 }
 
+/* Free the entire radix tree, invoking a free_callback function for each key's data. 
+ * An additional argument is passed to the free_callback function.*/
+ void raxFreeWithCallbackAndArgument(rax *rax, void (*free_callback)(void*, void*), void* argument) {
+    raxRecursiveFree(rax,rax->head,free_callback,argument);
+    assert(rax->numnodes == 0);
+    rax_free(rax);
+}
+
+/* Wrapper for the callback to adapt it for the context */
+void freeCallbackWrapper(void* data, void* argument) {
+    if (!argument) {
+        return;
+    }
+    void (*free_callback)(void*) = (void (*)(void*))argument;
+    free_callback(data);
+}
+
 /* Free a whole radix tree, calling the specified callback in order to
  * free the auxiliary data. */
 void raxFreeWithCallback(rax *rax, void (*free_callback)(void*)) {
-    raxRecursiveFree(rax,rax->head,free_callback);
-    assert(rax->numnodes == 0);
-    rax_free(rax);
+    raxFreeWithCallbackAndArgument(rax, freeCallbackWrapper, (void*)free_callback);
 }
 
 /* Free a whole radix tree. */
diff --git a/src/redis/rax.h b/src/redis/rax.h
index a25bc17e8f2b..0d2f1f8ae9c3 100644
--- a/src/redis/rax.h
+++ b/src/redis/rax.h
@@ -195,6 +195,7 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old);
 void *raxFind(rax *rax, unsigned char *s, size_t len);
 void raxFree(rax *rax);
 void raxFreeWithCallback(rax *rax, void (*free_callback)(void*));
+void raxFreeWithCallbackAndArgument(rax *rax, void (*free_callback)(void*, void*), void* argument);
 void raxStart(raxIterator *it, rax *rt);
 int raxSeek(raxIterator *it, const char *op, unsigned char *ele, size_t len);
 int raxNext(raxIterator *it);