Add IndexMap::try_insert and try_insert_full

src/inner/entry.rs

@@ -9,14 +9,9 @@ impl<'a, K, V> Entry<'a, K, V> {
     where
         K: Eq,
     {
-        let eq = equal(&key, &map.entries);
-        match map.indices.find_entry(hash.get(), eq) {
-            Ok(entry) => Entry::Occupied(OccupiedEntry {
-                bucket: entry.bucket_index(),
-                index: *entry.get(),
-                map,
-            }),
-            Err(_) => Entry::Vacant(VacantEntry { map, hash, key }),
+        match OccupiedEntry::find(map, hash, &key) {
+            Ok(entry) => Entry::Occupied(entry),
+            Err(map) => Entry::Vacant(VacantEntry { map, hash, key }),
         }
     }
 }
@@ -32,6 +27,25 @@ pub struct OccupiedEntry<'a, K, V> {
 }
 
 impl<'a, K, V> OccupiedEntry<'a, K, V> {
+    pub(crate) fn find(
+        map: &'a mut Core<K, V>,
+        hash: HashValue,
+        key: &K,
+    ) -> Result<Self, &'a mut Core<K, V>>
+    where
+        K: Eq,
+    {
+        let eq = equal(key, &map.entries);
+        match map.indices.find_entry(hash.get(), eq) {
+            Ok(entry) => Ok(Self {
+                bucket: entry.bucket_index(),
+                index: *entry.get(),
+                map,
+            }),
+            Err(_) => Err(map),
+        }
+    }
+
     /// Constructor for `RawEntryMut::from_hash`
     pub(crate) fn from_hash<F>(
         map: &'a mut Core<K, V>,

src/map.rs

@@ -468,6 +468,40 @@ where
         self.core.insert_full(hash, key, value)
     }
 
+    /// Tries to insert a key-value pair in the map, and returns a mutable reference to the value
+    /// in the new entry.
+    ///
+    /// If the map already had this key present, nothing is updated, and an error is returned with
+    /// the occupied entry and the given key and value.
+    pub fn try_insert(&mut self, key: K, value: V) -> Result<&mut V, OccupiedError<'_, K, V>> {
+        let hash = self.hash(&key);
+        match OccupiedEntry::find(&mut self.core, hash, &key) {
+            Ok(entry) => Err(OccupiedError { entry, key, value }),
+            Err(map) => Ok(map.insert_unique(hash, key, value).value_mut()),
+        }
+    }
+
+    /// Tries to insert a key-value pair in the map, and returns the index and references to the
+    /// key and value in the new entry.
+    ///
+    /// If the map already had this key present, nothing is updated, and an error is returned with
+    /// the occupied entry and the given key and value.
+    pub fn try_insert_full(
+        &mut self,
+        key: K,
+        value: V,
+    ) -> Result<(usize, &K, &mut V), OccupiedError<'_, K, V>> {
+        let hash = self.hash(&key);
+        match OccupiedEntry::find(&mut self.core, hash, &key) {
+            Ok(entry) => Err(OccupiedError { entry, key, value }),
+            Err(map) => {
+                let index = map.len();
+                let (key, value) = map.insert_unique(hash, key, value).ref_mut();
+                Ok((index, key, value))
+            }
+        }
+    }
+
     /// Insert a key-value pair in the map at its ordered position among sorted keys.
     ///
     /// This is equivalent to finding the position with
@@ -1890,3 +1924,27 @@ where
     S: BuildHasher,
 {
 }
+
+/// The error returned by [`try_insert`][IndexMap::try_insert] or
+/// [`try_insert_full`][IndexMap::try_insert_full] when the key already exists.
+///
+/// Contains the occupied entry and the key and value that were not inserted.
+#[non_exhaustive]
+pub struct OccupiedError<'a, K, V> {
+    /// The entry in the map that was already occupied.
+    pub entry: OccupiedEntry<'a, K, V>,
+    /// The key which was not inserted.
+    pub key: K,
+    /// The value which was not inserted.
+    pub value: V,
+}
+
+impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for OccupiedError<'_, K, V> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("OccupiedError")
+            .field("key", self.entry.key())
+            .field("old_value", self.entry.get())
+            .field("new_value", &self.value)
+            .finish_non_exhaustive()
+    }
+}

src/map/tests.rs

@@ -29,15 +29,15 @@ fn insert() {
 
     for (i, &elt) in insert.iter().enumerate() {
         assert_eq!(map.len(), i);
-        map.insert(elt, elt);
+        assert_eq!(map.insert(elt, elt), None);
         assert_eq!(map.len(), i + 1);
         assert_eq!(map.get(&elt), Some(&elt));
         assert_eq!(map[&elt], elt);
     }
     println!("{:?}", map);
 
     for &elt in &not_present {
-        assert!(map.get(&elt).is_none());
+        assert_eq!(map.get(&elt), None);
     }
 }
 
@@ -51,15 +51,15 @@ fn insert_full() {
         assert_eq!(map.len(), i);
         let (index, existing) = map.insert_full(elt, elt);
         assert_eq!(existing, None);
-        assert_eq!(Some(index), map.get_full(&elt).map(|x| x.0));
+        assert_eq!(map.get_full(&elt), Some((index, &elt, &elt)));
         assert_eq!(map.len(), i + 1);
     }
 
     let len = map.len();
     for &elt in &present {
         let (index, existing) = map.insert_full(elt, elt);
         assert_eq!(existing, Some(elt));
-        assert_eq!(Some(index), map.get_full(&elt).map(|x| x.0));
+        assert_eq!(map.get_full(&elt), Some((index, &elt, &elt)));
         assert_eq!(map.len(), len);
     }
 }
@@ -89,6 +89,48 @@ fn insert_2() {
     }
 }
 
+#[test]
+fn try_insert() {
+    let insert = [0, 4, 2, 12, 8, 7, 11, 5];
+    let not_present = [1, 3, 6, 9, 10];
+    let mut map = IndexMap::with_capacity(insert.len());
+
+    for (i, &elt) in insert.iter().enumerate() {
+        assert_eq!(map.len(), i);
+        assert!(map.try_insert(elt, elt).is_ok());
+        assert_eq!(map.len(), i + 1);
+        assert_eq!(map.get(&elt), Some(&elt));
+        assert_eq!(map[&elt], elt);
+    }
+    println!("{:?}", map);
+
+    for &elt in &not_present {
+        assert_eq!(map.get(&elt), None);
+    }
+}
+
+#[test]
+fn try_insert_full() {
+    let insert = vec![9, 2, 7, 1, 4, 6, 13];
+    let present = vec![1, 6, 2];
+    let mut map = IndexMap::with_capacity(insert.len());
+
+    for (i, &elt) in insert.iter().enumerate() {
+        assert_eq!(map.len(), i);
+        let (index, ..) = map.try_insert_full(elt, elt).unwrap();
+        assert_eq!(map.get_full(&elt), Some((index, &elt, &elt)));
+        assert_eq!(map.len(), i + 1);
+    }
+
+    let len = map.len();
+    for &elt in &present {
+        let entry = map.try_insert_full(elt, elt).unwrap_err().entry;
+        assert_eq!(entry.key(), &elt);
+        assert_eq!(Some(entry.index()), map.get_index_of(&elt));
+        assert_eq!(map.len(), len);
+    }
+}
+
 #[test]
 fn insert_order() {
     let insert = [0, 4, 2, 12, 8, 7, 11, 5, 3, 17, 19, 22, 23];