use core::borrow::Borrow; use core::hint; use core::ops::RangeBounds; use core::ptr; use super::node::{marker, ForceResult::*, Handle, NodeRef}; use super::search::SearchBound; use crate::alloc::Allocator; // `front` and `back` are always both `None` or both `Some`. pub struct LeafRange { front: Option, marker::Edge>>, back: Option, marker::Edge>>, } impl<'a, K: 'a, V: 'a> Clone for LeafRange, K, V> { fn clone(&self) -> Self { LeafRange { front: self.front.clone(), back: self.back.clone() } } } impl LeafRange { pub fn none() -> Self { LeafRange { front: None, back: None } } fn is_empty(&self) -> bool { self.front == self.back } /// Temporarily takes out another, immutable equivalent of the same range. pub fn reborrow(&self) -> LeafRange, K, V> { LeafRange { front: self.front.as_ref().map(|f| f.reborrow()), back: self.back.as_ref().map(|b| b.reborrow()), } } } impl<'a, K, V> LeafRange, K, V> { #[inline] pub fn next_checked(&mut self) -> Option<(&'a K, &'a V)> { self.perform_next_checked(|kv| kv.into_kv()) } #[inline] pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a V)> { self.perform_next_back_checked(|kv| kv.into_kv()) } } impl<'a, K, V> LeafRange, K, V> { #[inline] pub fn next_checked(&mut self) -> Option<(&'a K, &'a mut V)> { self.perform_next_checked(|kv| unsafe { ptr::read(kv) }.into_kv_valmut()) } #[inline] pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a mut V)> { self.perform_next_back_checked(|kv| unsafe { ptr::read(kv) }.into_kv_valmut()) } } impl LeafRange { /// If possible, extract some result from the following KV and move to the edge beyond it. fn perform_next_checked(&mut self, f: F) -> Option where F: Fn(&Handle, marker::KV>) -> R, { if self.is_empty() { None } else { super::mem::replace(self.front.as_mut().unwrap(), |front| { let kv = front.next_kv().ok().unwrap(); let result = f(&kv); (kv.next_leaf_edge(), Some(result)) }) } } /// If possible, extract some result from the preceding KV and move to the edge beyond it. fn perform_next_back_checked(&mut self, f: F) -> Option where F: Fn(&Handle, marker::KV>) -> R, { if self.is_empty() { None } else { super::mem::replace(self.back.as_mut().unwrap(), |back| { let kv = back.next_back_kv().ok().unwrap(); let result = f(&kv); (kv.next_back_leaf_edge(), Some(result)) }) } } } enum LazyLeafHandle { Root(NodeRef), // not yet descended Edge(Handle, marker::Edge>), } impl<'a, K: 'a, V: 'a> Clone for LazyLeafHandle, K, V> { fn clone(&self) -> Self { match self { LazyLeafHandle::Root(root) => LazyLeafHandle::Root(*root), LazyLeafHandle::Edge(edge) => LazyLeafHandle::Edge(*edge), } } } impl LazyLeafHandle { fn reborrow(&self) -> LazyLeafHandle, K, V> { match self { LazyLeafHandle::Root(root) => LazyLeafHandle::Root(root.reborrow()), LazyLeafHandle::Edge(edge) => LazyLeafHandle::Edge(edge.reborrow()), } } } // `front` and `back` are always both `None` or both `Some`. pub struct LazyLeafRange { front: Option>, back: Option>, } impl<'a, K: 'a, V: 'a> Clone for LazyLeafRange, K, V> { fn clone(&self) -> Self { LazyLeafRange { front: self.front.clone(), back: self.back.clone() } } } impl LazyLeafRange { pub fn none() -> Self { LazyLeafRange { front: None, back: None } } /// Temporarily takes out another, immutable equivalent of the same range. pub fn reborrow(&self) -> LazyLeafRange, K, V> { LazyLeafRange { front: self.front.as_ref().map(|f| f.reborrow()), back: self.back.as_ref().map(|b| b.reborrow()), } } } impl<'a, K, V> LazyLeafRange, K, V> { #[inline] pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { unsafe { self.init_front().unwrap().next_unchecked() } } #[inline] pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { unsafe { self.init_back().unwrap().next_back_unchecked() } } } impl<'a, K, V> LazyLeafRange, K, V> { #[inline] pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { unsafe { self.init_front().unwrap().next_unchecked() } } #[inline] pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { unsafe { self.init_back().unwrap().next_back_unchecked() } } } impl LazyLeafRange { fn take_front( &mut self, ) -> Option, marker::Edge>> { match self.front.take()? { LazyLeafHandle::Root(root) => Some(root.first_leaf_edge()), LazyLeafHandle::Edge(edge) => Some(edge), } } #[inline] pub unsafe fn deallocating_next_unchecked( &mut self, alloc: A, ) -> Handle, marker::KV> { debug_assert!(self.front.is_some()); let front = self.init_front().unwrap(); unsafe { front.deallocating_next_unchecked(alloc) } } #[inline] pub unsafe fn deallocating_next_back_unchecked( &mut self, alloc: A, ) -> Handle, marker::KV> { debug_assert!(self.back.is_some()); let back = self.init_back().unwrap(); unsafe { back.deallocating_next_back_unchecked(alloc) } } #[inline] pub fn deallocating_end(&mut self, alloc: A) { if let Some(front) = self.take_front() { front.deallocating_end(alloc) } } } impl LazyLeafRange { fn init_front( &mut self, ) -> Option<&mut Handle, marker::Edge>> { if let Some(LazyLeafHandle::Root(root)) = &self.front { self.front = Some(LazyLeafHandle::Edge(unsafe { ptr::read(root) }.first_leaf_edge())); } match &mut self.front { None => None, Some(LazyLeafHandle::Edge(edge)) => Some(edge), // SAFETY: the code above would have replaced it. Some(LazyLeafHandle::Root(_)) => unsafe { hint::unreachable_unchecked() }, } } fn init_back( &mut self, ) -> Option<&mut Handle, marker::Edge>> { if let Some(LazyLeafHandle::Root(root)) = &self.back { self.back = Some(LazyLeafHandle::Edge(unsafe { ptr::read(root) }.last_leaf_edge())); } match &mut self.back { None => None, Some(LazyLeafHandle::Edge(edge)) => Some(edge), // SAFETY: the code above would have replaced it. Some(LazyLeafHandle::Root(_)) => unsafe { hint::unreachable_unchecked() }, } } } impl NodeRef { /// Finds the distinct leaf edges delimiting a specified range in a tree. /// /// If such distinct edges exist, returns them in ascending order, meaning /// that a non-zero number of calls to `next_unchecked` on the `front` of /// the result and/or calls to `next_back_unchecked` on the `back` of the /// result will eventually reach the same edge. /// /// If there are no such edges, i.e., if the tree contains no key within /// the range, returns an empty `front` and `back`. /// /// # Safety /// Unless `BorrowType` is `Immut`, do not use the handles to visit the same /// KV twice. unsafe fn find_leaf_edges_spanning_range( self, range: R, ) -> LeafRange where Q: Ord, K: Borrow, R: RangeBounds, { match self.search_tree_for_bifurcation(&range) { Err(_) => LeafRange::none(), Ok(( node, lower_edge_idx, upper_edge_idx, mut lower_child_bound, mut upper_child_bound, )) => { let mut lower_edge = unsafe { Handle::new_edge(ptr::read(&node), lower_edge_idx) }; let mut upper_edge = unsafe { Handle::new_edge(node, upper_edge_idx) }; loop { match (lower_edge.force(), upper_edge.force()) { (Leaf(f), Leaf(b)) => return LeafRange { front: Some(f), back: Some(b) }, (Internal(f), Internal(b)) => { (lower_edge, lower_child_bound) = f.descend().find_lower_bound_edge(lower_child_bound); (upper_edge, upper_child_bound) = b.descend().find_upper_bound_edge(upper_child_bound); } _ => unreachable!("BTreeMap has different depths"), } } } } } } fn full_range( root1: NodeRef, root2: NodeRef, ) -> LazyLeafRange { LazyLeafRange { front: Some(LazyLeafHandle::Root(root1)), back: Some(LazyLeafHandle::Root(root2)), } } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { /// Finds the pair of leaf edges delimiting a specific range in a tree. /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. pub fn range_search(self, range: R) -> LeafRange, K, V> where Q: ?Sized + Ord, K: Borrow, R: RangeBounds, { // SAFETY: our borrow type is immutable. unsafe { self.find_leaf_edges_spanning_range(range) } } /// Finds the pair of leaf edges delimiting an entire tree. pub fn full_range(self) -> LazyLeafRange, K, V> { full_range(self, self) } } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { /// Splits a unique reference into a pair of leaf edges delimiting a specified range. /// The result are non-unique references allowing (some) mutation, which must be used /// carefully. /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. /// /// # Safety /// Do not use the duplicate handles to visit the same KV twice. pub fn range_search(self, range: R) -> LeafRange, K, V> where Q: ?Sized + Ord, K: Borrow, R: RangeBounds, { unsafe { self.find_leaf_edges_spanning_range(range) } } /// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree. /// The results are non-unique references allowing mutation (of values only), so must be used /// with care. pub fn full_range(self) -> LazyLeafRange, K, V> { // We duplicate the root NodeRef here -- we will never visit the same KV // twice, and never end up with overlapping value references. let self2 = unsafe { ptr::read(&self) }; full_range(self, self2) } } impl NodeRef { /// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree. /// The results are non-unique references allowing massively destructive mutation, so must be /// used with the utmost care. pub fn full_range(self) -> LazyLeafRange { // We duplicate the root NodeRef here -- we will never access it in a way // that overlaps references obtained from the root. let self2 = unsafe { ptr::read(&self) }; full_range(self, self2) } } impl Handle, marker::Edge> { /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV /// on the right side, which is either in the same leaf node or in an ancestor node. /// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node. pub fn next_kv( self, ) -> Result< Handle, marker::KV>, NodeRef, > { let mut edge = self.forget_node_type(); loop { edge = match edge.right_kv() { Ok(kv) => return Ok(kv), Err(last_edge) => match last_edge.into_node().ascend() { Ok(parent_edge) => parent_edge.forget_node_type(), Err(root) => return Err(root), }, } } } /// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV /// on the left side, which is either in the same leaf node or in an ancestor node. /// If the leaf edge is the first one in the tree, returns [`Result::Err`] with the root node. pub fn next_back_kv( self, ) -> Result< Handle, marker::KV>, NodeRef, > { let mut edge = self.forget_node_type(); loop { edge = match edge.left_kv() { Ok(kv) => return Ok(kv), Err(last_edge) => match last_edge.into_node().ascend() { Ok(parent_edge) => parent_edge.forget_node_type(), Err(root) => return Err(root), }, } } } } impl Handle, marker::Edge> { /// Given an internal edge handle, returns [`Result::Ok`] with a handle to the neighboring KV /// on the right side, which is either in the same internal node or in an ancestor node. /// If the internal edge is the last one in the tree, returns [`Result::Err`] with the root node. fn next_kv( self, ) -> Result< Handle, marker::KV>, NodeRef, > { let mut edge = self; loop { edge = match edge.right_kv() { Ok(internal_kv) => return Ok(internal_kv), Err(last_edge) => match last_edge.into_node().ascend() { Ok(parent_edge) => parent_edge, Err(root) => return Err(root), }, } } } } impl Handle, marker::Edge> { /// Given a leaf edge handle into a dying tree, returns the next leaf edge /// on the right side, and the key-value pair in between, if they exist. /// /// If the given edge is the last one in a leaf, this method deallocates /// the leaf, as well as any ancestor nodes whose last edge was reached. /// This implies that if no more key-value pair follows, the entire tree /// will have been deallocated and there is nothing left to return. /// /// # Safety /// - The given edge must not have been previously returned by counterpart /// `deallocating_next_back`. /// - The returned KV handle is only valid to access the key and value, /// and only valid until the next call to a `deallocating_` method. unsafe fn deallocating_next( self, alloc: A, ) -> Option<(Self, Handle, marker::KV>)> { let mut edge = self.forget_node_type(); loop { edge = match edge.right_kv() { Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)), Err(last_edge) => { match unsafe { last_edge.into_node().deallocate_and_ascend(alloc.clone()) } { Some(parent_edge) => parent_edge.forget_node_type(), None => return None, } } } } } /// Given a leaf edge handle into a dying tree, returns the next leaf edge /// on the left side, and the key-value pair in between, if they exist. /// /// If the given edge is the first one in a leaf, this method deallocates /// the leaf, as well as any ancestor nodes whose first edge was reached. /// This implies that if no more key-value pair follows, the entire tree /// will have been deallocated and there is nothing left to return. /// /// # Safety /// - The given edge must not have been previously returned by counterpart /// `deallocating_next`. /// - The returned KV handle is only valid to access the key and value, /// and only valid until the next call to a `deallocating_` method. unsafe fn deallocating_next_back( self, alloc: A, ) -> Option<(Self, Handle, marker::KV>)> { let mut edge = self.forget_node_type(); loop { edge = match edge.left_kv() { Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)), Err(last_edge) => { match unsafe { last_edge.into_node().deallocate_and_ascend(alloc.clone()) } { Some(parent_edge) => parent_edge.forget_node_type(), None => return None, } } } } } /// Deallocates a pile of nodes from the leaf up to the root. /// This is the only way to deallocate the remainder of a tree after /// `deallocating_next` and `deallocating_next_back` have been nibbling at /// both sides of the tree, and have hit the same edge. As it is intended /// only to be called when all keys and values have been returned, /// no cleanup is done on any of the keys or values. fn deallocating_end(self, alloc: A) { let mut edge = self.forget_node_type(); while let Some(parent_edge) = unsafe { edge.into_node().deallocate_and_ascend(alloc.clone()) } { edge = parent_edge.forget_node_type(); } } } impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { /// Moves the leaf edge handle to the next leaf edge and returns references to the /// key and value in between. /// /// # Safety /// There must be another KV in the direction travelled. unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { super::mem::replace(self, |leaf_edge| { let kv = leaf_edge.next_kv().ok().unwrap(); (kv.next_leaf_edge(), kv.into_kv()) }) } /// Moves the leaf edge handle to the previous leaf edge and returns references to the /// key and value in between. /// /// # Safety /// There must be another KV in the direction travelled. unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { super::mem::replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv().ok().unwrap(); (kv.next_back_leaf_edge(), kv.into_kv()) }) } } impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { /// Moves the leaf edge handle to the next leaf edge and returns references to the /// key and value in between. /// /// # Safety /// There must be another KV in the direction travelled. unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { let kv = super::mem::replace(self, |leaf_edge| { let kv = leaf_edge.next_kv().ok().unwrap(); (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv) }); // Doing this last is faster, according to benchmarks. kv.into_kv_valmut() } /// Moves the leaf edge handle to the previous leaf and returns references to the /// key and value in between. /// /// # Safety /// There must be another KV in the direction travelled. unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { let kv = super::mem::replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv().ok().unwrap(); (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv) }); // Doing this last is faster, according to benchmarks. kv.into_kv_valmut() } } impl Handle, marker::Edge> { /// Moves the leaf edge handle to the next leaf edge and returns the key and value /// in between, deallocating any node left behind while leaving the corresponding /// edge in its parent node dangling. /// /// # Safety /// - There must be another KV in the direction travelled. /// - That KV was not previously returned by counterpart /// `deallocating_next_back_unchecked` on any copy of the handles /// being used to traverse the tree. /// /// The only safe way to proceed with the updated handle is to compare it, drop it, /// or call this method or counterpart `deallocating_next_back_unchecked` again. unsafe fn deallocating_next_unchecked( &mut self, alloc: A, ) -> Handle, marker::KV> { super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next(alloc).unwrap() }) } /// Moves the leaf edge handle to the previous leaf edge and returns the key and value /// in between, deallocating any node left behind while leaving the corresponding /// edge in its parent node dangling. /// /// # Safety /// - There must be another KV in the direction travelled. /// - That leaf edge was not previously returned by counterpart /// `deallocating_next_unchecked` on any copy of the handles /// being used to traverse the tree. /// /// The only safe way to proceed with the updated handle is to compare it, drop it, /// or call this method or counterpart `deallocating_next_unchecked` again. unsafe fn deallocating_next_back_unchecked( &mut self, alloc: A, ) -> Handle, marker::KV> { super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next_back(alloc).unwrap() }) } } impl NodeRef { /// Returns the leftmost leaf edge in or underneath a node - in other words, the edge /// you need first when navigating forward (or last when navigating backward). #[inline] pub fn first_leaf_edge(self) -> Handle, marker::Edge> { let mut node = self; loop { match node.force() { Leaf(leaf) => return leaf.first_edge(), Internal(internal) => node = internal.first_edge().descend(), } } } /// Returns the rightmost leaf edge in or underneath a node - in other words, the edge /// you need last when navigating forward (or first when navigating backward). #[inline] pub fn last_leaf_edge(self) -> Handle, marker::Edge> { let mut node = self; loop { match node.force() { Leaf(leaf) => return leaf.last_edge(), Internal(internal) => node = internal.last_edge().descend(), } } } } pub enum Position { Leaf(NodeRef), Internal(NodeRef), InternalKV(Handle, marker::KV>), } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { /// Visits leaf nodes and internal KVs in order of ascending keys, and also /// visits internal nodes as a whole in a depth first order, meaning that /// internal nodes precede their individual KVs and their child nodes. pub fn visit_nodes_in_order(self, mut visit: F) where F: FnMut(Position, K, V>), { match self.force() { Leaf(leaf) => visit(Position::Leaf(leaf)), Internal(internal) => { visit(Position::Internal(internal)); let mut edge = internal.first_edge(); loop { edge = match edge.descend().force() { Leaf(leaf) => { visit(Position::Leaf(leaf)); match edge.next_kv() { Ok(kv) => { visit(Position::InternalKV(kv)); kv.right_edge() } Err(_) => return, } } Internal(internal) => { visit(Position::Internal(internal)); internal.first_edge() } } } } } } /// Calculates the number of elements in a (sub)tree. pub fn calc_length(self) -> usize { let mut result = 0; self.visit_nodes_in_order(|pos| match pos { Position::Leaf(node) => result += node.len(), Position::Internal(node) => result += node.len(), Position::InternalKV(_) => (), }); result } } impl Handle, marker::KV> { /// Returns the leaf edge closest to a KV for forward navigation. pub fn next_leaf_edge(self) -> Handle, marker::Edge> { match self.force() { Leaf(leaf_kv) => leaf_kv.right_edge(), Internal(internal_kv) => { let next_internal_edge = internal_kv.right_edge(); next_internal_edge.descend().first_leaf_edge() } } } /// Returns the leaf edge closest to a KV for backward navigation. pub fn next_back_leaf_edge( self, ) -> Handle, marker::Edge> { match self.force() { Leaf(leaf_kv) => leaf_kv.left_edge(), Internal(internal_kv) => { let next_internal_edge = internal_kv.left_edge(); next_internal_edge.descend().last_leaf_edge() } } } } impl NodeRef { /// Returns the leaf edge corresponding to the first point at which the /// given bound is true. pub fn lower_bound( self, mut bound: SearchBound<&Q>, ) -> Handle, marker::Edge> where Q: Ord, K: Borrow, { let mut node = self; loop { let (edge, new_bound) = node.find_lower_bound_edge(bound); match edge.force() { Leaf(edge) => return edge, Internal(edge) => { node = edge.descend(); bound = new_bound; } } } } /// Returns the leaf edge corresponding to the last point at which the /// given bound is true. pub fn upper_bound( self, mut bound: SearchBound<&Q>, ) -> Handle, marker::Edge> where Q: Ord, K: Borrow, { let mut node = self; loop { let (edge, new_bound) = node.find_upper_bound_edge(bound); match edge.force() { Leaf(edge) => return edge, Internal(edge) => { node = edge.descend(); bound = new_bound; } } } } }