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