fix: correctly update diagnostics when files are opened and closed

Basically, this tracks the changes to `subscriptions` we use when
issuing a publish_diagnostics.

Author: matklad

crates/rust-analyzer/src/main_loop.rs

@@ -406,25 +406,49 @@ impl GlobalState {
         }
 
         let state_changed = self.process_changes();
+        let memdocs_added_or_removed = self.mem_docs.take_changes();
 
-        if self.is_quiescent() && !was_quiescent {
-            for flycheck in &self.flycheck {
-                flycheck.update();
+        if self.is_quiescent() {
+            if !was_quiescent {
+                for flycheck in &self.flycheck {
+                    flycheck.update();
+                }
             }
-        }
 
-        if self.is_quiescent() && (!was_quiescent || state_changed) {
-            self.update_file_notifications_on_threadpool();
+            if !was_quiescent || state_changed {
+                // Ensure that only one cache priming task can run at a time
+                self.prime_caches_queue.request_op();
+                if self.prime_caches_queue.should_start_op() {
+                    self.task_pool.handle.spawn_with_sender({
+                        let snap = self.snapshot();
+                        move |sender| {
+                            let cb = |progress| {
+                                sender.send(Task::PrimeCaches(progress)).unwrap();
+                            };
+                            match snap.analysis.prime_caches(cb) {
+                                Ok(()) => (),
+                                Err(_canceled) => (),
+                            }
+                        }
+                    });
+                }
 
-            // Refresh semantic tokens if the client supports it.
-            if self.config.semantic_tokens_refresh() {
-                self.semantic_tokens_cache.lock().clear();
-                self.send_request::<lsp_types::request::SemanticTokensRefesh>((), |_, _| ());
+                // Refresh semantic tokens if the client supports it.
+                if self.config.semantic_tokens_refresh() {
+                    self.semantic_tokens_cache.lock().clear();
+                    self.send_request::<lsp_types::request::SemanticTokensRefesh>((), |_, _| ());
+                }
+
+                // Refresh code lens if the client supports it.
+                if self.config.code_lens_refresh() {
+                    self.send_request::<lsp_types::request::CodeLensRefresh>((), |_, _| ());
+                }
             }
 
-            // Refresh code lens if the client supports it.
-            if self.config.code_lens_refresh() {
-                self.send_request::<lsp_types::request::CodeLensRefresh>((), |_, _| ());
+            if !was_quiescent || state_changed || memdocs_added_or_removed {
+                if self.config.publish_diagnostics() {
+                    self.update_diagnostics()
+                }
             }
         }
 
@@ -586,42 +610,32 @@ impl GlobalState {
                     {
                         log::error!("duplicate DidOpenTextDocument: {}", path)
                     }
-                    let changed = this
-                        .vfs
+                    this.vfs
                         .write()
                         .0
                         .set_file_contents(path, Some(params.text_document.text.into_bytes()));
-
-                    // If the VFS contents are unchanged, update diagnostics, since `handle_event`
-                    // won't see any changes. This avoids missing diagnostics when opening a file.
-                    //
-                    // If the file *was* changed, `handle_event` will already recompute and send
-                    // diagnostics. We can't do it here, since the *current* file contents might be
-                    // unset in salsa, since the VFS change hasn't been applied to the database yet.
-                    if !changed {
-                        this.maybe_update_diagnostics();
-                    }
                 }
                 Ok(())
             })?
             .on::<lsp_types::notification::DidChangeTextDocument>(|this, params| {
                 if let Ok(path) = from_proto::vfs_path(&params.text_document.uri) {
-                    let doc = match this.mem_docs.get_mut(&path) {
-                        Some(doc) => doc,
+                    match this.mem_docs.get_mut(&path) {
+                        Some(doc) => {
+                            // The version passed in DidChangeTextDocument is the version after all edits are applied
+                            // so we should apply it before the vfs is notified.
+                            doc.version = params.text_document.version;
+                        }
                         None => {
                             log::error!("expected DidChangeTextDocument: {}", path);
                             return Ok(());
                         }
                     };
+
                     let vfs = &mut this.vfs.write().0;
                     let file_id = vfs.file_id(&path).unwrap();
                     let mut text = String::from_utf8(vfs.file_contents(file_id).to_vec()).unwrap();
                     apply_document_changes(&mut text, params.content_changes);
 
-                    // The version passed in DidChangeTextDocument is the version after all edits are applied
-                    // so we should apply it before the vfs is notified.
-                    doc.version = params.text_document.version;
-
                     vfs.set_file_contents(path.clone(), Some(text.into_bytes()));
                 }
                 Ok(())
@@ -696,53 +710,33 @@ impl GlobalState {
             .finish();
         Ok(())
     }
-    fn update_file_notifications_on_threadpool(&mut self) {
-        self.maybe_update_diagnostics();
-
-        // Ensure that only one cache priming task can run at a time
-        self.prime_caches_queue.request_op();
-        if self.prime_caches_queue.should_start_op() {
-            self.task_pool.handle.spawn_with_sender({
-                let snap = self.snapshot();
-                move |sender| {
-                    let cb = |progress| {
-                        sender.send(Task::PrimeCaches(progress)).unwrap();
-                    };
-                    match snap.analysis.prime_caches(cb) {
-                        Ok(()) => (),
-                        Err(_canceled) => (),
-                    }
-                }
-            });
-        }
-    }
-    fn maybe_update_diagnostics(&mut self) {
+
+    fn update_diagnostics(&mut self) {
         let subscriptions = self
             .mem_docs
             .iter()
             .map(|path| self.vfs.read().0.file_id(path).unwrap())
             .collect::<Vec<_>>();
 
         log::trace!("updating notifications for {:?}", subscriptions);
-        if self.config.publish_diagnostics() {
-            let snapshot = self.snapshot();
-            self.task_pool.handle.spawn(move || {
-                let diagnostics = subscriptions
-                    .into_iter()
-                    .filter_map(|file_id| {
-                        handlers::publish_diagnostics(&snapshot, file_id)
-                            .map_err(|err| {
-                                if !is_cancelled(&*err) {
-                                    log::error!("failed to compute diagnostics: {:?}", err);
-                                }
-                                ()
-                            })
-                            .ok()
-                            .map(|diags| (file_id, diags))
-                    })
-                    .collect::<Vec<_>>();
-                Task::Diagnostics(diagnostics)
-            })
-        }
+
+        let snapshot = self.snapshot();
+        self.task_pool.handle.spawn(move || {
+            let diagnostics = subscriptions
+                .into_iter()
+                .filter_map(|file_id| {
+                    handlers::publish_diagnostics(&snapshot, file_id)
+                        .map_err(|err| {
+                            if !is_cancelled(&*err) {
+                                log::error!("failed to compute diagnostics: {:?}", err);
+                            }
+                            ()
+                        })
+                        .ok()
+                        .map(|diags| (file_id, diags))
+                })
+                .collect::<Vec<_>>();
+            Task::Diagnostics(diagnostics)
+        })
     }
 }

crates/rust-analyzer/src/mem_docs.rs

@@ -1,5 +1,7 @@
 //! In-memory document information.
 
+use std::mem;
+
 use rustc_hash::FxHashMap;
 use vfs::VfsPath;
 
@@ -10,19 +12,22 @@ use vfs::VfsPath;
 #[derive(Default, Clone)]
 pub(crate) struct MemDocs {
     mem_docs: FxHashMap<VfsPath, DocumentData>,
+    added_or_removed: bool,
 }
 
 impl MemDocs {
     pub(crate) fn contains(&self, path: &VfsPath) -> bool {
         self.mem_docs.contains_key(path)
     }
     pub(crate) fn insert(&mut self, path: VfsPath, data: DocumentData) -> Result<(), ()> {
+        self.added_or_removed = true;
         match self.mem_docs.insert(path, data) {
             Some(_) => Err(()),
             None => Ok(()),
         }
     }
     pub(crate) fn remove(&mut self, path: &VfsPath) -> Result<(), ()> {
+        self.added_or_removed = true;
         match self.mem_docs.remove(path) {
             Some(_) => Ok(()),
             None => Err(()),
@@ -32,12 +37,16 @@ impl MemDocs {
         self.mem_docs.get(path)
     }
     pub(crate) fn get_mut(&mut self, path: &VfsPath) -> Option<&mut DocumentData> {
+        // NB: don't set `self.added_or_removed` here, as that purposefully only
+        // tracks changes to the key set.
         self.mem_docs.get_mut(path)
     }
-
     pub(crate) fn iter(&self) -> impl Iterator<Item = &VfsPath> {
         self.mem_docs.keys()
     }
+    pub(crate) fn take_changes(&mut self) -> bool {
+        mem::replace(&mut self.added_or_removed, false)
+    }
 }
 
 /// Information about a document that the Language Client