Add support for amortized resizes

Author: jonhoo

Cargo.toml

@@ -33,13 +33,20 @@ bench = false
 autocfg = "1"
 [dependencies]
 serde = { version = "1.0", optional = true, default-features = false }
-rayon = { version = "1.0", optional = true }
+rayon_ = { version = "1.0", optional = true, package = "rayon" }
+atone = { version = "0.3.1", optional = true }
 
 [dependencies.hashbrown]
 version = "0.8.1"
 default-features = false
 features = ["raw"]
 
+[dependencies.griddle]
+version = "0.3.1"
+default-features = false
+features = ["raw"]
+optional = true
+
 [dev-dependencies]
 itertools = "0.9"
 rand = {version = "0.7", features = ["small_rng"] }
@@ -50,7 +57,11 @@ fxhash = "0.2.1"
 
 [features]
 # Serialization with serde 1.0
-serde-1 = ["serde"]
+serde-1 = ["serde", "atone/serde", "griddle/serde"]
+rayon = ["rayon_", "atone/rayon", "griddle/rayon"]
+
+# Use griddle over hashbrown, and atone over Vec, for amortized resizes
+amortize = ["griddle", "atone"]
 
 # for testing only, of course
 test_low_transition_point = []

src/lib.rs

@@ -84,6 +84,11 @@ extern crate alloc;
 
 extern crate hashbrown;
 
