Merge pull request #202 from cuviper/safer-raw-access

Add more safe methods to RawTable

Author: Amanieu

src/map.rs

@@ -807,8 +807,8 @@ where
         Q: Hash + Eq,
     {
         // Avoid `Option::map` because it bloats LLVM IR.
-        match self.get_key_value(k) {
-            Some((_, v)) => Some(v),
+        match self.get_inner(k) {
+            Some(&(_, ref v)) => Some(v),
             None => None,
         }
     }
@@ -838,17 +838,23 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq,
     {
-        let hash = make_hash(&self.hash_builder, k);
         // Avoid `Option::map` because it bloats LLVM IR.
-        match self.table.find(hash, |x| k.eq(x.0.borrow())) {
-            Some(item) => unsafe {
-                let &(ref key, ref value) = item.as_ref();
-                Some((key, value))
-            },
+        match self.get_inner(k) {
+            Some(&(ref key, ref value)) => Some((key, value)),
             None => None,
         }
     }
 
+    #[inline]
+    fn get_inner<Q: ?Sized>(&self, k: &Q) -> Option<&(K, V)>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        let hash = make_hash(&self.hash_builder, k);
+        self.table.get(hash, |x| k.eq(x.0.borrow()))
+    }
+
     /// Returns the key-value pair corresponding to the supplied key, with a mutable reference to value.
     ///
     /// The supplied key may be any borrowed form of the map's key type, but
@@ -878,13 +884,9 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq,
     {
-        let hash = make_hash(&self.hash_builder, k);
         // Avoid `Option::map` because it bloats LLVM IR.
-        match self.table.find(hash, |x| k.eq(x.0.borrow())) {
-            Some(item) => unsafe {
-                let &mut (ref key, ref mut value) = item.as_mut();
-                Some((key, value))
-            },
+        match self.get_inner_mut(k) {
+            Some(&mut (ref key, ref mut value)) => Some((key, value)),
             None => None,
         }
     }
@@ -914,7 +916,7 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq,
     {
-        self.get(k).is_some()
+        self.get_inner(k).is_some()
     }
 
     /// Returns a mutable reference to the value corresponding to the key.
@@ -944,14 +946,23 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq,
     {
-        let hash = make_hash(&self.hash_builder, k);
         // Avoid `Option::map` because it bloats LLVM IR.
-        match self.table.find(hash, |x| k.eq(x.0.borrow())) {
-            Some(item) => Some(unsafe { &mut item.as_mut().1 }),
+        match self.get_inner_mut(k) {
+            Some(&mut (_, ref mut v)) => Some(v),
             None => None,
         }
     }
 
+    #[inline]
+    fn get_inner_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut (K, V)>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        let hash = make_hash(&self.hash_builder, k);
+        self.table.get_mut(hash, |x| k.eq(x.0.borrow()))
+    }
+
     /// Inserts a key-value pair into the map.
     ///
     /// If the map did not have this key present, [`None`] is returned.
@@ -979,16 +990,14 @@ where
     /// ```
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn insert(&mut self, k: K, v: V) -> Option<V> {
-        unsafe {
-            let hash = make_hash(&self.hash_builder, &k);
-            if let Some(item) = self.table.find(hash, |x| k.eq(&x.0)) {
-                Some(mem::replace(&mut item.as_mut().1, v))
-            } else {
-                let hash_builder = &self.hash_builder;
-                self.table
-                    .insert(hash, (k, v), |x| make_hash(hash_builder, &x.0));
-                None
-            }
+        let hash = make_hash(&self.hash_builder, &k);
+        if let Some((_, item)) = self.table.get_mut(hash, |x| k.eq(&x.0)) {
+            Some(mem::replace(item, v))
+        } else {
+            let hash_builder = &self.hash_builder;
+            self.table
+                .insert(hash, (k, v), |x| make_hash(hash_builder, &x.0));
+            None
         }
     }
 
@@ -1051,14 +1060,8 @@ where
         K: Borrow<Q>,
         Q: Hash + Eq,
     {
-        unsafe {
-            let hash = make_hash(&self.hash_builder, &k);
-            if let Some(item) = self.table.find(hash, |x| k.eq(x.0.borrow())) {
-                Some(self.table.remove(item))
-            } else {
-                None
-            }
-        }
+        let hash = make_hash(&self.hash_builder, &k);
+        self.table.remove_entry(hash, |x| k.eq(x.0.borrow()))
     }
 }
 
@@ -1590,11 +1593,8 @@ impl<'a, K, V, S> RawEntryBuilder<'a, K, V, S> {
     where
         F: FnMut(&K) -> bool,
     {
-        match self.map.table.find(hash, |(k, _)| is_match(k)) {
-            Some(item) => unsafe {
-                let &(ref key, ref value) = item.as_ref();
-                Some((key, value))
-            },
+        match self.map.table.get(hash, |(k, _)| is_match(k)) {
+            Some(&(ref key, ref value)) => Some((key, value)),
             None => None,
         }
     }
@@ -1961,11 +1961,10 @@ impl<'a, K, V, S> RawVacantEntryMut<'a, K, V, S> {
     where
         H: Fn(&K) -> u64,
     {
-        unsafe {
-            let elem = self.table.insert(hash, (key, value), |x| hasher(&x.0));
-            let &mut (ref mut k, ref mut v) = elem.as_mut();
-            (k, v)
-        }
+        let &mut (ref mut k, ref mut v) = self
+            .table
+            .insert_entry(hash, (key, value), |x| hasher(&x.0));
+        (k, v)
     }
 
     #[cfg_attr(feature = "inline-more", inline)]
@@ -2974,10 +2973,11 @@ impl<'a, K, V, S> VacantEntry<'a, K, V, S> {
         S: BuildHasher,
     {
         let hash_builder = &self.table.hash_builder;
-        let bucket = self.table.table.insert(self.hash, (self.key, value), |x| {
+        let table = &mut self.table.table;
+        let entry = table.insert_entry(self.hash, (self.key, value), |x| {
             make_hash(hash_builder, &x.0)
         });
-        unsafe { &mut bucket.as_mut().1 }
+        &mut entry.1
     }
 
     #[cfg_attr(feature = "inline-more", inline)]

src/raw/mod.rs

@@ -545,6 +545,20 @@ impl<T> RawTable<T> {
         item.drop();
     }
 
+    /// Finds and erases an element from the table, dropping it in place.
+    /// Returns true if an element was found.
+    #[cfg(feature = "raw")]
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn erase_entry(&mut self, hash: u64, eq: impl FnMut(&T) -> bool) -> bool {
+        // Avoid `Option::map` because it bloats LLVM IR.
+        if let Some(bucket) = self.find(hash, eq) {
+            unsafe { self.erase(bucket) };
+            true
+        } else {
+            false
+        }
+    }
+
     /// Removes an element from the table, returning it.
     #[cfg_attr(feature = "inline-more", inline)]
     #[allow(clippy::needless_pass_by_value)]
@@ -554,6 +568,16 @@ impl<T> RawTable<T> {
         item.read()
     }
 
+    /// Finds and removes an element from the table, returning it.
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn remove_entry(&mut self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<T> {
+        // Avoid `Option::map` because it bloats LLVM IR.
+        match self.find(hash, eq) {
+            Some(bucket) => Some(unsafe { self.remove(bucket) }),
+            None => None,
+        }
+    }
+
     /// Returns an iterator for a probe sequence on the table.
     ///
     /// This iterator never terminates, but is guaranteed to visit each bucket
@@ -910,7 +934,7 @@ impl<T> RawTable<T> {
         }
     }
 
-    /// Inserts a new element into the table.
+    /// Inserts a new element into the table, and returns its raw bucket.
     ///
     /// This does not check if the given element already exists in the table.
     #[cfg_attr(feature = "inline-more", inline)]
@@ -936,6 +960,14 @@ impl<T> RawTable<T> {
         }
     }
 
+    /// Inserts a new element into the table, and returns a mutable reference to it.
+    ///
+    /// This does not check if the given element already exists in the table.
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn insert_entry(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> &mut T {
+        unsafe { self.insert(hash, value, hasher).as_mut() }
+    }
+
     /// Inserts a new element into the table, without growing the table.
     ///
     /// There must be enough space in the table to insert the new element.
@@ -1001,6 +1033,26 @@ impl<T> RawTable<T> {
         }
     }
 
+    /// Gets a reference to an element in the table.
+    #[inline]
+    pub fn get(&self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<&T> {
+        // Avoid `Option::map` because it bloats LLVM IR.
+        match self.find(hash, eq) {
+            Some(bucket) => Some(unsafe { bucket.as_ref() }),
+            None => None,
+        }
+    }
+
+    /// Gets a mutable reference to an element in the table.
+    #[inline]
+    pub fn get_mut(&mut self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<&mut T> {
+        // Avoid `Option::map` because it bloats LLVM IR.
+        match self.find(hash, eq) {
+            Some(bucket) => Some(unsafe { bucket.as_mut() }),
+            None => None,
+        }
+    }
+
     /// Returns the number of elements the map can hold without reallocating.
     ///
     /// This number is a lower bound; the table might be able to hold