3 files changed, 31 insertions, 13 deletions

diff --git a/crates/rust-analyzer/src/global_state.rs b/crates/rust-analyzer/src/global_state.rs
index 94973b90a..80937dbc4 100644
--- a/crates/rust-analyzer/src/global_state.rs
+++ b/crates/rust-analyzer/src/global_state.rs
@@ -12,7 +12,7 @@ use parking_lot::RwLock;
 use ra_db::{CrateId, VfsPath};
 use ra_ide::{Analysis, AnalysisChange, AnalysisHost, FileId};
 use ra_project_model::{CargoWorkspace, ProcMacroClient, ProjectWorkspace, Target};
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashMap;
 
 use crate::{
     config::Config,
@@ -69,7 +69,7 @@ pub(crate) struct GlobalState {
     pub(crate) config: Config,
     pub(crate) analysis_host: AnalysisHost,
     pub(crate) diagnostics: DiagnosticCollection,
-    pub(crate) mem_docs: FxHashSet<VfsPath>,
+    pub(crate) mem_docs: FxHashMap<VfsPath, Option<i64>>,
     pub(crate) vfs: Arc<RwLock<(vfs::Vfs, FxHashMap<FileId, LineEndings>)>>,
     pub(crate) status: Status,
     pub(crate) source_root_config: SourceRootConfig,
@@ -84,6 +84,7 @@ pub(crate) struct GlobalStateSnapshot {
     pub(crate) analysis: Analysis,
     pub(crate) check_fixes: CheckFixes,
     pub(crate) latest_requests: Arc<RwLock<LatestRequests>>,
+    mem_docs: FxHashMap<VfsPath, Option<i64>>,
     vfs: Arc<RwLock<(vfs::Vfs, FxHashMap<FileId, LineEndings>)>>,
     pub(crate) workspaces: Arc<Vec<ProjectWorkspace>>,
 }
@@ -117,7 +118,7 @@ impl GlobalState {
             config,
             analysis_host,
             diagnostics: Default::default(),
-            mem_docs: FxHashSet::default(),
+            mem_docs: FxHashMap::default(),
             vfs: Arc::new(RwLock::new((vfs::Vfs::default(), FxHashMap::default()))),
             status: Status::default(),
             source_root_config: SourceRootConfig::default(),
@@ -183,6 +184,7 @@ impl GlobalState {
             vfs: Arc::clone(&self.vfs),
             latest_requests: Arc::clone(&self.latest_requests),
             check_fixes: Arc::clone(&self.diagnostics.check_fixes),
+            mem_docs: self.mem_docs.clone(),
         }
     }
 
@@ -255,6 +257,11 @@ impl GlobalStateSnapshot {
         self.vfs.read().1[&id]
     }
 
+    pub(crate) fn url_file_version(&self, url: &Url) -> Option<i64> {
+        let path = from_proto::vfs_path(&url).ok()?;
+        self.mem_docs.get(&path).copied()?
+    }
+
     pub(crate) fn anchored_path(&self, file_id: FileId, path: &str) -> Url {
         let mut base = self.vfs.read().0.file_path(file_id);
         base.pop();
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index bb7c4c0c6..00ae6b900 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -210,7 +210,7 @@ impl GlobalState {
                             let vfs = &mut self.vfs.write().0;
                             for (path, contents) in files {
                                 let path = VfsPath::from(path);
-                                if !self.mem_docs.contains(&path) {
+                                if !self.mem_docs.contains_key(&path) {
                                     vfs.set_file_contents(path, contents)
                                 }
                             }
@@ -299,7 +299,7 @@ impl GlobalState {
         if self.status == Status::Ready && (state_changed || prev_status == Status::Loading) {
             let subscriptions = self
                 .mem_docs
-                .iter()
+                .keys()
                 .map(|path| self.vfs.read().0.file_id(&path).unwrap())
                 .collect::<Vec<_>>();
 
@@ -310,8 +310,12 @@ impl GlobalState {
             for file_id in diagnostic_changes {
                 let url = file_id_to_url(&self.vfs.read().0, file_id);
                 let diagnostics = self.diagnostics.diagnostics_for(file_id).cloned().collect();
+                let version = from_proto::vfs_path(&url)
+                    .map(|path| self.mem_docs.get(&path).copied().flatten())
+                    .unwrap_or_default();
+
                 self.send_notification::<lsp_types::notification::PublishDiagnostics>(
-                    lsp_types::PublishDiagnosticsParams { uri: url, diagnostics, version: None },
+                    lsp_types::PublishDiagnosticsParams { uri: url, diagnostics, version },
                 );
             }
         }
@@ -400,7 +404,11 @@ impl GlobalState {
             })?
             .on::<lsp_types::notification::DidOpenTextDocument>(|this, params| {
                 if let Ok(path) = from_proto::vfs_path(&params.text_document.uri) {
-                    if !this.mem_docs.insert(path.clone()) {
+                    if this
+                        .mem_docs
+                        .insert(path.clone(), Some(params.text_document.version))
+                        .is_some()
+                    {
                         log::error!("duplicate DidOpenTextDocument: {}", path)
                     }
                     this.vfs
@@ -412,18 +420,20 @@ impl GlobalState {
             })?
             .on::<lsp_types::notification::DidChangeTextDocument>(|this, params| {
                 if let Ok(path) = from_proto::vfs_path(&params.text_document.uri) {
-                    assert!(this.mem_docs.contains(&path));
+                    assert!(this.mem_docs.contains_key(&path));
                     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);
-                    vfs.set_file_contents(path, Some(text.into_bytes()))
+                    vfs.set_file_contents(path.clone(), Some(text.into_bytes()));
+
+                    this.mem_docs.insert(path, params.text_document.version);
                 }
                 Ok(())
             })?
             .on::<lsp_types::notification::DidCloseTextDocument>(|this, params| {
                 if let Ok(path) = from_proto::vfs_path(&params.text_document.uri) {
-                    if !this.mem_docs.remove(&path) {
+                    if this.mem_docs.remove(&path).is_none() {
                         log::error!("orphan DidCloseTextDocument: {}", path)
                     }
                     if let Some(path) = path.as_path() {
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index fc94f28b9..c6935c029 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -480,9 +480,10 @@ pub(crate) fn url_from_abs_path(path: &Path) -> lsp_types::Url {
 pub(crate) fn versioned_text_document_identifier(
     snap: &GlobalStateSnapshot,
     file_id: FileId,
-    version: Option<i64>,
 ) -> lsp_types::VersionedTextDocumentIdentifier {
-    lsp_types::VersionedTextDocumentIdentifier { uri: url(snap, file_id), version }
+    let url = url(snap, file_id);
+    let version = snap.url_file_version(&url);
+    lsp_types::VersionedTextDocumentIdentifier { uri: url, version }
 }
 
 pub(crate) fn location(
@@ -571,7 +572,7 @@ pub(crate) fn snippet_text_document_edit(
     is_snippet: bool,
     source_file_edit: SourceFileEdit,
 ) -> Result<lsp_ext::SnippetTextDocumentEdit> {
-    let text_document = versioned_text_document_identifier(snap, source_file_edit.file_id, None);
+    let text_document = versioned_text_document_identifier(snap, source_file_edit.file_id);
     let line_index = snap.analysis.file_line_index(source_file_edit.file_id)?;
     let line_endings = snap.file_line_endings(source_file_edit.file_id);
     let edits = source_file_edit