Skip to content

Commit 2c3c891

Browse files
frank-kingfrank-king
committed
BTreeMap::entry: Avoid allocating if no insertion
1 parent 10dccdc commit 2c3c891

File tree

4 files changed

+86
-41
lines changed

4 files changed

+86
-41
lines changed

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

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
212212
let mut out_tree = clone_subtree(internal.first_edge().descend());
213213

214214
{
215-
let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
215+
let out_root = out_tree.root.as_mut().unwrap();
216216
let mut out_node = out_root.push_internal_level();
217217
let mut in_edge = internal.first_edge();
218218
while let Ok(kv) = in_edge.right_kv() {
@@ -278,11 +278,12 @@ where
278278

279279
fn replace(&mut self, key: K) -> Option<K> {
280280
let (map, dormant_map) = DormantMutRef::new(self);
281-
let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
281+
let root_node = map.root.get_or_insert_with(Root::new).borrow_mut();
282282
match root_node.search_tree::<K>(&key) {
283283
Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
284284
GoDown(handle) => {
285-
VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(());
285+
VacantEntry { key, handle: Some(handle), dormant_map, _marker: PhantomData }
286+
.insert(());
286287
None
287288
}
288289
}
@@ -1032,7 +1033,7 @@ impl<K, V> BTreeMap<K, V> {
10321033

10331034
let self_iter = mem::take(self).into_iter();
10341035
let other_iter = mem::take(other).into_iter();
1035-
let root = BTreeMap::ensure_is_owned(&mut self.root);
1036+
let root = self.root.get_or_insert_with(Root::new);
10361037
root.append_from_sorted_iters(self_iter, other_iter, &mut self.length)
10371038
}
10381039

@@ -1144,14 +1145,20 @@ impl<K, V> BTreeMap<K, V> {
11441145
where
11451146
K: Ord,
11461147
{
1147-
// FIXME(@porglezomp) Avoid allocating if we don't insert
11481148
let (map, dormant_map) = DormantMutRef::new(self);
1149-
let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
1150-
match root_node.search_tree(&key) {
1151-
Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }),
1152-
GoDown(handle) => {
1153-
Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData })
1154-
}
1149+
match map.root {
1150+
None => Vacant(VacantEntry { key, handle: None, dormant_map, _marker: PhantomData }),
1151+
Some(ref mut root) => match root.borrow_mut().search_tree(&key) {
1152+
Found(handle) => {
1153+
Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData })
1154+
}
1155+
GoDown(handle) => Vacant(VacantEntry {
1156+
key,
1157+
handle: Some(handle),
1158+
dormant_map,
1159+
_marker: PhantomData,
1160+
}),
1161+
},
11551162
}
11561163
}
11571164

@@ -2247,12 +2254,6 @@ impl<K, V> BTreeMap<K, V> {
22472254
pub const fn is_empty(&self) -> bool {
22482255
self.len() == 0
22492256
}
2250-
2251-
/// If the root node is the empty (non-allocated) root node, allocate our
2252-
/// own node. Is an associated function to avoid borrowing the entire BTreeMap.
2253-
fn ensure_is_owned(root: &mut Option<Root<K, V>>) -> &mut Root<K, V> {
2254-
root.get_or_insert_with(Root::new)
2255-
}
22562257
}
22572258

22582259
#[cfg(test)]

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

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ impl<K: Debug + Ord, V: Debug> Debug for Entry<'_, K, V> {
4040
#[stable(feature = "rust1", since = "1.0.0")]
4141
pub struct VacantEntry<'a, K: 'a, V: 'a> {
4242
pub(super) key: K,
43-
pub(super) handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
43+
/// `None` for a (empty) map without root
44+
pub(super) handle: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
4445
pub(super) dormant_map: DormantMutRef<'a, BTreeMap<K, V>>,
4546

