Skip to content

Commit

Permalink
BTreeMap: bring back the key slice for immutable lookup
Browse files Browse the repository at this point in the history
  • Loading branch information
ssomers committed Jan 20, 2021
1 parent a4cbb44 commit c0e8075
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 43 deletions.
43 changes: 18 additions & 25 deletions library/alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -238,13 +238,12 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
/// 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`.
Expand Down Expand Up @@ -344,26 +343,6 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
}
}

impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, 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<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
/// 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
Expand Down Expand Up @@ -421,6 +400,14 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, 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<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
Expand Down Expand Up @@ -987,7 +974,11 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marke

impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Immut<'a>, 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)
}
}

Expand All @@ -997,6 +988,7 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, 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() }
}
Expand All @@ -1010,6 +1002,7 @@ impl<'a, K, V, NodeType> Handle<NodeRef<marker::ValMut<'a>, K, V, NodeType>, mar

impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, 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 {
Expand Down
12 changes: 1 addition & 11 deletions library/alloc/src/collections/btree/node/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,17 +29,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, 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) => {
Expand Down
11 changes: 4 additions & 7 deletions library/alloc/src/collections/btree/search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,18 +71,15 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
Q: Ord,
K: Borrow<Q>,
{
// 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())
}
}

0 comments on commit c0e8075

Please sign in to comment.