diff --git a/src/tree.rs b/src/tree.rs index 7ec3c91..863e452 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -623,7 +623,6 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { } impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType>, marker::Value> { - /// SAFETY: the caller must ensure that `idx < node.len()` pub(super) unsafe fn into_value(self) -> &'a V { let leaf = self.node.into_leaf(); let v = leaf.value.as_ref().unwrap(); @@ -631,6 +630,14 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeTyp } } +impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType>, marker::Value> { + pub(super) unsafe fn value_mut(&mut self) -> &mut V { + let leaf = self.node.as_leaf_mut(); + let v = leaf.value.as_mut().unwrap(); + v + } +} + impl Handle, marker::Edge> { pub(super) fn forget_node_type( self, @@ -1053,6 +1060,11 @@ mod search { } } +enum HandleOrTree<'a, BorrowType, K, V, NodeType, HandleType> { + Handle(Handle, HandleType>), + Tree(borrow::DormantMutRef<'a, Tree>), +} + mod entry { use core::marker::PhantomData; @@ -1087,7 +1099,7 @@ mod entry { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("VacantEntry") .field("key", &self.key) - .field("handle", &self.handle) + // .field("handle", &self.handle) .finish() } } @@ -1107,15 +1119,13 @@ mod entry { { pub(super) key: Q, pub(super) handle: - Option, K, V, marker::LeafOrInternal>, marker::Edge>>, - pub(super) dormant_map: DormantMutRef<'a, Tree>, + super::HandleOrTree<'a, marker::Mut<'a>, K, V, marker::LeafOrInternal, marker::Edge>, pub(super) _marker: PhantomData<&'a mut (K, V)>, } pub struct OccupiedEntry<'a, K, V> { pub(super) handle: Handle, K, V, marker::LeafOrInternal>, marker::Value>, - pub(super) dormant_map: DormantMutRef<'a, Tree>, pub(super) _marker: PhantomData<&'a mut (K, V)>, } @@ -1123,6 +1133,9 @@ mod entry { pub fn get(&self) -> &V { unsafe { self.handle.reborrow().into_value() } } + pub fn get_mut(&mut self) -> &mut V { + unsafe { self.handle.value_mut() } + } } impl<'a, Q, K: Ord, V> Entry<'a, Q, K, V> @@ -1141,29 +1154,93 @@ mod entry { where Q: Iterator, { - pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V> { + pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> { + use super::HandleOrTree; let handle = match self.handle { // no root node yet - None => { - let tree = unsafe { self.dormant_map.reborrow() }; + HandleOrTree::Tree(mut tree) => { + // SAFETY: there are no nodes in the tree yet + let tree = unsafe { tree.reborrow() }; let root = tree.root.insert(NodeRef::new_internal()); unsafe { Handle::new_edge(root.borrow_mut(), 0).insert_recursing(self.key, value) } } - Some(handle) => unsafe { handle.insert_recursing(self.key, value) }, + HandleOrTree::Handle(handle) => unsafe { handle.insert_recursing(self.key, value) }, }; OccupiedEntry { handle, - dormant_map: self.dormant_map, _marker: PhantomData, } } } } +mod subtree { + use core::marker::PhantomData; + + use crate::tree::{HandleOrTree, search}; + + use super::{ + NodeRef, OnceAndIter, + entry::{Entry, OccupiedEntry, VacantEntry}, + marker, + }; + + pub struct SubtreeMut<'a, K, V> { + pub(super) root: NodeRef, K, V, marker::LeafOrInternal>, + pub(super) _marker: core::marker::PhantomData<&'a mut (K, V)>, + } + + impl<'a, K, V> SubtreeMut<'a, K, V> + where + K: Ord, + { + pub fn get(&self, mut key_seq: Q) -> Option<&V> + where + Q: Iterator, + { + let root = self.root.reborrow(); + match root.search_tree(&mut key_seq) { + search::SearchResult::Found(handle) => Some(unsafe { handle.into_value() }), + _ => { + // key not found + None + } + } + } + + pub fn entry(&'_ mut self, mut key_seq: Q) -> Entry<'_, OnceAndIter, K, V> + where + Q: Iterator, + { + use Entry::*; + + // SAFETY: this is actually our borrow? + let root = unsafe { self.root.reborrow_mut() }; + + match root.search_tree(&mut key_seq) { + search::SearchResult::Found(handle) => Occupied(OccupiedEntry { + handle, + _marker: PhantomData, + }), + search::SearchResult::GoDown(handle) => Vacant(VacantEntry { + key: key_seq.into(), + handle: HandleOrTree::Handle(handle.forget_node_type()), + _marker: PhantomData, + }), + search::SearchResult::Insert(key, handle) => Vacant(VacantEntry { + key: OnceAndIter::once(key, key_seq), + handle: HandleOrTree::Handle(handle), + _marker: PhantomData, + }), + } + } + } +} + mod borrow { use core::{marker::PhantomData, ptr::NonNull}; @@ -1361,89 +1438,51 @@ impl Tree where K: Ord, { - // pub fn closest_entry( - // &mut self, - // mut key_seq: Q, - // ) -> Option, K, V>> - // where - // Q: Iterator, - // { - // use borrow::DormantMutRef; - // use entry::Entry; + fn as_subtree_mut<'a>(&'a mut self) -> Option> { + let root = self + .root + .as_mut()? + .borrow_mut() + .dormant() + .as_leaf_or_internal(); - // let (tree, dormant) = DormantMutRef::new(self); + Some(subtree::SubtreeMut { + root: unsafe { root.awaken() }, + _marker: PhantomData, + }) + } - // match tree.root { - // None => None, - // Some(ref mut root) => match root.borrow_mut().forget_type().search_tree(&mut key_seq) { - // search::SearchResult::Found(handle) => { - // Some(Entry::Occupied(entry::OccupiedEntry { - // handle, - // dormant_map: dormant, - // _marker: PhantomData, - // })) - // } - // search::SearchResult::GoDown(handle) => match key_seq.next() { - // None => Some(Entry::Occupied(entry::OccupiedEntry { - // handle: handle.forget_node_type().into_value(), - // dormant_map: dormant, - // _marker: PhantomData, - // })), - // Some(key) => Some(Entry::Vacant(entry::VacantEntry { - // key: OnceAndIter::once(key, key_seq), - // handle: Some(handle.forget_node_type()), - // dormant_map: dormant, - // _marker: PhantomData, - // })), - // }, - // search::SearchResult::Insert(key, handle) => { - // Some(Entry::Vacant(entry::VacantEntry { - // key: OnceAndIter::once(key, key_seq), - // handle: Some(handle), - // dormant_map: dormant, - // _marker: PhantomData, - // })) - // } - // }, - // } - // } - - pub fn entry(&'_ mut self, mut key_seq: Q) -> entry::Entry<'_, OnceAndIter, K, V> + pub fn entry<'a, Q>(&'a mut self, key_seq: Q) -> entry::Entry<'a, OnceAndIter, K, V> where Q: Iterator, { use borrow::DormantMutRef; - use entry::{Entry::*, OccupiedEntry, VacantEntry}; + use entry::{Entry::*, VacantEntry}; let (tree, dormant) = DormantMutRef::new(self); - match tree.root { + let entry = match tree.as_subtree_mut() { + Some(mut subtree) => { + let entry = subtree.entry(key_seq); + // SAFETY: extending the lifetime is fine because we borrow the tree for 'a, + // and no references to the subtree are live after this. + // The same could be achieved using `dormant.reborrow()` a bunch + // of times while destructuring the entry. + unsafe { + core::mem::transmute::< + entry::Entry<'_, OnceAndIter, K, V>, + entry::Entry<'a, OnceAndIter, K, V>, + >(entry) + } + } None => Vacant(VacantEntry { key: key_seq.into(), - handle: None, - dormant_map: dormant, + handle: HandleOrTree::Tree(dormant), _marker: PhantomData, }), - Some(ref mut root) => match root.borrow_mut().forget_type().search_tree(&mut key_seq) { - search::SearchResult::Found(handle) => Occupied(OccupiedEntry { - handle, - dormant_map: dormant, - _marker: PhantomData, - }), - search::SearchResult::GoDown(handle) => Vacant(VacantEntry { - key: key_seq.into(), - handle: Some(handle.forget_node_type()), - dormant_map: dormant, - _marker: PhantomData, - }), - search::SearchResult::Insert(key, handle) => Vacant(VacantEntry { - key: OnceAndIter::once(key, key_seq), - handle: Some(handle), - dormant_map: dormant, - _marker: PhantomData, - }), - }, - } + }; + + entry } }