Add TreeMap::remove_between that can take abstract start and end points

Nathan Sobo created 2 years ago

This commit introduces a new adaptor trait for SeekTarget that works around
frustrating issues with lifetimes. It wraps the arguments in a newtype wrapper
that lives on the stack to avoid the lifetime getting extended to the caller
of the method.

This allows us to introduce a PathSuccessor object that can be passed as the
end argument of remove_between to remove a whole subtree.

Change summary

crates/project/src/worktree.rs  |   6 
crates/sum_tree/src/sum_tree.rs |   4 
crates/sum_tree/src/tree_map.rs | 190 +++++++++++++++-------------------
3 files changed, 90 insertions(+), 110 deletions(-)

Detailed changes

crates/project/src/worktree.rs

@@ -54,7 +54,7 @@ use std::{
     },
     time::{Duration, SystemTime},
 };
-use sum_tree::{Bias, Edit, SeekTarget, SumTree, TreeMap, TreeSet};
+use sum_tree::{Bias, Edit, SeekTarget, SumTree, TreeMap, TreeSet, PathDescendants};
 use util::{paths::HOME, ResultExt, TakeUntilExt, TryFutureExt};
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)]
@@ -3023,9 +3023,7 @@ impl BackgroundScanner {
             snapshot.repository_entries.update(&work_dir, |entry| {
                 entry
                     .worktree_statuses
-                    .remove_by(&repo_path, |stored_path| {
-                        stored_path.starts_with(&repo_path)
-                    })
+                    .remove_range(&repo_path, &PathDescendants(&repo_path))
             });
         }
 

crates/sum_tree/src/sum_tree.rs

@@ -5,7 +5,7 @@ use arrayvec::ArrayVec;
 pub use cursor::{Cursor, FilterCursor, Iter};
 use std::marker::PhantomData;
 use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
-pub use tree_map::{TreeMap, TreeSet};
+pub use tree_map::{TreeMap, TreeSet, PathDescendants};
 
 #[cfg(test)]
 const TREE_BASE: usize = 2;
@@ -47,7 +47,7 @@ impl<'a, T: Summary> Dimension<'a, T> for T {
 }
 
 pub trait SeekTarget<'a, S: Summary, D: Dimension<'a, S>>: fmt::Debug {
-    fn cmp(&self, cursor_location: &D, cx: &S::Context) -> Ordering;
+   fn cmp(&self, cursor_location: &D, cx: &S::Context) -> Ordering;
 }
 
 impl<'a, S: Summary, D: Dimension<'a, S> + Ord> SeekTarget<'a, S, D> for D {

crates/sum_tree/src/tree_map.rs

@@ -1,4 +1,8 @@
-use std::{cmp::Ordering, fmt::Debug};
+use std::{
+    cmp::Ordering,
+    fmt::Debug,
+    path::{Path, PathBuf},
+};
 
 use crate::{Bias, Dimension, Edit, Item, KeyedItem, SeekTarget, SumTree, Summary};
 
@@ -73,6 +77,17 @@ impl<K: Clone + Debug + Default + Ord, V: Clone + Debug> TreeMap<K, V> {
         removed
     }
 
