From 00b345fdfe426b00c7da0219585eda03c20a9171 Mon Sep 17 00:00:00 2001
From: Mikayla Maki <mikayla.c.maki@gmail.com>
Date: Wed, 10 May 2023 11:28:39 -0700
Subject: [PATCH] Use sum tree traversal to remove paths

---
 crates/project/src/worktree.rs  | 26 ++++++++-----
 crates/sum_tree/src/tree_map.rs | 68 +++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/crates/project/src/worktree.rs b/crates/project/src/worktree.rs
index 1fc4fcf5a8..ff9f7fde9a 100644
--- a/crates/project/src/worktree.rs
+++ b/crates/project/src/worktree.rs
@@ -7,7 +7,7 @@ use client::{proto, Client};
 use clock::ReplicaId;
 use collections::{HashMap, VecDeque};
 use fs::{
-    repository::{GitRepository, GitFileStatus, RepoPath},
+    repository::{GitFileStatus, GitRepository, RepoPath},
     Fs, LineEnding,
 };
 use futures::{
@@ -46,6 +46,7 @@ use std::{
     future::Future,
     mem,
     ops::{Deref, DerefMut},
+
     path::{Path, PathBuf},
     pin::Pin,
     sync::{
@@ -2896,12 +2897,13 @@ impl BackgroundScanner {
                 .git_repositories
                 .update(&work_dir_id, |entry| entry.scan_id = scan_id);
 
-            // TODO: Status Replace linear scan with smarter sum tree traversal
-            snapshot
-                .repository_entries
-                .update(&work_dir, |entry| entry.worktree_statuses.retain(|stored_path, _| {
-                    !stored_path.starts_with(&repo_path)
-                }));
+            snapshot.repository_entries.update(&work_dir, |entry| {
+                entry
+                    .worktree_statuses
+                    .remove_from_while(&repo_path, |stored_path, _| {
+                        stored_path.starts_with(&repo_path)
+                    })
+            });
         }
 
         Some(())
@@ -2958,9 +2960,9 @@ impl BackgroundScanner {
                 .git_repositories
                 .update(&work_dir_id, |entry| entry.scan_id = scan_id);
 
-            snapshot
-                .repository_entries
-                .update(&work_dir, |entry| entry.worktree_statuses.insert(repo_path, status));
+            snapshot.repository_entries.update(&work_dir, |entry| {
+                entry.worktree_statuses.insert(repo_path, status)
+            });
         }
 
         Some(())
@@ -3948,6 +3950,8 @@ mod tests {
         std::fs::remove_file(work_dir.join(B_TXT)).unwrap();
         std::fs::remove_dir_all(work_dir.join("c")).unwrap();
 
+        dbg!(git_status(&repo));
+
         tree.flush_fs_events(cx).await;
 
         // Check that non-repo behavior is tracked
@@ -3955,6 +3959,8 @@ mod tests {
             let snapshot = tree.snapshot();
             let (_, repo) = snapshot.repository_entries.iter().next().unwrap();
 
+            dbg!(&repo.worktree_statuses);
+
             assert_eq!(repo.worktree_statuses.iter().count(), 0);
             assert_eq!(repo.worktree_statuses.get(&Path::new(A_TXT).into()), None);
             assert_eq!(repo.worktree_statuses.get(&Path::new(B_TXT).into()), None);
diff --git a/crates/sum_tree/src/tree_map.rs b/crates/sum_tree/src/tree_map.rs
index ab37d2577a..359137d439 100644
--- a/crates/sum_tree/src/tree_map.rs
+++ b/crates/sum_tree/src/tree_map.rs
@@ -82,6 +82,36 @@ 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 remove_from_while<F>(&mut self, from: &K, mut f: F)
+    where F: FnMut(&K, &V) -> bool
+    {
+        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
+        let from_key = MapKeyRef(Some(from));
+        let mut new_tree = cursor.slice(&from_key, Bias::Left, &());
+        while let Some(item) = cursor.item() {
+            if !f(&item.key, &item.value) {
+                break;
+            }
+            cursor.next(&());
+        }
+        new_tree.push_tree(cursor.suffix(&()), &());
+        drop(cursor);
+        self.0 = new_tree;
+    }
+
+
     pub fn update<F, T>(&mut self, key: &K, f: F) -> Option<T>
     where
         F: FnOnce(&mut V) -> T,
@@ -272,4 +302,42 @@ mod tests {
         map.retain(|key, _| *key % 2 == 0);
         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_from_while() {
+        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_from_while(&"ba", |key, _| key.starts_with(&"ba"));
+
+        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));
+    }
 }