Rollup merge of #67686 - ssomers:keys_start_slasher, r=Mark-Simulacrum

Simplify NodeHeader by avoiding slices in BTreeMaps with shared roots

Simplify a complicated piece of code that creates slices of keys in node leaves.
This commit is contained in:
Mazdak Farrokhzad 2020-01-21 19:42:17 +01:00 committed by GitHub
commit d532a04a1c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 24 additions and 65 deletions

View File

@ -1968,7 +1968,7 @@ fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeBounds<Q>>(
(i, false) => i,
},
(_, Unbounded) => 0,
(true, Included(_)) => min_node.keys().len(),
(true, Included(_)) => min_node.len(),
(true, Excluded(_)) => 0,
};
@ -1987,9 +1987,9 @@ fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeBounds<Q>>(
}
(i, false) => i,
},
(_, Unbounded) => max_node.keys().len(),
(_, Unbounded) => max_node.len(),
(true, Included(_)) => 0,
(true, Excluded(_)) => max_node.keys().len(),
(true, Excluded(_)) => max_node.len(),
};
if !diverged {

View File

@ -54,10 +54,8 @@
/// `NodeHeader` because we do not want unnecessary padding between `len` and the keys.
/// Crucially, `NodeHeader` can be safely transmuted to different K and V. (This is exploited
/// by `as_header`.)
/// See `into_key_slice` for an explanation of K2. K2 cannot be safely transmuted around
/// because the size of `NodeHeader` depends on its alignment!
#[repr(C)]
struct NodeHeader<K, V, K2 = ()> {
struct NodeHeader<K, V> {
/// We use `*const` as opposed to `*mut` so as to be covariant in `K` and `V`.
/// This either points to an actual node or is null.
parent: *const InternalNode<K, V>,
@ -72,9 +70,6 @@ struct NodeHeader<K, V, K2 = ()> {
/// This next to `parent_idx` to encourage the compiler to join `len` and
/// `parent_idx` into the same 32-bit word, reducing space overhead.
len: u16,
/// See `into_key_slice`.
keys_start: [K2; 0],
}
#[repr(C)]
struct LeafNode<K, V> {
@ -128,7 +123,7 @@ unsafe impl Sync for NodeHeader<(), ()> {}
// We use just a header in order to save space, since no operation on an empty tree will
// ever take a pointer past the first key.
static EMPTY_ROOT_NODE: NodeHeader<(), ()> =
NodeHeader { parent: ptr::null(), parent_idx: MaybeUninit::uninit(), len: 0, keys_start: [] };
NodeHeader { parent: ptr::null(), parent_idx: MaybeUninit::uninit(), len: 0 };
/// The underlying representation of internal nodes. As with `LeafNode`s, these should be hidden
/// behind `BoxedNode`s to prevent dropping uninitialized keys and values. Any pointer to an
@ -390,14 +385,13 @@ pub fn is_shared_root(&self) -> bool {
}
/// Borrows a view into the keys stored in the node.
/// Works on all possible nodes, including the shared root.
pub fn keys(&self) -> &[K] {
/// The caller must ensure that the node is not the shared root.
pub unsafe fn keys(&self) -> &[K] {
self.reborrow().into_key_slice()
}
/// Borrows a view into the values stored in the node.
/// The caller must ensure that the node is not the shared root.
/// This function is not public, so doesn't have to support shared roots like `keys` does.
fn vals(&self) -> &[V] {
self.reborrow().into_val_slice()
}
@ -515,7 +509,6 @@ fn as_leaf_mut(&mut self) -> *mut LeafNode<K, V> {
}
/// The caller must ensure that the node is not the shared root.
/// This function is not public, so doesn't have to support shared roots like `keys` does.
fn keys_mut(&mut self) -> &mut [K] {
unsafe { self.reborrow_mut().into_key_slice_mut() }
}
@ -527,48 +520,11 @@ fn vals_mut(&mut self) -> &mut [V] {
}
impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
fn into_key_slice(self) -> &'a [K] {
// We have to be careful here because we might be pointing to the shared root.
// In that case, we must not create an `&LeafNode`. We could just return
// an empty slice whenever the length is 0 (this includes the shared root),
// but we want to avoid that run-time check.
// Instead, we create a slice pointing into the node whenever possible.
// We can sometimes do this even for the shared root, as the slice will be
// empty and `NodeHeader` contains an empty `keys_start` array.
// We cannot *always* do this because:
// - `keys_start` is not correctly typed because we want `NodeHeader`'s size to
// not depend on the alignment of `K` (needed because `as_header` should be safe).
// For this reason, `NodeHeader` has this `K2` parameter (that's usually `()`
// and hence just adds a size-0-align-1 field, not affecting layout).
// If the correctly typed header is more highly aligned than the allocated header,
// we cannot transmute safely.
// - Even if we can transmute, the offset of a correctly typed `keys_start` might
// be different and outside the bounds of the allocated header!
// So we do an alignment check and a size check first, that will be evaluated
// at compile-time, and only do any run-time check in the rare case that
// the compile-time checks signal danger.
if (mem::align_of::<NodeHeader<K, V, K>>() > mem::align_of::<NodeHeader<K, V>>()
|| mem::size_of::<NodeHeader<K, V, K>>() != mem::size_of::<NodeHeader<K, V>>())
&& self.is_shared_root()
{
&[]
} else {
// If we are a `LeafNode<K, V>`, we can always transmute to
// `NodeHeader<K, V, K>` and `keys_start` always has the same offset
// as the actual `keys`.
// Thanks to the checks above, we know that we can transmute to
// `NodeHeader<K, V, K>` and that `keys_start` will be
// in-bounds of some allocation even if this is the shared root!
// (We might be one-past-the-end, but that is allowed by LLVM.)
// Thus we can use `NodeHeader<K, V, K>`
// to compute the pointer where the keys start.
// This entire hack will become unnecessary once
// <https://github.com/rust-lang/rfcs/pull/2582> lands, then we can just take a raw
// pointer to the `keys` field of `*const InternalNode<K, V>`.
let header = self.as_header() as *const _ as *const NodeHeader<K, V, K>;
let keys = unsafe { &(*header).keys_start as *const _ as *const K };
unsafe { slice::from_raw_parts(keys, self.len()) }
}
/// The caller must ensure that the node is not the shared root.
unsafe fn into_key_slice(self) -> &'a [K] {
debug_assert!(!self.is_shared_root());
// We cannot be the shared root, so `as_leaf` is okay.
slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().keys), self.len())
}
/// The caller must ensure that the node is not the shared root.
@ -578,9 +534,10 @@ fn into_val_slice(self) -> &'a [V] {
unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&self.as_leaf().vals), self.len()) }
}
/// The caller must ensure that the node is not the shared root.
fn into_slices(self) -> (&'a [K], &'a [V]) {
let k = unsafe { ptr::read(&self) };
(k.into_key_slice(), self.into_val_slice())
(unsafe { k.into_key_slice() }, self.into_val_slice())
}
}

View File

@ -61,16 +61,18 @@ pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
{
// This function is defined over all borrow types (immutable, mutable, owned),
// and may be called on the shared root in each case.
// Crucially, we use `keys()` here, i.e., we work with immutable data.
// `keys_mut()` does not support the shared root, so we cannot use it.
// Using `keys()` is fine here even if BorrowType is mutable, as all we return
// is an index -- not a reference.
for (i, k) in node.keys().iter().enumerate() {
match key.cmp(k.borrow()) {
Ordering::Greater => {}
Ordering::Equal => return (i, true),
Ordering::Less => return (i, false),
let len = node.len();
if len > 0 {
let keys = unsafe { node.keys() }; // safe because a non-empty node cannot be the shared root
for (i, k) in keys.iter().enumerate() {
match key.cmp(k.borrow()) {
Ordering::Greater => {}
Ordering::Equal => return (i, true),
Ordering::Less => return (i, false),
}
}
}
(node.keys().len(), false)
(len, false)
}