In-place growth for HashMap.

Author: pczarn

src/libstd/collections/hash/map.rs

@@ -23,7 +23,7 @@ use iter::{Iterator, ExactSizeIterator, IntoIterator, IteratorExt, FromIterator,
 use marker::Sized;
 use mem::{self, replace};
 use num::{Int, UnsignedInt};
-use ops::{Drop, FnMut, Index, IndexMut};
+use ops::{Deref, DerefMut, Drop, FnMut, Index, IndexMut};
 use option::Option::{self, Some, None};
 use rand::{self, Rng};
 use result::Result::{self, Ok, Err};
@@ -445,7 +445,7 @@ fn robin_hood<'a, K: 'a, V: 'a>(bucket: FullBucketMut<'a, K, V>,
 
 // Performs insertion with relaxed requirements.
 // The caller should ensure that invariants of Robin Hood linear probing hold.
-fn insert_hashed_ordered<M, K, V>(arg: M, h: SafeHash, k: K, v: V) -> M
+fn insert_hashed_ordered<M: Put, K, V>(arg: M, h: SafeHash, k: K, v: V) -> M
     where RawTable<K, V>: BorrowFromMut<M>
 {
     let table = TableRef(arg);
@@ -652,12 +652,20 @@ impl<K, V, S> HashMap<K, V, S>
         }
 
         // Grow the table.
-        let mut destination = RawTable::new(new_capacity);
+        let is_inplace = self.table.grow_inplace(new_capacity);
+
+        let mut destination = if is_inplace {
+            // Resizing in-place.
+            None
+        } else {
+            // Borrow self.table in both branches to satisfy the checker.
+            Some(RawTable::new(new_capacity))
+        };
 
         // Iterate over `old_capacity` buckets, which constitute half of
         // the table which was resized in-place, or the entire
         // `old_table`.
-        let mut bucket = Bucket::at_index(&mut self.table, 0).ok().unwrap();
+        let mut bucket = Bucket::at_index(&mut self.table, 0).ok().unwrap().iter_to(old_capacity);
 
         // "So a few of the first shall be last: for many be called,
         // but few chosen."
@@ -703,20 +711,35 @@ impl<K, V, S> HashMap<K, V, S>
         //    ^ exit once table.size == 0
         let idx_end = bucket.index() + old_capacity;
 
-        while bucket.index() != idx_end {
-            bucket = match bucket.peek() {
-                Full(bucket) => {
-                    let h = *bucket.read().0;
-                    let (b, k, v) = bucket.take();
-                    insert_hashed_ordered(&mut destination, h, k, v);
-                    b.into_bucket()
-                }
-                Empty(b) => b.into_bucket()
-            };
-            bucket.next(); // wraps at old_capacity
-        }
+        if let Some(mut dest) = destination {
+            while bucket.index() != idx_end {
+                bucket = match bucket.peek() {
+                    Full(bucket) => {
+                        let h = *bucket.read().0;
+                        let (b, k, v) = bucket.take();
+                        insert_hashed_ordered(&mut dest, h, k, v);
+                        b.into_bucket()
+                    }
+                    Empty(b) => b.into_bucket()
+                };
+                bucket.next(); // wraps at old_capacity
+            }
 
-        replace(bucket.into_table(), destination);
+            replace(bucket.into_table(), dest);
+        } else {
+            while bucket.index() != idx_end {
+                bucket = match bucket.peek() {
+                    Full(bucket) => {
+                        let h = *bucket.read().0;
+                        let (b, k, v) = bucket.take();
+                        // Resizing in-place.
+                        insert_hashed_ordered(b.into_bucket(), h, k, v)
+                    }
+                    Empty(b) => b.into_bucket()
+                };
+                bucket.next(); // wraps at old_capacity
+            }
+        }
     }
 
     /// Shrinks the capacity of the map as much as possible. It will drop

src/libstd/collections/hash/table.rs

@@ -22,9 +22,9 @@ use mem::{self, min_align_of, size_of};
 use num::{Int, UnsignedInt};
 use ops::{Deref, DerefMut, Drop};
 use option::Option::{self, Some, None};
-use ptr::{self, Unique, PtrExt, copy_nonoverlapping_memory, zero_memory};
+use ptr::{self, Unique, PtrExt, copy_memory, copy_nonoverlapping_memory, zero_memory};
 use result::Result::{self, Ok, Err};
-use rt::heap::{allocate, deallocate, EMPTY};
+use rt::heap::{EMPTY, allocate, deallocate, reallocate_inplace};
 use collections::hash_state::HashState;
 use core::nonzero::NonZero;
 
