From ccdf8971bb29c4cbd8de63f4e45fd38284890ea0 Mon Sep 17 00:00:00 2001
From: Xiangpeng Hao <haoxiangpeng123@gmail.com>
Date: Wed, 20 Nov 2024 23:29:43 -0600
Subject: [PATCH] fix leak

---
 src/node_ptr.rs |  7 -------
 src/tree.rs     | 31 ++++++++++++-------------------
 2 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/src/node_ptr.rs b/src/node_ptr.rs
index d81ca4f..0bf5198 100644
--- a/src/node_ptr.rs
+++ b/src/node_ptr.rs
@@ -2,7 +2,6 @@ use std::ptr::NonNull;
 
 use crate::{
     base_node::{BaseNode, Node},
-    node_256::Node256,
     utils::KeyTracker,
 };
 
@@ -49,12 +48,6 @@ impl NodePtr {
         }
     }
 
-    pub(crate) fn from_root(ptr: NonNull<Node256>) -> Self {
-        Self {
-            sub_node: unsafe { std::mem::transmute::<NonNull<Node256>, NonNull<BaseNode>>(ptr) },
-        }
-    }
-
     fn from_node_new(ptr: NonNull<BaseNode>) -> Self {
         Self { sub_node: ptr }
     }
diff --git a/src/tree.rs b/src/tree.rs
index 55241ec..cc0e9ba 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -33,30 +33,23 @@ impl<const K_LEN: usize> Default for RawCongee<K_LEN> {
 
 impl<const K_LEN: usize, A: Allocator + Clone> Drop for RawCongee<K_LEN, A> {
     fn drop(&mut self) {
-        let mut sub_nodes = vec![(NodePtr::from_root(self.root), 0)];
+        let mut sub_nodes = vec![(unsafe { std::mem::transmute(self.root) }, 0)];
 
         while let Some((node, level)) = sub_nodes.pop() {
-            match node.downcast::<K_LEN>(level) {
-                PtrType::Payload(_) => {
-                    continue;
-                }
-                PtrType::SubNode(sub_node) => {
-                    let node_lock = BaseNode::read_lock(sub_node).unwrap();
-                    let children = node_lock.as_ref().get_children(0, 255);
-                    for (_k, n) in children {
-                        match n.downcast::<K_LEN>(level) {
-                            PtrType::Payload(_) => {}
-                            PtrType::SubNode(sub_sub_node) => {
-                                let node_lock = BaseNode::read_lock(sub_sub_node).unwrap();
-                                sub_nodes.push((n, node_lock.as_ref().prefix().len()));
-                            }
-                        }
-                    }
-                    unsafe {
-                        BaseNode::drop_node(sub_node, self.allocator.clone());
+            let node_lock = BaseNode::read_lock(node).unwrap();
+            let children = node_lock.as_ref().get_children(0, 255);
+            for (_k, n) in children {
+                match n.downcast::<K_LEN>(level) {
+                    PtrType::Payload(_) => {}
+                    PtrType::SubNode(sub_sub_node) => {
+                        let node_lock = BaseNode::read_lock(sub_sub_node).unwrap();
+                        sub_nodes.push((sub_sub_node, node_lock.as_ref().prefix().len()));
                     }
                 }
             }
+            unsafe {
+                BaseNode::drop_node(node, self.allocator.clone());
+            }
         }
 
         // see this: https://github.com/XiangpengHao/congee/issues/20