Add support for amortized resizes

Author: jonhoo

Cargo.toml

@@ -34,13 +34,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"] }
@@ -52,6 +59,10 @@ fxhash = "0.2.1"
 [features]
 # Serialization with serde 1.0
 serde-1 = ["serde"]
+rayon = ["rayon_", "atone/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

@@ -86,9 +86,11 @@ extern crate alloc;
 extern crate std;
 
 #[cfg(not(has_std))]
+#[cfg_attr(feature = "amortize", allow(unused_imports))]
 use alloc::vec::{self, Vec};
 
 #[cfg(has_std)]
+#[cfg_attr(feature = "amortize", allow(unused_imports))]
 use std::vec::{self, Vec};
 
 #[macro_use]
@@ -113,6 +115,11 @@ pub use crate::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)]
@@ -182,10 +189,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

@@ -8,13 +8,12 @@ pub use crate::mutable_keys::MutableKeys;
 #[cfg(feature = "rayon")]
 pub use crate::rayon::map as rayon;
 
-use crate::vec::{self, Vec};
+use crate::EntryVec;
 use ::core::cmp::Ordering;
 use ::core::fmt;
 use ::core::hash::{BuildHasher, Hash, Hasher};
 use ::core::iter::FromIterator;
 use ::core::ops::{Index, IndexMut, RangeFull};
-use ::core::slice::{Iter as SliceIter, IterMut as SliceIterMut};
 
 #[cfg(has_std)]
 use std::collections::hash_map::RandomState;
@@ -101,23 +100,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);
     }
@@ -618,6 +617,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));
         });
     }
@@ -635,6 +636,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));
         });
     }
@@ -648,7 +651,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(),
         }
@@ -724,7 +731,7 @@ impl<K, V, S> IndexMap<K, V, S> {
 /// [`keys`]: struct.IndexMap.html#method.keys
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct Keys<'a, K, V> {
-    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> {
@@ -768,7 +775,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, V> {
-    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> {
@@ -812,7 +819,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, V> {
-    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> {
@@ -841,7 +848,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, V> {
-    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> {
@@ -885,7 +892,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, V> {
-    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> {
@@ -914,7 +921,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: vec::IntoIter<Bucket<K, V>>,
+    pub(crate) iter: <EntryVec<Bucket<K, V>> as IntoIterator>::IntoIter,
 }
 
 impl<K, V> Iterator for IntoIter<K, V> {
@@ -936,6 +943,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()
@@ -950,7 +962,10 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
 /// [`drain`]: struct.IndexMap.html#method.drain
 /// [`IndexMap`]: struct.IndexMap.html
 pub struct Drain<'a, K, V> {
-    pub(crate) iter: vec::Drain<'a, Bucket<K, V>>,
+    #[cfg(not(feature = "amortize"))]
+    pub(crate) iter: crate::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> {
@@ -1307,7 +1322,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

@@ -9,28 +9,35 @@
 
 mod raw;
 
+#[cfg(feature = "amortize")]
+use griddle::raw::RawTable;
+#[cfg(not(feature = "amortize"))]
 use hashbrown::raw::RawTable;
 
-use crate::vec::{Drain, Vec};
 use core::cmp;
 use core::fmt;
 use core::mem::replace;
 use core::ops::RangeFull;
 
 use crate::equivalent::Equivalent;
 use crate::util::enumerate;
+use crate::EntryVec;
 use crate::{Bucket, Entries, HashValue};
+#[cfg(feature = "amortize")]
+use atone::vc::Drain;
+#[cfg(not(feature = "amortize"))]
+use std::vec::Drain;
 
 /// 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()
 }
 
@@ -40,8 +47,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 }
     }
@@ -74,23 +87,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();
@@ -102,15 +115,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),
         }
     }
 
@@ -218,7 +231,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

@@ -5,8 +5,14 @@
 use super::{Entry, Equivalent, HashValue, IndexMapCore, VacantEntry};
 use core::fmt;
 use core::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 +111,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`