From 2c3c891df0c4f0c29caee10c7289502ca801ff98 Mon Sep 17 00:00:00 2001 From: Frank King Date: Wed, 9 Mar 2022 22:26:39 +0800 Subject: [PATCH] BTreeMap::entry: Avoid allocating if no insertion --- library/alloc/src/collections/btree/map.rs | 35 ++++++++------- .../alloc/src/collections/btree/map/entry.rs | 38 ++++++++++------ .../alloc/src/collections/btree/map/tests.rs | 45 ++++++++++++++++--- library/alloc/src/collections/btree/node.rs | 9 ++-- 4 files changed, 86 insertions(+), 41 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 7890c1040f0..07b1970a65d 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -212,7 +212,7 @@ impl Clone for BTreeMap { let mut out_tree = clone_subtree(internal.first_edge().descend()); { - let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root); + let out_root = out_tree.root.as_mut().unwrap(); let mut out_node = out_root.push_internal_level(); let mut in_edge = internal.first_edge(); while let Ok(kv) = in_edge.right_kv() { @@ -278,11 +278,12 @@ where fn replace(&mut self, key: K) -> Option { let (map, dormant_map) = DormantMutRef::new(self); - let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut(); + let root_node = map.root.get_or_insert_with(Root::new).borrow_mut(); match root_node.search_tree::(&key) { Found(mut kv) => Some(mem::replace(kv.key_mut(), key)), GoDown(handle) => { - VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(()); + VacantEntry { key, handle: Some(handle), dormant_map, _marker: PhantomData } + .insert(()); None } } @@ -1032,7 +1033,7 @@ impl BTreeMap { let self_iter = mem::take(self).into_iter(); let other_iter = mem::take(other).into_iter(); - let root = BTreeMap::ensure_is_owned(&mut self.root); + let root = self.root.get_or_insert_with(Root::new); root.append_from_sorted_iters(self_iter, other_iter, &mut self.length) } @@ -1144,14 +1145,20 @@ impl BTreeMap { where K: Ord, { - // FIXME(@porglezomp) Avoid allocating if we don't insert let (map, dormant_map) = DormantMutRef::new(self); - let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut(); - match root_node.search_tree(&key) { - Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }), - GoDown(handle) => { - Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData }) - } + match map.root { + None => Vacant(VacantEntry { key, handle: None, dormant_map, _marker: PhantomData }), + Some(ref mut root) => match root.borrow_mut().search_tree(&key) { + Found(handle) => { + Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }) + } + GoDown(handle) => Vacant(VacantEntry { + key, + handle: Some(handle), + dormant_map, + _marker: PhantomData, + }), + }, } } @@ -2247,12 +2254,6 @@ impl BTreeMap { pub const fn is_empty(&self) -> bool { self.len() == 0 } - - /// If the root node is the empty (non-allocated) root node, allocate our - /// own node. Is an associated function to avoid borrowing the entire BTreeMap. - fn ensure_is_owned(root: &mut Option>) -> &mut Root { - root.get_or_insert_with(Root::new) - } } #[cfg(test)] diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs index 66608d09082..cacd06b5df1 100644 --- a/library/alloc/src/collections/btree/map/entry.rs +++ b/library/alloc/src/collections/btree/map/entry.rs @@ -40,7 +40,8 @@ impl Debug for Entry<'_, K, V> { #[stable(feature = "rust1", since = "1.0.0")] pub struct VacantEntry<'a, K: 'a, V: 'a> { pub(super) key: K, - pub(super) handle: Handle, K, V, marker::Leaf>, marker::Edge>, + /// `None` for a (empty) map without root + pub(super) handle: Option, K, V, marker::Leaf>, marker::Edge>>, pub(super) dormant_map: DormantMutRef<'a, BTreeMap>, // Be invariant in `K` and `V` @@ -312,22 +313,33 @@ impl<'a, K: Ord, V> VacantEntry<'a, K, V> { /// ``` #[stable(feature = "rust1", since = "1.0.0")] pub fn insert(self, value: V) -> &'a mut V { - let out_ptr = match self.handle.insert_recursing(self.key, value) { - (None, val_ptr) => { - // SAFETY: We have consumed self.handle and the handle returned. - let map = unsafe { self.dormant_map.awaken() }; - map.length += 1; - val_ptr - } - (Some(ins), val_ptr) => { - drop(ins.left); + let out_ptr = match self.handle { + None => { // SAFETY: We have consumed self.handle and the reference returned. let map = unsafe { self.dormant_map.awaken() }; - let root = map.root.as_mut().unwrap(); - root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right); - map.length += 1; + let mut root = NodeRef::new_leaf(); + let val_ptr = root.borrow_mut().push(self.key, value) as *mut V; + map.root = Some(root.forget_type()); + map.length = 1; val_ptr } + Some(handle) => match handle.insert_recursing(self.key, value) { + (None, val_ptr) => { + // SAFETY: We have consumed self.handle and the handle returned. + let map = unsafe { self.dormant_map.awaken() }; + map.length += 1; + val_ptr + } + (Some(ins), val_ptr) => { + drop(ins.left); + // SAFETY: We have consumed self.handle and the reference returned. + let map = unsafe { self.dormant_map.awaken() }; + let root = map.root.as_mut().unwrap(); + root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right); + map.length += 1; + val_ptr + } + }, }; // Now that we have finished growing the tree using borrowed references, // dereference the pointer to a part of it, that we picked up along the way. diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 65468d5fe57..7fe584adb23 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -115,8 +115,9 @@ impl BTreeMap { K: Ord, { let iter = mem::take(self).into_iter(); - let root = BTreeMap::ensure_is_owned(&mut self.root); - root.bulk_push(iter, &mut self.length); + if !iter.is_empty() { + self.root.insert(Root::new()).bulk_push(iter, &mut self.length); + } } } @@ -914,7 +915,7 @@ mod test_drain_filter { fn empty() { let mut map: BTreeMap = BTreeMap::new(); map.drain_filter(|_, _| unreachable!("there's nothing to decide on")); - assert!(map.is_empty()); + assert_eq!(map.height(), None); map.check(); } @@ -1410,7 +1411,7 @@ fn test_clear() { assert_eq!(map.len(), len); map.clear(); map.check(); - assert!(map.is_empty()); + assert_eq!(map.height(), None); } } @@ -1789,7 +1790,7 @@ fn test_occupied_entry_key() { let mut a = BTreeMap::new(); let key = "hello there"; let value = "value goes here"; - assert!(a.is_empty()); + assert_eq!(a.height(), None); a.insert(key, value); assert_eq!(a.len(), 1); assert_eq!(a[key], value); @@ -1809,9 +1810,9 @@ fn test_vacant_entry_key() { let key = "hello there"; let value = "value goes here"; - assert!(a.is_empty()); + assert_eq!(a.height(), None); match a.entry(key) { - Occupied(_) => panic!(), + Occupied(_) => unreachable!(), Vacant(e) => { assert_eq!(key, *e.key()); e.insert(value); @@ -1822,6 +1823,36 @@ fn test_vacant_entry_key() { a.check(); } +#[test] +fn test_vacant_entry_no_insert() { + let mut a = BTreeMap::<&str, ()>::new(); + let key = "hello there"; + + // Non-allocated + assert_eq!(a.height(), None); + match a.entry(key) { + Occupied(_) => unreachable!(), + Vacant(e) => assert_eq!(key, *e.key()), + } + // Ensures the tree has no root. + assert_eq!(a.height(), None); + a.check(); + + // Allocated but still empty + a.insert(key, ()); + a.remove(&key); + assert_eq!(a.height(), Some(0)); + assert!(a.is_empty()); + match a.entry(key) { + Occupied(_) => unreachable!(), + Vacant(e) => assert_eq!(key, *e.key()), + } + // Ensures the allocated root is not changed. + assert_eq!(a.height(), Some(0)); + assert!(a.is_empty()); + a.check(); +} + #[test] fn test_first_last_entry() { let mut a = BTreeMap::new(); diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 44f5bc850b8..b5f0edf6b33 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -213,7 +213,7 @@ unsafe impl Send for NodeRef unsafe impl Send for NodeRef {} impl NodeRef { - fn new_leaf() -> Self { + pub fn new_leaf() -> Self { Self::from_new_leaf(LeafNode::new()) } @@ -619,15 +619,16 @@ impl NodeRef { } impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Leaf> { - /// Adds a key-value pair to the end of the node. - pub fn push(&mut self, key: K, val: V) { + /// Adds a key-value pair to the end of the node, and returns + /// the mutable reference of the inserted value. + pub fn push(&mut self, key: K, val: V) -> &mut V { let len = self.len_mut(); let idx = usize::from(*len); assert!(idx < CAPACITY); *len += 1; unsafe { self.key_area_mut(idx).write(key); - self.val_area_mut(idx).write(val); + self.val_area_mut(idx).write(val) } } }