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

   1 use core::fmt::{self, Debug};
   2 use core::marker::PhantomData;
   3 use core::mem;
   4
   5 use crate::alloc::{Allocator, Global};
   6
   7 use super::super::borrow::DormantMutRef;
   8 use super::super::node::{marker, Handle, NodeRef};
   9 use super::BTreeMap;
  10
  11 use Entry::*;
  12
  13 /// A view into a single entry in a map, which may either be vacant or occupied.
  14 ///
  15 /// This `enum` is constructed from the [`entry`] method on [`BTreeMap`].
  16 ///
  17 /// [`entry`]: BTreeMap::entry
  18 #[stable(feature = "rust1", since = "1.0.0")]
  19 #[cfg_attr(not(test), rustc_diagnostic_item = "BTreeEntry")]
  20 pub enum Entry<
  21     'a,
  22     K: 'a,
  23     V: 'a,
  24     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  25 > {
  26     /// A vacant entry.
  27     #[stable(feature = "rust1", since = "1.0.0")]
  28     Vacant(#[stable(feature = "rust1", since = "1.0.0")] VacantEntry<'a, K, V, A>),
  29
  30     /// An occupied entry.
  31     #[stable(feature = "rust1", since = "1.0.0")]
  32     Occupied(#[stable(feature = "rust1", since = "1.0.0")] OccupiedEntry<'a, K, V, A>),
  33 }
  34
  35 #[stable(feature = "debug_btree_map", since = "1.12.0")]
  36 impl<K: Debug + Ord, V: Debug, A: Allocator> Debug for Entry<'_, K, V, A> {
  37     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  38         match *self {
  39             Vacant(ref v) => f.debug_tuple("Entry").field(v).finish(),
  40             Occupied(ref o) => f.debug_tuple("Entry").field(o).finish(),
  41         }
  42     }
  43 }
  44
  45 /// A view into a vacant entry in a `BTreeMap`.
  46 /// It is part of the [`Entry`] enum.
  47 #[stable(feature = "rust1", since = "1.0.0")]
  48 pub struct VacantEntry<
  49     'a,
  50     K,
  51     V,
  52     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  53 > {
  54     pub(super) key: K,
  55     /// `None` for a (empty) map without root
  56     pub(super) handle: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
  57     pub(super) dormant_map: DormantMutRef<'a, BTreeMap<K, V, A>>,
  58
  59     pub(super) alloc: &'a A,
  60
  61     // Be invariant in `K` and `V`
  62     pub(super) _marker: PhantomData<&'a mut (K, V)>,
  63 }
  64
  65 #[stable(feature = "debug_btree_map", since = "1.12.0")]
  66 impl<K: Debug + Ord, V, A: Allocator> Debug for VacantEntry<'_, K, V, A> {
  67     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  68         f.debug_tuple("VacantEntry").field(self.key()).finish()
  69     }
  70 }
  71
  72 /// A view into an occupied entry in a `BTreeMap`.
  73 /// It is part of the [`Entry`] enum.
  74 #[stable(feature = "rust1", since = "1.0.0")]
  75 pub struct OccupiedEntry<
  76     'a,
  77     K,
  78     V,
  79     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
  80 > {
  81     pub(super) handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>,
  82     pub(super) dormant_map: DormantMutRef<'a, BTreeMap<K, V, A>>,
  83
  84     pub(super) alloc: &'a A,
  85
  86     // Be invariant in `K` and `V`
  87     pub(super) _marker: PhantomData<&'a mut (K, V)>,
  88 }
  89
  90 #[stable(feature = "debug_btree_map", since = "1.12.0")]
  91 impl<K: Debug + Ord, V: Debug, A: Allocator> Debug for OccupiedEntry<'_, K, V, A> {
  92     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  93         f.debug_struct("OccupiedEntry").field("key", self.key()).field("value", self.get()).finish()
  94     }
  95 }
  96
  97 /// The error returned by [`try_insert`](BTreeMap::try_insert) when the key already exists.
  98 ///
  99 /// Contains the occupied entry, and the value that was not inserted.
 100 #[unstable(feature = "map_try_insert", issue = "82766")]
 101 pub struct OccupiedError<'a, K: 'a, V: 'a, A: Allocator = Global> {
 102     /// The entry in the map that was already occupied.
 103     pub entry: OccupiedEntry<'a, K, V, A>,
 104     /// The value which was not inserted, because the entry was already occupied.
 105     pub value: V,
 106 }
 107
 108 #[unstable(feature = "map_try_insert", issue = "82766")]
 109 impl<K: Debug + Ord, V: Debug, A: Allocator> Debug for OccupiedError<'_, K, V, A> {
 110     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 111         f.debug_struct("OccupiedError")
 112             .field("key", self.entry.key())
 113             .field("old_value", self.entry.get())
 114             .field("new_value", &self.value)
 115             .finish()
 116     }
 117 }
 118
 119 #[unstable(feature = "map_try_insert", issue = "82766")]
 120 impl<'a, K: Debug + Ord, V: Debug, A: Allocator> fmt::Display for OccupiedError<'a, K, V, A> {
 121     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 122         write!(
 123             f,
 124             "failed to insert {:?}, key {:?} already exists with value {:?}",
 125             self.value,
 126             self.entry.key(),
 127             self.entry.get(),
 128         )
 129     }
 130 }
 131
 132 impl<'a, K: Ord, V, A: Allocator> Entry<'a, K, V, A> {
 133     /// Ensures a value is in the entry by inserting the default if empty, and returns
 134     /// a mutable reference to the value in the entry.
 135     ///
 136     /// # Examples
 137     ///
 138     /// ```
 139     /// use std::collections::BTreeMap;
 140     ///
 141     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 142     /// map.entry("poneyland").or_insert(12);
 143     ///
 144     /// assert_eq!(map["poneyland"], 12);
 145     /// ```
 146     #[stable(feature = "rust1", since = "1.0.0")]
 147     pub fn or_insert(self, default: V) -> &'a mut V {
 148         match self {
 149             Occupied(entry) => entry.into_mut(),
 150             Vacant(entry) => entry.insert(default),
 151         }
 152     }
 153
 154     /// Ensures a value is in the entry by inserting the result of the default function if empty,
 155     /// and returns a mutable reference to the value in the entry.
 156     ///
 157     /// # Examples
 158     ///
 159     /// ```
 160     /// use std::collections::BTreeMap;
 161     ///
 162     /// let mut map: BTreeMap<&str, String> = BTreeMap::new();
 163     /// let s = "hoho".to_string();
 164     ///
 165     /// map.entry("poneyland").or_insert_with(|| s);
 166     ///
 167     /// assert_eq!(map["poneyland"], "hoho".to_string());
 168     /// ```
 169     #[stable(feature = "rust1", since = "1.0.0")]
 170     pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V {
 171         match self {
 172             Occupied(entry) => entry.into_mut(),
 173             Vacant(entry) => entry.insert(default()),
 174         }
 175     }
 176
 177     /// Ensures a value is in the entry by inserting, if empty, the result of the default function.
 178     /// This method allows for generating key-derived values for insertion by providing the default
 179     /// function a reference to the key that was moved during the `.entry(key)` method call.
 180     ///
 181     /// The reference to the moved key is provided so that cloning or copying the key is
 182     /// unnecessary, unlike with `.or_insert_with(|| ... )`.
 183     ///
 184     /// # Examples
 185     ///
 186     /// ```
 187     /// use std::collections::BTreeMap;
 188     ///
 189     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 190     ///
 191     /// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
 192     ///
 193     /// assert_eq!(map["poneyland"], 9);
 194     /// ```
 195     #[inline]
 196     #[stable(feature = "or_insert_with_key", since = "1.50.0")]
 197     pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
 198         match self {
 199             Occupied(entry) => entry.into_mut(),
 200             Vacant(entry) => {
 201                 let value = default(entry.key());
 202                 entry.insert(value)
 203             }
 204         }
 205     }
 206
 207     /// Returns a reference to this entry's key.
 208     ///
 209     /// # Examples
 210     ///
 211     /// ```
 212     /// use std::collections::BTreeMap;
 213     ///
 214     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 215     /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
 216     /// ```
 217     #[stable(feature = "map_entry_keys", since = "1.10.0")]
 218     pub fn key(&self) -> &K {
 219         match *self {
 220             Occupied(ref entry) => entry.key(),
 221             Vacant(ref entry) => entry.key(),
 222         }
 223     }
 224
 225     /// Provides in-place mutable access to an occupied entry before any
 226     /// potential inserts into the map.
 227     ///
 228     /// # Examples
 229     ///
 230     /// ```
 231     /// use std::collections::BTreeMap;
 232     ///
 233     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 234     ///
 235     /// map.entry("poneyland")
 236     ///    .and_modify(|e| { *e += 1 })
 237     ///    .or_insert(42);
 238     /// assert_eq!(map["poneyland"], 42);
 239     ///
 240     /// map.entry("poneyland")
 241     ///    .and_modify(|e| { *e += 1 })
 242     ///    .or_insert(42);
 243     /// assert_eq!(map["poneyland"], 43);
 244     /// ```
 245     #[stable(feature = "entry_and_modify", since = "1.26.0")]
 246     pub fn and_modify<F>(self, f: F) -> Self
 247     where
 248         F: FnOnce(&mut V),
 249     {
 250         match self {
 251             Occupied(mut entry) => {
 252                 f(entry.get_mut());
 253                 Occupied(entry)
 254             }
 255             Vacant(entry) => Vacant(entry),
 256         }
 257     }
 258 }
 259
 260 impl<'a, K: Ord, V: Default, A: Allocator> Entry<'a, K, V, A> {
 261     #[stable(feature = "entry_or_default", since = "1.28.0")]
 262     /// Ensures a value is in the entry by inserting the default value if empty,
 263     /// and returns a mutable reference to the value in the entry.
 264     ///
 265     /// # Examples
 266     ///
 267     /// ```
 268     /// use std::collections::BTreeMap;
 269     ///
 270     /// let mut map: BTreeMap<&str, Option<usize>> = BTreeMap::new();
 271     /// map.entry("poneyland").or_default();
 272     ///
 273     /// assert_eq!(map["poneyland"], None);
 274     /// ```
 275     pub fn or_default(self) -> &'a mut V {
 276         match self {
 277             Occupied(entry) => entry.into_mut(),
 278             Vacant(entry) => entry.insert(Default::default()),
 279         }
 280     }
 281 }
 282
 283 impl<'a, K: Ord, V, A: Allocator> VacantEntry<'a, K, V, A> {
 284     /// Gets a reference to the key that would be used when inserting a value
 285     /// through the VacantEntry.
 286     ///
 287     /// # Examples
 288     ///
 289     /// ```
 290     /// use std::collections::BTreeMap;
 291     ///
 292     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 293     /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
 294     /// ```
 295     #[stable(feature = "map_entry_keys", since = "1.10.0")]
 296     pub fn key(&self) -> &K {
 297         &self.key
 298     }
 299
 300     /// Take ownership of the key.
 301     ///
 302     /// # Examples
 303     ///
 304     /// ```
 305     /// use std::collections::BTreeMap;
 306     /// use std::collections::btree_map::Entry;
 307     ///
 308     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 309     ///
 310     /// if let Entry::Vacant(v) = map.entry("poneyland") {
 311     ///     v.into_key();
 312     /// }
 313     /// ```
 314     #[stable(feature = "map_entry_recover_keys2", since = "1.12.0")]
 315     pub fn into_key(self) -> K {
 316         self.key
 317     }
 318
 319     /// Sets the value of the entry with the `VacantEntry`'s key,
 320     /// and returns a mutable reference to it.
 321     ///
 322     /// # Examples
 323     ///
 324     /// ```
 325     /// use std::collections::BTreeMap;
 326     /// use std::collections::btree_map::Entry;
 327     ///
 328     /// let mut map: BTreeMap<&str, u32> = BTreeMap::new();
 329     ///
 330     /// if let Entry::Vacant(o) = map.entry("poneyland") {
 331     ///     o.insert(37);
 332     /// }
 333     /// assert_eq!(map["poneyland"], 37);
 334     /// ```
 335     #[stable(feature = "rust1", since = "1.0.0")]
 336     pub fn insert(self, value: V) -> &'a mut V {
 337         let out_ptr = match self.handle {
 338             None => {
 339                 // SAFETY: There is no tree yet so no reference to it exists.
 340                 let map = unsafe { self.dormant_map.awaken() };
 341                 let mut root = NodeRef::new_leaf(self.alloc);
 342                 let val_ptr = root.borrow_mut().push(self.key, value) as *mut V;
 343                 map.root = Some(root.forget_type());
 344                 map.length = 1;
 345                 val_ptr
 346             }
 347             Some(handle) => match handle.insert_recursing(self.key, value, self.alloc) {
 348                 (None, val_ptr) => {
 349                     // SAFETY: We have consumed self.handle.
 350                     let map = unsafe { self.dormant_map.awaken() };
 351                     map.length += 1;
 352                     val_ptr
 353                 }
 354                 (Some(ins), val_ptr) => {
 355                     drop(ins.left);
 356                     // SAFETY: We have consumed self.handle and dropped the
 357                     // remaining reference to the tree, ins.left.
 358                     let map = unsafe { self.dormant_map.awaken() };
 359                     let root = map.root.as_mut().unwrap(); // same as ins.left
 360                     root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right);
 361                     map.length += 1;
 362                     val_ptr
 363                 }
 364             },
 365         };
 366         // Now that we have finished growing the tree using borrowed references,
 367         // dereference the pointer to a part of it, that we picked up along the way.
 368         unsafe { &mut *out_ptr }
 369     }
 370 }
 371
 372 impl<'a, K: Ord, V, A: Allocator> OccupiedEntry<'a, K, V, A> {
 373     /// Gets a reference to the key in the entry.
 374     ///
 375     /// # Examples
 376     ///
 377     /// ```
 378     /// use std::collections::BTreeMap;
 379     ///
 380     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 381     /// map.entry("poneyland").or_insert(12);
 382     /// assert_eq!(map.entry("poneyland").key(), &"poneyland");
 383     /// ```
 384     #[must_use]
 385     #[stable(feature = "map_entry_keys", since = "1.10.0")]
 386     pub fn key(&self) -> &K {
 387         self.handle.reborrow().into_kv().0
 388     }
 389
 390     /// Take ownership of the key and value from the map.
 391     ///
 392     /// # Examples
 393     ///
 394     /// ```
 395     /// use std::collections::BTreeMap;
 396     /// use std::collections::btree_map::Entry;
 397     ///
 398     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 399     /// map.entry("poneyland").or_insert(12);
 400     ///
 401     /// if let Entry::Occupied(o) = map.entry("poneyland") {
 402     ///     // We delete the entry from the map.
 403     ///     o.remove_entry();
 404     /// }
 405     ///
 406     /// // If now try to get the value, it will panic:
 407     /// // println!("{}", map["poneyland"]);
 408     /// ```
 409     #[stable(feature = "map_entry_recover_keys2", since = "1.12.0")]
 410     pub fn remove_entry(self) -> (K, V) {
 411         self.remove_kv()
 412     }
 413
 414     /// Gets a reference to the value in the entry.
 415     ///
 416     /// # Examples
 417     ///
 418     /// ```
 419     /// use std::collections::BTreeMap;
 420     /// use std::collections::btree_map::Entry;
 421     ///
 422     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 423     /// map.entry("poneyland").or_insert(12);
 424     ///
 425     /// if let Entry::Occupied(o) = map.entry("poneyland") {
 426     ///     assert_eq!(o.get(), &12);
 427     /// }
 428     /// ```
 429     #[must_use]
 430     #[stable(feature = "rust1", since = "1.0.0")]
 431     pub fn get(&self) -> &V {
 432         self.handle.reborrow().into_kv().1
 433     }
 434
 435     /// Gets a mutable reference to the value in the entry.
 436     ///
 437     /// If you need a reference to the `OccupiedEntry` that may outlive the
 438     /// destruction of the `Entry` value, see [`into_mut`].
 439     ///
 440     /// [`into_mut`]: OccupiedEntry::into_mut
 441     ///
 442     /// # Examples
 443     ///
 444     /// ```
 445     /// use std::collections::BTreeMap;
 446     /// use std::collections::btree_map::Entry;
 447     ///
 448     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 449     /// map.entry("poneyland").or_insert(12);
 450     ///
 451     /// assert_eq!(map["poneyland"], 12);
 452     /// if let Entry::Occupied(mut o) = map.entry("poneyland") {
 453     ///     *o.get_mut() += 10;
 454     ///     assert_eq!(*o.get(), 22);
 455     ///
 456     ///     // We can use the same Entry multiple times.
 457     ///     *o.get_mut() += 2;
 458     /// }
 459     /// assert_eq!(map["poneyland"], 24);
 460     /// ```
 461     #[stable(feature = "rust1", since = "1.0.0")]
 462     pub fn get_mut(&mut self) -> &mut V {
 463         self.handle.kv_mut().1
 464     }
 465
 466     /// Converts the entry into a mutable reference to its value.
 467     ///
 468     /// If you need multiple references to the `OccupiedEntry`, see [`get_mut`].
 469     ///
 470     /// [`get_mut`]: OccupiedEntry::get_mut
 471     ///
 472     /// # Examples
 473     ///
 474     /// ```
 475     /// use std::collections::BTreeMap;
 476     /// use std::collections::btree_map::Entry;
 477     ///
 478     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 479     /// map.entry("poneyland").or_insert(12);
 480     ///
 481     /// assert_eq!(map["poneyland"], 12);
 482     /// if let Entry::Occupied(o) = map.entry("poneyland") {
 483     ///     *o.into_mut() += 10;
 484     /// }
 485     /// assert_eq!(map["poneyland"], 22);
 486     /// ```
 487     #[must_use = "`self` will be dropped if the result is not used"]
 488     #[stable(feature = "rust1", since = "1.0.0")]
 489     pub fn into_mut(self) -> &'a mut V {
 490         self.handle.into_val_mut()
 491     }
 492
 493     /// Sets the value of the entry with the `OccupiedEntry`'s key,
 494     /// and returns the entry's old value.
 495     ///
 496     /// # Examples
 497     ///
 498     /// ```
 499     /// use std::collections::BTreeMap;
 500     /// use std::collections::btree_map::Entry;
 501     ///
 502     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 503     /// map.entry("poneyland").or_insert(12);
 504     ///
 505     /// if let Entry::Occupied(mut o) = map.entry("poneyland") {
 506     ///     assert_eq!(o.insert(15), 12);
 507     /// }
 508     /// assert_eq!(map["poneyland"], 15);
 509     /// ```
 510     #[stable(feature = "rust1", since = "1.0.0")]
 511     pub fn insert(&mut self, value: V) -> V {
 512         mem::replace(self.get_mut(), value)
 513     }
 514
 515     /// Takes the value of the entry out of the map, and returns it.
 516     ///
 517     /// # Examples
 518     ///
 519     /// ```
 520     /// use std::collections::BTreeMap;
 521     /// use std::collections::btree_map::Entry;
 522     ///
 523     /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
 524     /// map.entry("poneyland").or_insert(12);
 525     ///
 526     /// if let Entry::Occupied(o) = map.entry("poneyland") {
 527     ///     assert_eq!(o.remove(), 12);
 528     /// }
 529     /// // If we try to get "poneyland"'s value, it'll panic:
 530     /// // println!("{}", map["poneyland"]);
 531     /// ```
 532     #[stable(feature = "rust1", since = "1.0.0")]
 533     pub fn remove(self) -> V {
 534         self.remove_kv().1
 535     }
 536
 537     // Body of `remove_entry`, probably separate because the name reflects the returned pair.
 538     pub(super) fn remove_kv(self) -> (K, V) {
 539         let mut emptied_internal_root = false;
 540         let (old_kv, _) =
 541             self.handle.remove_kv_tracking(|| emptied_internal_root = true, self.alloc);
 542         // SAFETY: we consumed the intermediate root borrow, `self.handle`.
 543         let map = unsafe { self.dormant_map.awaken() };
 544         map.length -= 1;
 545         if emptied_internal_root {
 546             let root = map.root.as_mut().unwrap();
 547             root.pop_internal_level(&*self.alloc);
 548         }
 549         old_kv
 550     }
 551 }