+    pub fn remove_range(&mut self, start: &impl MapSeekTarget<K>, end: &impl MapSeekTarget<K>) {
+        let start = MapSeekTargetAdaptor(start);
+        let end = MapSeekTargetAdaptor(end);
+        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
+        let mut new_tree = cursor.slice(&start, Bias::Left, &());
+        cursor.seek(&end, Bias::Left, &());
+        new_tree.push_tree(cursor.suffix(&()), &());
+        drop(cursor);
+        self.0 = new_tree;
+    }
+
     /// Returns the key-value pair with the greatest key less than or equal to the given key.
     pub fn closest(&self, key: &K) -> Option<(&K, &V)> {
         let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
@@ -82,17 +97,6 @@ impl<K: Clone + Debug + Default + Ord, V: Clone + Debug> TreeMap<K, V> {
         cursor.item().map(|item| (&item.key, &item.value))
     }
 
-    pub fn remove_between(&mut self, from: &K, until: &K) {
-        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
-        let from_key = MapKeyRef(Some(from));
-        let mut new_tree = cursor.slice(&from_key, Bias::Left, &());
-        let until_key = MapKeyRef(Some(until));
-        cursor.seek_forward(&until_key, Bias::Left, &());
-        new_tree.push_tree(cursor.suffix(&()), &());
-        drop(cursor);
-        self.0 = new_tree;
-    }
-
     pub fn iter_from<'a>(&'a self, from: &'a K) -> impl Iterator<Item = (&K, &V)> + '_ {
         let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
         let from_key = MapKeyRef(Some(from));
@@ -160,46 +164,43 @@ impl<K: Clone + Debug + Default + Ord, V: Clone + Debug> TreeMap<K, V> {
 
         self.0.edit(edits, &());
     }
+}
 
-    pub fn remove_by<F>(&mut self, key: &K, f: F)
-    where
-        F: Fn(&K) -> bool,
-    {
-        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
-        let key = MapKeyRef(Some(key));
-        let mut new_tree = cursor.slice(&key, Bias::Left, &());
-        let until = RemoveByTarget(key, &f);
-        cursor.seek_forward(&until, Bias::Right, &());
-        new_tree.push_tree(cursor.suffix(&()), &());
-        drop(cursor);
-        self.0 = new_tree;
+#[derive(Debug)]
+struct MapSeekTargetAdaptor<'a, T>(&'a T);
+
+impl<'a, K: Debug + Clone + Default + Ord, T: MapSeekTarget<K>>
+    SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>> for MapSeekTargetAdaptor<'_, T>
+{
+    fn cmp(&self, cursor_location: &MapKeyRef<K>, _: &()) -> Ordering {
+        MapSeekTarget::cmp(self.0, cursor_location)
     }
 }
 
-struct RemoveByTarget<'a, K>(MapKeyRef<'a, K>, &'a dyn Fn(&K) -> bool);
+pub trait MapSeekTarget<K>: Debug {
+    fn cmp(&self, cursor_location: &MapKeyRef<K>) -> Ordering;
+}
 
-impl<'a, K: Debug> Debug for RemoveByTarget<'a, K> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("RemoveByTarget")
-            .field("key", &self.0)
-            .field("F", &"<...>")
-            .finish()
+impl<K: Debug + Ord> MapSeekTarget<K> for K {
+    fn cmp(&self, cursor_location: &MapKeyRef<K>) -> Ordering {
+        if let Some(key) = &cursor_location.0 {
+            self.cmp(key)
+        } else {
+            Ordering::Greater
+        }
     }
 }
 
-impl<'a, K: Debug + Clone + Default + Ord> SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>>
-    for RemoveByTarget<'_, K>
-{
-    fn cmp(
-        &self,
-        cursor_location: &MapKeyRef<'a, K>,
-        _cx: &<MapKey<K> as Summary>::Context,
-    ) -> Ordering {
-        if let Some(cursor_location) = cursor_location.0 {
-            if (self.1)(cursor_location) {
-                Ordering::Equal
+#[derive(Debug)]
+pub struct PathDescendants<'a>(&'a Path);
+
+impl MapSeekTarget<PathBuf> for PathDescendants<'_> {
+    fn cmp(&self, cursor_location: &MapKeyRef<PathBuf>) -> Ordering {
+        if let Some(key) = &cursor_location.0 {
+            if key.starts_with(&self.0) {
+                Ordering::Greater
             } else {
-                self.0 .0.unwrap().cmp(cursor_location)
+                self.0.cmp(key)
             }
         } else {
             Ordering::Greater
@@ -266,7 +267,7 @@ where
     K: Clone + Debug + Default + Ord,
 {
     fn cmp(&self, cursor_location: &MapKeyRef<K>, _: &()) -> Ordering {
-        self.0.cmp(&cursor_location.0)
+        Ord::cmp(&self.0, &cursor_location.0)
     }
 }
 
@@ -353,66 +354,6 @@ mod tests {
         assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&4, &"d"), (&6, &"f")]);
     }
 
-    #[test]
-    fn test_remove_between() {
-        let mut map = TreeMap::default();
-
-        map.insert("a", 1);
-        map.insert("b", 2);
-        map.insert("baa", 3);
-        map.insert("baaab", 4);
-        map.insert("c", 5);
-
-        map.remove_between(&"ba", &"bb");
-
-        assert_eq!(map.get(&"a"), Some(&1));
-        assert_eq!(map.get(&"b"), Some(&2));
-        assert_eq!(map.get(&"baaa"), None);
-        assert_eq!(map.get(&"baaaab"), None);
-        assert_eq!(map.get(&"c"), Some(&5));
-    }
-
-    #[test]
-    fn test_remove_by() {
-        let mut map = TreeMap::default();
-
-        map.insert("a", 1);
-        map.insert("aa", 1);
-        map.insert("b", 2);
-        map.insert("baa", 3);
-        map.insert("baaab", 4);
-        map.insert("c", 5);
-        map.insert("ca", 6);
-
-        map.remove_by(&"ba", |key| key.starts_with("ba"));
-
-        assert_eq!(map.get(&"a"), Some(&1));
-        assert_eq!(map.get(&"aa"), Some(&1));
-        assert_eq!(map.get(&"b"), Some(&2));
-        assert_eq!(map.get(&"baaa"), None);
-        assert_eq!(map.get(&"baaaab"), None);
-        assert_eq!(map.get(&"c"), Some(&5));
-        assert_eq!(map.get(&"ca"), Some(&6));
-
-        map.remove_by(&"c", |key| key.starts_with("c"));
-
-        assert_eq!(map.get(&"a"), Some(&1));
-        assert_eq!(map.get(&"aa"), Some(&1));
-        assert_eq!(map.get(&"b"), Some(&2));
-        assert_eq!(map.get(&"c"), None);
-        assert_eq!(map.get(&"ca"), None);
-
-        map.remove_by(&"a", |key| key.starts_with("a"));
-
-        assert_eq!(map.get(&"a"), None);
-        assert_eq!(map.get(&"aa"), None);
-        assert_eq!(map.get(&"b"), Some(&2));
-
-        map.remove_by(&"b", |key| key.starts_with("b"));
-
-        assert_eq!(map.get(&"b"), None);
-    }
-
     #[test]
     fn test_iter_from() {
         let mut map = TreeMap::default();
@@ -461,4 +402,45 @@ mod tests {
         assert_eq!(map.get(&"c"), Some(&3));
         assert_eq!(map.get(&"d"), Some(&4));
     }
+
+    #[test]
+    fn test_remove_between_and_path_successor() {
+        let mut map = TreeMap::default();
+
+        map.insert(PathBuf::from("a"), 1);
+        map.insert(PathBuf::from("a/a"), 1);
+        map.insert(PathBuf::from("b"), 2);
+        map.insert(PathBuf::from("b/a/a"), 3);
+        map.insert(PathBuf::from("b/a/a/a/b"), 4);
+        map.insert(PathBuf::from("c"), 5);
+        map.insert(PathBuf::from("c/a"), 6);
+
+        map.remove_range(&PathBuf::from("b/a"), &PathDescendants(&PathBuf::from("b/a")));
+
+        assert_eq!(map.get(&PathBuf::from("a")), Some(&1));
+        assert_eq!(map.get(&PathBuf::from("a/a")), Some(&1));
+        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
+        assert_eq!(map.get(&PathBuf::from("b/a/a")), None);
+        assert_eq!(map.get(&PathBuf::from("b/a/a/a/b")), None);
+        assert_eq!(map.get(&PathBuf::from("c")), Some(&5));
+        assert_eq!(map.get(&PathBuf::from("c/a")), Some(&6));
+
+        map.remove_range(&PathBuf::from("c"), &PathDescendants(&PathBuf::from("c")));
+
+        assert_eq!(map.get(&PathBuf::from("a")), Some(&1));
+        assert_eq!(map.get(&PathBuf::from("a/a")), Some(&1));
+        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
+        assert_eq!(map.get(&PathBuf::from("c")), None);
+        assert_eq!(map.get(&PathBuf::from("c/a")), None);
+
+        map.remove_range(&PathBuf::from("a"), &PathDescendants(&PathBuf::from("a")));
+
+        assert_eq!(map.get(&PathBuf::from("a")), None);
+        assert_eq!(map.get(&PathBuf::from("a/a")), None);
+        assert_eq!(map.get(&PathBuf::from("b")), Some(&2));
+
+        map.remove_range(&PathBuf::from("b"), &PathDescendants(&PathBuf::from("b")));
+
+        assert_eq!(map.get(&PathBuf::from("b")), None);
+    }
 }