Add methods to replace the key at a given index

src/map.rs

@@ -652,6 +652,37 @@ where
         }
     }
 
+    /// Replaces the key at the given index. The new key does not need to be
+    /// equivalent to the one it is replacing, but it must be unique to the rest
+    /// of the map.
+    ///
+    /// Returns `Ok(old_key)` if successful, or `Err((other_index, key))` if an
+    /// equivalent key already exists at a different index. The map will be
+    /// unchanged in the error case.
+    ///
+    /// Direct indexing can be used to change the corresponding value: simply
+    /// `map[index] = value`, or `mem::replace(&mut map[index], value)` to
+    /// retrieve the old value as well.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    #[track_caller]
+    pub fn replace_index(&mut self, index: usize, key: K) -> Result<K, (usize, K)> {
+        // If there's a direct match, we don't even need to hash it.
+        let entry = &mut self.as_entries_mut()[index];
+        if key == entry.key {
+            return Ok(mem::replace(&mut entry.key, key));
+        }
+
+        let hash = self.hash(&key);
+        if let Some(i) = self.core.get_index_of(hash, &key) {
+            debug_assert_ne!(i, index);
+            return Err((i, key));
+        }
+        Ok(self.core.replace_index_unique(index, hash, key))
+    }
+
     /// Get the given key’s corresponding entry in the map for insertion and/or
     /// in-place manipulation.
     ///

src/map/core.rs

@@ -385,6 +385,13 @@ impl<K, V> IndexMapCore<K, V> {
         }
     }
 
+    /// Replaces the key at the given index,
+    /// *without* checking whether it already exists.
+    #[track_caller]
+    pub(crate) fn replace_index_unique(&mut self, index: usize, hash: HashValue, key: K) -> K {
+        self.borrow_mut().replace_index_unique(index, hash, key).0
+    }
+
     /// Remove an entry by shifting all entries that follow it
     pub(crate) fn shift_remove_full<Q>(&mut self, hash: HashValue, key: &Q) -> Option<(usize, K, V)>
     where
@@ -561,6 +568,28 @@ impl<'a, K, V> RefMut<'a, K, V> {
         OccupiedEntry::new(self.entries, entry)
     }
 
+    /// Replaces the key at the given index,
+    /// *without* checking whether it already exists.
+    #[track_caller]
+    fn replace_index_unique(
+        self,
+        index: usize,
+        hash: HashValue,
+        key: K,
+    ) -> (K, OccupiedEntry<'a, K, V>) {
+        // NB: This removal and insertion isn't "no grow" (with unreachable hasher)
+        // because hashbrown's tombstones might force a resize anyway.
+        erase_index(self.indices, self.entries[index].hash, index);
+        let table_entry = self
+            .indices
+            .insert_unique(hash.get(), index, get_hash(&self.entries));
+
+        let entry = &mut self.entries[index];
+        entry.hash = hash;
+        let old_key = mem::replace(&mut entry.key, key);
+        (old_key, OccupiedEntry::new(self.entries, table_entry))
+    }
+
     /// Insert a key-value pair in `entries` at a particular index,
     /// *without* checking whether it already exists.
     fn shift_insert_unique(&mut self, index: usize, hash: HashValue, key: K, value: V) {

src/map/core/entry.rs

@@ -414,6 +414,17 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
             .shift_insert_unique(index, self.hash, self.key, value);
         &mut self.map.entries[index].value
     }
+
+    /// Replaces the key at the given index with this entry's key, returning the
+    /// old key and an `OccupiedEntry` for that index.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    #[track_caller]
+    pub fn replace_index(self, index: usize) -> (K, OccupiedEntry<'a, K, V>) {
+        self.map.replace_index_unique(index, self.hash, self.key)
+    }
 }
 
 impl<K: fmt::Debug, V> fmt::Debug for VacantEntry<'_, K, V> {

src/set.rs

@@ -550,6 +550,22 @@ where
         }
     }
 
+    /// Replaces the value at the given index. The new value does not need to be
+    /// equivalent to the one it is replacing, but it must be unique to the rest
+    /// of the set.
+    ///
+    /// Returns `Ok(old_value)` if successful, or `Err((other_index, value))` if
+    /// an equivalent value already exists at a different index. The set will be
+    /// unchanged in the error case.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    #[track_caller]
+    pub fn replace_index(&mut self, index: usize, value: T) -> Result<T, (usize, T)> {
+        self.map.replace_index(index, value)
+    }
+
     /// Return an iterator over the values that are in `self` but not `other`.
     ///
     /// Values are produced in the same order that they appear in `self`.

tests/quick.rs

@@ -130,6 +130,60 @@ quickcheck_limit! {
         true
     }
 
+    fn replace_index(insert: Vec<u8>, index: u8, new_key: u8) -> TestResult {
+        if insert.is_empty() {
+            return TestResult::discard();
+        }
+        let mut map = IndexMap::new();
+        for &key in &insert {
+            map.insert(key, ());
+        }
+        let mut index = usize::from(index);
+        if index < map.len() {
+            match map.replace_index(index, new_key) {
+                Ok(old_key) => {
+                    assert!(old_key == new_key || !map.contains_key(&old_key));
+                }
+                Err((i, key)) => {
+                    assert_eq!(key, new_key);
+                    index = i;
+                }
+            }
+            assert_eq!(map.get_index_of(&new_key), Some(index));
+            assert_eq!(map.get_index(index), Some((&new_key, &())));
+            TestResult::passed()
+        } else {
+            TestResult::must_fail(move || map.replace_index(index, new_key))
+        }
+    }
+
+    fn vacant_replace_index(insert: Vec<u8>, index: u8, new_key: u8) -> TestResult {
+        if insert.is_empty() {
+            return TestResult::discard();
+        }
+        let mut map = IndexMap::new();
+        for &key in &insert {
+            map.insert(key, ());
+        }
+        let index = usize::from(index);
+        if let Some((&old_key, &())) = map.get_index(index) {
+            match map.entry(new_key) {
+                Entry::Occupied(_) => return TestResult::discard(),
+                Entry::Vacant(entry) => {
+                    let (replaced_key, entry) = entry.replace_index(index);
+                    assert_eq!(old_key, replaced_key);
+                    assert_eq!(*entry.key(), new_key);
+                }
+            };
+            assert!(!map.contains_key(&old_key));
+            assert_eq!(map.get_index_of(&new_key), Some(index));
+            assert_eq!(map.get_index(index), Some((&new_key, &())));
+            TestResult::passed()
+        } else {
+            TestResult::must_fail(move || map.replace_index(index, new_key))
+        }
+    }
+
     fn pop(insert: Vec<u8>) -> bool {
         let mut map = IndexMap::new();
         for &key in &insert {