diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index f8729c33c6713..92f02fb60168b 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -5,7 +5,6 @@ use core::hash::{Hash, Hasher}; use core::iter::{FromIterator, FusedIterator, Peekable}; use core::marker::PhantomData; use core::mem::{self, ManuallyDrop}; -use core::ops::Bound::{Excluded, Included, Unbounded}; use core::ops::{Index, RangeBounds}; use core::ptr; @@ -408,8 +407,8 @@ impl fmt::Debug for Range<'_, K, V> { /// [`range_mut`]: BTreeMap::range_mut #[stable(feature = "btree_range", since = "1.17.0")] pub struct RangeMut<'a, K: 'a, V: 'a> { - front: Option, K, V, marker::Leaf>, marker::Edge>>, - back: Option, K, V, marker::Leaf>, marker::Edge>>, + front: Option, K, V, marker::Leaf>, marker::Edge>>, + back: Option, K, V, marker::Leaf>, marker::Edge>>, // Be invariant in `K` and `V` _marker: PhantomData<&'a mut (K, V)>, @@ -999,7 +998,7 @@ impl BTreeMap { R: RangeBounds, { if let Some(root) = &self.root { - let (f, b) = range_search(root.node_as_ref(), range); + let (f, b) = root.node_as_ref().range_search(range); Range { front: Some(f), back: Some(b) } } else { @@ -1045,7 +1044,7 @@ impl BTreeMap { R: RangeBounds, { if let Some(root) = &mut self.root { - let (f, b) = range_search(root.node_as_mut(), range); + let (f, b) = root.node_as_valmut().range_search(range); RangeMut { front: Some(f), back: Some(b), _marker: PhantomData } } else { @@ -1478,7 +1477,7 @@ impl IntoIterator for BTreeMap { fn into_iter(self) -> IntoIter { let mut me = ManuallyDrop::new(self); if let Some(root) = me.root.take() { - let (f, b) = full_range_search(root.into_ref()); + let (f, b) = root.into_ref().full_range(); IntoIter { front: Some(f), back: Some(b), length: me.length } } else { @@ -1942,7 +1941,7 @@ impl<'a, K, V> RangeMut<'a, K, V> { self.front == self.back } - unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) { + unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { unsafe { unwrap_unchecked(self.front.as_mut()).next_unchecked() } } } @@ -1963,7 +1962,7 @@ impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> { impl FusedIterator for RangeMut<'_, K, V> {} impl<'a, K, V> RangeMut<'a, K, V> { - unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) { + unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { unsafe { unwrap_unchecked(self.back.as_mut()).next_back_unchecked() } } } @@ -2073,119 +2072,6 @@ where } } -/// Finds the leaf edges delimiting a specified range in or underneath a node. -fn range_search>( - root: NodeRef, - range: R, -) -> ( - Handle, marker::Edge>, - Handle, marker::Edge>, -) -where - Q: Ord, - K: Borrow, -{ - match (range.start_bound(), range.end_bound()) { - (Excluded(s), Excluded(e)) if s == e => { - panic!("range start and end are equal and excluded in BTreeMap") - } - (Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => { - panic!("range start is greater than range end in BTreeMap") - } - _ => {} - }; - - // We duplicate the root NodeRef here -- we will never access it in a way - // that overlaps references obtained from the root. - let mut min_node = unsafe { ptr::read(&root) }; - let mut max_node = root; - let mut min_found = false; - let mut max_found = false; - - loop { - let front = match (min_found, range.start_bound()) { - (false, Included(key)) => match search::search_node(min_node, key) { - Found(kv) => { - min_found = true; - kv.left_edge() - } - GoDown(edge) => edge, - }, - (false, Excluded(key)) => match search::search_node(min_node, key) { - Found(kv) => { - min_found = true; - kv.right_edge() - } - GoDown(edge) => edge, - }, - (true, Included(_)) => min_node.last_edge(), - (true, Excluded(_)) => min_node.first_edge(), - (_, Unbounded) => min_node.first_edge(), - }; - - let back = match (max_found, range.end_bound()) { - (false, Included(key)) => match search::search_node(max_node, key) { - Found(kv) => { - max_found = true; - kv.right_edge() - } - GoDown(edge) => edge, - }, - (false, Excluded(key)) => match search::search_node(max_node, key) { - Found(kv) => { - max_found = true; - kv.left_edge() - } - GoDown(edge) => edge, - }, - (true, Included(_)) => max_node.first_edge(), - (true, Excluded(_)) => max_node.last_edge(), - (_, Unbounded) => max_node.last_edge(), - }; - - if front.partial_cmp(&back) == Some(Ordering::Greater) { - panic!("Ord is ill-defined in BTreeMap range"); - } - match (front.force(), back.force()) { - (Leaf(f), Leaf(b)) => { - return (f, b); - } - (Internal(min_int), Internal(max_int)) => { - min_node = min_int.descend(); - max_node = max_int.descend(); - } - _ => unreachable!("BTreeMap has different depths"), - }; - } -} - -/// Equivalent to `range_search(k, v, ..)` without the `Ord` bound. -fn full_range_search( - root: NodeRef, -) -> ( - Handle, marker::Edge>, - Handle, marker::Edge>, -) { - // We duplicate the root NodeRef here -- we will never access it in a way - // that overlaps references obtained from the root. - let mut min_node = unsafe { ptr::read(&root) }; - let mut max_node = root; - loop { - let front = min_node.first_edge(); - let back = max_node.last_edge(); - match (front.force(), back.force()) { - (Leaf(f), Leaf(b)) => { - return (f, b); - } - (Internal(min_int), Internal(max_int)) => { - min_node = min_int.descend(); - max_node = max_int.descend(); - } - _ => unreachable!("BTreeMap has different depths"), - }; - } -} - impl BTreeMap { /// Gets an iterator over the entries of the map, sorted by key. /// @@ -2211,7 +2097,7 @@ impl BTreeMap { #[stable(feature = "rust1", since = "1.0.0")] pub fn iter(&self) -> Iter<'_, K, V> { if let Some(root) = &self.root { - let (f, b) = full_range_search(root.node_as_ref()); + let (f, b) = root.node_as_ref().full_range(); Iter { range: Range { front: Some(f), back: Some(b) }, length: self.length } } else { @@ -2243,7 +2129,7 @@ impl BTreeMap { #[stable(feature = "rust1", since = "1.0.0")] pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { if let Some(root) = &mut self.root { - let (f, b) = full_range_search(root.node_as_mut()); + let (f, b) = root.node_as_valmut().full_range(); IterMut { range: RangeMut { front: Some(f), back: Some(b), _marker: PhantomData }, @@ -2826,7 +2712,7 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter if stole_from_left && at_leaf { // SAFETY: This is safe since we just added an element to our node. unsafe { - pos.next_unchecked(); + pos.move_next_unchecked(); } } break; diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index b7b66ac7ceccd..376060b3143de 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -1,10 +1,210 @@ +use core::borrow::Borrow; +use core::cmp::Ordering; use core::intrinsics; use core::mem; +use core::ops::Bound::{Excluded, Included, Unbounded}; +use core::ops::RangeBounds; use core::ptr; use super::node::{marker, ForceResult::*, Handle, NodeRef}; +use super::search::{self, SearchResult}; use super::unwrap_unchecked; +/// Finds the leaf edges delimiting a specified range in or underneath a node. +fn range_search( + root1: NodeRef, + root2: NodeRef, + range: R, +) -> ( + Handle, marker::Edge>, + Handle, marker::Edge>, +) +where + Q: ?Sized + Ord, + K: Borrow, + R: RangeBounds, +{ + match (range.start_bound(), range.end_bound()) { + (Excluded(s), Excluded(e)) if s == e => { + panic!("range start and end are equal and excluded in BTreeMap") + } + (Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => { + panic!("range start is greater than range end in BTreeMap") + } + _ => {} + }; + + let mut min_node = root1; + let mut max_node = root2; + let mut min_found = false; + let mut max_found = false; + + loop { + let front = match (min_found, range.start_bound()) { + (false, Included(key)) => match search::search_node(min_node, key) { + SearchResult::Found(kv) => { + min_found = true; + kv.left_edge() + } + SearchResult::GoDown(edge) => edge, + }, + (false, Excluded(key)) => match search::search_node(min_node, key) { + SearchResult::Found(kv) => { + min_found = true; + kv.right_edge() + } + SearchResult::GoDown(edge) => edge, + }, + (true, Included(_)) => min_node.last_edge(), + (true, Excluded(_)) => min_node.first_edge(), + (_, Unbounded) => min_node.first_edge(), + }; + + let back = match (max_found, range.end_bound()) { + (false, Included(key)) => match search::search_node(max_node, key) { + SearchResult::Found(kv) => { + max_found = true; + kv.right_edge() + } + SearchResult::GoDown(edge) => edge, + }, + (false, Excluded(key)) => match search::search_node(max_node, key) { + SearchResult::Found(kv) => { + max_found = true; + kv.left_edge() + } + SearchResult::GoDown(edge) => edge, + }, + (true, Included(_)) => max_node.first_edge(), + (true, Excluded(_)) => max_node.last_edge(), + (_, Unbounded) => max_node.last_edge(), + }; + + if front.partial_cmp(&back) == Some(Ordering::Greater) { + panic!("Ord is ill-defined in BTreeMap range"); + } + match (front.force(), back.force()) { + (Leaf(f), Leaf(b)) => { + return (f, b); + } + (Internal(min_int), Internal(max_int)) => { + min_node = min_int.descend(); + max_node = max_int.descend(); + } + _ => unreachable!("BTreeMap has different depths"), + }; + } +} + +/// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound. +fn full_range( + root1: NodeRef, + root2: NodeRef, +) -> ( + Handle, marker::Edge>, + Handle, marker::Edge>, +) { + let mut min_node = root1; + let mut max_node = root2; + loop { + let front = min_node.first_edge(); + let back = max_node.last_edge(); + match (front.force(), back.force()) { + (Leaf(f), Leaf(b)) => { + return (f, b); + } + (Internal(min_int), Internal(max_int)) => { + min_node = min_int.descend(); + max_node = max_int.descend(); + } + _ => unreachable!("BTreeMap has different depths"), + }; + } +} + +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { + /// Creates a pair of leaf edges delimiting a specified range in or underneath a node. + pub fn range_search( + self, + range: R, + ) -> ( + Handle, K, V, marker::Leaf>, marker::Edge>, + Handle, K, V, marker::Leaf>, marker::Edge>, + ) + where + Q: ?Sized + Ord, + K: Borrow, + R: RangeBounds, + { + range_search(self, self, range) + } + + /// Returns (self.first_leaf_edge(), self.last_leaf_edge()), but more efficiently. + pub fn full_range( + self, + ) -> ( + Handle, K, V, marker::Leaf>, marker::Edge>, + Handle, K, V, marker::Leaf>, marker::Edge>, + ) { + 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. + pub fn range_search( + self, + range: R, + ) -> ( + Handle, K, V, marker::Leaf>, marker::Edge>, + Handle, K, V, marker::Leaf>, marker::Edge>, + ) + where + Q: ?Sized + Ord, + K: Borrow, + R: RangeBounds, + { + // 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) }; + range_search(self, self2, 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, + ) -> ( + Handle, K, V, marker::Leaf>, marker::Edge>, + Handle, K, V, marker::Leaf>, marker::Edge>, + ) { + // 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, + ) -> ( + Handle, marker::Edge>, + Handle, marker::Edge>, + ) { + // 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. @@ -75,12 +275,13 @@ impl Handle, marke macro_rules! def_next_kv_uncheched_dealloc { { unsafe fn $name:ident : $adjacent_kv:ident } => { /// Given a leaf edge handle into an owned tree, returns a handle to the next KV, - /// while deallocating any node left behind. - /// Unsafe for two reasons: - /// - The caller must ensure that the leaf edge is not the last one in the tree. - /// - The node pointed at by the given handle, and its ancestors, may be deallocated, - /// while the reference to those nodes in the surviving ancestors is left dangling; - /// thus using the returned handle to navigate further is dangerous. + /// while deallocating any node left behind yet leaving the corresponding edge + /// in its parent node dangling. + /// + /// # Safety + /// - The leaf edge must not be the last one in the direction travelled. + /// - The node carrying the next KV returned must not have been deallocated by a + /// previous call on any handle obtained for this tree. unsafe fn $name ( leaf_edge: Handle, marker::Edge>, ) -> Handle, marker::KV> { @@ -103,6 +304,15 @@ macro_rules! def_next_kv_uncheched_dealloc { def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv} def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv} +/// This replaces the value behind the `v` unique reference by calling the +/// relevant function. +/// +/// If a panic occurs in the `change` closure, the entire process will be aborted. +#[inline] +fn take_mut(v: &mut T, change: impl FnOnce(T) -> T) { + replace(v, |value| (change(value), ())) +} + /// This replaces the value behind the `v` unique reference by calling the /// relevant function, and returns a result obtained along the way. /// @@ -128,7 +338,9 @@ fn replace(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R { 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. - /// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree. + /// + /// # Safety + /// There must be another KV in the direction travelled. pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { replace(self, |leaf_edge| { let kv = leaf_edge.next_kv(); @@ -139,7 +351,9 @@ impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Ed /// Moves the leaf edge handle to the previous leaf edge and returns references to the /// key and value in between. - /// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree. + /// + /// # Safety + /// There must be another KV in the direction travelled. pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv(); @@ -149,53 +363,69 @@ impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Ed } } -impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { +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. - /// Unsafe for two reasons: - /// - The caller must ensure that the leaf edge is not the last one in the tree. - /// - Using the updated handle may well invalidate the returned references. - pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) { + /// The returned references might be invalidated when the updated handle is used again. + /// + /// # Safety + /// There must be another KV in the direction travelled. + pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { let kv = replace(self, |leaf_edge| { let kv = leaf_edge.next_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) }; (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv) }); // Doing the descend (and perhaps another move) invalidates the references - // returned by `into_kv_mut`, so we have to do this last. - kv.into_kv_mut() + // returned by `into_kv_valmut`, so we have to do this last. + kv.into_kv_valmut() } /// Moves the leaf edge handle to the previous leaf and returns references to the /// key and value in between. - /// Unsafe for two reasons: - /// - The caller must ensure that the leaf edge is not the first one in the tree. - /// - Using the updated handle may well invalidate the returned references. - pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) { + /// The returned references might be invalidated when the updated handle is used again. + /// + /// # Safety + /// There must be another KV in the direction travelled. + pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { let kv = replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) }; (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv) }); // Doing the descend (and perhaps another move) invalidates the references - // returned by `into_kv_mut`, so we have to do this last. - kv.into_kv_mut() + // returned by `into_kv_valmut`, so we have to do this last. + kv.into_kv_valmut() + } +} + +impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { + /// Moves the leaf edge handle to the next leaf edge. + /// + /// # Safety + /// There must be another KV in the direction travelled. + pub unsafe fn move_next_unchecked(&mut self) { + take_mut(self, |leaf_edge| { + let kv = leaf_edge.next_kv(); + let kv = unsafe { unwrap_unchecked(kv.ok()) }; + kv.next_leaf_edge() + }) } } impl Handle, marker::Edge> { /// Moves the leaf edge handle to the next leaf edge and returns the key and value - /// in between, while deallocating any node left behind. - /// Unsafe for two reasons: - /// - The caller must ensure that the leaf edge is not the last one in the tree - /// and is not a handle previously resulting from counterpart `next_back_unchecked`. - /// - Further use of the updated leaf edge handle is very dangerous. In particular, - /// if the leaf edge is the last edge of a node, that node and possibly ancestors - /// will be deallocated, while the reference to those nodes in the surviving ancestor - /// is left dangling. - /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// call this method again subject to both preconditions listed in the first point, - /// or call counterpart `next_back_unchecked` subject to its preconditions. + /// 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 `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, + /// call this method again subject to its safety conditions, or call counterpart + /// `next_back_unchecked` subject to its safety conditions. pub unsafe fn next_unchecked(&mut self) -> (K, V) { replace(self, |leaf_edge| { let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) }; @@ -205,18 +435,18 @@ impl Handle, marker::Edge> { }) } - /// Moves the leaf edge handle to the previous leaf edge and returns the key - /// and value in between, while deallocating any node left behind. - /// Unsafe for two reasons: - /// - The caller must ensure that the leaf edge is not the first one in the tree - /// and is not a handle previously resulting from counterpart `next_unchecked`. - /// - Further use of the updated leaf edge handle is very dangerous. In particular, - /// if the leaf edge is the first edge of a node, that node and possibly ancestors - /// will be deallocated, while the reference to those nodes in the surviving ancestor - /// is left dangling. - /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// call this method again subject to both preconditions listed in the first point, - /// or call counterpart `next_unchecked` subject to its preconditions. + /// 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 `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, + /// call this method again subject to its safety conditions, or call counterpart + /// `next_unchecked` subject to its safety conditions. pub unsafe fn next_back_unchecked(&mut self) -> (K, V) { replace(self, |leaf_edge| { let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) }; diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index acc2ae73572ba..772bdf357de78 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -186,6 +186,15 @@ impl Root { } } + pub fn node_as_valmut(&mut self) -> NodeRef, K, V, marker::LeafOrInternal> { + NodeRef { + height: self.height, + node: self.node.as_ptr(), + root: ptr::null(), + _marker: PhantomData, + } + } + pub fn into_ref(self) -> NodeRef { NodeRef { height: self.height, @@ -253,9 +262,12 @@ impl Root { /// A reference to a node. /// /// This type has a number of parameters that controls how it acts: -/// - `BorrowType`: This can be `Immut<'a>` or `Mut<'a>` for some `'a` or `Owned`. +/// - `BorrowType`: This can be `Immut<'a>`, `Mut<'a>` or `ValMut<'a>' for some `'a` +/// or `Owned`. /// When this is `Immut<'a>`, the `NodeRef` acts roughly like `&'a Node`, /// when this is `Mut<'a>`, the `NodeRef` acts roughly like `&'a mut Node`, +/// when this is `ValMut<'a>`, the `NodeRef` acts as immutable with respect +/// to keys and tree structure, but allows mutable references to values, /// and when this is `Owned`, the `NodeRef` acts roughly like `Box`. /// - `K` and `V`: These control what types of things are stored in the nodes. /// - `Type`: This can be `Leaf`, `Internal`, or `LeafOrInternal`. When this is @@ -282,6 +294,7 @@ unsafe impl Sync for NodeRef Send for NodeRef, K, V, Type> {} unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef, K, V, Type> {} +unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef, K, V, Type> {} unsafe impl Send for NodeRef {} impl NodeRef { @@ -515,6 +528,22 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { } } +impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { + /// Same as the marker::Mut method, but far more dangerous because ValMut-based iterators: + /// - have front and back handles often refering to the same node, so `self` is not unique; + /// - hand out mutable references to parts of these slices to the public. + fn into_slices_mut(self) -> (&'a [K], &'a mut [V]) { + let len = self.len(); + let leaf = self.node.as_ptr(); + // SAFETY: The keys and values of a node must always be initialized up to length. + let keys = unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&(*leaf).keys), len) }; + let vals = unsafe { + slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).vals), len) + }; + (keys, vals) + } +} + impl<'a, K, V> NodeRef, K, V, marker::Leaf> { /// Adds a key/value pair to the end of the node. pub fn push(&mut self, key: K, val: V) { @@ -1053,11 +1082,13 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType> let vals = self.node.into_val_slice_mut(); unsafe { vals.get_unchecked_mut(self.idx) } } +} - pub fn into_kv_mut(self) -> (&'a mut K, &'a mut V) { +impl<'a, K, V, NodeType> Handle, K, V, NodeType>, marker::KV> { + pub fn into_kv_valmut(self) -> (&'a K, &'a mut V) { unsafe { let (keys, vals) = self.node.into_slices_mut(); - (keys.get_unchecked_mut(self.idx), vals.get_unchecked_mut(self.idx)) + (keys.get_unchecked(self.idx), vals.get_unchecked_mut(self.idx)) } } } @@ -1558,6 +1589,7 @@ pub mod marker { pub enum Owned {} pub struct Immut<'a>(PhantomData<&'a ()>); pub struct Mut<'a>(PhantomData<&'a mut ()>); + pub struct ValMut<'a>(PhantomData<&'a mut ()>); pub enum KV {} pub enum Edge {}