diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 8ab3f58c1adba..57fad4217b6c6 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -238,13 +238,12 @@ impl NodeRef { /// such restrictions: /// - For each type parameter, we can only define a method either generically /// or for one particular type. For example, we cannot define a method like -/// `key_at` generically for all `BorrowType`, because we want it to return -/// `&'a K` for most choices of `BorrowType`, but plain `K` for `Owned`. -/// We cannot define `key_at` once for all types that carry a lifetime. +/// `into_kv` generically for all `BorrowType`, or once for all types that +/// carry a lifetime, because we want it to return `&'a` references. /// Therefore, we define it only for the least powerful type `Immut<'a>`. /// - We cannot get implicit coercion from say `Mut<'a>` to `Immut<'a>`. /// Therefore, we have to explicitly call `reborrow` on a more powerfull -/// `NodeRef` in order to reach a method like `key_at`. +/// `NodeRef` in order to reach a method like `into_kv`. /// /// All methods on `NodeRef` that return some kind of reference, either: /// - Take `self` by value, and return the lifetime carried by `BorrowType`. @@ -344,26 +343,6 @@ impl NodeRef { } } -impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { - /// Exposes one of the keys stored in the node. - /// - /// # Safety - /// The node has more than `idx` initialized elements. - pub unsafe fn key_at(self, idx: usize) -> &'a K { - debug_assert!(idx < self.len()); - unsafe { self.into_leaf().keys.get_unchecked(idx).assume_init_ref() } - } - - /// Exposes one of the values stored in the node. - /// - /// # Safety - /// The node has more than `idx` initialized elements. - unsafe fn val_at(self, idx: usize) -> &'a V { - debug_assert!(idx < self.len()); - unsafe { self.into_leaf().vals.get_unchecked(idx).assume_init_ref() } - } -} - impl NodeRef { /// Finds the parent of the current node. Returns `Ok(handle)` if the current /// node actually has a parent, where `handle` points to the edge of the parent @@ -421,6 +400,14 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { // SAFETY: there can be no mutable references into this tree borrowed as `Immut`. unsafe { &*ptr } } + + /// Borrows a view into the keys stored in the node. + pub fn keys(&self) -> &[K] { + let leaf = self.into_leaf(); + unsafe { + MaybeUninit::slice_assume_init_ref(leaf.keys.get_unchecked(..usize::from(leaf.len))) + } + } } impl NodeRef { @@ -987,7 +974,11 @@ impl Handle, marke impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType>, marker::KV> { pub fn into_kv(self) -> (&'a K, &'a V) { - (unsafe { self.node.key_at(self.idx) }, unsafe { self.node.val_at(self.idx) }) + debug_assert!(self.idx < self.node.len()); + let leaf = self.node.into_leaf(); + let k = unsafe { leaf.keys.get_unchecked(self.idx).assume_init_ref() }; + let v = unsafe { leaf.vals.get_unchecked(self.idx).assume_init_ref() }; + (k, v) } } @@ -997,6 +988,7 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType> } pub fn into_val_mut(self) -> &'a mut V { + debug_assert!(self.idx < self.node.len()); let leaf = self.node.into_leaf_mut(); unsafe { leaf.vals.get_unchecked_mut(self.idx).assume_init_mut() } } @@ -1010,6 +1002,7 @@ impl<'a, K, V, NodeType> Handle, K, V, NodeType>, mar impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType>, marker::KV> { pub fn kv_mut(&mut self) -> (&mut K, &mut V) { + debug_assert!(self.idx < self.node.len()); // We cannot call separate key and value methods, because calling the second one // invalidates the reference returned by the first. unsafe { diff --git a/library/alloc/src/collections/btree/node/tests.rs b/library/alloc/src/collections/btree/node/tests.rs index 48ce9f2bd89c8..11433cd845bee 100644 --- a/library/alloc/src/collections/btree/node/tests.rs +++ b/library/alloc/src/collections/btree/node/tests.rs @@ -29,17 +29,7 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> navigate::Position::Leaf(leaf) => { let depth = self.height(); let indent = " ".repeat(depth); - result += &format!("\n{}", indent); - if leaf.len() == 0 { - result += "(empty node)"; - } else { - for idx in 0..leaf.len() { - if idx > 0 { - result += ", "; - } - result += &format!("{:?}", unsafe { leaf.key_at(idx) }); - } - } + result += &format!("\n{}{:?}", indent, leaf.keys()); } navigate::Position::Internal(_) => {} navigate::Position::InternalKV(kv) => { diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs index f62eae3f4d5a2..d8bebed758f1a 100644 --- a/library/alloc/src/collections/btree/search.rs +++ b/library/alloc/src/collections/btree/search.rs @@ -71,18 +71,15 @@ impl NodeRef { Q: Ord, K: Borrow, { - // This function is defined over all borrow types (immutable, mutable, owned). - // Using `keys_at()` is fine here even if BorrowType is mutable, as all we return - // is an index -- not a reference. - let len = self.len(); - for i in 0..len { - let k = unsafe { self.reborrow().key_at(i) }; + let node = self.reborrow(); + let keys = node.keys(); + for (i, k) in keys.iter().enumerate() { match key.cmp(k.borrow()) { Ordering::Greater => {} Ordering::Equal => return IndexResult::KV(i), Ordering::Less => return IndexResult::Edge(i), } } - IndexResult::Edge(len) + IndexResult::Edge(keys.len()) } }