fix leak

XiangpengHao · Nov 21, 2024 · ccdf897 · ccdf897
1 parent b708afc
commit ccdf897
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 26 deletions.
diff --git 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
@@ -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