Skip to content

Commit

Permalink
Rollup merge of rust-lang#75200 - ssomers:btree_valmut, r=Mark-Simula…
Browse files Browse the repository at this point in the history 
…crum

 BTreeMap: introduce marker::ValMut and reserve Mut for unique access

The mutable BTreeMap iterators (apart from `DrainFilter`) are double-ended, meaning they have to rely on a front and a back handle that each represent a reference into the tree. Reserve a type category `marker::ValMut` for them, so that we guarantee that they cannot reach operations on handles with borrow type `marker::Mut`and that these operations can assume unique access to the tree.

Including rust-lang#75195, benchmarks report no genuine change:
```
benchcmp old new --threshold 5
 name                                 old ns/iter  new ns/iter  diff ns/iter   diff %  speedup
 btree::map::iter_100                 3,333        3,023                -310   -9.30%   x 1.10
 btree::map::range_unbounded_vs_iter  36,624       31,569             -5,055  -13.80%   x 1.16
```

r? @Mark-Simulacrum
  • Loading branch information
@Dylan-DPC
Dylan-DPC authored Sep 4, 2020
2 parents 42d896a + e5f9d7f commit 57672a2
Show file tree
Hide file tree
Showing 3 changed files with 319 additions and 171 deletions.
134 changes: 10 additions & 124 deletions library/alloc/src/collections/btree/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -408,8 +407,8 @@ impl<K: fmt::Debug, V: fmt::Debug> 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<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
back: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
front: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,
back: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,

// Be invariant in `K` and `V`
_marker: PhantomData<&'a mut (K, V)>,
Expand Down Expand Up @@ -999,7 +998,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
R: RangeBounds<T>,
{
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 {
Expand Down Expand Up @@ -1045,7 +1044,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
R: RangeBounds<T>,
{
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 {
Expand Down Expand Up @@ -1478,7 +1477,7 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
fn into_iter(self) -> IntoIter<K, V> {
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 {
Expand Down Expand Up @@ -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() }
}
}
Expand All @@ -1963,7 +1962,7 @@ impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
impl<K, V> 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() }
}
}
Expand Down Expand Up @@ -2073,119 +2072,6 @@ where
}
}

/// Finds the leaf edges delimiting a specified range in or underneath a node.
fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeBounds<Q>>(
root: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R,
) -> (
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
)
where
Q: Ord,
K: Borrow<Q>,
{
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<BorrowType, K, V>(
root: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
) -> (
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, 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<K, V> BTreeMap<K, V> {
/// Gets an iterator over the entries of the map, sorted by key.
///
Expand All @@ -2211,7 +2097,7 @@ impl<K, V> BTreeMap<K, V> {
#[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 {
Expand Down Expand Up @@ -2243,7 +2129,7 @@ impl<K, V> BTreeMap<K, V> {
#[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 },
Expand Down Expand Up @@ -2826,7 +2712,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, 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;
Expand Down
Loading

0 comments on commit 57672a2

Please sign in to comment.