From 216485e4488a7636803a94cb460ec8e77442c781 Mon Sep 17 00:00:00 2001 From: janis Date: Thu, 18 Sep 2025 19:29:00 +0200 Subject: [PATCH] impl drop for free --- src/tree.rs | 273 +++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 249 insertions(+), 24 deletions(-) diff --git a/src/tree.rs b/src/tree.rs index 5df10a6..a648d90 100644 --- a/src/tree.rs +++ b/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 { parent: Option>, parent_idx: MaybeUninit, @@ -484,7 +507,7 @@ impl NodeRef { } 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 NodeRef { 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, K, V> { NodeRef { @@ -651,7 +678,8 @@ impl NodeRef { impl<'a, K, V> NodeRef, 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, 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 NodeRef { // 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 Handle, marker::Edge> { pub(super) fn into_value(self) -> Handle, marker::Value> { unsafe { Handle::new_value(self.node) } } + + pub(crate) fn right_edge( + self, + ) -> Result, 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, 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, K, V>, marker::Edge> { @@ -1304,6 +1352,12 @@ pub struct Tree { _marker: PhantomData>, } +impl Default for Tree { + fn default() -> Self { + Self::new() + } +} + impl<'a, K: core::fmt::Debug + 'a, V: core::fmt::Debug + 'a> core::fmt::Debug for &'a Tree { 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 { + pub(crate) enum LeafHandle { Root(NodeRef), Edge(Handle, marker::Edge>), } @@ -1500,8 +1554,85 @@ mod range { } } + impl Handle, marker::Edge> { + unsafe fn deallocating_next( + self, + ) -> Option, 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, 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 LeafRange { - fn init_front(&mut self) -> Option<&mut Handle, marker::Edge>> { + pub(crate) fn init_front( + &mut self, + ) -> Option<&mut Handle, 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 LeafRange { + 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( + start: Option>, + end: Option>, + ) -> LeafRange { + 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(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( + v: &mut T, + change: impl FnOnce(T) -> Result<(T, R), T>, +) -> Option { + 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 { + range: range::LeafRange, +} + +impl IntoIter { + fn new(tree: Tree) -> 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 Drop for IntoIter { + fn drop(&mut self) { + while let Some(_) = self.range.dying_next() {} + } +} + +impl Drop for Tree { + fn drop(&mut self) { + drop(unsafe { IntoIter::new(core::ptr::read(self)) }); + } } #[cfg(test)] mod tests { + use super::*; + fn build_tree() -> Tree { + 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(); } }