+#[cfg(feature = "amortize")]
+extern crate atone;
+#[cfg(feature = "amortize")]
+extern crate griddle;
+
 #[cfg(not(has_std))]
 pub(crate) mod std {
     pub use core::*;
@@ -121,6 +126,11 @@ pub use set::IndexSet;
 
 // shared private items
 
+#[cfg(feature = "amortize")]
+type EntryVec<T> = atone::Vc<T>;
+#[cfg(not(feature = "amortize"))]
+type EntryVec<T> = Vec<T>;
+
 /// Hash value newtype. Not larger than usize, since anything larger
 /// isn't used for selecting position anyway.
 #[derive(Clone, Copy, Debug, PartialEq)]
@@ -190,10 +200,10 @@ impl<K, V> Bucket<K, V> {
 
 trait Entries {
     type Entry;
-    fn into_entries(self) -> Vec<Self::Entry>;
-    fn as_entries(&self) -> &[Self::Entry];
-    fn as_entries_mut(&mut self) -> &mut [Self::Entry];
+    fn into_entries(self) -> EntryVec<Self::Entry>;
+    fn as_entries(&self) -> &EntryVec<Self::Entry>;
+    fn as_entries_mut(&mut self) -> &mut EntryVec<Self::Entry>;
     fn with_entries<F>(&mut self, f: F)
     where
-        F: FnOnce(&mut [Self::Entry]);
+        F: FnOnce(&mut EntryVec<Self::Entry>);
 }

src/map.rs

@@ -26,6 +26,7 @@ use std::fmt;
 use self::core::IndexMapCore;
 use equivalent::Equivalent;
 use util::third;
+use EntryVec;
 use {Bucket, Entries, HashValue};
 
 pub use self::core::{Entry, OccupiedEntry, VacantEntry};
@@ -105,23 +106,23 @@ impl<K, V, S> Entries for IndexMap<K, V, S> {
     type Entry = Bucket<K, V>;
 
     #[inline]
-    fn into_entries(self) -> Vec<Self::Entry> {
+    fn into_entries(self) -> EntryVec<Self::Entry> {
         self.core.into_entries()
     }
 
     #[inline]
-    fn as_entries(&self) -> &[Self::Entry] {
+    fn as_entries(&self) -> &EntryVec<Self::Entry> {
         self.core.as_entries()
     }
 
     #[inline]
-    fn as_entries_mut(&mut self) -> &mut [Self::Entry] {
+    fn as_entries_mut(&mut self) -> &mut EntryVec<Self::Entry> {
         self.core.as_entries_mut()
     }
 
     fn with_entries<F>(&mut self, f: F)
     where
-        F: FnOnce(&mut [Self::Entry]),
+        F: FnOnce(&mut EntryVec<Self::Entry>),
     {
         self.core.with_entries(f);
     }
@@ -622,6 +623,8 @@ where
         K: Ord,
     {
         self.with_entries(|entries| {
+            #[cfg(feature = "amortize")]
+            let entries = entries.make_contiguous();
             entries.sort_by(|a, b| Ord::cmp(&a.key, &b.key));
         });
     }
@@ -639,6 +642,8 @@ where
         F: FnMut(&K, &V, &K, &V) -> Ordering,
     {
         self.with_entries(move |entries| {
+            #[cfg(feature = "amortize")]
+            let entries = entries.make_contiguous();
             entries.sort_by(move |a, b| cmp(&a.key, &a.value, &b.key, &b.value));
         });
     }
@@ -652,7 +657,11 @@ where
         F: FnMut(&K, &V, &K, &V) -> Ordering,
     {
         let mut entries = self.into_entries();
-        entries.sort_by(move |a, b| cmp(&a.key, &a.value, &b.key, &b.value));
+        {
+            #[cfg(feature = "amortize")]
+            let entries = entries.make_contiguous();
+            entries.sort_by(move |a, b| cmp(&a.key, &a.value, &b.key, &b.value));
+        }
         IntoIter {
             iter: entries.into_iter(),
         }
@@ -720,10 +729,6 @@ impl<K, V, S> IndexMap<K, V, S> {
     }
 }
 
-use std::slice::Iter as SliceIter;
-use std::slice::IterMut as SliceIterMut;
-use std::vec::IntoIter as VecIntoIter;
-
 /// An iterator over the keys of a `IndexMap`.
 ///
 /// This `struct` is created by the [`keys`] method on [`IndexMap`]. See its
@@ -732,7 +737,7 @@ use std::vec::IntoIter as VecIntoIter;
 /// [`keys`]: struct.IndexMap.html#method.keys
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct Keys<'a, K: 'a, V: 'a> {
-    pub(crate) iter: SliceIter<'a, Bucket<K, V>>,
+    pub(crate) iter: <&'a EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<'a, K, V> Iterator for Keys<'a, K, V> {
@@ -776,7 +781,7 @@ impl<'a, K: fmt::Debug, V> fmt::Debug for Keys<'a, K, V> {
 /// [`values`]: struct.IndexMap.html#method.values
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct Values<'a, K: 'a, V: 'a> {
-    iter: SliceIter<'a, Bucket<K, V>>,
+    iter: <&'a EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<'a, K, V> Iterator for Values<'a, K, V> {
@@ -820,7 +825,7 @@ impl<'a, K, V: fmt::Debug> fmt::Debug for Values<'a, K, V> {
 /// [`values_mut`]: struct.IndexMap.html#method.values_mut
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct ValuesMut<'a, K: 'a, V: 'a> {
-    iter: SliceIterMut<'a, Bucket<K, V>>,
+    iter: <&'a mut EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
@@ -849,7 +854,7 @@ impl<'a, K, V> ExactSizeIterator for ValuesMut<'a, K, V> {
 /// [`iter`]: struct.IndexMap.html#method.iter
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct Iter<'a, K: 'a, V: 'a> {
-    iter: SliceIter<'a, Bucket<K, V>>,
+    iter: <&'a EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<'a, K, V> Iterator for Iter<'a, K, V> {
@@ -893,7 +898,7 @@ impl<'a, K: fmt::Debug, V: fmt::Debug> fmt::Debug for Iter<'a, K, V> {
 /// [`iter_mut`]: struct.IndexMap.html#method.iter_mut
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct IterMut<'a, K: 'a, V: 'a> {
-    iter: SliceIterMut<'a, Bucket<K, V>>,
+    iter: <&'a mut EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<'a, K, V> Iterator for IterMut<'a, K, V> {
@@ -922,7 +927,7 @@ impl<'a, K, V> ExactSizeIterator for IterMut<'a, K, V> {
 /// [`into_iter`]: struct.IndexMap.html#method.into_iter
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct IntoIter<K, V> {
-    pub(crate) iter: VecIntoIter<Bucket<K, V>>,
+    pub(crate) iter: <EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<K, V> Iterator for IntoIter<K, V> {
@@ -944,6 +949,11 @@ impl<K, V> ExactSizeIterator for IntoIter<K, V> {
 }
 
 impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
+    #[cfg(feature = "amortize")]
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("IntoIter").finish()
+    }
+    #[cfg(not(feature = "amortize"))]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let iter = self.iter.as_slice().iter().map(Bucket::refs);
         f.debug_list().entries(iter).finish()
@@ -962,7 +972,10 @@ where
     K: 'a,
     V: 'a,
 {
+    #[cfg(not(feature = "amortize"))]
     pub(crate) iter: ::std::vec::Drain<'a, Bucket<K, V>>,
+    #[cfg(feature = "amortize")]
+    pub(crate) iter: atone::vc::Drain<'a, Bucket<K, V>>,
 }
 
 impl<'a, K, V> Iterator for Drain<'a, K, V> {
@@ -1320,7 +1333,9 @@ mod tests {
             assert_eq!(map.get(&i), Some(&(i * i)));
             map.shrink_to_fit();
             assert_eq!(map.len(), i + 1);
-            assert_eq!(map.capacity(), i + 1);
+            if !cfg!(feature = "amortize") {
+                assert_eq!(map.capacity(), i + 1);
+            }
             assert_eq!(map.get(&i), Some(&(i * i)));
         }
     }

src/map/core.rs

@@ -12,28 +12,36 @@ mod raw;
 #[cfg(not(has_std))]
 use std::vec::Vec;
 
+#[cfg(feature = "amortize")]
+use griddle::raw::RawTable;
+#[cfg(not(feature = "amortize"))]
 use hashbrown::raw::RawTable;
 
 use std::cmp;
 use std::fmt;
 use std::mem::replace;
 use std::ops::RangeFull;
+
+#[cfg(feature = "amortize")]
+use atone::vc::Drain;
+#[cfg(not(feature = "amortize"))]
 use std::vec::Drain;
 
 use equivalent::Equivalent;
 use util::enumerate;
+use EntryVec;
 use {Bucket, Entries, HashValue};
 
 /// Core of the map that does not depend on S
 pub(crate) struct IndexMapCore<K, V> {
     /// indices mapping from the entry hash to its index.
     indices: RawTable<usize>,
     /// entries is a dense vec of entries in their order.
-    entries: Vec<Bucket<K, V>>,
+    entries: EntryVec<Bucket<K, V>>,
 }
 
 #[inline(always)]
-fn get_hash<K, V>(entries: &[Bucket<K, V>]) -> impl Fn(&usize) -> u64 + '_ {
+fn get_hash<K, V>(entries: &EntryVec<Bucket<K, V>>) -> impl Fn(&usize) -> u64 + '_ {
     move |&i| entries[i].hash.get()
 }
 
@@ -43,8 +51,14 @@ where
     V: Clone,
 {
     fn clone(&self) -> Self {
+        #[cfg(feature = "amortize")]
+        let indices = {
+            let hasher = get_hash(&self.entries);
+            self.indices.clone_with_hasher(hasher)
+        };
+        #[cfg(not(feature = "amortize"))]
         let indices = self.indices.clone();
-        let mut entries = Vec::with_capacity(indices.capacity());
+        let mut entries = EntryVec::with_capacity(indices.capacity());
         entries.clone_from(&self.entries);
         IndexMapCore { indices, entries }
     }
@@ -77,23 +91,23 @@ impl<K, V> Entries for IndexMapCore<K, V> {
     type Entry = Bucket<K, V>;
 
     #[inline]
-    fn into_entries(self) -> Vec<Self::Entry> {
+    fn into_entries(self) -> EntryVec<Self::Entry> {
         self.entries
     }
 
     #[inline]
-    fn as_entries(&self) -> &[Self::Entry] {
+    fn as_entries(&self) -> &EntryVec<Self::Entry> {
         &self.entries
     }
 
     #[inline]
-    fn as_entries_mut(&mut self) -> &mut [Self::Entry] {
+    fn as_entries_mut(&mut self) -> &mut EntryVec<Self::Entry> {
         &mut self.entries
     }
 
     fn with_entries<F>(&mut self, f: F)
     where
-        F: FnOnce(&mut [Self::Entry]),
+        F: FnOnce(&mut EntryVec<Self::Entry>),
     {
         f(&mut self.entries);
         self.rebuild_hash_table();
@@ -105,15 +119,15 @@ impl<K, V> IndexMapCore<K, V> {
     pub(crate) fn new() -> Self {
         IndexMapCore {
             indices: RawTable::new(),
-            entries: Vec::new(),
+            entries: EntryVec::new(),
         }
     }
 
     #[inline]
     pub(crate) fn with_capacity(n: usize) -> Self {
         IndexMapCore {
             indices: RawTable::with_capacity(n),
-            entries: Vec::with_capacity(n),
+            entries: EntryVec::with_capacity(n),
         }
     }
 
@@ -221,7 +235,11 @@ impl<K, V> IndexMapCore<K, V> {
         debug_assert!(self.indices.capacity() >= self.entries.len());
         for (i, entry) in enumerate(&self.entries) {
             // We should never have to reallocate, so there's no need for a real hasher.
+            #[cfg(not(feature = "amortize"))]
             self.indices.insert_no_grow(entry.hash.get(), i);
+            #[cfg(feature = "amortize")]
+            self.indices
+                .insert_no_grow(entry.hash.get(), i, |_| unreachable!());
         }
     }
 }

src/map/core/raw.rs

@@ -3,10 +3,17 @@
 //! mostly in dealing with its bucket "pointers".
 
 use super::{Entry, Equivalent, HashValue, IndexMapCore, VacantEntry};
-use hashbrown::raw::RawTable;
 use std::fmt;
 use std::mem::replace;
 
+#[cfg(feature = "amortize")]
+use griddle::raw::RawTable;
+#[cfg(not(feature = "amortize"))]
+use hashbrown::raw::RawTable;
+
+#[cfg(feature = "amortize")]
+type RawBucket = griddle::raw::Bucket<usize>;
+#[cfg(not(feature = "amortize"))]
 type RawBucket = hashbrown::raw::Bucket<usize>;
 
 pub(super) struct DebugIndices<'a>(pub &'a RawTable<usize>);
@@ -105,6 +112,9 @@ impl<K, V> IndexMapCore<K, V> {
         // correct indices that point to the entries that followed the removed entry.
         // use a heuristic between a full sweep vs. a `find()` for every shifted item.
         let raw_capacity = self.indices.buckets();
+        #[cfg(feature = "amortize")]
+        let shifted_entries = self.entries.range(index..);
+        #[cfg(not(feature = "amortize"))]
         let shifted_entries = &self.entries[index..];
         if shifted_entries.len() > raw_capacity / 2 {
             // shift all indices greater than `index`