diff --git a/src/map.rs b/src/map.rs index 872d4547..dd7b6f0b 100644 --- a/src/map.rs +++ b/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 { + // 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. /// diff --git a/src/map/core.rs b/src/map/core.rs index cd0a966c..b7c55b66 100644 --- a/src/map/core.rs +++ b/src/map/core.rs @@ -385,6 +385,13 @@ impl IndexMapCore { } } + /// 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(&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) { diff --git a/src/map/core/entry.rs b/src/map/core/entry.rs index 39a20d5f..b4aba740 100644 --- a/src/map/core/entry.rs +++ b/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 fmt::Debug for VacantEntry<'_, K, V> { diff --git a/src/set.rs b/src/set.rs index 3e54463e..65fe224d 100644 --- a/src/set.rs +++ b/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 { + 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`. diff --git a/tests/quick.rs b/tests/quick.rs index e9682115..99bce6dc 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -130,6 +130,60 @@ quickcheck_limit! { true } + fn replace_index(insert: Vec, 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, 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) -> bool { let mut map = IndexMap::new(); for &key in &insert {