From d65964b0733c6aa41a4e3ac1d1259083f6f6bc86 Mon Sep 17 00:00:00 2001 From: ltdk Date: Tue, 27 Jan 2026 14:15:44 -0500 Subject: [PATCH] Make HashSet::replace work without raw table internals --- src/map.rs | 208 ++++++++++++++++++++++++++++++++++++++++++++++++----- src/set.rs | 23 +++--- 2 files changed, 205 insertions(+), 26 deletions(-) diff --git a/src/map.rs b/src/map.rs index f009f4132..2154ec3e1 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1796,8 +1796,13 @@ where /// ``` #[cfg_attr(feature = "inline-more", inline)] pub fn insert(&mut self, k: K, v: V) -> Option { - let hash = make_hash::(&self.hash_builder, &k); - match self.find_or_find_insert_index(hash, &k) { + let hash = make_hash(&self.hash_builder, &k); + let equivalent = equivalent_key(&k); + let hasher = make_hasher(&self.hash_builder); + match self + .table + .find_or_find_insert_index(hash, equivalent, hasher) + { Ok(bucket) => Some(mem::replace(unsafe { &mut bucket.as_mut().1 }, v)), Err(index) => { unsafe { @@ -1808,22 +1813,6 @@ where } } - #[cfg_attr(feature = "inline-more", inline)] - pub(crate) fn find_or_find_insert_index( - &mut self, - hash: u64, - key: &Q, - ) -> Result, usize> - where - Q: Equivalent + ?Sized, - { - self.table.find_or_find_insert_index( - hash, - equivalent_key(key), - make_hasher(&self.hash_builder), - ) - } - /// Insert a key-value pair into the map without checking /// if the key already exists in the map. /// @@ -3779,6 +3768,97 @@ impl<'a, K, V, S, A: Allocator> OccupiedEntry<'a, K, V, S, A> { unsafe { &self.elem.as_ref().0 } } + /// Replaces the key in the entry with a new one. + /// + /// # Panics + /// + /// This method panics if `key` is not equivalent to the key in the entry. + /// + /// # Examples + /// + /// ``` + /// use hashbrown::hash_map::{Entry, HashMap}; + /// + /// let mut map: HashMap<&str, u32> = HashMap::new(); + /// + /// let old_key = "poneyland"; + /// let new_key = Box::leak(old_key.to_owned().into_boxed_str()); + /// map.entry(old_key).or_insert(12); + /// match map.entry("poneyland") { + /// Entry::Vacant(_) => panic!(), + /// Entry::Occupied(mut entry) => { + /// let replaced = entry.replace_key(new_key); + /// assert!(std::ptr::eq(replaced, old_key)); + /// assert!(std::ptr::eq(*entry.key(), new_key)); + /// }, + /// } + /// + /// # // appease miri; no memory leaks here! + /// # drop(map); + /// # unsafe { + /// # Box::from_raw(new_key); + /// # } + /// ``` + #[cfg_attr(feature = "inline-more", inline)] + pub fn replace_key(&mut self, key: K) -> K + where + K: Equivalent, + { + assert!( + self.key().equivalent(&key), + "replaced key is not equivalent to the one in the entry" + ); + + // SAFETY: We verified that the keys were equivalent. + unsafe { self.replace_key_unchecked(key) } + } + + /// Replaces the key in the entry with a new one, without checking the + /// equivalence of the key. + /// + /// # Safety + /// + /// This operation is safe if you replace the key with an equivalent one. + /// + /// Additionally, this operation (and following operations) are guaranteed + /// to not violate memory safety. + /// + /// However this operation is still unsafe because the resulting `HashMap` + /// may be passed to unsafe code which does expect the map to behave + /// correctly. If the map has keys at unexpected positions inside it, + /// future operations may panic, loop forever, or return unexpected results, + /// potentially violating memory safety. + /// + /// # Examples + /// + /// ``` + /// use hashbrown::hash_map::{Entry, HashMap}; + /// + /// let mut map: HashMap<&str, u32> = HashMap::new(); + /// + /// let old_key = "poneyland"; + /// let new_key = Box::leak(old_key.to_owned().into_boxed_str()); + /// map.entry(old_key).or_insert(12); + /// match map.entry("poneyland") { + /// Entry::Vacant(_) => panic!(), + /// Entry::Occupied(mut entry) => { + /// let replaced = unsafe { entry.replace_key_unchecked(new_key) }; + /// assert!(std::ptr::eq(replaced, old_key)); + /// assert!(std::ptr::eq(*entry.key(), new_key)); + /// }, + /// } + /// + /// # // appease miri; no memory leaks here! + /// # drop(map); + /// # unsafe { + /// # Box::from_raw(new_key); + /// # } + /// ``` + #[cfg_attr(feature = "inline-more", inline)] + pub unsafe fn replace_key_unchecked(&mut self, key: K) -> K { + mem::replace(unsafe { &mut self.elem.as_mut().0 }, key) + } + /// Take the ownership of the key and value from the map. /// Keeps the allocated memory for reuse. /// @@ -4504,6 +4584,50 @@ impl<'map, 'key, K, Q: ?Sized, V, S, A: Allocator> VacantEntryRef<'map, 'key, K, self.insert_entry_with_key(key, value).into_mut() } + /// Sets the key and value of the entry and returns a mutable reference to + /// the inserted value, without checking the equivalence of the key. + /// + /// See [`insert_with_key`](Self::insert_with_key) for more information. + /// + /// # Safety + /// + /// This operation is safe if the keys are equivalent. + /// + /// Additionally, this operation (and following operations) are guaranteed + /// to not violate memory safety. + /// + /// However this operation is still unsafe because the resulting `HashMap` + /// may be passed to unsafe code which does expect the map to behave + /// correctly. If the map has keys at unexpected positions inside it, + /// future operations may panic, loop forever, or return unexpected results, + /// potentially violating memory safety. + /// + /// # Example + /// + /// ``` + /// use hashbrown::hash_map::EntryRef; + /// use hashbrown::HashMap; + /// + /// let mut map = HashMap::<(String, String), char>::new(); + /// let k = ("c".to_string(), "C".to_string()); + /// let v = match map.entry_ref(&k) { + /// // SAFETY: We trust the `Clone` implementation to return an equivalent value + /// EntryRef::Vacant(r) => unsafe { r.insert_with_key_unchecked(k.clone(), 'c') }, + /// // In this branch we avoid the clone. + /// EntryRef::Occupied(r) => r.into_mut(), + /// }; + /// assert_eq!(*v, 'c'); + /// ``` + #[cfg_attr(feature = "inline-more", inline)] + pub unsafe fn insert_with_key_unchecked(self, key: K, value: V) -> &'map mut V + where + K: Hash, + S: BuildHasher, + { + // SAFETY: Guaranteed by caller. + unsafe { self.insert_entry_with_key_unchecked(key, value) }.into_mut() + } + /// Sets the value of the entry with the [`VacantEntryRef`]'s key, /// and returns an [`OccupiedEntry`]. /// @@ -4578,6 +4702,54 @@ impl<'map, 'key, K, Q: ?Sized, V, S, A: Allocator> VacantEntryRef<'map, 'key, K, (self.key).equivalent(&key), "key used for Entry creation is not equivalent to the one used for insertion" ); + // SAFETY: We checked equivalence first. + unsafe { self.insert_entry_with_key_unchecked(key, value) } + } + + /// Sets the key and value of the entry and returns an [`OccupiedEntry`], + /// without checking the equivalence of the key. + /// + /// See [`insert_entry_with_key`](Self::insert_entry_with_key) for more information. + /// + /// # Safety + /// + /// This operation is safe if the keys are equivalent. + /// + /// Additionally, this operation (and following operations) are guaranteed + /// to not violate memory safety. + /// + /// However this operation is still unsafe because the resulting `HashMap` + /// may be passed to unsafe code which does expect the map to behave + /// correctly. If the map has keys at unexpected positions inside it, + /// future operations may panic, loop forever, or return unexpected results, + /// potentially violating memory safety. + /// + /// # Example + /// + /// ``` + /// use hashbrown::hash_map::EntryRef; + /// use hashbrown::HashMap; + /// + /// let mut map = HashMap::<(String, String), char>::new(); + /// let k = ("c".to_string(), "C".to_string()); + /// let r = match map.entry_ref(&k) { + /// // SAFETY: We trust the `Clone` implementation to return an equivalent key + /// EntryRef::Vacant(r) => unsafe { r.insert_entry_with_key_unchecked(k.clone(), 'c') }, + /// // In this branch we avoid the clone. + /// EntryRef::Occupied(r) => r, + /// }; + /// assert_eq!(r.get(), &'c'); + /// ``` + #[cfg_attr(feature = "inline-more", inline)] + pub unsafe fn insert_entry_with_key_unchecked( + self, + key: K, + value: V, + ) -> OccupiedEntry<'map, K, V, S, A> + where + K: Hash, + S: BuildHasher, + { let elem = self.table.table.insert( self.hash, (key, value), diff --git a/src/set.rs b/src/set.rs index 81d61c893..366a85af6 100644 --- a/src/set.rs +++ b/src/set.rs @@ -1,9 +1,9 @@ use crate::{Equivalent, TryReserveError}; +use core::cell::UnsafeCell; +use core::fmt; use core::hash::{BuildHasher, Hash}; use core::iter::{Chain, FusedIterator}; use core::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Sub, SubAssign}; -use core::{fmt, mem}; -use map::make_hash; use super::map::{self, HashMap, Keys}; use crate::DefaultHashBuilder; @@ -1112,12 +1112,18 @@ where /// ``` #[cfg_attr(feature = "inline-more", inline)] pub fn replace(&mut self, value: T) -> Option { - let hash = make_hash(&self.map.hash_builder, &value); - match self.map.find_or_find_insert_index(hash, &value) { - Ok(bucket) => Some(mem::replace(unsafe { &mut bucket.as_mut().0 }, value)), - Err(index) => { + let value = UnsafeCell::new(value); + // SAFETY: We know the key is no longer accessed after the initial check. + match self.map.entry_ref(unsafe { &*value.get() }) { + map::EntryRef::Occupied(mut entry) => { + // SAFETY: We know the key will not be accessed any more, and + // that the key is equivalent to the one in the entry. + Some(unsafe { entry.replace_key_unchecked(value.into_inner()) }) + } + map::EntryRef::Vacant(entry) => { + // SAFETY: A value is equivalent to itself. unsafe { - self.map.table.insert_at_index(hash, index, (value, ())); + entry.insert_with_key_unchecked(value.into_inner(), ()); } None } @@ -2529,8 +2535,9 @@ fn assert_covariance() { #[cfg(test)] mod test_set { - use super::{Equivalent, HashSet, make_hash}; + use super::{Equivalent, HashSet}; use crate::DefaultHashBuilder; + use crate::map::make_hash; use std::vec::Vec; #[test]