Rework how we send diagnostics to client.

The previous way of sending from the thread pool suffered from stale
diagnostics due to being canceled before we could clear the old ones.

The key change is moving to sending diagnostics from the main loop
thread, but doing all the hard work in the thread pool. This should
provide the best of both worlds, with little to no of the downsides.

This should hopefully fix a lot of issues, but we'll need testing in
each individual issue to be sure.

5 files changed, 140 insertions, 70 deletions

diff --git a/crates/ra_lsp_server/src/diagnostics.rs b/crates/ra_lsp_server/src/diagnostics.rs
new file mode 100644
index 000000000..ea08bce24
--- /dev/null
+++ b/crates/ra_lsp_server/src/diagnostics.rs
@@ -0,0 +1,85 @@
+//! Book keeping for keeping diagnostics easily in sync with the client.
+use lsp_types::{CodeActionOrCommand, Diagnostic, Range};
+use ra_ide::FileId;
+use std::{collections::HashMap, sync::Arc};
+
+pub type CheckFixes = Arc<HashMap<FileId, Vec<Fix>>>;
+
+#[derive(Debug, Default, Clone)]
+pub struct DiagnosticCollection {
+    pub native: HashMap<FileId, Vec<Diagnostic>>,
+    pub check: HashMap<FileId, Vec<Diagnostic>>,
+    pub check_fixes: CheckFixes,
+}
+
+#[derive(Debug, Clone)]
+pub struct Fix {
+    pub range: Range,
+    pub action: CodeActionOrCommand,
+}
+
+#[derive(Debug)]
+pub enum DiagnosticTask {
+    ClearCheck,
+    AddCheck(FileId, Diagnostic, Vec<CodeActionOrCommand>),
+    SetNative(FileId, Vec<Diagnostic>),
+}
+
+impl DiagnosticCollection {
+    pub fn clear_check(&mut self) -> Vec<FileId> {
+        Arc::make_mut(&mut self.check_fixes).clear();
+        self.check.drain().map(|(key, _value)| key).collect()
+    }
+
+    pub fn add_check_diagnostic(
+        &mut self,
+        file_id: FileId,
+        diagnostic: Diagnostic,
+        fixes: Vec<CodeActionOrCommand>,
+    ) {
+        let diagnostics = self.check.entry(file_id).or_default();
+        for existing_diagnostic in diagnostics.iter() {
+            if are_diagnostics_equal(&existing_diagnostic, &diagnostic) {
+                return;
+            }
+        }
+
+        let check_fixes = Arc::make_mut(&mut self.check_fixes);
+        check_fixes
+            .entry(file_id)
+            .or_default()
+            .extend(fixes.into_iter().map(|action| Fix { range: diagnostic.range, action }));
+        diagnostics.push(diagnostic);
+    }
+
+    pub fn set_native_diagnostics(&mut self, file_id: FileId, diagnostics: Vec<Diagnostic>) {
+        self.native.insert(file_id, diagnostics);
+    }
+
+    pub fn diagnostics_for(&self, file_id: FileId) -> impl Iterator<Item = &Diagnostic> {
+        let native = self.native.get(&file_id).into_iter().flatten();
+        let check = self.check.get(&file_id).into_iter().flatten();
+        native.chain(check)
+    }
+
+    pub fn handle_task(&mut self, task: DiagnosticTask) -> Vec<FileId> {
+        match task {
+            DiagnosticTask::ClearCheck => self.clear_check(),
+            DiagnosticTask::AddCheck(file_id, diagnostic, fixes) => {
+                self.add_check_diagnostic(file_id, diagnostic, fixes);
+                vec![file_id]
+            }
+            DiagnosticTask::SetNative(file_id, diagnostics) => {
+                self.set_native_diagnostics(file_id, diagnostics);
+                vec![file_id]
+            }
+        }
+    }
+}
+
+fn are_diagnostics_equal(left: &Diagnostic, right: &Diagnostic) -> bool {
+    left.source == right.source
+        && left.severity == right.severity
+        && left.range == right.range
+        && left.message == right.message
+}
diff --git a/crates/ra_lsp_server/src/lib.rs b/crates/ra_lsp_server/src/lib.rs
index 2ca149fd5..1208c1343 100644
--- a/crates/ra_lsp_server/src/lib.rs
+++ b/crates/ra_lsp_server/src/lib.rs
@@ -29,6 +29,7 @@ mod markdown;
 pub mod req;
 mod config;
 mod world;
