Add VecMap::entry_ref (#52601)

Eric Holk created

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

Change summary

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(-)

Detailed changes

crates/collections/src/vecmap.rs 🔗

@@ -45,6 +45,20 @@ impl<K: Eq, V> VecMap<K, V> {
             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<K, V>,
     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<F>(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<F>(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<K, V>,
+    key: &'key K,
+}

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<Cell<usize>>,
+}
+
+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<String, i32> = 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<String, i32> = 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<CountedKey, i32> = 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<CountedKey, i32> = 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<String, String> = 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<_>>(),
+        vec![(&"hello".to_string(), &"HELLO".to_string())]
+    );
+}
+
+#[test]
+fn test_entry_ref_or_insert_with_not_called_when_occupied() {
+    let mut map: VecMap<String, i32> = 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<_>>(), vec![(&key, &1)]);
+}
+
+#[test]
+fn test_entry_ref_or_insert_default() {
+    let mut map: VecMap<String, i32> = VecMap::new();
+    map.entry_ref(&"a".to_string()).or_insert_default();
+    assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&"a".to_string(), &0)]);
+}
+
+#[test]
+fn test_entry_ref_key() {
+    let mut map: VecMap<String, i32> = 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<String, i32> = VecMap::new();
+    let key = "a".to_string();
+    let value = map.entry_ref(&key).or_insert(0);
+    *value = 5;
+    assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&key, &5)]);
+}

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) {