Implement placement-in protocol for HashMap

Charlie Fan · F001 · commit 584c79888b8e · 2017-03-11T21:08:18.000+08:00
diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs
@@ -19,8 +19,9 @@ use fmt::{self, Debug};
 use hash::{Hash, Hasher, BuildHasher, SipHasher13};
 use iter::{FromIterator, FusedIterator};
 use mem::{self, replace};
-use ops::{Deref, Index};
+use ops::{Deref, Index, InPlace, Place, Placer};
 use rand::{self, Rng};
+use ptr;
 
 use super::table::{self, Bucket, EmptyBucket, FullBucket, FullBucketMut, RawTable, SafeHash};
 use super::table::BucketState::{Empty, Full};
@@ -483,7 +484,7 @@ fn robin_hood<'a, K: 'a, V: 'a>(bucket: FullBucketMut<'a, K, V>,
                                 mut hash: SafeHash,
                                 mut key: K,
                                 mut val: V)
-                                -> &'a mut V {
+                                -> FullBucketMut<'a, K, V> {
     let start_index = bucket.index();
     let size = bucket.table().size();
     // Save the *starting point*.
@@ -515,7 +516,7 @@ fn robin_hood<'a, K: 'a, V: 'a>(bucket: FullBucketMut<'a, K, V>,
                     // bucket, which is a FullBucket on top of a
                     // FullBucketMut, into just one FullBucketMut. The "table"
                     // refers to the inner FullBucketMut in this context.
-                    return bucket.into_table().into_mut_refs().1;
+                    return bucket.into_table();
                 }
                 Full(bucket) => bucket,
             };
@@ -1818,6 +1819,80 @@ impl<'a, K, V> fmt::Debug for Drain<'a, K, V>
     }
 }
 
+/// A place for insertion to a `Entry`.
+///
+/// See [`HashMap::entry`](struct.HashMap.html#method.entry) for details.
+#[must_use = "places do nothing unless written to with `<-` syntax"]
+#[unstable(feature = "collection_placement",
+           reason = "struct name and placement protocol is subject to change",
+           issue = "30172")]
+pub struct EntryPlace<'a, K: 'a, V: 'a> {
+    bucket: FullBucketMut<'a, K, V>,
+}
+
+#[unstable(feature = "collection_placement",
+           reason = "struct name and placement protocol is subject to change",
+           issue = "30172")]
+impl<'a, K: 'a + Debug, V: 'a + Debug> Debug for EntryPlace<'a, K, V> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("EntryPlace")
+            .field("key", self.bucket.read().0)
+            .field("value", self.bucket.read().1)
+            .finish()
+    }
+}
+
+#[unstable(feature = "collection_placement",
+           reason = "struct name and placement protocol is subject to change",
+           issue = "30172")]
+impl<'a, K, V> Drop for EntryPlace<'a, K, V> {
+    fn drop(&mut self) {
+        // Inplacement insertion failed. Only key need to drop.
+        // The value is failed to insert into map.
+        unsafe { self.bucket.remove_key() };
+    }
+}
+
+#[unstable(feature = "collection_placement",
+           reason = "placement protocol is subject to change",
+           issue = "30172")]
+impl<'a, K, V> Placer<V> for Entry<'a, K, V> {
+    type Place = EntryPlace<'a, K, V>;
+
+    fn make_place(self) -> EntryPlace<'a, K, V> {
+        let b = match self {
+            Occupied(mut o) => {
+                unsafe { ptr::drop_in_place(o.elem.read_mut().1); }
+                o.elem
+            }
+            Vacant(v) => {
+                unsafe { v.insert_key() }
+            }
+        };
+        EntryPlace { bucket: b }
+    }
+}
+
+#[unstable(feature = "collection_placement",
+           reason = "placement protocol is subject to change",
+           issue = "30172")]
+impl<'a, K, V> Place<V> for EntryPlace<'a, K, V> {
+    fn pointer(&mut self) -> *mut V {
+        self.bucket.read_mut().1
+    }
+}
+
+#[unstable(feature = "collection_placement",
+           reason = "placement protocol is subject to change",
+           issue = "30172")]
+impl<'a, K, V> InPlace<V> for EntryPlace<'a, K, V> {
+    type Owner = ();
+
+    unsafe fn finalize(self) {
+        mem::forget(self);
+    }
+}
+
 impl<'a, K, V> Entry<'a, K, V> {
     #[stable(feature = "rust1", since = "1.0.0")]
     /// Ensures a value is in the entry by inserting the default if empty, and returns
@@ -2108,7 +2183,7 @@ impl<'a, K: 'a, V: 'a> VacantEntry<'a, K, V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn insert(self, value: V) -> &'a mut V {
-        match self.elem {
+        let b = match self.elem {
             NeqElem(mut bucket, disp) => {
                 if disp >= DISPLACEMENT_THRESHOLD {
                     bucket.table_mut().set_tag(true);
@@ -2119,7 +2194,28 @@ impl<'a, K: 'a, V: 'a> VacantEntry<'a, K, V> {
                 if disp >= DISPLACEMENT_THRESHOLD {
                     bucket.table_mut().set_tag(true);
                 }
-                bucket.put(self.hash, self.key, value).into_mut_refs().1
+                bucket.put(self.hash, self.key, value)
+            },
+        };
+        b.into_mut_refs().1
+    }
+
+    // Only used for InPlacement insert. Avoid unnecessary value copy.
+    // The value remains uninitialized.
+    unsafe fn insert_key(self) -> FullBucketMut<'a, K, V> {
+        match self.elem {
+            NeqElem(mut bucket, disp) => {
+                if disp >= DISPLACEMENT_THRESHOLD {
+                    bucket.table_mut().set_tag(true);
+                }
+                let uninit = mem::uninitialized();
+                robin_hood(bucket, disp, self.hash, self.key, uninit)
+            },
+            NoElem(mut bucket, disp) => {
+                if disp >= DISPLACEMENT_THRESHOLD {
+                    bucket.table_mut().set_tag(true);
+                }
+                bucket.put_key(self.hash, self.key)
             },
         }
     }
@@ -2392,6 +2488,7 @@ mod test_map {
     use super::RandomState;
     use cell::RefCell;
     use rand::{thread_rng, Rng};
+    use panic;
 
     #[test]
     fn test_zero_capacities() {
@@ -3265,4 +3362,57 @@ mod test_map {
         }
         panic!("Adaptive early resize failed");
     }
+
+    #[test]
+    fn test_placement_in() {
+        let mut map = HashMap::new();
+        map.extend((0..10).map(|i| (i, i)));
+
+        map.entry(100) <- 100;
+        assert_eq!(map[&100], 100);
+
+        map.entry(0) <- 10;
+        assert_eq!(map[&0], 10);
+
+        assert_eq!(map.len(), 11);
+    }
+
+    #[test]
+    fn test_placement_panic() {
+        let mut map = HashMap::new();
+        map.extend((0..10).map(|i| (i, i)));
+
+        fn mkpanic() -> usize { panic!() }
+
+        // modify existing key
+        // when panic happens, previous key is removed.
+        let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { map.entry(0) <- mkpanic(); }));
+        assert_eq!(map.len(), 9);
+        assert!(!map.contains_key(&0));
+
+        // add new key
+        let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { map.entry(100) <- mkpanic(); }));
+        assert_eq!(map.len(), 9);
+        assert!(!map.contains_key(&100));
+    }
+
+    #[test]
+    fn test_placement_drop() {
+        // correctly drop
+        struct TestV<'a>(&'a mut bool);
+        impl<'a> Drop for TestV<'a> {
+            fn drop(&mut self) {
+                if !*self.0 { panic!("value double drop!"); } // no double drop
+                *self.0 = false;
+            }
+        }
+
+        fn makepanic<'a>() -> TestV<'a> { panic!() }
+
+        let mut can_drop = true;
+        let mut hm = HashMap::new();
+        hm.insert(0, TestV(&mut can_drop));
+        let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { hm.entry(0) <- makepanic(); }));
+        assert_eq!(hm.len(), 0);
+    }
 }
