Merge pull request #294 from cuviper/splice

Add splicing iterators

Author: cuviper

src/map.rs

@@ -14,7 +14,7 @@ mod tests;
 
 pub use self::core::{Entry, IndexedEntry, OccupiedEntry, VacantEntry};
 pub use self::iter::{
-    Drain, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, Values, ValuesMut,
+    Drain, IntoIter, IntoKeys, IntoValues, Iter, IterMut, Keys, Splice, Values, ValuesMut,
 };
 pub use self::slice::Slice;
 pub use crate::mutable_keys::MutableKeys;
@@ -418,6 +418,44 @@ where
         let hash = self.hash(&key);
         self.core.entry(hash, key)
     }
+
+    /// Creates a splicing iterator that replaces the specified range in the map
+    /// with the given `replace_with` key-value iterator and yields the removed
+    /// items. `replace_with` does not need to be the same length as `range`.
+    ///
+    /// The `range` is removed even if the iterator is not consumed until the
+    /// end. It is unspecified how many elements are removed from the map if the
+    /// `Splice` value is leaked.
+    ///
+    /// The input iterator `replace_with` is only consumed when the `Splice`
+    /// value is dropped. If a key from the iterator matches an existing entry
+    /// in the map (outside of `range`), then the value will be updated in that
+    /// position. Otherwise, the new key-value pair will be inserted in the
+    /// replaced `range`.
+    ///
+    /// ***Panics*** if the starting point is greater than the end point or if
+    /// the end point is greater than the length of the map.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use indexmap::IndexMap;
+    ///
+    /// let mut map = IndexMap::from([(0, '_'), (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')]);
+    /// let new = [(5, 'E'), (4, 'D'), (3, 'C'), (2, 'B'), (1, 'A')];
+    /// let removed: Vec<_> = map.splice(2..4, new).collect();
+    ///
+    /// // 1 and 4 got new values, while 5, 3, and 2 were newly inserted.
+    /// assert!(map.into_iter().eq([(0, '_'), (1, 'A'), (5, 'E'), (3, 'C'), (2, 'B'), (4, 'D')]));
+    /// assert_eq!(removed, &[(2, 'b'), (3, 'c')]);
+    /// ```
+    pub fn splice<R, I>(&mut self, range: R, replace_with: I) -> Splice<'_, I::IntoIter, K, V, S>
+    where
+        R: RangeBounds<usize>,
+        I: IntoIterator<Item = (K, V)>,
+    {
+        Splice::new(self, range, replace_with.into_iter())
+    }
 }
 
 impl<K, V, S> IndexMap<K, V, S>

src/map/core.rs

@@ -11,7 +11,7 @@ mod raw;
 
 use hashbrown::raw::RawTable;
 
-use crate::vec::{Drain, Vec};
+use crate::vec::{self, Vec};
 use crate::TryReserveError;
 use core::fmt;
 use core::mem;
@@ -160,7 +160,7 @@ impl<K, V> IndexMapCore<K, V> {
         }
     }
 