4647
// Be invariant in `K` and `V`
@@ -312,22 +313,33 @@ impl<'a, K: Ord, V> VacantEntry<'a, K, V> {
312313
/// ```
313314
#[stable(feature = "rust1", since = "1.0.0")]
314315
pub fn insert(self, value: V) -> &'a mut V {
315-
let out_ptr = match self.handle.insert_recursing(self.key, value) {
316-
(None, val_ptr) => {
317-
// SAFETY: We have consumed self.handle and the handle returned.
318-
let map = unsafe { self.dormant_map.awaken() };
319-
map.length += 1;
320-
val_ptr
321-
}
322-
(Some(ins), val_ptr) => {
323-
drop(ins.left);
316+
let out_ptr = match self.handle {
317+
None => {
324318
// SAFETY: We have consumed self.handle and the reference returned.
325319
let map = unsafe { self.dormant_map.awaken() };
326-
let root = map.root.as_mut().unwrap();
327-
root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right);
328-
map.length += 1;
320+
let mut root = NodeRef::new_leaf();
321+
let val_ptr = root.borrow_mut().push(self.key, value) as *mut V;
322+
map.root = Some(root.forget_type());
323+
map.length = 1;
329324
val_ptr
330325
}
326+
Some(handle) => match handle.insert_recursing(self.key, value) {
327+
(None, val_ptr) => {
328+
// SAFETY: We have consumed self.handle and the handle returned.
329+
let map = unsafe { self.dormant_map.awaken() };
330+
map.length += 1;
331+
val_ptr
332+
}
333+
(Some(ins), val_ptr) => {
334+
drop(ins.left);
335+
// SAFETY: We have consumed self.handle and the reference returned.
336+
let map = unsafe { self.dormant_map.awaken() };
337+
let root = map.root.as_mut().unwrap();
338+
root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right);
339+
map.length += 1;
340+
val_ptr
341+
}
342+
},
331343
};
332344
// Now that we have finished growing the tree using borrowed references,
333345
// dereference the pointer to a part of it, that we picked up along the way.

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

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,9 @@ impl<K, V> BTreeMap<K, V> {
115115
K: Ord,
116116
{
117117
let iter = mem::take(self).into_iter();
118-
let root = BTreeMap::ensure_is_owned(&mut self.root);
119-
root.bulk_push(iter, &mut self.length);
118+
if !iter.is_empty() {
119+
self.root.insert(Root::new()).bulk_push(iter, &mut self.length);
120+
}
120121
}
121122
}
122123

@@ -914,7 +915,7 @@ mod test_drain_filter {
914915
fn empty() {
915916
let mut map: BTreeMap<i32, i32> = BTreeMap::new();
916917
map.drain_filter(|_, _| unreachable!("there's nothing to decide on"));
917-
assert!(map.is_empty());
918+
assert_eq!(map.height(), None);
918919
map.check();
919920
}
920921

@@ -1410,7 +1411,7 @@ fn test_clear() {
14101411
assert_eq!(map.len(), len);
14111412
map.clear();
14121413
map.check();
1413-
assert!(map.is_empty());
1414+
assert_eq!(map.height(), None);
14141415
}
14151416
}
14161417

@@ -1789,7 +1790,7 @@ fn test_occupied_entry_key() {
17891790
let mut a = BTreeMap::new();
17901791
let key = "hello there";
17911792
let value = "value goes here";
1792-
assert!(a.is_empty());
1793+
assert_eq!(a.height(), None);
17931794
a.insert(key, value);
17941795
assert_eq!(a.len(), 1);
17951796
assert_eq!(a[key], value);
@@ -1809,9 +1810,9 @@ fn test_vacant_entry_key() {
18091810
let key = "hello there";
18101811
let value = "value goes here";
18111812

1812-
assert!(a.is_empty());
1813+
assert_eq!(a.height(), None);
18131814
match a.entry(key) {
1814-
Occupied(_) => panic!(),
1815+
Occupied(_) => unreachable!(),
18151816
Vacant(e) => {
18161817
assert_eq!(key, *e.key());
18171818
e.insert(value);
@@ -1822,6 +1823,36 @@ fn test_vacant_entry_key() {
18221823
a.check();
18231824
}
18241825

1826+
#[test]
1827+
fn test_vacant_entry_no_insert() {
1828+
let mut a = BTreeMap::<&str, ()>::new();
1829+
let key = "hello there";
1830+
1831+
// Non-allocated
1832+
assert_eq!(a.height(), None);
1833+
match a.entry(key) {
1834+
Occupied(_) => unreachable!(),
1835+
Vacant(e) => assert_eq!(key, *e.key()),
1836+
}
1837+
// Ensures the tree has no root.
1838+
assert_eq!(a.height(), None);
1839+
a.check();
1840+
1841+
// Allocated but still empty
1842+
a.insert(key, ());
1843+
a.remove(&key);
1844+
assert_eq!(a.height(), Some(0));
1845+
assert!(a.is_empty());
1846+
match a.entry(key) {
1847+
Occupied(_) => unreachable!(),
1848+
Vacant(e) => assert_eq!(key, *e.key()),
1849+
}
1850+
// Ensures the allocated root is not changed.
1851+
assert_eq!(a.height(), Some(0));
1852+
assert!(a.is_empty());
1853+
a.check();
1854+
}
1855+
18251856
#[test]
18261857
fn test_first_last_entry() {
18271858
let mut a = BTreeMap::new();

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type>
213213
unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}
214214

215215
impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
216-
fn new_leaf() -> Self {
216+
pub fn new_leaf() -> Self {
217217
Self::from_new_leaf(LeafNode::new())
218218
}
219219

@@ -619,15 +619,16 @@ impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
619619
}
620620

621621
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
622-
/// Adds a key-value pair to the end of the node.
623-
pub fn push(&mut self, key: K, val: V) {
622+
/// Adds a key-value pair to the end of the node, and returns
623+
/// the mutable reference of the inserted value.
624+
pub fn push(&mut self, key: K, val: V) -> &mut V {
624625
let len = self.len_mut();
625626
let idx = usize::from(*len);
626627
assert!(idx < CAPACITY);
627628
*len += 1;
628629
unsafe {
629630
self.key_area_mut(idx).write(key);
630-
self.val_area_mut(idx).write(val);
631+
self.val_area_mut(idx).write(val)
631632
}
632633
}
633634
}

0 commit comments

Comments
 (0)