@@ -199,15 +199,15 @@ impl<K, V, M, S> Borrow<RawTable<K, V>> for Bucket<K, V, M, S>
     where M: Borrow<RawTable<K, V>>
 {
     fn borrow(&self) -> &RawTable<K, V> {
-        bucket.table.0.borrow()
+        self.table.0.borrow()
     }
 }
 
 impl<K, V, M, S> BorrowMut<RawTable<K, V>> for Bucket<K, V, M, S>
     where M: Borrow<RawTable<K, V>>
 {
     fn borrow_mut(&mut self) -> &mut RawTable<K, V> {
-        bucket.table.0.borrow_mut()
+        self.table.0.borrow_mut()
     }
 }
 
@@ -271,6 +271,14 @@ impl<K, V, M> Bucket<K, V, M> where M: Borrow<RawTable<K, V>> {
         Bucket::at_index(table, 0).map(|b| b.into_iter()).unwrap_or_else(|b| b.into_iter())
     }
 
+    /// Narrows down the range of iteration, which must be a power of 2.
+    pub fn iter_to(mut self, limit: usize) -> Bucket<K, V, M> {
+        assert!(limit <= self.table.capacity());
+        assert!(limit.is_power_of_two());
+        self.capacity = limit;
+        self
+    }
+
     /// Reads a bucket at a given index, returning an enum indicating whether
     /// it's initialized or not. You need to match on this enum to get
     /// the appropriate types to call most of the other functions in
@@ -479,7 +487,7 @@ impl<K, V> RawTable<K, V> {
                 RawTable {
                     capacity: capacity,
                     size:     0,
-                    middle:   Unique((hashes as *mut (K, V)).offset(capacity as isize)),
+                    middle:   Unique::new((hashes as *mut (K, V)).offset(capacity as isize)),
                 }
             };
 
@@ -514,6 +522,47 @@ impl<K, V> RawTable<K, V> {
         }
     }
 
+    pub fn grow_inplace(&mut self, capacity: uint) -> bool {
+        assert!(capacity.is_power_of_two());
+        assert!(capacity >= self.capacity);
+
+        if self.middle.ptr.is_null() {
+            return false;
+        }
+
+        let new_size = checked_size_generic::<K, V>(capacity);
+
+        unsafe {
+            let ptr = self.middle.ptr.offset(-(self.capacity as isize)) as *mut u8;
+            let is_inplace = reallocate_inplace(ptr,
+                                                size_generic::<K, V>(self.capacity),
+                                                new_size,
+                                                align::<K, V>()) >= new_size;
+
+            if is_inplace {
+                let cap_diff = (capacity - self.capacity) as isize;
+                let hashes = self.middle.ptr.offset(cap_diff) as *mut Option<SafeHash>;
+                // Copy the array of hashes. Maybe it's already in cache.
+                if size_of::<(K, V)>() >= size_of::<Option<SafeHash>>() {
+                    // The regions of memory occupied by old and new hash arrays are disjoint.
+                    // before: [KVKVKVKV|h h h h ]
+                    // after:  [KVKVKVKV|KVKVKVKV|h h h h h h h h ]
+                    copy_nonoverlapping_memory(hashes, self.middle.ptr as *const _, self.capacity);
+                } else {
+                    // before: [KVKVKVKV|h   h   |h   h   ]
+                    // after:  [KVKVKVKV|KVKVKVKV|h   h   h   h   h   h   h   h   ]
+                    copy_memory(hashes, self.middle.ptr as *const _, self.capacity);
+                }
+                zero_memory(hashes.offset(self.capacity as int), capacity - self.capacity);
+
+                self.middle = Unique::new(self.middle.ptr.offset(cap_diff));
+                self.capacity = capacity;
+            }
+
+            return is_inplace;
+        }
+    }
+
     /// The hashtable's capacity, similar to a vector's.
     pub fn capacity(&self) -> usize {
         self.capacity
@@ -570,13 +619,13 @@ fn align<K, V>() -> usize {
 /// A newtyped RawBucket. Not copyable.
 pub struct RawFullBucket<K, V, M>(RawBucket<K, V>);
 
-impl<'t, K, V, M: 't> RawFullBucket<K, V, M> where RawTable<K, V>: BorrowFrom<M> {
+impl<'t, K, V, M: 't> RawFullBucket<K, V, M> where M: Borrow<RawTable<K, V>> {
     pub fn into_refs(self) -> (&'t K, &'t V) {
         unsafe { (&(*self.0.kval).0, &(*self.0.kval).1) }
     }
 }
 
-impl<'t, K, V, M: 't> RawFullBucket<K, V, M> where RawTable<K, V>: BorrowFromMut<M> {
+impl<'t, K, V, M: 't> RawFullBucket<K, V, M> where M: BorrowMut<RawTable<K, V>> {
     pub fn into_mut_refs(self) -> (&'t mut K, &'t mut V) {
         unsafe { (&mut (*self.0.kval).0, &mut (*self.0.kval).1) }
     }