-    pub(crate) fn drain<R>(&mut self, range: R) -> Drain<'_, Bucket<K, V>>
+    pub(crate) fn drain<R>(&mut self, range: R) -> vec::Drain<'_, Bucket<K, V>>
     where
         R: RangeBounds<usize>,
     {
@@ -192,6 +192,28 @@ impl<K, V> IndexMapCore<K, V> {
         Self { indices, entries }
     }
 
+    pub(crate) fn split_splice<R>(&mut self, range: R) -> (Self, vec::IntoIter<Bucket<K, V>>)
+    where
+        R: RangeBounds<usize>,
+    {
+        let range = simplify_range(range, self.len());
+        self.erase_indices(range.start, self.entries.len());
+        let entries = self.entries.split_off(range.end);
+        let drained = self.entries.split_off(range.start);
+
+        let mut indices = RawTable::with_capacity(entries.len());
+        raw::insert_bulk_no_grow(&mut indices, &entries);
+        (Self { indices, entries }, drained.into_iter())
+    }
+
+    /// Append from another map without checking whether items already exist.
+    pub(crate) fn append_unchecked(&mut self, other: &mut Self) {
+        self.reserve(other.len());
+        raw::insert_bulk_no_grow(&mut self.indices, &other.entries);
+        self.entries.append(&mut other.entries);
+        other.indices.clear();
+    }
+
     /// Reserve capacity for `additional` more key-value pairs.
     pub(crate) fn reserve(&mut self, additional: usize) {
         self.indices.reserve(additional, get_hash(&self.entries));

src/map/iter.rs

@@ -1,9 +1,11 @@
+use super::core::IndexMapCore;
 use super::{Bucket, Entries, IndexMap, Slice};
 
 use alloc::vec::{self, Vec};
 use core::fmt;
+use core::hash::{BuildHasher, Hash};
 use core::iter::FusedIterator;
-use core::ops::Index;
+use core::ops::{Index, RangeBounds};
 use core::slice;
 
 impl<'a, K, V, S> IntoIterator for &'a IndexMap<K, V, S> {
@@ -580,3 +582,132 @@ impl<K, V> Default for IntoValues<K, V> {
         }
     }
 }
+
+/// A splicing iterator for `IndexMap`.
+///
+/// This `struct` is created by [`IndexMap::splice()`].
+/// See its documentation for more.
+pub struct Splice<'a, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    map: &'a mut IndexMap<K, V, S>,
+    tail: IndexMapCore<K, V>,
+    drain: vec::IntoIter<Bucket<K, V>>,
+    replace_with: I,
+}
+
+impl<'a, I, K, V, S> Splice<'a, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    pub(super) fn new<R>(map: &'a mut IndexMap<K, V, S>, range: R, replace_with: I) -> Self
+    where
+        R: RangeBounds<usize>,
+    {
+        let (tail, drain) = map.core.split_splice(range);
+        Self {
+            map,
+            tail,
+            drain,
+            replace_with,
+        }
+    }
+}
+
+impl<I, K, V, S> Drop for Splice<'_, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    fn drop(&mut self) {
+        // Finish draining unconsumed items. We don't strictly *have* to do this
+        // manually, since we already split it into separate memory, but it will
+        // match the drop order of `vec::Splice` items this way.
+        let _ = self.drain.nth(usize::MAX);
+
+        // Now insert all the new items. If a key matches an existing entry, it
+        // keeps the original position and only replaces the value, like `insert`.
+        while let Some((key, value)) = self.replace_with.next() {
+            // Since the tail is disjoint, we can try to update it first,
+            // or else insert (update or append) the primary map.
+            let hash = self.map.hash(&key);
+            if let Some(i) = self.tail.get_index_of(hash, &key) {
+                self.tail.as_entries_mut()[i].value = value;
+            } else {
+                self.map.core.insert_full(hash, key, value);
+            }
+        }
+
+        // Finally, re-append the tail
+        self.map.core.append_unchecked(&mut self.tail);
+    }
+}
+
+impl<I, K, V, S> Iterator for Splice<'_, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.drain.next().map(Bucket::key_value)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.drain.size_hint()
+    }
+}
+
+impl<I, K, V, S> DoubleEndedIterator for Splice<'_, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.drain.next_back().map(Bucket::key_value)
+    }
+}
+
+impl<I, K, V, S> ExactSizeIterator for Splice<'_, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+    fn len(&self) -> usize {
+        self.drain.len()
+    }
+}
+
+impl<I, K, V, S> FusedIterator for Splice<'_, I, K, V, S>
+where
+    I: Iterator<Item = (K, V)>,
+    K: Hash + Eq,
+    S: BuildHasher,
+{
+}
+
+impl<'a, I, K, V, S> fmt::Debug for Splice<'a, I, K, V, S>
+where
+    I: fmt::Debug + Iterator<Item = (K, V)>,
+    K: fmt::Debug + Hash + Eq,
+    V: fmt::Debug,
+    S: BuildHasher,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Follow `vec::Splice` in only printing the drain and replacement
+        f.debug_struct("Splice")
+            .field("drain", &self.drain)
+            .field("replace_with", &self.replace_with)
+            .finish()
+    }
+}

src/set.rs

@@ -6,7 +6,9 @@ mod slice;
 #[cfg(test)]
 mod tests;
 
-pub use self::iter::{Difference, Drain, Intersection, IntoIter, Iter, SymmetricDifference, Union};
+pub use self::iter::{
+    Difference, Drain, Intersection, IntoIter, Iter, Splice, SymmetricDifference, Union,
+};
 pub use self::slice::Slice;
 
 #[cfg(feature = "rayon")]
@@ -423,6 +425,43 @@ where
     {
         Union::new(self, other)
     }
+
+    /// Creates a splicing iterator that replaces the specified range in the set
+    /// with the given `replace_with` iterator and yields the removed items.
+    /// `replace_with` does not need to be the same length as `range`.
+    ///
+    /// The `range` is removed even if the iterator is not consumed until the
+    /// end. It is unspecified how many elements are removed from the set if the
+    /// `Splice` value is leaked.
+    ///
+    /// The input iterator `replace_with` is only consumed when the `Splice`
+    /// value is dropped. If a value from the iterator matches an existing entry
+    /// in the set (outside of `range`), then the original will be unchanged.
+    /// Otherwise, the new value will be inserted in the replaced `range`.
+    ///
+    /// ***Panics*** if the starting point is greater than the end point or if
+    /// the end point is greater than the length of the set.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use indexmap::IndexSet;
+    ///
+    /// let mut set = IndexSet::from([0, 1, 2, 3, 4]);
+    /// let new = [5, 4, 3, 2, 1];
+    /// let removed: Vec<_> = set.splice(2..4, new).collect();
+    ///
+    /// // 1 and 4 kept their positions, while 5, 3, and 2 were newly inserted.
+    /// assert!(set.into_iter().eq([0, 1, 5, 3, 2, 4]));
+    /// assert_eq!(removed, &[2, 3]);
+    /// ```
+    pub fn splice<R, I>(&mut self, range: R, replace_with: I) -> Splice<'_, I::IntoIter, T, S>
+    where
+        R: RangeBounds<usize>,
+        I: IntoIterator<Item = T>,
+    {
+        Splice::new(self, range, replace_with.into_iter())
+    }
 }
 
 impl<T, S> IndexSet<T, S>