fix(lsp): handle mbc documents properly (#9151)

Co-authored-by: Ryan Dahl <ry@tinyclouds.org>

1 files changed, 93 insertions, 86 deletions

diff --git a/cli/lsp/diagnostics.rs b/cli/lsp/diagnostics.rs
index 6632620ec..36e7b093c 100644
--- a/cli/lsp/diagnostics.rs
+++ b/cli/lsp/diagnostics.rs
@@ -4,7 +4,6 @@ use super::analysis::get_lint_references;
 use super::analysis::references_to_diagnostics;
 use super::analysis::ResolvedDependency;
 use super::language_server::StateSnapshot;
-use super::memory_cache::FileId;
 use super::tsc;
 
 use crate::diagnostics;
@@ -13,7 +12,7 @@ use crate::media_type::MediaType;
 use deno_core::error::custom_error;
 use deno_core::error::AnyError;
 use deno_core::serde_json;
-use deno_core::serde_json::Value;
+use deno_core::ModuleSpecifier;
 use lspower::lsp_types;
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -28,43 +27,47 @@ pub enum DiagnosticSource {
 
 #[derive(Debug, Default, Clone)]
 pub struct DiagnosticCollection {
-  map: HashMap<(FileId, DiagnosticSource), Vec<lsp_types::Diagnostic>>,
-  versions: HashMap<FileId, i32>,
-  changes: HashSet<FileId>,
+  map: HashMap<(ModuleSpecifier, DiagnosticSource), Vec<lsp_types::Diagnostic>>,
+  versions: HashMap<ModuleSpecifier, i32>,
+  changes: HashSet<ModuleSpecifier>,
 }
 
 impl DiagnosticCollection {
   pub fn set(
     &mut self,
-    file_id: FileId,
+    specifier: ModuleSpecifier,
     source: DiagnosticSource,
     version: Option<i32>,
     diagnostics: Vec<lsp_types::Diagnostic>,
   ) {
-    self.map.insert((file_id, source), diagnostics);
+    self.map.insert((specifier.clone(), source), diagnostics);
     if let Some(version) = version {
-      self.versions.insert(file_id, version);
+      self.versions.insert(specifier.clone(), version);
     }
-    self.changes.insert(file_id);
+    self.changes.insert(specifier);
   }
 
   pub fn diagnostics_for(
     &self,
-    file_id: FileId,
-    source: DiagnosticSource,
+    specifier: &ModuleSpecifier,
+    source: &DiagnosticSource,
   ) -> impl Iterator<Item = &lsp_types::Diagnostic> {
-    self.map.get(&(file_id, source)).into_iter().flatten()
+    self
+      .map
+      .get(&(specifier.clone(), source.clone()))
+      .into_iter()
+      .flatten()
   }
 
-  pub fn get_version(&self, file_id: &FileId) -> Option<i32> {
-    self.versions.get(file_id).cloned()
+  pub fn get_version(&self, specifier: &ModuleSpecifier) -> Option<i32> {
+    self.versions.get(specifier).cloned()
   }
 
-  pub fn invalidate(&mut self, file_id: &FileId) {
-    self.versions.remove(file_id);
+  pub fn invalidate(&mut self, specifier: &ModuleSpecifier) {
+    self.versions.remove(specifier);
   }
 
-  pub fn take_changes(&mut self) -> Option<HashSet<FileId>> {
+  pub fn take_changes(&mut self) -> Option<HashSet<ModuleSpecifier>> {
     if self.changes.is_empty() {
       return None;
     }
@@ -72,7 +75,8 @@ impl DiagnosticCollection {
   }
 }
 
-pub type DiagnosticVec = Vec<(FileId, Option<i32>, Vec<lsp_types::Diagnostic>)>;
+pub type DiagnosticVec =
+  Vec<(ModuleSpecifier, Option<i32>, Vec<lsp_types::Diagnostic>)>;
 
 pub async fn generate_lint_diagnostics(
   state_snapshot: StateSnapshot,
@@ -81,25 +85,24 @@ pub async fn generate_lint_diagnostics(
   tokio::task::spawn_blocking(move || {
     let mut diagnostic_list = Vec::new();
 
-    let file_cache = state_snapshot.file_cache.lock().unwrap();
-    for (specifier, doc_data) in state_snapshot.doc_data.iter() {
-      let file_id = file_cache.lookup(specifier).unwrap();
-      let version = doc_data.version;
-      let current_version = diagnostic_collection.get_version(&file_id);
+    let documents = state_snapshot.documents.lock().unwrap();
+    for specifier in documents.open_specifiers() {
+      let version = documents.version(specifier);
+      let current_version = diagnostic_collection.get_version(specifier);
       if version != current_version {
         let media_type = MediaType::from(specifier);
-        if let Ok(source_code) = file_cache.get_contents(file_id) {
+        if let Ok(Some(source_code)) = documents.content(specifier) {
           if let Ok(references) =
             get_lint_references(specifier, &media_type, &source_code)
           {
             if !references.is_empty() {
               diagnostic_list.push((
-                file_id,
+                specifier.clone(),
                 version,
                 references_to_diagnostics(references),
               ));
             } else {
-              diagnostic_list.push((file_id, version, Vec::new()));
+              diagnostic_list.push((specifier.clone(), version, Vec::new()));
             }
           }
         } else {
@@ -154,7 +157,7 @@ fn to_lsp_range(
   }
 }
 
-type TsDiagnostics = Vec<diagnostics::Diagnostic>;
+type TsDiagnostics = HashMap<String, Vec<diagnostics::Diagnostic>>;
 
 fn get_diagnostic_message(diagnostic: &diagnostics::Diagnostic) -> String {
   if let Some(message) = diagnostic.message_text.clone() {
@@ -197,65 +200,70 @@ fn to_lsp_related_information(
 }
 
 fn ts_json_to_diagnostics(
-  value: Value,
-) -> Result<Vec<lsp_types::Diagnostic>, AnyError> {
-  let ts_diagnostics: TsDiagnostics = serde_json::from_value(value)?;
-  Ok(
-    ts_diagnostics
-      .iter()
-      .filter_map(|d| {
-        if let (Some(start), Some(end)) = (&d.start, &d.end) {
-          Some(lsp_types::Diagnostic {
-            range: to_lsp_range(start, end),
-            severity: Some((&d.category).into()),
-            code: Some(lsp_types::NumberOrString::Number(d.code as i32)),
-            code_description: None,
-            source: Some("deno-ts".to_string()),
-            message: get_diagnostic_message(d),
-            related_information: to_lsp_related_information(
-              &d.related_information,
-            ),
-            tags: match d.code {
-              // These are codes that indicate the variable is unused.
-              6133 | 6192 | 6196 => {
-                Some(vec![lsp_types::DiagnosticTag::Unnecessary])
-              }
-              _ => None,
-            },
-            data: None,
-          })
-        } else {
-          None
-        }
-      })
-      .collect(),
-  )
+  diagnostics: &[diagnostics::Diagnostic],
+) -> Vec<lsp_types::Diagnostic> {
+  diagnostics
+    .iter()
+    .filter_map(|d| {
+      if let (Some(start), Some(end)) = (&d.start, &d.end) {
+        Some(lsp_types::Diagnostic {
+          range: to_lsp_range(start, end),
+          severity: Some((&d.category).into()),
+          code: Some(lsp_types::NumberOrString::Number(d.code as i32)),
+          code_description: None,
+          source: Some("deno-ts".to_string()),
+          message: get_diagnostic_message(d),
+          related_information: to_lsp_related_information(
+            &d.related_information,
+          ),
+          tags: match d.code {
+            // These are codes that indicate the variable is unused.
+            6133 | 6192 | 6196 => {
+              Some(vec![lsp_types::DiagnosticTag::Unnecessary])
+            }
+            _ => None,
+          },
+          data: None,
+        })
+      } else {
+        None
+      }
+    })
+    .collect()
 }
 
 pub async fn generate_ts_diagnostics(
-  ts_server: &tsc::TsServer,
-  diagnostic_collection: &DiagnosticCollection,
   state_snapshot: StateSnapshot,
+  diagnostic_collection: DiagnosticCollection,
+  ts_server: &tsc::TsServer,
 ) -> Result<DiagnosticVec, AnyError> {
   let mut diagnostics = Vec::new();
-  let state_snapshot_ = state_snapshot.clone();
-  for (specifier, doc_data) in state_snapshot_.doc_data.iter() {
-    let file_id = {
-      // TODO(lucacasonato): this is highly inefficient
-      let file_cache = state_snapshot_.file_cache.lock().unwrap();
-      file_cache.lookup(specifier).unwrap()
-    };
-    let version = doc_data.version;
-    let current_version = diagnostic_collection.get_version(&file_id);
-    if version != current_version {
-      let req = tsc::RequestMethod::GetDiagnostics(specifier.clone());
-      let ts_diagnostics = ts_json_to_diagnostics(
-        ts_server.request(state_snapshot.clone(), req).await?,
-      )?;
-      diagnostics.push((file_id, version, ts_diagnostics));
+  let mut specifiers = Vec::new();
+  {
+    let documents = state_snapshot.documents.lock().unwrap();
+    for specifier in documents.open_specifiers() {
+      let version = documents.version(specifier);
+      let current_version = diagnostic_collection.get_version(specifier);
+      if version != current_version {
+        specifiers.push(specifier.clone());
+      }
+    }
+  }
+  if !specifiers.is_empty() {
+    let req = tsc::RequestMethod::GetDiagnostics(specifiers);
+    let res = ts_server.request(state_snapshot.clone(), req).await?;
+    let ts_diagnostic_map: TsDiagnostics = serde_json::from_value(res)?;
+    for (specifier_str, ts_diagnostics) in ts_diagnostic_map.iter() {
+      let specifier = ModuleSpecifier::resolve_url(specifier_str)?;
+      let version =
+        state_snapshot.documents.lock().unwrap().version(&specifier);
+      diagnostics.push((
+        specifier,
+        version,
+        ts_json_to_diagnostics(ts_diagnostics),
+      ));
     }
   }
-
   Ok(diagnostics)
 }
 
@@ -266,19 +274,18 @@ pub async fn generate_dependency_diagnostics(
   tokio::task::spawn_blocking(move || {
     let mut diagnostics = Vec::new();
 
-    let file_cache = state_snapshot.file_cache.lock().unwrap();
     let mut sources = if let Ok(sources) = state_snapshot.sources.lock() {
       sources
     } else {
       return Err(custom_error("Deadlock", "deadlock locking sources"));
     };
-    for (specifier, doc_data) in state_snapshot.doc_data.iter() {
-      let file_id = file_cache.lookup(specifier).unwrap();
-      let version = doc_data.version;
-      let current_version = diagnostic_collection.get_version(&file_id);
+    let documents = state_snapshot.documents.lock().unwrap();
+    for specifier in documents.open_specifiers() {
+      let version = documents.version(specifier);
+      let current_version = diagnostic_collection.get_version(specifier);
       if version != current_version {
         let mut diagnostic_list = Vec::new();
-        if let Some(dependencies) = &doc_data.dependencies {
+        if let Some(dependencies) = documents.dependencies(specifier) {
           for (_, dependency) in dependencies.iter() {
             if let (Some(code), Some(range)) = (
               &dependency.maybe_code,
@@ -299,7 +306,7 @@ pub async fn generate_dependency_diagnostics(
                   })
                 }
                 ResolvedDependency::Resolved(specifier) => {
-                  if !(state_snapshot.doc_data.contains_key(&specifier) || sources.contains(&specifier)) {
+                  if !(documents.contains(&specifier) || sources.contains(&specifier)) {
                     let is_local = specifier.as_url().scheme() == "file";
                     diagnostic_list.push(lsp_types::Diagnostic {
                       range: *range,
@@ -322,7 +329,7 @@ pub async fn generate_dependency_diagnostics(
             }
           }
         }
-        diagnostics.push((file_id, version, diagnostic_list))
+        diagnostics.push((specifier.clone(), version, diagnostic_list))
       }
     }