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),
         }
     }