Skip to content

Commit

Permalink
Add BTreeMap::try_insert and btree_map::OccupiedError.
Browse files Browse the repository at this point in the history
  • Loading branch information
m-ou-se committed Mar 4, 2021
1 parent 7f32f62 commit 09cbcdc
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 1 deletion.
36 changes: 35 additions & 1 deletion library/alloc/src/collections/btree/map.rs
Expand Up @@ -14,7 +14,7 @@ use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::SearchResult::*;

mod entry;
pub use entry::{Entry, OccupiedEntry, VacantEntry};
pub use entry::{Entry, OccupiedEntry, OccupiedError, VacantEntry};
use Entry::*;

/// Minimum number of elements in nodes that are not a root.
Expand Down Expand Up @@ -836,6 +836,40 @@ impl<K, V> BTreeMap<K, V> {
}
}

/// Tries to insert a key-value pair into the map, and returns
/// a mutable reference to the value in the entry.
///
/// If the map already had this key present, nothing is updated, and
/// an error containing the occupied entry and the value is returned.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(map_try_insert)]
///
/// use std::collections::BTreeMap;
///
/// let mut map = BTreeMap::new();
/// assert_eq!(map.try_insert(37, "a").unwrap(), &"a");
///
/// let err = map.try_insert(37, "b").unwrap_err();
/// assert_eq!(err.entry.key(), &37);
/// assert_eq!(err.entry.get(), &"a");
/// assert_eq!(err.value, "b");
/// ```
#[unstable(feature = "map_try_insert", issue = "none")]
pub fn try_insert(&mut self, key: K, value: V) -> Result<&mut V, OccupiedError<'_, K, V>>
where
K: Ord,
{
match self.entry(key) {
Occupied(entry) => Err(OccupiedError { entry, value }),
Vacant(entry) => Ok(entry.insert(value)),
}
}

/// Removes a key from the map, returning the value at the key if the key
/// was previously in the map.
///
Expand Down
21 changes: 21 additions & 0 deletions library/alloc/src/collections/btree/map/entry.rs
Expand Up @@ -71,6 +71,27 @@ impl<K: Debug + Ord, V: Debug> Debug for OccupiedEntry<'_, K, V> {
}
}

/// The error returned by [`try_insert`](BTreeMap::try_insert) when the key already exists.
///
/// Contains the occupied entry, and the value that was not inserted.
#[unstable(feature = "map_try_insert", issue = "none")]
pub struct OccupiedError<'a, K: 'a, V: 'a> {
/// The entry in the map that was already occupied.
pub entry: OccupiedEntry<'a, K, V>,
/// The value which was not inserted, because the entry was already occupied.
pub value: V,
}

#[unstable(feature = "map_try_insert", issue = "none")]
impl<K: Debug + Ord, V: Debug> Debug for OccupiedError<'_, K, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OccupiedError")
.field("entry", &self.entry)
.field("value", &self.value)
.finish()
}
}

impl<'a, K: Ord, V> Entry<'a, K, V> {
/// Ensures a value is in the entry by inserting the default if empty, and returns
/// a mutable reference to the value in the entry.
Expand Down

0 comments on commit 09cbcdc

Please sign in to comment.