Merge pull request #310 from cuviper/shift_insert

Add more indexed methods to entries

src/map.rs

@@ -26,6 +26,7 @@ pub use crate::rayon::map as rayon;
 use ::core::cmp::Ordering;
 use ::core::fmt;
 use ::core::hash::{BuildHasher, Hash, Hasher};
+use ::core::mem;
 use ::core::ops::{Index, IndexMut, RangeBounds};
 use alloc::boxed::Box;
 use alloc::vec::Vec;
@@ -382,14 +383,14 @@ where
     ///
     /// If an equivalent key already exists in the map: the key remains and
     /// retains in its place in the order, its corresponding value is updated
-    /// with `value` and the older value is returned inside `Some(_)`.
+    /// with `value`, and the older value is returned inside `Some(_)`.
     ///
     /// If no equivalent key existed in the map: the new key-value pair is
     /// inserted, last in order, and `None` is returned.
     ///
     /// Computes in **O(1)** time (amortized average).
     ///
-    /// See also [`entry`][Self::entry] if you you want to insert *or* modify,
+    /// See also [`entry`][Self::entry] if you want to insert *or* modify,
     /// or [`insert_full`][Self::insert_full] if you need to get the index of
     /// the corresponding key-value pair.
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
@@ -400,19 +401,48 @@ where
     ///
     /// If an equivalent key already exists in the map: the key remains and
     /// retains in its place in the order, its corresponding value is updated
-    /// with `value` and the older value is returned inside `(index, Some(_))`.
+    /// with `value`, and the older value is returned inside `(index, Some(_))`.
     ///
     /// If no equivalent key existed in the map: the new key-value pair is
     /// inserted, last in order, and `(index, None)` is returned.
     ///
     /// Computes in **O(1)** time (amortized average).
     ///
-    /// See also [`entry`][Self::entry] if you you want to insert *or* modify.
+    /// See also [`entry`][Self::entry] if you want to insert *or* modify.
     pub fn insert_full(&mut self, key: K, value: V) -> (usize, Option<V>) {
         let hash = self.hash(&key);
         self.core.insert_full(hash, key, value)
     }
 
+    /// Insert a key-value pair in the map at the given index.
+    ///
+    /// If an equivalent key already exists in the map: the key remains and
+    /// is moved to the new position in the map, its corresponding value is updated
+    /// with `value`, and the older value is returned inside `Some(_)`.
+    ///
+    /// If no equivalent key existed in the map: the new key-value pair is
+    /// inserted at the given index, and `None` is returned.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    ///
+    /// See also [`entry`][Self::entry] if you want to insert *or* modify,
+    /// perhaps only using the index for new entries with [`VacantEntry::shift_insert`].
+    pub fn shift_insert(&mut self, index: usize, key: K, value: V) -> Option<V> {
+        match self.entry(key) {
+            Entry::Occupied(mut entry) => {
+                let old = mem::replace(entry.get_mut(), value);
+                entry.move_index(index);
+                Some(old)
+            }
+            Entry::Vacant(entry) => {
+                entry.shift_insert(index, value);
+                None
+            }
+        }
+    }
+
     /// Get the given key’s corresponding entry in the map for insertion and/or
     /// in-place manipulation.
     ///
@@ -1053,6 +1083,8 @@ impl<K, V, S> IndexMap<K, V, S> {
     /// Swaps the position of two key-value pairs in the map.
     ///
     /// ***Panics*** if `a` or `b` are out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
     pub fn swap_indices(&mut self, a: usize, b: usize) {
         self.core.swap_indices(a, b)
     }

src/map/core.rs

@@ -311,6 +311,17 @@ impl<K, V> IndexMapCore<K, V> {
         self.entries.push(Bucket { hash, key, value });
     }
 
+    /// Insert a key-value pair in `entries` at a particular index,
+    /// *without* checking whether it already exists.
+    fn insert_entry(&mut self, index: usize, hash: HashValue, key: K, value: V) {
+        if self.entries.len() == self.entries.capacity() {
+            // Reserve our own capacity synced to the indices,
+            // rather than letting `Vec::insert` just double it.
+            self.reserve_entries(1);
+        }
+        self.entries.insert(index, Bucket { hash, key, value });
+    }
+
     /// Return the index in `entries` where an equivalent key can be found
     pub(crate) fn get_index_of<Q>(&self, hash: HashValue, key: &Q) -> Option<usize>
     where
@@ -361,6 +372,29 @@ impl<K, V> IndexMapCore<K, V> {
         }
     }
 
