Re-index project when a worktree is registered

Antonio Scandurra and Kyle Caverly created

Co-Authored-By: Kyle Caverly <kyle@zed.dev>

Change summary

crates/semantic_index/src/semantic_index.rs       | 69 +++++-----------
crates/semantic_index/src/semantic_index_tests.rs |  7 -
2 files changed, 24 insertions(+), 52 deletions(-)

Detailed changes

crates/semantic_index/src/semantic_index.rs 🔗

@@ -34,7 +34,6 @@ use util::{
     paths::EMBEDDINGS_DIR,
     ResultExt, TryFutureExt,
 };
-use workspace::WorkspaceCreated;
 
 const SEMANTIC_INDEX_VERSION: usize = 9;
 const BACKGROUND_INDEXING_DELAY: Duration = Duration::from_secs(5 * 60);
@@ -57,24 +56,6 @@ pub fn init(
         return;
     }
 
-    cx.subscribe_global::<WorkspaceCreated, _>({
-        move |event, cx| {
-            let Some(semantic_index) = SemanticIndex::global(cx) else {
-                return;
-            };
-            let workspace = &event.0;
-            if let Some(workspace) = workspace.upgrade(cx) {
-                let project = workspace.read(cx).project().clone();
-                if project.read(cx).is_local() {
-                    semantic_index.update(cx, |index, cx| {
-                        index.register_project(project, cx);
-                    });
-                }
-            }
-        }
-    })
-    .detach();
-
     cx.spawn(move |mut cx| async move {
         let semantic_index = SemanticIndex::new(
             fs,
@@ -426,31 +407,6 @@ impl SemanticIndex {
         }
     }
 
-    fn register_project(
-        &mut self,
-        project: ModelHandle<Project>,
-        cx: &mut ModelContext<Self>,
-    ) -> &mut ProjectState {
-        if !self.projects.contains_key(&project.downgrade()) {
-            log::trace!("Registering Project for Semantic Index");
-
-            let subscription = cx.subscribe(&project, |this, project, event, cx| match event {
-                project::Event::WorktreeAdded | project::Event::WorktreeRemoved(_) => {
-                    this.project_worktrees_changed(project.clone(), cx);
-                }
-                project::Event::WorktreeUpdatedEntries(worktree_id, changes) => {
-                    this.project_entries_changed(project, *worktree_id, changes.clone(), cx);
-                }
-                _ => {}
-            });
-            self.projects
-                .insert(project.downgrade(), ProjectState::new(subscription));
-            self.project_worktrees_changed(project.clone(), cx);
-        }
-
-        self.projects.get_mut(&project.downgrade()).unwrap()
-    }
-
     fn register_worktree(
         &mut self,
         project: ModelHandle<Project>,
@@ -542,11 +498,14 @@ impl SemanticIndex {
                         anyhow::Ok(changed_paths)
                     })
                     .await?;
-                this.update(&mut cx, |this, _| {
+                this.update(&mut cx, |this, cx| {
                     let project_state = this
                         .projects
                         .get_mut(&project)
                         .ok_or_else(|| anyhow!("project not registered"))?;
+                    let project = project
+                        .upgrade(cx)
+                        .ok_or_else(|| anyhow!("project was dropped"))?;
 
                     if let Some(WorktreeState::Registering(state)) =
                         project_state.worktrees.remove(&worktree_id)
@@ -560,6 +519,7 @@ impl SemanticIndex {
                             changed_paths,
                         }),
                     );
+                    this.index_project(project, cx);
 
                     anyhow::Ok(())
                 })?;
@@ -762,7 +722,24 @@ impl SemanticIndex {
     }
 
     pub fn index_project(&mut self, project: ModelHandle<Project>, cx: &mut ModelContext<Self>) {
-        let project_state = self.register_project(project.clone(), cx);
+        if !self.projects.contains_key(&project.downgrade()) {
+            log::trace!("Registering Project for Semantic Index");
+
+            let subscription = cx.subscribe(&project, |this, project, event, cx| match event {
+                project::Event::WorktreeAdded | project::Event::WorktreeRemoved(_) => {
+                    this.project_worktrees_changed(project.clone(), cx);
+                }
+                project::Event::WorktreeUpdatedEntries(worktree_id, changes) => {
+                    this.project_entries_changed(project, *worktree_id, changes.clone(), cx);
+                }
+                _ => {}
+            });
+            self.projects
+                .insert(project.downgrade(), ProjectState::new(subscription));
+            self.project_worktrees_changed(project.clone(), cx);
+        }
+
+        let project_state = self.projects.get_mut(&project.downgrade()).unwrap();
 
         let mut pending_files = Vec::new();
         let mut files_to_delete = Vec::new();

crates/semantic_index/src/semantic_index_tests.rs 🔗

@@ -87,14 +87,9 @@ async fn test_semantic_index(deterministic: Arc<Deterministic>, cx: &mut TestApp
 
     let project = Project::test(fs.clone(), ["/the-root".as_ref()], cx).await;
 
-    semantic_index.update(cx, |store, cx| {
-        store.register_project(project.clone(), cx);
-    });
-    deterministic.run_until_parked();
-
+    semantic_index.update(cx, |store, cx| store.index_project(project.clone(), cx));
     let pending_file_count =
         semantic_index.read_with(cx, |index, _| index.pending_file_count(&project).unwrap());
-    semantic_index.update(cx, |store, cx| store.index_project(project.clone(), cx));
     deterministic.run_until_parked();
     assert_eq!(*pending_file_count.borrow(), 3);
     deterministic.advance_clock(EMBEDDING_QUEUE_FLUSH_TIMEOUT);