+mod diagnostics;
 
 pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
 pub use crate::{
diff --git a/crates/ra_lsp_server/src/main_loop.rs b/crates/ra_lsp_server/src/main_loop.rs
index 508fe08c0..12961ba37 100644
--- a/crates/ra_lsp_server/src/main_loop.rs
+++ b/crates/ra_lsp_server/src/main_loop.rs
@@ -17,16 +17,17 @@ use std::{
 use crossbeam_channel::{select, unbounded, RecvError, Sender};
 use lsp_server::{Connection, ErrorCode, Message, Notification, Request, RequestId, Response};
 use lsp_types::{ClientCapabilities, NumberOrString};
-use ra_cargo_watch::{CheckOptions, CheckTask};
+use ra_cargo_watch::{url_from_path_with_drive_lowercasing, CheckOptions, CheckTask};
 use ra_ide::{Canceled, FeatureFlags, FileId, LibraryData, SourceRootId};
 use ra_prof::profile;
-use ra_vfs::{VfsTask, Watch};
+use ra_vfs::{VfsFile, VfsTask, Watch};
 use relative_path::RelativePathBuf;
 use rustc_hash::FxHashSet;
 use serde::{de::DeserializeOwned, Serialize};
 use threadpool::ThreadPool;
 
 use crate::{
+    diagnostics::DiagnosticTask,
     main_loop::{
         pending_requests::{PendingRequest, PendingRequests},
         subscriptions::Subscriptions,
@@ -254,6 +255,7 @@ pub fn main_loop(
 enum Task {
     Respond(Response),
     Notify(Notification),
+    Diagnostic(DiagnosticTask),
 }
 
 enum Event {
@@ -359,7 +361,7 @@ fn loop_turn(
             world_state.maybe_collect_garbage();
             loop_state.in_flight_libraries -= 1;
         }
-        Event::CheckWatcher(task) => on_check_task(pool, task, world_state, task_sender)?,
+        Event::CheckWatcher(task) => on_check_task(task, world_state, task_sender)?,
         Event::Msg(msg) => match msg {
             Message::Request(req) => on_request(
                 world_state,
@@ -464,6 +466,7 @@ fn on_task(
         Task::Notify(n) => {
             msg_sender.send(n.into()).unwrap();
         }
+        Task::Diagnostic(task) => on_diagnostic_task(task, msg_sender, state),
     }
 }
 
@@ -621,23 +624,26 @@ fn on_notification(
 }
 
 fn on_check_task(
-    pool: &ThreadPool,
     task: CheckTask,
     world_state: &mut WorldState,
     task_sender: &Sender<Task>,
 ) -> Result<()> {
-    let urls = match task {
+    match task {
         CheckTask::ClearDiagnostics => {
-            let state = Arc::get_mut(&mut world_state.check_watcher.state)
-                .expect("couldn't get check watcher state as mutable");
-            state.clear()
+            task_sender.send(Task::Diagnostic(DiagnosticTask::ClearCheck))?;
         }
 
-        CheckTask::AddDiagnostic(url, diagnostic) => {
-            let state = Arc::get_mut(&mut world_state.check_watcher.state)
-                .expect("couldn't get check watcher state as mutable");
-            state.add_diagnostic_with_fixes(url.clone(), diagnostic);
-            vec![url]
+        CheckTask::AddDiagnostic { url, diagnostic, fixes } => {
+            let path = url.to_file_path().map_err(|()| format!("invalid uri: {}", url))?;
+            let file_id = world_state
+                .vfs
+                .read()
+                .path2file(&path)
+                .map(|it| FileId(it.0))
+                .ok_or_else(|| format!("unknown file: {}", path.to_string_lossy()))?;
+
+            task_sender
+                .send(Task::Diagnostic(DiagnosticTask::AddCheck(file_id, diagnostic, fixes)))?;
         }
 
         CheckTask::Status(progress) => {
@@ -647,31 +653,30 @@ fn on_check_task(
             };
             let not = notification_new::<req::Progress>(params);
             task_sender.send(Task::Notify(not)).unwrap();
-            Vec::new()
         }
     };
 
-    let subscriptions = urls
-        .into_iter()
-        .map(|url| {
-            let path = url.to_file_path().map_err(|()| format!("invalid uri: {}", url))?;
-            Ok(world_state.vfs.read().path2file(&path).map(|it| FileId(it.0)))
-        })
-        .filter_map(|res| res.transpose())
-        .collect::<Result<Vec<_>>>()?;
+    Ok(())
+}
 
-    // We manually send a diagnostic update when the watcher asks
-    // us to, to avoid the issue of having to change the file to
-    // receive updated diagnostics.
-    update_file_notifications_on_threadpool(
-        pool,
-        world_state.snapshot(),
-        false,
-        task_sender.clone(),
-        subscriptions,
-    );
+fn on_diagnostic_task(task: DiagnosticTask, msg_sender: &Sender<Message>, state: &mut WorldState) {
+    let subscriptions = state.diagnostics.handle_task(task);
 
-    Ok(())
+    for file_id in subscriptions {
+        let path = state.vfs.read().file2path(VfsFile(file_id.0));
+        let uri = match url_from_path_with_drive_lowercasing(&path) {
+            Ok(uri) => uri,
+            Err(err) => {
+                log::error!("Couldn't convert path to url ({}): {:?}", err, path.to_string_lossy());
+                continue;
+            }
+        };
+
+        let diagnostics = state.diagnostics.diagnostics_for(file_id).cloned().collect();
+        let params = req::PublishDiagnosticsParams { uri, diagnostics, version: None };
+        let not = notification_new::<req::PublishDiagnostics>(params);
+        msg_sender.send(not.into()).unwrap();
+    }
 }
 
 struct PoolDispatcher<'a> {
@@ -819,9 +824,8 @@ fn update_file_notifications_on_threadpool(
                             log::error!("failed to compute diagnostics: {:?}", e);
                         }
                     }
-                    Ok(params) => {
-                        let not = notification_new::<req::PublishDiagnostics>(params);
-                        task_sender.send(Task::Notify(not)).unwrap();
+                    Ok(task) => {
+                        task_sender.send(Task::Diagnostic(task)).unwrap();
                     }
                 }
             }
diff --git a/crates/ra_lsp_server/src/main_loop/handlers.rs b/crates/ra_lsp_server/src/main_loop/handlers.rs
index 9aa1e7eea..839bca342 100644
--- a/crates/ra_lsp_server/src/main_loop/handlers.rs
+++ b/crates/ra_lsp_server/src/main_loop/handlers.rs
@@ -33,6 +33,7 @@ use crate::{
         to_call_hierarchy_item, to_location, Conv, ConvWith, FoldConvCtx, MapConvWith, TryConvWith,
         TryConvWithToVec,
     },
+    diagnostics::DiagnosticTask,
     req::{self, Decoration, InlayHint, InlayHintsParams, InlayKind},
     world::WorldSnapshot,
     LspError, Result,
@@ -656,6 +657,7 @@ pub fn handle_code_action(
         .filter(|(diag_range, _fix)| diag_range.intersection(&range).is_some())
         .map(|(_range, fix)| fix);
 
+    // TODO: When done, we won't need this, only the one pulling from world.diagnostics
     for source_edit in fixes_from_diagnostics {
         let title = source_edit.label.clone();
         let edit = source_edit.try_conv_with(&world)?;
@@ -676,28 +678,12 @@ pub fn handle_code_action(
         res.push(action.into());
     }
 
-    for fix in world.check_watcher.fixes_for(&params.text_document.uri).into_iter().flatten() {
-        let fix_range = fix.location.range.conv_with(&line_index);
+    for fix in world.check_fixes.get(&file_id).into_iter().flatten() {
+        let fix_range = fix.range.conv_with(&line_index);
         if fix_range.intersection(&range).is_none() {
             continue;
         }
-
-        let edit = {
-            let edits = vec![TextEdit::new(fix.location.range, fix.replacement.clone())];
-            let mut edit_map = std::collections::HashMap::new();
-            edit_map.insert(fix.location.uri.clone(), edits);
-            WorkspaceEdit::new(edit_map)
-        };
-
-        let action = CodeAction {
-            title: fix.title.clone(),
-            kind: Some("quickfix".to_string()),
-            diagnostics: Some(fix.diagnostics.clone()),
-            edit: Some(edit),
-            command: None,
-            is_preferred: None,
-        };
-        res.push(action.into());
+        res.push(fix.action.clone());
     }
 
     for assist in world.analysis().assists(FileRange { file_id, range })?.into_iter() {
@@ -875,14 +861,10 @@ pub fn handle_document_highlight(
     ))
 }
 
-pub fn publish_diagnostics(
-    world: &WorldSnapshot,
-    file_id: FileId,
-) -> Result<req::PublishDiagnosticsParams> {
+pub fn publish_diagnostics(world: &WorldSnapshot, file_id: FileId) -> Result<DiagnosticTask> {
     let _p = profile("publish_diagnostics");
-    let uri = world.file_id_to_uri(file_id)?;
     let line_index = world.analysis().file_line_index(file_id)?;
-    let mut diagnostics: Vec<Diagnostic> = world
+    let diagnostics: Vec<Diagnostic> = world
         .analysis()
         .diagnostics(file_id)?
         .into_iter()
@@ -896,10 +878,7 @@ pub fn publish_diagnostics(
             tags: None,
         })
         .collect();
-    if let Some(check_diags) = world.check_watcher.diagnostics_for(&uri) {
-        diagnostics.extend(check_diags.iter().cloned());
-    }
-    Ok(req::PublishDiagnosticsParams { uri, diagnostics, version: None })
+    Ok(DiagnosticTask::SetNative(file_id, diagnostics))
 }
 
 pub fn publish_decorations(
diff --git a/crates/ra_lsp_server/src/world.rs b/crates/ra_lsp_server/src/world.rs
index 3059ef9ec..1ee02b47c 100644
--- a/crates/ra_lsp_server/src/world.rs
+++ b/crates/ra_lsp_server/src/world.rs
@@ -12,9 +12,7 @@ use crossbeam_channel::{unbounded, Receiver};
 use lsp_server::ErrorCode;
 use lsp_types::Url;
 use parking_lot::RwLock;
-use ra_cargo_watch::{
-    url_from_path_with_drive_lowercasing, CheckOptions, CheckState, CheckWatcher,
-};
+use ra_cargo_watch::{url_from_path_with_drive_lowercasing, CheckOptions, CheckWatcher};
 use ra_ide::{
     Analysis, AnalysisChange, AnalysisHost, CrateGraph, FeatureFlags, FileId, LibraryData,
     SourceRootId,
@@ -25,6 +23,7 @@ use ra_vfs_glob::{Glob, RustPackageFilterBuilder};
 use relative_path::RelativePathBuf;
 
 use crate::{
+    diagnostics::{CheckFixes, DiagnosticCollection},
     main_loop::pending_requests::{CompletedRequest, LatestRequests},
     LspError, Result,
 };
@@ -55,6 +54,7 @@ pub struct WorldState {
     pub task_receiver: Receiver<VfsTask>,
     pub latest_requests: Arc<RwLock<LatestRequests>>,
     pub check_watcher: CheckWatcher,
+    pub diagnostics: DiagnosticCollection,
 }
 
 /// An immutable snapshot of the world's state at a point in time.
@@ -63,7 +63,7 @@ pub struct WorldSnapshot {
     pub workspaces: Arc<Vec<ProjectWorkspace>>,
     pub analysis: Analysis,
     pub latest_requests: Arc<RwLock<LatestRequests>>,
-    pub check_watcher: CheckState,
+    pub check_fixes: CheckFixes,
     vfs: Arc<RwLock<Vfs>>,
 }
 
@@ -159,6 +159,7 @@ impl WorldState {
             task_receiver,
             latest_requests: Default::default(),
             check_watcher,
+            diagnostics: Default::default(),
         }
     }
 
@@ -220,7 +221,7 @@ impl WorldState {
             analysis: self.analysis_host.analysis(),
             vfs: Arc::clone(&self.vfs),
             latest_requests: Arc::clone(&self.latest_requests),
-            check_watcher: (*self.check_watcher.state).clone(),
+            check_fixes: Arc::clone(&self.diagnostics.check_fixes),
         }
     }