impl drop for free

src/tree.rs

@@ -13,18 +13,40 @@
 // tree.find('i') -> None
 
 use alloc::boxed::Box;
-use core::{marker::PhantomData, mem::MaybeUninit, ptr::NonNull};
+use core::{
+    marker::PhantomData,
+    mem::{ManuallyDrop, MaybeUninit},
+    ptr::NonNull,
+};
 
 use crate::tree::subtree::Subtree;
 
 mod marker {
+    #![allow(dead_code)]
     use core::marker::PhantomData;
 
+    pub trait NodeType {}
+
+    #[derive(Debug)]
+    pub struct Internal;
+    #[derive(Debug)]
+    pub struct Leaf;
+    #[derive(Debug)]
+    pub struct LeafOrInternal;
+
+    impl NodeType for Internal {}
+    impl NodeType for Leaf {}
+    impl NodeType for LeafOrInternal {}
+
+    pub trait HandleType {}
     #[derive(Debug)]
     pub struct Edge;
     #[derive(Debug)]
     pub struct Value;
 
+    impl HandleType for Edge {}
+    impl HandleType for Value {}
+
     #[derive(Debug)]
     pub struct Owned;
     #[derive(Debug)]
@@ -54,6 +76,7 @@ mod marker {
 }
 
 #[derive(Debug)]
+#[repr(C)]
 pub(crate) struct LeafNode<K, V> {
     parent: Option<BoxedNode<K, V>>,
     parent_idx: MaybeUninit<u16>,
@@ -484,7 +507,7 @@ impl<BorrowType, K, V> NodeRef<BorrowType, K, V> {
     }
 
     pub(super) fn capacity(&self) -> usize {
-        unsafe { usize::from((*Self::as_leaf_ptr(self)).len) }
+        unsafe { usize::from((*Self::as_leaf_ptr(self)).capacity) }
     }
 }
 
@@ -493,6 +516,10 @@ impl<BorrowType, K, V> NodeRef<BorrowType, K, V> {
         unsafe { (*Self::as_leaf_ptr(self)).len == 0 }
     }
 
+    pub(crate) fn has_descendants(&self) -> bool {
+        self.as_leaf().len > 0
+    }
+
     /// Temporarily takes out another, immutable reference to the same node.
     pub(super) fn reborrow(&self) -> NodeRef<marker::Immut<'_>, K, V> {
         NodeRef {
@@ -651,7 +678,8 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V> {
 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V> {
     pub(super) fn grow_node(&mut self) {
         // grow the node
-        let Some(new_capacity) = self.capacity_mut().checked_mul(2) else {
+        let capacity = self.capacity();
+        let Some(new_capacity) = capacity.checked_mul(2) else {
             panic!("Node capacity overflow");
         };
 
@@ -679,8 +707,8 @@ impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V> {
 
                 // dealloc old keys and edges.
                 // This doesn't drop because the keys and edges were moved.
-                _ = Box::from_non_null(old_keys);
-                _ = Box::from_non_null(old_edges);
+                _ = Box::<[_]>::from_non_null(NonNull::slice_from_raw_parts(old_keys, capacity));
+                _ = Box::<[_]>::from_non_null(NonNull::slice_from_raw_parts(old_edges, capacity));
             }
         }
 
@@ -716,6 +744,7 @@ impl<K, V> NodeRef<marker::Owned, K, V> {
         // insert new key and child node.
         // SAFETY: we just grew the allocations.
         unsafe {
+            assert!(new_len == parent.idx + 1);
             slice_insert(parent.node.key_area_mut(..new_len), parent.idx, key);
             slice_insert(parent.node.edge_area_mut(..new_len), parent.idx, self.node);
         }
@@ -733,6 +762,25 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V>, marker::Edge> {
     pub(super) fn into_value(self) -> Handle<NodeRef<BorrowType, K, V>, marker::Value> {
         unsafe { Handle::new_value(self.node) }
     }
+
+    pub(crate) fn right_edge(
+        self,
+    ) -> Result<Handle<NodeRef<BorrowType, K, V>, marker::Edge>, Self> {
+        let len = self.node.len();
+        if self.idx + 1 < len {
+            Ok(unsafe { Handle::new_edge(self.node, self.idx + 1) })
+        } else {
+            Err(self)
+        }
+    }
+
+    pub(crate) fn left_edge(self) -> Result<Handle<NodeRef<BorrowType, K, V>, marker::Edge>, Self> {
+        if self.idx > 0 {
+            Ok(unsafe { Handle::new_edge(self.node, self.idx - 1) })
+        } else {
+            Err(self)
+        }
+    }
 }
 
 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V>, marker::Edge> {
@@ -1304,6 +1352,12 @@ pub struct Tree<K, V> {
     _marker: PhantomData<alloc::boxed::Box<(K, V)>>,
 }
 
+impl<K, V> Default for Tree<K, V> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<'a, K: core::fmt::Debug + 'a, V: core::fmt::Debug + 'a> core::fmt::Debug for &'a Tree<K, V> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         let mut map = f.debug_map();
@@ -1472,7 +1526,7 @@ mod range {
 
     use crate::tree::{Handle, NodeRef};
 
-    enum LeafHandle<BorrowType, K, V> {
+    pub(crate) enum LeafHandle<BorrowType, K, V> {
         Root(NodeRef<BorrowType, K, V>),
         Edge(Handle<NodeRef<BorrowType, K, V>, marker::Edge>),
     }
@@ -1500,8 +1554,85 @@ mod range {
         }
     }
 
+    impl<K, V> Handle<NodeRef<marker::Dying, K, V>, marker::Edge> {
+        unsafe fn deallocating_next(
+            self,
+        ) -> Option<Handle<NodeRef<marker::Dying, K, V>, marker::Edge>> {
+            let mut edge = self;
+            loop {
+                // if this node has descendants, we want to drop them first.
+                if edge.node.has_descendants() {
+                    return Some(edge.descend().first_edge());
+                }
+
+                let mut last_edge = edge;
+                edge = loop {
+                    // no more descendants: deallocate this node and ascend.
+                    // after ascending, check if we need to go right and drop
+                    // more descendants of the parent.
+                    // otherwise, continue deallocating and ascending
+                    match last_edge.node.deallocate_and_ascend() {
+                        Some(parent) => match parent.right_edge() {
+                            Ok(next) => break next,
+                            Err(last) => last_edge = last,
+                        },
+                        None => return None,
+                    }
+                };
+            }
+        }
+
+        unsafe fn dellocating_next_back(
+            self,
+        ) -> Option<Handle<NodeRef<marker::Dying, K, V>, marker::Edge>> {
+            let mut edge = self;
+            loop {
+                if edge.node.has_descendants() {
+                    match edge.left_edge() {
+                        Ok(next) => return Some(next),
+                        Err(node) => edge = node,
+                    }
+                } else {
+                    match edge.node.deallocate_and_ascend() {
+                        Some(parent) => edge = parent,
+                        None => return None,
+                    }
+                }
+            }
+        }
+
+        pub(crate) unsafe fn deallocating_next_unchecked(&mut self) {
+            super::replace(self, |edge| {
+                (unsafe { edge.deallocating_next().unwrap() }, ())
+            })
+        }
+
+        pub(crate) unsafe fn try_deallocating_next(&mut self) -> bool {
+            super::maybe_replace(self, |edge| unsafe {
+                let err = core::ptr::read(&edge);
+                edge.deallocating_next().map(|e| (e, ())).ok_or(err)
+            })
+            .is_some()
+        }
+
+        unsafe fn deallocating_next_back_unchecked(&mut self) {
+            super::replace(self, |edge| {
+                (unsafe { edge.dellocating_next_back().unwrap() }, ())
+            })
+        }
+
+        pub(crate) fn deallocating_end(self) {
+            let mut edge = self;
+            while let Some(parent) = edge.node.deallocate_and_ascend() {
+                edge = parent;
+            }
+        }
+    }
+
     impl<BorrowType: marker::BorrowType, K, V> LeafRange<BorrowType, K, V> {
-        fn init_front(&mut self) -> Option<&mut Handle<NodeRef<BorrowType, K, V>, marker::Edge>> {
+        pub(crate) fn init_front(
+            &mut self,
+        ) -> Option<&mut Handle<NodeRef<BorrowType, K, V>, marker::Edge>> {
             if let Some(LeafHandle::Root(root)) = &self.front {
                 self.front = Some(LeafHandle::Edge(
                     unsafe { ptr::read(root) }.first_leaf_edge(),
@@ -1511,6 +1642,7 @@ mod range {
             match &mut self.front {
                 None => None,
                 Some(LeafHandle::Edge(edge)) => Some(edge),
+                // if it was root, we've just replaced it with an edge.
                 Some(LeafHandle::Root(_)) => unreachable!(),
             }
         }
@@ -1525,33 +1657,126 @@ mod range {
             match &mut self.back {
                 None => None,
                 Some(LeafHandle::Edge(edge)) => Some(edge),
+                // if it was root, we've just replaced it with an edge.
                 Some(LeafHandle::Root(_)) => unreachable!(),
             }
         }
     }
+
+    impl<K, V> LeafRange<marker::Dying, K, V> {
+        pub(crate) fn dying_next(&mut self) -> Option<()> {
+            let front = self.init_front().unwrap();
+            unsafe { front.try_deallocating_next().then_some(()) }
+        }
+    }
+
+    pub(crate) fn full_range<BorrowType, K, V>(
+        start: Option<NodeRef<BorrowType, K, V>>,
+        end: Option<NodeRef<BorrowType, K, V>>,
+    ) -> LeafRange<BorrowType, K, V> {
+        LeafRange {
+            front: start.map(LeafHandle::Root),
+            back: end.map(LeafHandle::Root),
+        }
+    }
+}
+/// This replaces the value behind the `v` unique reference by calling the
+/// relevant function, and returns a result obtained along the way.
+///
+/// If a panic occurs in the `change` closure, the entire process will be aborted.
+#[inline]
+pub(super) fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
+    use core::{mem, ptr};
+    struct PanicGuard;
+    impl Drop for PanicGuard {
+        fn drop(&mut self) {
+            panic!()
+        }
+    }
+    let guard = PanicGuard;
+    let value = unsafe { ptr::read(v) };
+    let (new_value, ret) = change(value);
+    unsafe {
+        ptr::write(v, new_value);
+    }
+    mem::forget(guard);
+    ret
+}
+
+pub(crate) fn maybe_replace<T, R>(
+    v: &mut T,
+    change: impl FnOnce(T) -> Result<(T, R), T>,
+) -> Option<R> {
+    use core::{mem, ptr};
+    struct PanicGuard;
+    impl Drop for PanicGuard {
+        fn drop(&mut self) {
+            panic!()
+        }
+    }
+    let guard = PanicGuard;
+    let value = unsafe { ptr::read(v) };
+    let (new_value, ret) = match change(value) {
+        Err(old_value) => (old_value, None),
+        Ok((new_value, ret)) => (new_value, Some(ret)),
+    };
+    unsafe {
+        ptr::write(v, new_value);
+    }
+    mem::forget(guard);
+    ret
+}
+
+struct IntoIter<K, V> {
+    range: range::LeafRange<marker::Dying, K, V>,
+}
+
+impl<K, V> IntoIter<K, V> {
+    fn new(tree: Tree<K, V>) -> Self {
+        let mut tree = ManuallyDrop::new(tree);
+        if let Some(root) = tree.root.take() {
+            let root = root.into_dying();
+            let root2 = unsafe { core::ptr::read(&root) };
+
+            Self {
+                range: range::full_range(Some(root), Some(root2)),
+            }
+        } else {
+            Self {
+                range: range::full_range(None, None),
+            }
+        }
+    }
+}
+
+impl<K, V> Drop for IntoIter<K, V> {
+    fn drop(&mut self) {
+        while let Some(_) = self.range.dying_next() {}
+    }
+}
+
+impl<K, V> Drop for Tree<K, V> {
+    fn drop(&mut self) {
+        drop(unsafe { IntoIter::new(core::ptr::read(self)) });
+    }
 }
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
 
+    fn build_tree() -> Tree<char, i32> {
+        let mut tree = Tree::new();
+        tree.entry("asdf".chars()).or_insert(1);
+        tree.entry("asd".chars()).or_insert(2);
+        tree.entry("asdg".chars()).or_insert(3);
+
+        tree
+    }
+
     #[test]
-    fn promote_leaf() {
-        #[derive(Debug, PartialEq, Eq)]
-        struct Test(&'static str);
-
-        impl Drop for Test {
-            fn drop(&mut self) {
-                std::eprintln!("Dropping: {}", self.0);
-            }
-        }
-
-        let mut leaf = NodeRef::<_, (), Test>::new();
-        leaf.borrow_mut().as_leaf_mut().value = Some(Test("test"));
-        let mut root = NodeRef::new();
-
-        let mut leaf = leaf.reparent(unsafe { Handle::new_edge(root.borrow_mut(), 0) }, ());
-
-        assert_eq!(leaf.as_leaf_mut().value, Some(Test("test")));
+    fn drop_tree() {
+        let tree = build_tree();
     }
 }