diff options
Diffstat (limited to 'library/alloc/src')
| -rw-r--r-- | library/alloc/src/alloc.rs | 8 | ||||
| -rw-r--r-- | library/alloc/src/collections/btree/map/entry.rs | 105 | ||||
| -rw-r--r-- | library/alloc/src/collections/btree/map/tests.rs | 48 | ||||
| -rw-r--r-- | library/alloc/src/collections/btree/node.rs | 20 |
4 files changed, 144 insertions, 37 deletions
diff --git a/library/alloc/src/alloc.rs b/library/alloc/src/alloc.rs index 98402a46201..04b7315e650 100644 --- a/library/alloc/src/alloc.rs +++ b/library/alloc/src/alloc.rs @@ -61,7 +61,7 @@ pub use std::alloc::Global; /// of the allocator registered with the `#[global_allocator]` attribute /// if there is one, or the `std` crate’s default. /// -/// This function is expected to be deprecated in favor of the `alloc` method +/// This function is expected to be deprecated in favor of the `allocate` method /// of the [`Global`] type when it and the [`Allocator`] trait become stable. /// /// # Safety @@ -106,7 +106,7 @@ pub unsafe fn alloc(layout: Layout) -> *mut u8 { /// of the allocator registered with the `#[global_allocator]` attribute /// if there is one, or the `std` crate’s default. /// -/// This function is expected to be deprecated in favor of the `dealloc` method +/// This function is expected to be deprecated in favor of the `deallocate` method /// of the [`Global`] type when it and the [`Allocator`] trait become stable. /// /// # Safety @@ -125,7 +125,7 @@ pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) { /// of the allocator registered with the `#[global_allocator]` attribute /// if there is one, or the `std` crate’s default. /// -/// This function is expected to be deprecated in favor of the `realloc` method +/// This function is expected to be deprecated in favor of the `grow` and `shrink` methods /// of the [`Global`] type when it and the [`Allocator`] trait become stable. /// /// # Safety @@ -145,7 +145,7 @@ pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 /// of the allocator registered with the `#[global_allocator]` attribute /// if there is one, or the `std` crate’s default. /// -/// This function is expected to be deprecated in favor of the `alloc_zeroed` method +/// This function is expected to be deprecated in favor of the `allocate_zeroed` method /// of the [`Global`] type when it and the [`Allocator`] trait become stable. /// /// # Safety diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs index 75bb86916a8..0da6af54bc2 100644 --- a/library/alloc/src/collections/btree/map/entry.rs +++ b/library/alloc/src/collections/btree/map/entry.rs @@ -269,6 +269,31 @@ impl<'a, K: Ord, V, A: Allocator + Clone> Entry<'a, K, V, A> { Vacant(entry) => Vacant(entry), } } + + /// Sets the value of the entry, and returns an `OccupiedEntry`. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_entry_insert)] + /// use std::collections::BTreeMap; + /// + /// let mut map: BTreeMap<&str, String> = BTreeMap::new(); + /// let entry = map.entry("poneyland").insert_entry("hoho".to_string()); + /// + /// assert_eq!(entry.key(), &"poneyland"); + /// ``` + #[inline] + #[unstable(feature = "btree_entry_insert", issue = "65225")] + pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V, A> { + match self { + Occupied(mut entry) => { + entry.insert(value); + entry + } + Vacant(entry) => entry.insert_entry(value), + } + } } impl<'a, K: Ord, V: Default, A: Allocator + Clone> Entry<'a, K, V, A> { @@ -348,41 +373,61 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> { /// ``` #[stable(feature = "rust1", since = "1.0.0")] #[rustc_confusables("push", "put")] - pub fn insert(mut self, value: V) -> &'a mut V { - let out_ptr = match self.handle { + pub fn insert(self, value: V) -> &'a mut V { + self.insert_entry(value).into_mut() + } + + /// Sets the value of the entry with the `VacantEntry`'s key, + /// and returns an `OccupiedEntry`. + /// + /// # Examples + /// + /// ``` + /// #![feature(btree_entry_insert)] + /// use std::collections::BTreeMap; + /// use std::collections::btree_map::Entry; + /// + /// let mut map: BTreeMap<&str, u32> = BTreeMap::new(); + /// + /// if let Entry::Vacant(o) = map.entry("poneyland") { + /// let entry = o.insert_entry(37); + /// assert_eq!(entry.get(), &37); + /// } + /// assert_eq!(map["poneyland"], 37); + /// ``` + #[unstable(feature = "btree_entry_insert", issue = "65225")] + pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V, A> { + let handle = match self.handle { None => { // SAFETY: There is no tree yet so no reference to it exists. - let map = unsafe { self.dormant_map.awaken() }; - let mut root = NodeRef::new_leaf(self.alloc.clone()); - let val_ptr = root.borrow_mut().push(self.key, value); - map.root = Some(root.forget_type()); - map.length = 1; - val_ptr - } - Some(handle) => { - let new_handle = - handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| { - drop(ins.left); - // SAFETY: Pushing a new root node doesn't invalidate - // handles to existing nodes. - let map = unsafe { self.dormant_map.reborrow() }; - let root = map.root.as_mut().unwrap(); // same as ins.left - root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right) - }); - - // Get the pointer to the value - let val_ptr = new_handle.into_val_mut(); - - // SAFETY: We have consumed self.handle. - let map = unsafe { self.dormant_map.awaken() }; - map.length += 1; - val_ptr + let map = unsafe { self.dormant_map.reborrow() }; + let root = map.root.insert(NodeRef::new_leaf(self.alloc.clone()).forget_type()); + // SAFETY: We *just* created the root as a leaf, and we're + // stacking the new handle on the original borrow lifetime. + unsafe { + let mut leaf = root.borrow_mut().cast_to_leaf_unchecked(); + leaf.push_with_handle(self.key, value) + } } + Some(handle) => handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| { + drop(ins.left); + // SAFETY: Pushing a new root node doesn't invalidate + // handles to existing nodes. + let map = unsafe { self.dormant_map.reborrow() }; + let root = map.root.as_mut().unwrap(); // same as ins.left + root.push_internal_level(self.alloc.clone()).push(ins.kv.0, ins.kv.1, ins.right) + }), }; - // 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. - unsafe { &mut *out_ptr } + // SAFETY: modifying the length doesn't invalidate handles to existing nodes. + unsafe { self.dormant_map.reborrow().length += 1 }; + + OccupiedEntry { + handle: handle.forget_node_type(), + dormant_map: self.dormant_map, + alloc: self.alloc, + _marker: PhantomData, + } } } diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index db16d82be7d..5975134382e 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2272,6 +2272,54 @@ fn test_into_iter_drop_leak_height_0() { #[test] #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] +fn test_into_iter_drop_leak_kv_panic_in_key() { + let a_k = CrashTestDummy::new(0); + let a_v = CrashTestDummy::new(1); + let b_k = CrashTestDummy::new(2); + let b_v = CrashTestDummy::new(3); + let c_k = CrashTestDummy::new(4); + let c_v = CrashTestDummy::new(5); + let mut map = BTreeMap::new(); + map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never)); + map.insert(b_k.spawn(Panic::InDrop), b_v.spawn(Panic::Never)); + map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never)); + + catch_unwind(move || drop(map.into_iter())).unwrap_err(); + + assert_eq!(a_k.dropped(), 1); + assert_eq!(a_v.dropped(), 1); + assert_eq!(b_k.dropped(), 1); + assert_eq!(b_v.dropped(), 1); + assert_eq!(c_k.dropped(), 1); + assert_eq!(c_v.dropped(), 1); +} + +#[test] +#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] +fn test_into_iter_drop_leak_kv_panic_in_val() { + let a_k = CrashTestDummy::new(0); + let a_v = CrashTestDummy::new(1); + let b_k = CrashTestDummy::new(2); + let b_v = CrashTestDummy::new(3); + let c_k = CrashTestDummy::new(4); + let c_v = CrashTestDummy::new(5); + let mut map = BTreeMap::new(); + map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never)); + map.insert(b_k.spawn(Panic::Never), b_v.spawn(Panic::InDrop)); + map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never)); + + catch_unwind(move || drop(map.into_iter())).unwrap_err(); + + assert_eq!(a_k.dropped(), 1); + assert_eq!(a_v.dropped(), 1); + assert_eq!(b_k.dropped(), 1); + assert_eq!(b_v.dropped(), 1); + assert_eq!(c_k.dropped(), 1); + assert_eq!(c_v.dropped(), 1); +} + +#[test] +#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")] fn test_into_iter_drop_leak_height_1() { let size = MIN_INSERTS_HEIGHT_1; for panic_point in vec![0, 1, size - 2, size - 1] { diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 2a853ef4216..0c93eff0d20 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -739,7 +739,7 @@ impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> { impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> { /// Unsafely asserts to the compiler the static information that this node is a `Leaf`. - unsafe fn cast_to_leaf_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> { + pub unsafe fn cast_to_leaf_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> { debug_assert!(self.height == 0); NodeRef { height: self.height, node: self.node, _marker: PhantomData } } @@ -1173,11 +1173,25 @@ impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV> /// The node that the handle refers to must not yet have been deallocated. #[inline] pub unsafe fn drop_key_val(mut self) { + // Run the destructor of the value even if the destructor of the key panics. + struct Dropper<'a, T>(&'a mut MaybeUninit<T>); + impl<T> Drop for Dropper<'_, T> { + #[inline] + fn drop(&mut self) { + unsafe { + self.0.assume_init_drop(); + } + } + } + debug_assert!(self.idx < self.node.len()); let leaf = self.node.as_leaf_dying(); unsafe { - leaf.keys.get_unchecked_mut(self.idx).assume_init_drop(); - leaf.vals.get_unchecked_mut(self.idx).assume_init_drop(); + let key = leaf.keys.get_unchecked_mut(self.idx); + let val = leaf.vals.get_unchecked_mut(self.idx); + let _guard = Dropper(val); + key.assume_init_drop(); + // dropping the guard will drop the value } } } |