+    fn insert_unique(&mut self, hash: HashValue, key: K, value: V) -> usize {
+        let i = self.indices.len();
+        self.indices.insert(hash.get(), i, get_hash(&self.entries));
+        debug_assert_eq!(i, self.entries.len());
+        self.push_entry(hash, key, value);
+        i
+    }
+
+    fn shift_insert_unique(&mut self, index: usize, hash: HashValue, key: K, value: V) {
+        let end = self.indices.len();
+        assert!(index <= end);
+        // Increment others first so we don't have duplicate indices.
+        self.increment_indices(index, end);
+        let entries = &*self.entries;
+        self.indices.insert(hash.get(), index, move |&i| {
+            // Adjust for the incremented indices to find hashes.
+            debug_assert_ne!(i, index);
+            let i = if i < index { i } else { i - 1 };
+            entries[i].hash.get()
+        });
+        self.insert_entry(index, hash, key, value);
+    }
+
     /// 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

src/map/core/entry.rs

@@ -1,5 +1,5 @@
 use super::raw::RawTableEntry;
-use super::{get_hash, IndexMapCore};
+use super::IndexMapCore;
 use crate::HashValue;
 use core::{fmt, mem};
 
@@ -237,6 +237,36 @@ impl<'a, K, V> OccupiedEntry<'a, K, V> {
         let (map, index) = self.raw.remove_index();
         map.shift_remove_finish(index)
     }
+
+    /// Moves the position of the entry to a new index
+    /// by shifting all other entries in-between.
+    ///
+    /// This is equivalent to [`IndexMap::move_index`][`crate::IndexMap::move_index`]
+    /// coming `from` the current [`.index()`][Self::index].
+    ///
+    /// * If `self.index() < to`, the other pairs will shift down while the targeted pair moves up.
+    /// * If `self.index() > to`, the other pairs will shift up while the targeted pair moves down.
+    ///
+    /// ***Panics*** if `to` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn move_index(self, to: usize) {
+        let (map, index) = self.raw.into_inner();
+        map.move_index(index, to);
+    }
+
+    /// Swaps the position of entry with another.
+    ///
+    /// This is equivalent to [`IndexMap::swap_indices`][`crate::IndexMap::swap_indices`]
+    /// with the current [`.index()`][Self::index] as one of the two being swapped.
+    ///
+    /// ***Panics*** if the `other` index is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    pub fn swap_indices(self, other: usize) {
+        let (map, index) = self.raw.into_inner();
+        map.swap_indices(index, other)
+    }
 }
 
 impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for OccupiedEntry<'_, K, V> {
@@ -275,13 +305,22 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
     /// Inserts the entry's key and the given value into the map, and returns a mutable reference
     /// to the value.
     pub fn insert(self, value: V) -> &'a mut V {
-        let i = self.index();
         let Self { map, hash, key } = self;
-        map.indices.insert(hash.get(), i, get_hash(&map.entries));
-        debug_assert_eq!(i, map.entries.len());
-        map.push_entry(hash, key, value);
+        let i = map.insert_unique(hash, key, value);
         &mut map.entries[i].value
     }
+
+    /// Inserts the entry's key and the given value into the map at the given index,
+    /// shifting others to the right, and returns a mutable reference to the value.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn shift_insert(self, index: usize, value: V) -> &'a mut V {
+        let Self { map, hash, key } = self;
+        map.shift_insert_unique(index, hash, key, value);
+        &mut map.entries[index].value
+    }
 }
 
 impl<K: fmt::Debug, V> fmt::Debug for VacantEntry<'_, K, V> {
@@ -383,6 +422,34 @@ impl<'a, K, V> IndexedEntry<'a, K, V> {
     pub fn shift_remove(self) -> V {
         self.shift_remove_entry().1
     }
+
+    /// Moves the position of the entry to a new index
+    /// by shifting all other entries in-between.
+    ///
+    /// This is equivalent to [`IndexMap::move_index`][`crate::IndexMap::move_index`]
+    /// coming `from` the current [`.index()`][Self::index].
+    ///
+    /// * If `self.index() < to`, the other pairs will shift down while the targeted pair moves up.
+    /// * If `self.index() > to`, the other pairs will shift up while the targeted pair moves down.
+    ///
+    /// ***Panics*** if `to` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn move_index(self, to: usize) {
+        self.map.move_index(self.index, to);
+    }
+
+    /// Swaps the position of entry with another.
+    ///
+    /// This is equivalent to [`IndexMap::swap_indices`][`crate::IndexMap::swap_indices`]
+    /// with the current [`.index()`][Self::index] as one of the two being swapped.
+    ///
+    /// ***Panics*** if the `other` index is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    pub fn swap_indices(self, other: usize) {
+        self.map.swap_indices(self.index, other)
+    }
 }
 
 impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IndexedEntry<'_, K, V> {

src/map/core/raw.rs

@@ -144,4 +144,10 @@ impl<'a, K, V> RawTableEntry<'a, K, V> {
         let (index, _slot) = unsafe { self.map.indices.remove(self.raw_bucket) };
         (self.map, index)
     }
+
+    /// Take no action, just return the index and the original map reference.
+    pub(super) fn into_inner(self) -> (&'a mut IndexMapCore<K, V>, usize) {
+        let index = self.index();
+        (self.map, index)
+    }
 }

