Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit e5f9d7f

Browse files
committedAug 19, 2020
BTreeMap: introduce marker::ValMut and reserve marker::Mut for unique access
1 parent 443e177 commit e5f9d7f

File tree

3 files changed

+319
-171
lines changed

3 files changed

+319
-171
lines changed
 

‎library/alloc/src/collections/btree/map.rs

Lines changed: 10 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ use core::hash::{Hash, Hasher};
55
use core::iter::{FromIterator, FusedIterator, Peekable};
66
use core::marker::PhantomData;
77
use core::mem::{self, ManuallyDrop};
8-
use core::ops::Bound::{Excluded, Included, Unbounded};
98
use core::ops::{Index, RangeBounds};
109
use core::ptr;
1110

@@ -408,8 +407,8 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Range<'_, K, V> {
408407
/// [`range_mut`]: BTreeMap::range_mut
409408
#[stable(feature = "btree_range", since = "1.17.0")]
410409
pub struct RangeMut<'a, K: 'a, V: 'a> {
411-
front: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
412-
back: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
410+
front: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,
411+
back: Option<Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>>,
413412

414413
// Be invariant in `K` and `V`
415414
_marker: PhantomData<&'a mut (K, V)>,
@@ -999,7 +998,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
999998
R: RangeBounds<T>,
1000999
{
10011000
if let Some(root) = &self.root {
1002-
let (f, b) = range_search(root.node_as_ref(), range);
1001+
let (f, b) = root.node_as_ref().range_search(range);
10031002

10041003
Range { front: Some(f), back: Some(b) }
10051004
} else {
@@ -1045,7 +1044,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
10451044
R: RangeBounds<T>,
10461045
{
10471046
if let Some(root) = &mut self.root {
1048-
let (f, b) = range_search(root.node_as_mut(), range);
1047+
let (f, b) = root.node_as_valmut().range_search(range);
10491048

10501049
RangeMut { front: Some(f), back: Some(b), _marker: PhantomData }
10511050
} else {
@@ -1478,7 +1477,7 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
14781477
fn into_iter(self) -> IntoIter<K, V> {
14791478
let mut me = ManuallyDrop::new(self);
14801479
if let Some(root) = me.root.take() {
1481-
let (f, b) = full_range_search(root.into_ref());
1480+
let (f, b) = root.into_ref().full_range();
14821481

14831482
IntoIter { front: Some(f), back: Some(b), length: me.length }
14841483
} else {
@@ -1942,7 +1941,7 @@ impl<'a, K, V> RangeMut<'a, K, V> {
19421941
self.front == self.back
19431942
}
19441943

1945-
unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
1944+
unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
19461945
unsafe { unwrap_unchecked(self.front.as_mut()).next_unchecked() }
19471946
}
19481947
}
@@ -1963,7 +1962,7 @@ impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
19631962
impl<K, V> FusedIterator for RangeMut<'_, K, V> {}
19641963

19651964
impl<'a, K, V> RangeMut<'a, K, V> {
1966-
unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
1965+
unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
19671966
unsafe { unwrap_unchecked(self.back.as_mut()).next_back_unchecked() }
19681967
}
19691968
}
@@ -2073,119 +2072,6 @@ where
20732072
}
20742073
}
20752074

2076-
/// Finds the leaf edges delimiting a specified range in or underneath a node.
2077-
fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeBounds<Q>>(
2078-
root: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
2079-
range: R,
2080-
) -> (
2081-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
2082-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
2083-
)
2084-
where
2085-
Q: Ord,
2086-
K: Borrow<Q>,
2087-
{
2088-
match (range.start_bound(), range.end_bound()) {
2089-
(Excluded(s), Excluded(e)) if s == e => {
2090-
panic!("range start and end are equal and excluded in BTreeMap")
2091-
}
2092-
(Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => {
2093-
panic!("range start is greater than range end in BTreeMap")
2094-
}
2095-
_ => {}
2096-
};
2097-
2098-
// We duplicate the root NodeRef here -- we will never access it in a way
2099-
// that overlaps references obtained from the root.
2100-
let mut min_node = unsafe { ptr::read(&root) };
2101-
let mut max_node = root;
2102-
let mut min_found = false;
2103-
let mut max_found = false;
2104-
2105-
loop {
2106-
let front = match (min_found, range.start_bound()) {
2107-
(false, Included(key)) => match search::search_node(min_node, key) {
2108-
Found(kv) => {
2109-
min_found = true;
2110-
kv.left_edge()
2111-
}
2112-
GoDown(edge) => edge,
2113-
},
2114-
(false, Excluded(key)) => match search::search_node(min_node, key) {
2115-
Found(kv) => {
2116-
min_found = true;
2117-
kv.right_edge()
2118-
}
2119-
GoDown(edge) => edge,
2120-
},
2121-
(true, Included(_)) => min_node.last_edge(),
2122-
(true, Excluded(_)) => min_node.first_edge(),
2123-
(_, Unbounded) => min_node.first_edge(),
2124-
};
2125-
2126-
let back = match (max_found, range.end_bound()) {
2127-
(false, Included(key)) => match search::search_node(max_node, key) {
2128-
Found(kv) => {
2129-
max_found = true;
2130-
kv.right_edge()
2131-
}
2132-
GoDown(edge) => edge,
2133-
},
2134-
(false, Excluded(key)) => match search::search_node(max_node, key) {
2135-
Found(kv) => {
2136-
max_found = true;
2137-
kv.left_edge()
2138-
}
2139-
GoDown(edge) => edge,
2140-
},
2141-
(true, Included(_)) => max_node.first_edge(),
2142-
(true, Excluded(_)) => max_node.last_edge(),
2143-
(_, Unbounded) => max_node.last_edge(),
2144-
};
2145-
2146-
if front.partial_cmp(&back) == Some(Ordering::Greater) {
2147-
panic!("Ord is ill-defined in BTreeMap range");
2148-
}
2149-
match (front.force(), back.force()) {
2150-
(Leaf(f), Leaf(b)) => {
2151-
return (f, b);
2152-
}
2153-
(Internal(min_int), Internal(max_int)) => {
2154-
min_node = min_int.descend();
2155-
max_node = max_int.descend();
2156-
}
2157-
_ => unreachable!("BTreeMap has different depths"),
2158-
};
2159-
}
2160-
}
2161-
2162-
/// Equivalent to `range_search(k, v, ..)` without the `Ord` bound.
2163-
fn full_range_search<BorrowType, K, V>(
2164-
root: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
2165-
) -> (
2166-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
2167-
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
2168-
) {
2169-
// We duplicate the root NodeRef here -- we will never access it in a way
2170-
// that overlaps references obtained from the root.
2171-
let mut min_node = unsafe { ptr::read(&root) };
2172-
let mut max_node = root;
2173-
loop {
2174-
let front = min_node.first_edge();
2175-
let back = max_node.last_edge();
2176-
match (front.force(), back.force()) {
2177-
(Leaf(f), Leaf(b)) => {
2178-
return (f, b);
2179-
}
2180-
(Internal(min_int), Internal(max_int)) => {
2181-
min_node = min_int.descend();
2182-
max_node = max_int.descend();
2183-
}
2184-
_ => unreachable!("BTreeMap has different depths"),
2185-
};
2186-
}
2187-
}
2188-
21892075
impl<K, V> BTreeMap<K, V> {
21902076
/// Gets an iterator over the entries of the map, sorted by key.
21912077
///
@@ -2211,7 +2097,7 @@ impl<K, V> BTreeMap<K, V> {
22112097
#[stable(feature = "rust1", since = "1.0.0")]
22122098
pub fn iter(&self) -> Iter<'_, K, V> {
22132099
if let Some(root) = &self.root {
2214-
let (f, b) = full_range_search(root.node_as_ref());
2100+
let (f, b) = root.node_as_ref().full_range();
22152101

22162102
Iter { range: Range { front: Some(f), back: Some(b) }, length: self.length }
22172103
} else {
@@ -2243,7 +2129,7 @@ impl<K, V> BTreeMap<K, V> {
22432129
#[stable(feature = "rust1", since = "1.0.0")]
22442130
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
22452131
if let Some(root) = &mut self.root {
2246-
let (f, b) = full_range_search(root.node_as_mut());
2132+
let (f, b) = root.node_as_valmut().full_range();
22472133

22482134
IterMut {
22492135
range: RangeMut { front: Some(f), back: Some(b), _marker: PhantomData },
@@ -2826,7 +2712,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
28262712
if stole_from_left && at_leaf {
28272713
// SAFETY: This is safe since we just added an element to our node.
28282714
unsafe {
2829-
pos.next_unchecked();
2715+
pos.move_next_unchecked();
28302716
}
28312717
}
28322718
break;

‎library/alloc/src/collections/btree/navigate.rs

Lines changed: 274 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,210 @@
1+
use core::borrow::Borrow;
2+
use core::cmp::Ordering;
13
use core::intrinsics;
24
use core::mem;
5+
use core::ops::Bound::{Excluded, Included, Unbounded};
6+
use core::ops::RangeBounds;
37
use core::ptr;
48

59
use super::node::{marker, ForceResult::*, Handle, NodeRef};
10+
use super::search::{self, SearchResult};
611
use super::unwrap_unchecked;
712

13+
/// Finds the leaf edges delimiting a specified range in or underneath a node.
14+
fn range_search<BorrowType, K, V, Q, R>(
15+
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
16+
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
17+
range: R,
18+
) -> (
19+
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
20+
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
21+
)
22+
where
23+
Q: ?Sized + Ord,
24+
K: Borrow<Q>,
25+
R: RangeBounds<Q>,
26+
{
27+
match (range.start_bound(), range.end_bound()) {
28+
(Excluded(s), Excluded(e)) if s == e => {
29+
panic!("range start and end are equal and excluded in BTreeMap")
30+
}
31+
(Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => {
32+
panic!("range start is greater than range end in BTreeMap")
33+
}
34+
_ => {}
35+
};
36+
37+
let mut min_node = root1;
38+
let mut max_node = root2;
39+
let mut min_found = false;
40+
let mut max_found = false;
41+
42+
loop {
43+
let front = match (min_found, range.start_bound()) {
44+
(false, Included(key)) => match search::search_node(min_node, key) {
45+
SearchResult::Found(kv) => {
46+
min_found = true;
47+
kv.left_edge()
48+
}
49+
SearchResult::GoDown(edge) => edge,
50+
},
51+
(false, Excluded(key)) => match search::search_node(min_node, key) {
52+
SearchResult::Found(kv) => {
53+
min_found = true;
54+
kv.right_edge()
55+
}
56+
SearchResult::GoDown(edge) => edge,
57+
},
58+
(true, Included(_)) => min_node.last_edge(),
59+
(true, Excluded(_)) => min_node.first_edge(),
60+
(_, Unbounded) => min_node.first_edge(),
61+
};
62+
63+
let back = match (max_found, range.end_bound()) {
64+
(false, Included(key)) => match search::search_node(max_node, key) {
65+
SearchResult::Found(kv) => {
66+
max_found = true;
67+
kv.right_edge()
68+
}
69+
SearchResult::GoDown(edge) => edge,
70+
},
71+
(false, Excluded(key)) => match search::search_node(max_node, key) {
72+
SearchResult::Found(kv) => {
73+
max_found = true;
74+
kv.left_edge()
75+
}
76+
SearchResult::GoDown(edge) => edge,
77+
},
78+
(true, Included(_)) => max_node.first_edge(),
79+
(true, Excluded(_)) => max_node.last_edge(),
80+
(_, Unbounded) => max_node.last_edge(),
81+
};
82+
83+
if front.partial_cmp(&back) == Some(Ordering::Greater) {
84+
panic!("Ord is ill-defined in BTreeMap range");
85+
}
86+
match (front.force(), back.force()) {
87+
(Leaf(f), Leaf(b)) => {
88+
return (f, b);
89+
}
90+
(Internal(min_int), Internal(max_int)) => {
91+
min_node = min_int.descend();
92+
max_node = max_int.descend();
93+
}
94+
_ => unreachable!("BTreeMap has different depths"),
95+
};
96+
}
97+
}
98+
99+
/// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound.
100+
fn full_range<BorrowType, K, V>(
101+
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
102+
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
103+
) -> (
104+
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
105+
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
106+
) {
107+
let mut min_node = root1;
108+
let mut max_node = root2;
109+
loop {
110+
let front = min_node.first_edge();
111+
let back = max_node.last_edge();
112+
match (front.force(), back.force()) {
113+
(Leaf(f), Leaf(b)) => {
114+
return (f, b);
115+
}
116+
(Internal(min_int), Internal(max_int)) => {
117+
min_node = min_int.descend();
118+
max_node = max_int.descend();
119+
}
120+
_ => unreachable!("BTreeMap has different depths"),
121+
};
122+
}
123+
}
124+
125+
impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
126+
/// Creates a pair of leaf edges delimiting a specified range in or underneath a node.
127+
pub fn range_search<Q, R>(
128+
self,
129+
range: R,
130+
) -> (
131+
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
132+
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
133+
)
134+
where
135+
Q: ?Sized + Ord,
136+
K: Borrow<Q>,
137+
R: RangeBounds<Q>,
138+
{
139+
range_search(self, self, range)
140+
}
141+
142+
/// Returns (self.first_leaf_edge(), self.last_leaf_edge()), but more efficiently.
143+
pub fn full_range(
144+
self,
145+
) -> (
146+
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
147+
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
148+
) {
149+
full_range(self, self)
150+
}
151+
}
152+
153+
impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal> {
154+
/// Splits a unique reference into a pair of leaf edges delimiting a specified range.
155+
/// The result are non-unique references allowing (some) mutation, which must be used
156+
/// carefully.
157+
pub fn range_search<Q, R>(
158+
self,
159+
range: R,
160+
) -> (
161+
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
162+
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
163+
)
164+
where
165+
Q: ?Sized + Ord,
166+
K: Borrow<Q>,
167+
R: RangeBounds<Q>,
168+
{
169+
// We duplicate the root NodeRef here -- we will never visit the same KV
170+
// twice, and never end up with overlapping value references.
171+
let self2 = unsafe { ptr::read(&self) };
172+
range_search(self, self2, range)
173+
}
174+
175+
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
176+
/// The results are non-unique references allowing mutation (of values only), so must be used
177+
/// with care.
178+
pub fn full_range(
179+
self,
180+
) -> (
181+
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
182+
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
183+
) {
184+
// We duplicate the root NodeRef here -- we will never visit the same KV
185+
// twice, and never end up with overlapping value references.
186+
let self2 = unsafe { ptr::read(&self) };
187+
full_range(self, self2)
188+
}
189+
}
190+
191+
impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
192+
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
193+
/// The results are non-unique references allowing massively destructive mutation, so must be
194+
/// used with the utmost care.
195+
pub fn full_range(
196+
self,
197+
) -> (
198+
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
199+
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
200+
) {
201+
// We duplicate the root NodeRef here -- we will never access it in a way
202+
// that overlaps references obtained from the root.
203+
let self2 = unsafe { ptr::read(&self) };
204+
full_range(self, self2)
205+
}
206+
}
207+
8208
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
9209
/// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
10210
/// on the right side, which is either in the same leaf node or in an ancestor node.
@@ -75,12 +275,13 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marke
75275
macro_rules! def_next_kv_uncheched_dealloc {
76276
{ unsafe fn $name:ident : $adjacent_kv:ident } => {
77277
/// Given a leaf edge handle into an owned tree, returns a handle to the next KV,
78-
/// while deallocating any node left behind.
79-
/// Unsafe for two reasons:
80-
/// - The caller must ensure that the leaf edge is not the last one in the tree.
81-
/// - The node pointed at by the given handle, and its ancestors, may be deallocated,
82-
/// while the reference to those nodes in the surviving ancestors is left dangling;
83-
/// thus using the returned handle to navigate further is dangerous.
278+
/// while deallocating any node left behind yet leaving the corresponding edge
279+
/// in its parent node dangling.
280+
///
281+
/// # Safety
282+
/// - The leaf edge must not be the last one in the direction travelled.
283+
/// - The node carrying the next KV returned must not have been deallocated by a
284+
/// previous call on any handle obtained for this tree.
84285
unsafe fn $name <K, V>(
85286
leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
86287
) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
@@ -103,6 +304,15 @@ macro_rules! def_next_kv_uncheched_dealloc {
103304
def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
104305
def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
105306

307+
/// This replaces the value behind the `v` unique reference by calling the
308+
/// relevant function.
309+
///
310+
/// If a panic occurs in the `change` closure, the entire process will be aborted.
311+
#[inline]
312+
fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
313+
replace(v, |value| (change(value), ()))
314+
}
315+
106316
/// This replaces the value behind the `v` unique reference by calling the
107317
/// relevant function, and returns a result obtained along the way.
108318
///
@@ -128,7 +338,9 @@ fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
128338
impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
129339
/// Moves the leaf edge handle to the next leaf edge and returns references to the
130340
/// key and value in between.
131-
/// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree.
341+
///
342+
/// # Safety
343+
/// There must be another KV in the direction travelled.
132344
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
133345
replace(self, |leaf_edge| {
134346
let kv = leaf_edge.next_kv();
@@ -139,7 +351,9 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
139351

140352
/// Moves the leaf edge handle to the previous leaf edge and returns references to the
141353
/// key and value in between.
142-
/// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree.
354+
///
355+
/// # Safety
356+
/// There must be another KV in the direction travelled.
143357
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
144358
replace(self, |leaf_edge| {
145359
let kv = leaf_edge.next_back_kv();
@@ -149,53 +363,69 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
149363
}
150364
}
151365

152-
impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
366+
impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge> {
153367
/// Moves the leaf edge handle to the next leaf edge and returns references to the
154368
/// key and value in between.
155-
/// Unsafe for two reasons:
156-
/// - The caller must ensure that the leaf edge is not the last one in the tree.
157-
/// - Using the updated handle may well invalidate the returned references.
158-
pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
369+
/// The returned references might be invalidated when the updated handle is used again.
370+
///
371+
/// # Safety
372+
/// There must be another KV in the direction travelled.
373+
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
159374
let kv = replace(self, |leaf_edge| {
160375
let kv = leaf_edge.next_kv();
161376
let kv = unsafe { unwrap_unchecked(kv.ok()) };
162377
(unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
163378
});
164379
// Doing the descend (and perhaps another move) invalidates the references
165-
// returned by `into_kv_mut`, so we have to do this last.
166-
kv.into_kv_mut()
380+
// returned by `into_kv_valmut`, so we have to do this last.
381+
kv.into_kv_valmut()
167382
}
168383

169384
/// Moves the leaf edge handle to the previous leaf and returns references to the
170385
/// key and value in between.
171-
/// Unsafe for two reasons:
172-
/// - The caller must ensure that the leaf edge is not the first one in the tree.
173-
/// - Using the updated handle may well invalidate the returned references.
174-
pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
386+
/// The returned references might be invalidated when the updated handle is used again.
387+
///
388+
/// # Safety
389+
/// There must be another KV in the direction travelled.
390+
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
175391
let kv = replace(self, |leaf_edge| {
176392
let kv = leaf_edge.next_back_kv();
177393
let kv = unsafe { unwrap_unchecked(kv.ok()) };
178394
(unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
179395
});
180396
// Doing the descend (and perhaps another move) invalidates the references
181-
// returned by `into_kv_mut`, so we have to do this last.
182-
kv.into_kv_mut()
397+
// returned by `into_kv_valmut`, so we have to do this last.
398+
kv.into_kv_valmut()
399+
}
400+
}
401+
402+
impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
403+
/// Moves the leaf edge handle to the next leaf edge.
404+
///
405+
/// # Safety
406+
/// There must be another KV in the direction travelled.
407+
pub unsafe fn move_next_unchecked(&mut self) {
408+
take_mut(self, |leaf_edge| {
409+
let kv = leaf_edge.next_kv();
410+
let kv = unsafe { unwrap_unchecked(kv.ok()) };
411+
kv.next_leaf_edge()
412+
})
183413
}
184414
}
185415

186416
impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
187417
/// Moves the leaf edge handle to the next leaf edge and returns the key and value
188-
/// in between, while deallocating any node left behind.
189-
/// Unsafe for two reasons:
190-
/// - The caller must ensure that the leaf edge is not the last one in the tree
191-
/// and is not a handle previously resulting from counterpart `next_back_unchecked`.
192-
/// - Further use of the updated leaf edge handle is very dangerous. In particular,
193-
/// if the leaf edge is the last edge of a node, that node and possibly ancestors
194-
/// will be deallocated, while the reference to those nodes in the surviving ancestor
195-
/// is left dangling.
196-
/// The only safe way to proceed with the updated handle is to compare it, drop it,
197-
/// call this method again subject to both preconditions listed in the first point,
198-
/// or call counterpart `next_back_unchecked` subject to its preconditions.
418+
/// in between, deallocating any node left behind while leaving the corresponding
419+
/// edge in its parent node dangling.
420+
///
421+
/// # Safety
422+
/// - There must be another KV in the direction travelled.
423+
/// - That KV was not previously returned by counterpart `next_back_unchecked`
424+
/// on any copy of the handles being used to traverse the tree.
425+
///
426+
/// The only safe way to proceed with the updated handle is to compare it, drop it,
427+
/// call this method again subject to its safety conditions, or call counterpart
428+
/// `next_back_unchecked` subject to its safety conditions.
199429
pub unsafe fn next_unchecked(&mut self) -> (K, V) {
200430
replace(self, |leaf_edge| {
201431
let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
@@ -205,18 +435,18 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
205435
})
206436
}
207437

208-
/// Moves the leaf edge handle to the previous leaf edge and returns the key
209-
/// and value in between, while deallocating any node left behind.
210-
/// Unsafe for two reasons:
211-
/// - The caller must ensure that the leaf edge is not the first one in the tree
212-
/// and is not a handle previously resulting from counterpart `next_unchecked`.
213-
/// - Further use of the updated leaf edge handle is very dangerous. In particular,
214-
/// if the leaf edge is the first edge of a node, that node and possibly ancestors
215-
/// will be deallocated, while the reference to those nodes in the surviving ancestor
216-
/// is left dangling.
217-
/// The only safe way to proceed with the updated handle is to compare it, drop it,
218-
/// call this method again subject to both preconditions listed in the first point,
219-
/// or call counterpart `next_unchecked` subject to its preconditions.
438+
/// Moves the leaf edge handle to the previous leaf edge and returns the key and value
439+
/// in between, deallocating any node left behind while leaving the corresponding
440+
/// edge in its parent node dangling.
441+
///
442+
/// # Safety
443+
/// - There must be another KV in the direction travelled.
444+
/// - That leaf edge was not previously returned by counterpart `next_unchecked`
445+
/// on any copy of the handles being used to traverse the tree.
446+
///
447+
/// The only safe way to proceed with the updated handle is to compare it, drop it,
448+
/// call this method again subject to its safety conditions, or call counterpart
449+
/// `next_unchecked` subject to its safety conditions.
220450
pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
221451
replace(self, |leaf_edge| {
222452
let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };

‎library/alloc/src/collections/btree/node.rs

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,15 @@ impl<K, V> Root<K, V> {
186186
}
187187
}
188188

189+
pub fn node_as_valmut(&mut self) -> NodeRef<marker::ValMut<'_>, K, V, marker::LeafOrInternal> {
190+
NodeRef {
191+
height: self.height,
192+
node: self.node.as_ptr(),
193+
root: ptr::null(),
194+
_marker: PhantomData,
195+
}
196+
}
197+
189198
pub fn into_ref(self) -> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
190199
NodeRef {
191200
height: self.height,
@@ -253,9 +262,12 @@ impl<K, V> Root<K, V> {
253262
/// A reference to a node.
254263
///
255264
/// This type has a number of parameters that controls how it acts:
256-
/// - `BorrowType`: This can be `Immut<'a>` or `Mut<'a>` for some `'a` or `Owned`.
265+
/// - `BorrowType`: This can be `Immut<'a>`, `Mut<'a>` or `ValMut<'a>' for some `'a`
266+
/// or `Owned`.
257267
/// When this is `Immut<'a>`, the `NodeRef` acts roughly like `&'a Node`,
258268
/// when this is `Mut<'a>`, the `NodeRef` acts roughly like `&'a mut Node`,
269+
/// when this is `ValMut<'a>`, the `NodeRef` acts as immutable with respect
270+
/// to keys and tree structure, but allows mutable references to values,
259271
/// and when this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`.
260272
/// - `K` and `V`: These control what types of things are stored in the nodes.
261273
/// - `Type`: This can be `Leaf`, `Internal`, or `LeafOrInternal`. When this is
@@ -282,6 +294,7 @@ unsafe impl<BorrowType, K: Sync, V: Sync, Type> Sync for NodeRef<BorrowType, K,
282294

283295
unsafe impl<'a, K: Sync + 'a, V: Sync + 'a, Type> Send for NodeRef<marker::Immut<'a>, K, V, Type> {}
284296
unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::Mut<'a>, K, V, Type> {}
297+
unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::ValMut<'a>, K, V, Type> {}
285298
unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type> {}
286299

287300
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
@@ -515,6 +528,22 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
515528
}
516529
}
517530

531+
impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::ValMut<'a>, K, V, Type> {
532+
/// Same as the marker::Mut method, but far more dangerous because ValMut-based iterators:
533+
/// - have front and back handles often refering to the same node, so `self` is not unique;
534+
/// - hand out mutable references to parts of these slices to the public.
535+
fn into_slices_mut(self) -> (&'a [K], &'a mut [V]) {
536+
let len = self.len();
537+
let leaf = self.node.as_ptr();
538+
// SAFETY: The keys and values of a node must always be initialized up to length.
539+
let keys = unsafe { slice::from_raw_parts(MaybeUninit::first_ptr(&(*leaf).keys), len) };
540+
let vals = unsafe {
541+
slice::from_raw_parts_mut(MaybeUninit::first_ptr_mut(&mut (*leaf).vals), len)
542+
};
543+
(keys, vals)
544+
}
545+
}
546+
518547
impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
519548
/// Adds a key/value pair to the end of the node.
520549
pub fn push(&mut self, key: K, val: V) {
@@ -1053,11 +1082,13 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
10531082
let vals = self.node.into_val_slice_mut();
10541083
unsafe { vals.get_unchecked_mut(self.idx) }
10551084
}
1085+
}
10561086

1057-
pub fn into_kv_mut(self) -> (&'a mut K, &'a mut V) {
1087+
impl<'a, K, V, NodeType> Handle<NodeRef<marker::ValMut<'a>, K, V, NodeType>, marker::KV> {
1088+
pub fn into_kv_valmut(self) -> (&'a K, &'a mut V) {
10581089
unsafe {
10591090
let (keys, vals) = self.node.into_slices_mut();
1060-
(keys.get_unchecked_mut(self.idx), vals.get_unchecked_mut(self.idx))
1091+
(keys.get_unchecked(self.idx), vals.get_unchecked_mut(self.idx))
10611092
}
10621093
}
10631094
}
@@ -1558,6 +1589,7 @@ pub mod marker {
15581589
pub enum Owned {}
15591590
pub struct Immut<'a>(PhantomData<&'a ()>);
15601591
pub struct Mut<'a>(PhantomData<&'a mut ()>);
1592+
pub struct ValMut<'a>(PhantomData<&'a mut ()>);
15611593

15621594
pub enum KV {}
15631595
pub enum Edge {}

0 commit comments

Comments
 (0)
Please sign in to comment.