diff --git a/src/map.rs b/src/map.rs index 4ee24d5f..c1bc2660 100644 --- a/src/map.rs +++ b/src/map.rs @@ -32,7 +32,7 @@ use alloc::vec::Vec; #[cfg(feature = "std")] use std::collections::hash_map::RandomState; -use self::core::IndexMapCore; +use self::core::{EntryAt, IndexMapCore}; use crate::util::{third, try_simplify_range}; use crate::{Bucket, Entries, Equivalent, HashValue, TryReserveError}; @@ -423,6 +423,18 @@ where self.core.entry(hash, key) } + /// Get an entry in the map by index for in-place manipulation. + /// + /// Valid indices are *0 <= index < self.len()* + /// + /// Computes in **O(1)** time. + pub fn entry_at(&mut self, index: usize) -> Option> { + if index >= self.len() { + return None; + } + Some(EntryAt::new(&mut self.core, index)) + } + /// Return `true` if an equivalent to `key` exists in the map. /// /// Computes in **O(1)** time (average). diff --git a/src/map/core.rs b/src/map/core.rs index 4a78035c..06f15f61 100644 --- a/src/map/core.rs +++ b/src/map/core.rs @@ -754,9 +754,128 @@ impl fmt::Debug for VacantEntry<'_, K, V> { } } +/// A view into an occupied entry in a `IndexMap` obtained by index. +/// +/// This struct is constructed from the `entry_at` method on `IndexMap`. +// SAFETY: We have a mutable reference to the map, which keeps the index +// valid and pointing to the correct entry. +pub struct EntryAt<'a, K, V> { + map: &'a mut IndexMapCore, + index: usize, +} + +impl<'a, K, V> EntryAt<'a, K, V> { + pub(crate) fn new(map: &'a mut IndexMapCore, index: usize) -> Self { + Self { map, index } + } + + /// Gets a reference to the entry's key in the map. + pub fn key(&self) -> &K { + &self.map.entries[self.index].key + } + + /// Gets a reference to the entry's value in the map. + pub fn get(&self) -> &V { + &self.map.entries[self.index].value + } + + /// Gets a mutable reference to the entry's value in the map. + /// + /// If you need a reference which may outlive the destruction of the + /// `Entry` value, see `into_mut`. + pub fn get_mut(&mut self) -> &mut V { + &mut self.map.entries[self.index].value + } + + /// Sets the value of the entry to `value`, and returns the entry's old value. + pub fn insert(&mut self, value: V) -> V { + mem::replace(self.get_mut(), value) + } + + /// Return the index of the key-value pair + #[inline] + pub fn index(&self) -> usize { + self.index + } + + /// Converts into a mutable reference to the entry's value in the map, + /// with a lifetime bound to the map itself. + pub fn into_mut(self) -> &'a mut V { + &mut self.map.entries[self.index].value + } + + /// Remove and return the key, value pair stored in the map for this entry + /// + /// Like `Vec::swap_remove`, the pair is removed by swapping it with the + /// last element of the map and popping it off. **This perturbs + /// the position of what used to be the last element!** + /// + /// Computes in **O(1)** time (average). + pub fn swap_remove_entry(self) -> (K, V) { + self.map.swap_remove_index(self.index).unwrap() + } + + /// Remove and return the key, value pair stored in the map for this entry + /// + /// Like `Vec::remove`, the pair is removed by shifting all of the + /// elements that follow it, preserving their relative order. + /// **This perturbs the index of all of those elements!** + /// + /// Computes in **O(n)** time (average). + pub fn shift_remove_entry(self) -> (K, V) { + self.map.shift_remove_index(self.index).unwrap() + } + + /// Remove the key, value pair stored in the map for this entry, and return the value. + /// + /// **NOTE:** This is equivalent to `.swap_remove()`. + pub fn remove(self) -> V { + self.swap_remove() + } + + /// Remove the key, value pair stored in the map for this entry, and return the value. + /// + /// Like `Vec::swap_remove`, the pair is removed by swapping it with the + /// last element of the map and popping it off. **This perturbs + /// the position of what used to be the last element!** + /// + /// Computes in **O(1)** time (average). + pub fn swap_remove(self) -> V { + self.swap_remove_entry().1 + } + + /// Remove the key, value pair stored in the map for this entry, and return the value. + /// + /// Like `Vec::remove`, the pair is removed by shifting all of the + /// elements that follow it, preserving their relative order. + /// **This perturbs the index of all of those elements!** + /// + /// Computes in **O(n)** time (average). + pub fn shift_remove(self) -> V { + self.shift_remove_entry().1 + } + + /// Remove and return the key, value pair stored in the map for this entry + /// + /// **NOTE:** This is equivalent to `.swap_remove_entry()`. + pub fn remove_entry(self) -> (K, V) { + self.swap_remove_entry() + } +} + +impl fmt::Debug for EntryAt<'_, K, V> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct(stringify!(OccupiedEntry)) + .field("key", self.key()) + .field("value", self.get()) + .finish() + } +} + #[test] fn assert_send_sync() { fn assert_send_sync() {} assert_send_sync::>(); assert_send_sync::>(); + assert_send_sync::>(); } diff --git a/src/map/tests.rs b/src/map/tests.rs index 3d2315d3..b11811d0 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -359,6 +359,36 @@ fn occupied_entry_key() { } } +#[test] +fn entry_at() { + let mut map = IndexMap::new(); + + assert!(map.entry_at(0).is_none()); + + map.insert(0, "0"); + map.insert(1, "1"); + map.insert(2, "2"); + map.insert(3, "3"); + + assert!(map.entry_at(4).is_none()); + + { + let e = map.entry_at(1).unwrap(); + assert_eq!(*e.key(), 1); + assert_eq!(*e.get(), "1"); + assert_eq!(e.remove(), "1"); + } + + { + let mut e = map.entry_at(1).unwrap(); + assert_eq!(*e.key(), 3); + assert_eq!(*e.get(), "3"); + assert_eq!(e.insert("4"), "3"); + } + + assert_eq!(*map.get(&3).unwrap(), "4"); +} + #[test] fn keys() { let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];