src/map/core/raw_entry_v1.rs

@@ -10,7 +10,7 @@
 //! `IndexMap` without such an opt-in trait.
 
 use super::raw::RawTableEntry;
-use super::{get_hash, IndexMapCore};
+use super::IndexMapCore;
 use crate::{Equivalent, HashValue, IndexMap};
 use core::fmt;
 use core::hash::{BuildHasher, Hash, Hasher};
@@ -539,6 +539,36 @@ impl<'a, K, V, S> RawOccupiedEntryMut<'a, K, V, S> {
         let (map, index) = self.raw.remove_index();
         map.shift_remove_finish(index)
     }
+
+    /// Moves the position of the entry to a new index
+    /// by shifting all other entries in-between.
+    ///
+    /// This is equivalent to [`IndexMap::move_index`]
+    /// coming `from` the current [`.index()`][Self::index].
+    ///
+    /// * If `self.index() < to`, the other pairs will shift down while the targeted pair moves up.
+    /// * If `self.index() > to`, the other pairs will shift up while the targeted pair moves down.
+    ///
+    /// ***Panics*** if `to` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn move_index(self, to: usize) {
+        let (map, index) = self.raw.into_inner();
+        map.move_index(index, to);
+    }
+
+    /// Swaps the position of entry with another.
+    ///
+    /// This is equivalent to [`IndexMap::swap_indices`]
+    /// with the current [`.index()`][Self::index] as one of the two being swapped.
+    ///
+    /// ***Panics*** if the `other` index is out of bounds.
+    ///
+    /// Computes in **O(1)** time (average).
+    pub fn swap_indices(self, other: usize) {
+        let (map, index) = self.raw.into_inner();
+        map.swap_indices(index, other)
+    }
 }
 
 /// A view into a vacant raw entry in an [`IndexMap`].
@@ -575,13 +605,43 @@ impl<'a, K, V, S> RawVacantEntryMut<'a, K, V, S> {
     /// Inserts the given key and value into the map with the provided hash,
     /// and returns mutable references to them.
     pub fn insert_hashed_nocheck(self, hash: u64, key: K, value: V) -> (&'a mut K, &'a mut V) {
-        let i = self.index();
-        let map = self.map;
         let hash = HashValue(hash as usize);
-        map.indices.insert(hash.get(), i, get_hash(&map.entries));
-        debug_assert_eq!(i, map.entries.len());
-        map.push_entry(hash, key, value);
-        map.entries[i].muts()
+        let i = self.map.insert_unique(hash, key, value);
+        self.map.entries[i].muts()
+    }
+
+    /// Inserts the given key and value into the map at the given index,
+    /// shifting others to the right, and returns mutable references to them.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn shift_insert(self, index: usize, key: K, value: V) -> (&'a mut K, &'a mut V)
+    where
+        K: Hash,
+        S: BuildHasher,
+    {
+        let mut h = self.hash_builder.build_hasher();
+        key.hash(&mut h);
+        self.shift_insert_hashed_nocheck(index, h.finish(), key, value)
+    }
+
+    /// Inserts the given key and value into the map with the provided hash
+    /// at the given index, and returns mutable references to them.
+    ///
+    /// ***Panics*** if `index` is out of bounds.
+    ///
+    /// Computes in **O(n)** time (average).
+    pub fn shift_insert_hashed_nocheck(
+        self,
+        index: usize,
+        hash: u64,
+        key: K,
+        value: V,
+    ) -> (&'a mut K, &'a mut V) {
+        let hash = HashValue(hash as usize);
+        self.map.shift_insert_unique(index, hash, key, value);
+        self.map.entries[index].muts()
     }
 }
 

src/map/tests.rs

@@ -108,6 +108,33 @@ fn insert_order() {
     }
 }
 
+#[test]
+fn shift_insert() {
+    let insert = [0, 4, 2, 12, 8, 7, 11, 5, 3, 17, 19, 22, 23];
+    let mut map = IndexMap::new();
+
+    for &elt in &insert {
+        map.shift_insert(0, elt, ());
+    }
+
+    assert_eq!(map.keys().count(), map.len());
+    assert_eq!(map.keys().count(), insert.len());
+    for (a, b) in insert.iter().rev().zip(map.keys()) {
+        assert_eq!(a, b);
+    }
+    for (i, k) in (0..insert.len()).zip(map.keys()) {
+        assert_eq!(map.get_index(i).unwrap().0, k);
+    }
+
+    // "insert" that moves an existing entry
+    map.shift_insert(0, insert[0], ());
+    assert_eq!(map.keys().count(), insert.len());
+    assert_eq!(insert[0], map.keys()[0]);
+    for (a, b) in insert[1..].iter().rev().zip(map.keys().skip(1)) {
+        assert_eq!(a, b);
+    }
+}
+
 #[test]
 fn grow() {
     let insert = [0, 4, 2, 12, 8, 7, 11];