diff --git a/src/libstd/collections/hash/table.rs b/src/libstd/collections/hash/table.rs
@@ -506,6 +506,22 @@ impl<K, V, M> EmptyBucket<K, V, M>
             table: self.table,
         }
     }
+
+    /// Puts given key, remain value uninitialized.
+    /// It is only used for inplacement insertion.
+    pub unsafe fn put_key(mut self, hash: SafeHash, key: K) -> FullBucket<K, V, M> {
+        *self.raw.hash = hash.inspect();
+        let pair_mut = self.raw.pair as *mut (K, V);
+        ptr::write(&mut (*pair_mut).0, key);
+
+        self.table.borrow_table_mut().size += 1;
+
+        FullBucket {
+            raw: self.raw,
+            idx: self.idx,
+            table: self.table,
+        }
+    }
 }
 
 impl<K, V, M: Deref<Target = RawTable<K, V>>> FullBucket<K, V, M> {
@@ -581,6 +597,17 @@ impl<'t, K, V> FullBucket<K, V, &'t mut RawTable<K, V>> {
             v)
         }
     }
+
+    /// Remove this bucket's `key` from the hashtable.
+    /// Only used for inplacement insertion.
+    /// NOTE: `Value` is uninitialized when this function is called, don't try to drop the `Value`.
+    pub unsafe fn remove_key(&mut self) {
+        self.table.size -= 1;
+
+        *self.raw.hash = EMPTY_BUCKET;
+        let pair_mut = self.raw.pair as *mut (K, V);
+        ptr::drop_in_place(&mut (*pair_mut).0); // only drop key
+    }
 }
 
 // This use of `Put` is misleading and restrictive, but safe and sufficient for our use cases
diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs
@@ -281,6 +281,7 @@
 #![feature(panic_unwind)]
 #![feature(peek)]
 #![feature(placement_in_syntax)]
+#![feature(placement_new_protocol)]
 #![feature(prelude_import)]
 #![feature(pub_restricted)]
 #![feature(rand)]
