Checkpoint

crates/project/src/project_tree.rs

@@ -8,9 +8,8 @@ mod toolchain_tree;
 
 use std::{
     borrow::Borrow,
-    collections::{btree_map::Entry as TreeEntry, hash_map::Entry, BTreeMap},
+    collections::{hash_map::Entry, BTreeMap},
     ops::ControlFlow,
-    path::Path,
     sync::Arc,
 };
 
@@ -28,13 +27,11 @@ use crate::{
 };
 
 pub(crate) use server_tree::LanguageServerTree;
-type IsRoot = bool;
 
 struct WorktreeRoots {
-    _roots: RootPathTrie<LanguageServerName>,
-    roots: BTreeMap<Arc<Path>, BTreeMap<LanguageServerName, IsRoot>>,
+    roots: RootPathTrie<LanguageServerName>,
     worktree_store: Model<WorktreeStore>,
-    worktree_subscription: Subscription,
+    _worktree_subscription: Subscription,
 }
 
 impl WorktreeRoots {
@@ -44,16 +41,16 @@ impl WorktreeRoots {
         cx: &mut AppContext,
     ) -> Model<Self> {
         cx.new_model(|cx| Self {
-            _roots: RootPathTrie::new(),
-            roots: Default::default(),
+            roots: RootPathTrie::new(),
             worktree_store,
-            worktree_subscription: cx.subscribe(&worktree, |this: &mut Self, _, event, cx| {
+            _worktree_subscription: cx.subscribe(&worktree, |this: &mut Self, _, event, cx| {
                 match event {
                     WorktreeEvent::UpdatedEntries(changes) => {
                         for (path, _, kind) in changes.iter() {
                             match kind {
                                 worktree::PathChange::Removed => {
-                                    this.roots.remove(path);
+                                    let path = TriePath::from(path.as_ref());
+                                    this.roots.remove(&path);
                                 }
                                 _ => {}
                             }
@@ -65,7 +62,8 @@ impl WorktreeRoots {
                         else {
                             return;
                         };
-                        this.roots.remove(&entry.path);
+                        let path = TriePath::from(entry.path.as_ref());
+                        this.roots.remove(&path);
                     }
                 }
             }),
@@ -150,11 +148,11 @@ impl ProjectTree {
 
         let key = TriePath::from(&*path);
 
-        let mut roots = worktree_roots.update(cx, |this, _| {
-            this._roots.walk(&key, &mut |path, labels| {
-                for label in labels {
+        worktree_roots.update(cx, |this, _| {
+            this.roots.walk(&key, &mut |path, labels| {
+                for (label, presence) in labels {
                     if let Some(slot) = roots.get_mut(label) {
-                        debug_assert_eq!(slot, &mut None, "For a given path to a root of a worktree there should be at most project root");
+                        debug_assert_eq!(slot, &mut None, "For a given path to a root of a worktree there should be at most project root of {label:?} kind");
                         let _ = slot.insert(ProjectPath {
                             worktree_id,
                             path: path.clone(),
@@ -168,7 +166,6 @@ impl ProjectTree {
                     ControlFlow::Continue(())
                 }
             });
-            roots
         });
 
         // for ancestor in path.ancestors().skip(1) {

crates/project/src/project_tree/path_trie.rs

@@ -13,12 +13,17 @@ use std::{
 /// a known root at `python/project` and the unexplored part is `subdir/another_subdir` - we need to run a scan on these 2 directories
 pub(super) struct RootPathTrie<Label> {
     path_component: Arc<OsStr>,
-    labels: BTreeSet<Label>,
+    labels: BTreeMap<Label, LabelPresence>,
     children: BTreeMap<Arc<OsStr>, RootPathTrie<Label>>,
 }
 
+#[derive(Clone, Copy, Debug, PartialOrd, PartialEq, Ord, Eq)]
+pub(super) enum LabelPresence {
+    Present,
+    KnownAbsent,
+}
 impl<Label: Ord> RootPathTrie<Label> {
-    pub(crate) fn new() -> Self {
+    pub(super) fn new() -> Self {
         Self::new_with_key(Arc::from(OsStr::new("")))
     }
     fn new_with_key(path_component: Arc<OsStr>) -> Self {
@@ -28,7 +33,7 @@ impl<Label: Ord> RootPathTrie<Label> {
             children: Default::default(),
         }
     }
-    pub(crate) fn insert(&mut self, path: &TriePath, value: Label) {
+    pub(super) fn insert(&mut self, path: &TriePath, value: Label, presence: LabelPresence) {
         let mut current = self;
 
         for key in path.0.iter() {
@@ -39,12 +44,16 @@ impl<Label: Ord> RootPathTrie<Label> {
                 Entry::Occupied(occupied_entry) => occupied_entry.into_mut(),
             };
         }
-        current.labels.insert(value);
+        let _previous_value = current.labels.insert(value, presence);
+        debug_assert_eq!(_previous_value, None);
     }
-    pub(crate) fn walk<'a>(
+    pub(super) fn walk<'a>(
         &'a self,
         path: &TriePath,
-        callback: &mut dyn for<'b> FnMut(&'b Arc<Path>, &'a BTreeSet<Label>) -> ControlFlow<()>,
+        callback: &mut dyn for<'b> FnMut(
+            &'b Arc<Path>,
+            &'a BTreeMap<Label, LabelPresence>,
+        ) -> ControlFlow<()>,
     ) {
         let mut current = self;
         let tmp_path = Arc::from(Path::new(""));
@@ -63,11 +72,41 @@ impl<Label: Ord> RootPathTrie<Label> {
             (callback)(&tmp_path, &current.labels);
         }
     }
+
+    pub(super) fn remove(&mut self, path: &TriePath) {
+        let mut current = &mut *self;
+        // Tracks how many nodes (starting from the leaf, going upwards) can be removed
+        let mut consecutive_node_chain = 0;
+        for path in path.0.iter() {
+            if current.children.len() > 1 {
+                consecutive_node_chain = 0;
+            } else if current.children.len() == 1 {
+                consecutive_node_chain += 1;
+            }
+            current = match current.children.get_mut(path) {
+                Some(child) => child,
+                None => return,
+            };
+        }
+        // Now walk the tree again, this time iterating only up to the root of the consecutive node chain.
+        let consecutive_chain_start = path.0.len() - consecutive_node_chain;
+        let mut current = self;
+        for path in path.0[..consecutive_chain_start].iter() {
+            current = match current.children.get_mut(path) {
+                Some(child) => child,
+                None => unreachable!(),
+            };
+        }
+        current
+            .children
+            .remove(&path.0[consecutive_chain_start])
+            .expect("The removal to succeed");
+    }
 }
 
 /// [TriePath] is a [Path] preprocessed for amortizing the cost of doing multiple lookups in distinct [RootPathTrie]s.
 #[derive(Clone)]
-pub(crate) struct TriePath(Arc<[Arc<OsStr>]>);
+pub(super) struct TriePath(Arc<[Arc<OsStr>]>);
 
 impl From<&Path> for TriePath {
     fn from(value: &Path) -> Self {