From 54a95e717d4ed1589c9b0e38648f65424defc74b Mon Sep 17 00:00:00 2001 From: Eric Holk Date: Fri, 27 Mar 2026 18:06:56 -0700 Subject: [PATCH] Add VecMap::entry_ref (#52601) Thinking more on #52596, I realized the `entry` method needs a key by value, which meant we were always cloning a path list even if it was already in the map. This PR adds an `entry_ref` that takes the key by reference and delays cloning the key until we know we're going to be inserting. Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Release Notes: - N/A --- crates/collections/src/vecmap.rs | 73 +++++++++++++ crates/collections/src/vecmap_tests.rs | 115 ++++++++++++++++++++ crates/sidebar/src/project_group_builder.rs | 2 +- 3 files changed, 189 insertions(+), 1 deletion(-) diff --git a/crates/collections/src/vecmap.rs b/crates/collections/src/vecmap.rs index 97972e7bf9f5ae43957648b8a44b10e3e45bc32f..bec6596b924742daf4e1da3831f1182557875d61 100644 --- a/crates/collections/src/vecmap.rs +++ b/crates/collections/src/vecmap.rs @@ -45,6 +45,20 @@ impl VecMap { None => Entry::Vacant(VacantEntry { map: self, key }), } } + + /// Like [`Self::entry`] but takes its key by reference instead of by value. + /// + /// This can be helpful if you have a key where cloning is expensive, as we + /// can avoid cloning the key until a value is inserted under that entry. + pub fn entry_ref<'a, 'k>(&'a mut self, key: &'k K) -> EntryRef<'k, 'a, K, V> { + match self.keys.iter().position(|k| k == key) { + Some(index) => EntryRef::Occupied(OccupiedEntry { + key: &self.keys[index], + value: &mut self.values[index], + }), + None => EntryRef::Vacant(VacantEntryRef { map: self, key }), + } + } } pub struct Iter<'a, K, V> { @@ -117,3 +131,62 @@ pub struct VacantEntry<'a, K, V> { map: &'a mut VecMap, key: K, } + +pub enum EntryRef<'key, 'map, K, V> { + Occupied(OccupiedEntry<'map, K, V>), + Vacant(VacantEntryRef<'key, 'map, K, V>), +} + +impl<'key, 'map, K, V> EntryRef<'key, 'map, K, V> { + pub fn key(&self) -> &K { + match self { + EntryRef::Occupied(entry) => entry.key, + EntryRef::Vacant(entry) => entry.key, + } + } +} + +impl<'key, 'map, K, V> EntryRef<'key, 'map, K, V> +where + K: Clone, +{ + pub fn or_insert_with_key(self, default: F) -> &'map mut V + where + F: FnOnce(&K) -> V, + { + match self { + EntryRef::Occupied(entry) => entry.value, + EntryRef::Vacant(entry) => { + entry.map.values.push(default(entry.key)); + entry.map.keys.push(entry.key.clone()); + match entry.map.values.last_mut() { + Some(value) => value, + None => unreachable!("vec empty after pushing to it"), + } + } + } + } + + pub fn or_insert_with(self, default: F) -> &'map mut V + where + F: FnOnce() -> V, + { + self.or_insert_with_key(|_| default()) + } + + pub fn or_insert(self, value: V) -> &'map mut V { + self.or_insert_with_key(|_| value) + } + + pub fn or_insert_default(self) -> &'map mut V + where + V: Default, + { + self.or_insert_with_key(|_| Default::default()) + } +} + +pub struct VacantEntryRef<'key, 'map, K, V> { + map: &'map mut VecMap, + key: &'key K, +} diff --git a/crates/collections/src/vecmap_tests.rs b/crates/collections/src/vecmap_tests.rs index 9fc7d430f3422374a7662bb476cbd99dd09d9a43..1f698f8331cc5044f23c19603005e253f8a81ef3 100644 --- a/crates/collections/src/vecmap_tests.rs +++ b/crates/collections/src/vecmap_tests.rs @@ -94,3 +94,118 @@ fn test_multiple_entries_independent() { map.entry(1).or_insert(99); assert_eq!(map.iter().count(), 3); } + +// entry_ref tests + +use std::cell::Cell; +use std::rc::Rc; + +#[derive(PartialEq, Eq)] +struct CountedKey { + value: String, + clone_count: Rc>, +} + +impl Clone for CountedKey { + fn clone(&self) -> Self { + self.clone_count.set(self.clone_count.get() + 1); + CountedKey { + value: self.value.clone(), + clone_count: self.clone_count.clone(), + } + } +} + +#[test] +fn test_entry_ref_vacant_or_insert() { + let mut map: VecMap = VecMap::new(); + let key = "a".to_string(); + let value = map.entry_ref(&key).or_insert(1); + assert_eq!(*value, 1); + assert_eq!(map.iter().count(), 1); +} + +#[test] +fn test_entry_ref_occupied_or_insert_keeps_existing() { + let mut map: VecMap = VecMap::new(); + map.entry_ref(&"a".to_string()).or_insert(1); + let value = map.entry_ref(&"a".to_string()).or_insert(99); + assert_eq!(*value, 1); + assert_eq!(map.iter().count(), 1); +} + +#[test] +fn test_entry_ref_key_not_cloned_when_occupied() { + let clone_count = Rc::new(Cell::new(0)); + let key = CountedKey { + value: "a".to_string(), + clone_count: clone_count.clone(), + }; + + let mut map: VecMap = VecMap::new(); + map.entry_ref(&key).or_insert(1); + let clones_after_insert = clone_count.get(); + + // Looking up an existing key must not clone it. + map.entry_ref(&key).or_insert(99); + assert_eq!(clone_count.get(), clones_after_insert); +} + +#[test] +fn test_entry_ref_key_cloned_exactly_once_on_vacant_insert() { + let clone_count = Rc::new(Cell::new(0)); + let key = CountedKey { + value: "a".to_string(), + clone_count: clone_count.clone(), + }; + + let mut map: VecMap = VecMap::new(); + map.entry_ref(&key).or_insert(1); + assert_eq!(clone_count.get(), 1); +} + +#[test] +fn test_entry_ref_or_insert_with_key() { + let mut map: VecMap = VecMap::new(); + let key = "hello".to_string(); + map.entry_ref(&key).or_insert_with_key(|k| k.to_uppercase()); + assert_eq!( + map.iter().collect::>(), + vec![(&"hello".to_string(), &"HELLO".to_string())] + ); +} + +#[test] +fn test_entry_ref_or_insert_with_not_called_when_occupied() { + let mut map: VecMap = VecMap::new(); + let key = "a".to_string(); + map.entry_ref(&key).or_insert(1); + map.entry_ref(&key) + .or_insert_with(|| panic!("should not be called")); + assert_eq!(map.iter().collect::>(), vec![(&key, &1)]); +} + +#[test] +fn test_entry_ref_or_insert_default() { + let mut map: VecMap = VecMap::new(); + map.entry_ref(&"a".to_string()).or_insert_default(); + assert_eq!(map.iter().collect::>(), vec![(&"a".to_string(), &0)]); +} + +#[test] +fn test_entry_ref_key() { + let mut map: VecMap = VecMap::new(); + let key = "a".to_string(); + assert_eq!(*map.entry_ref(&key).key(), key); + map.entry_ref(&key).or_insert(1); + assert_eq!(*map.entry_ref(&key).key(), key); +} + +#[test] +fn test_entry_ref_mut_ref_can_be_updated() { + let mut map: VecMap = VecMap::new(); + let key = "a".to_string(); + let value = map.entry_ref(&key).or_insert(0); + *value = 5; + assert_eq!(map.iter().collect::>(), vec![(&key, &5)]); +} diff --git a/crates/sidebar/src/project_group_builder.rs b/crates/sidebar/src/project_group_builder.rs index 8b4a0862dcd269310eb571f4db6703ed0e508fef..318dfac0a839e28ceb27c6036b87e6a13d9bc992 100644 --- a/crates/sidebar/src/project_group_builder.rs +++ b/crates/sidebar/src/project_group_builder.rs @@ -120,7 +120,7 @@ impl ProjectGroupBuilder { } fn project_group_entry(&mut self, name: &ProjectGroupName) -> &mut ProjectGroup { - self.project_groups.entry(name.clone()).or_insert_default() + self.project_groups.entry_ref(name).or_insert_default() } fn add_mapping(&mut self, work_directory: &Path, original_repo: &Path) {