From 318aa263d01d38286edfca0572943d3e6125cb03 Mon Sep 17 00:00:00 2001
From: Nayeem Rahman <nayeemrmn99@gmail.com>
Date: Wed, 15 May 2024 14:31:07 +0100
Subject: [PATCH] fix(lsp): respect types dependencies for tsc roots

---
 cli/lsp/documents.rs           |  2 +-
 cli/lsp/tsc.rs                 | 28 +++++++++++++++++++++++-----
 tests/integration/lsp_tests.rs | 28 ++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/cli/lsp/documents.rs b/cli/lsp/documents.rs
index 8c553b03d28707..8a6e8b65b568ee 100644
--- a/cli/lsp/documents.rs
+++ b/cli/lsp/documents.rs
@@ -1280,7 +1280,7 @@ impl Documents {
     self.dirty = false;
   }
 
-  fn resolve_dependency(
+  pub fn resolve_dependency(
     &self,
     specifier: &ModuleSpecifier,
     referrer: &ModuleSpecifier,
diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 6999b40eff4805..b41a8fb1930fc9 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -67,6 +67,7 @@ use regex::Captures;
 use regex::Regex;
 use serde_repr::Deserialize_repr;
 use serde_repr::Serialize_repr;
+use std::borrow::Cow;
 use std::cell::RefCell;
 use std::cmp;
 use std::collections::HashMap;
@@ -4261,12 +4262,12 @@ fn op_script_names(state: &mut OpState) -> Vec<String> {
     // ensure this is first so it resolves the node types first
     let specifier = "asset:///node_types.d.ts";
     result.push(specifier.to_string());
-    seen.insert(specifier);
+    seen.insert(Cow::Borrowed(specifier));
   }
 
   // inject these next because they're global
   for specifier in state.state_snapshot.resolver.graph_import_specifiers() {
-    if seen.insert(specifier.as_str()) {
+    if seen.insert(Cow::Borrowed(specifier.as_str())) {
       result.push(specifier.to_string());
     }
   }
@@ -4278,10 +4279,27 @@ fn op_script_names(state: &mut OpState) -> Vec<String> {
     .documents(DocumentsFilter::AllDiagnosable);
   for doc in &docs {
     let specifier = doc.specifier();
-    if seen.insert(specifier.as_str())
-      && (doc.is_open() || specifier.scheme() == "file")
+    let is_open = doc.is_open();
+    if seen.insert(Cow::Borrowed(specifier.as_str()))
+      && (is_open || specifier.scheme() == "file")
     {
-      result.push(specifier.to_string());
+      let types_specifier = (|| {
+        let documents = &state.state_snapshot.documents;
+        let types = doc.maybe_types_dependency().maybe_specifier()?;
+        let (types, _) = documents.resolve_dependency(types, specifier)?;
+        let types_doc = documents.get(&types)?;
+        Some(types_doc.specifier().clone())
+      })();
+      // If there is a types dep, use that as the root instead. But if the doc
+      // is open, include both as roots.
+      if let Some(types_specifier) = &types_specifier {
+        if seen.insert(Cow::Owned(types_specifier.to_string())) {
+          result.push(types_specifier.to_string());
+        }
+      }
+      if types_specifier.is_none() || is_open {
+        result.push(specifier.to_string());
+      }
     }
   }
 
diff --git a/tests/integration/lsp_tests.rs b/tests/integration/lsp_tests.rs
index 47fefeafea309c..6aec33e9969cac 100644
--- a/tests/integration/lsp_tests.rs
+++ b/tests/integration/lsp_tests.rs
@@ -8954,6 +8954,34 @@ fn lsp_diagnostics_deno_types() {
   client.shutdown();
 }
 
+#[test]
+fn lsp_root_with_global_reference_types() {
+  let context = TestContextBuilder::new()
+    .use_http_server()
+    .use_temp_cwd()
+    .build();
+  let temp_dir = context.temp_dir();
+  let file = source_file(
+    temp_dir.path().join("file.ts"),
+    "import 'http://localhost:4545/subdir/foo_types.d.ts'; Foo.bar;",
+  );
+  let file2 = source_file(
+    temp_dir.path().join("file2.ts"),
+    r#"/// <reference types="http://localhost:4545/subdir/foo_types.d.ts" />"#,
+  );
+  let mut client = context.new_lsp_command().build();
+  client.initialize_default();
+  client.write_request(
+    "workspace/executeCommand",
+    json!({
+      "command": "deno.cache",
+      "arguments": [[], file2.uri()],
+    }),
+  );
+  let diagnostics = client.did_open_file(&file);
+  assert_eq!(json!(diagnostics.all()), json!([]));
+}
+
 #[test]
 fn lsp_diagnostics_refresh_dependents() {
   let context = TestContextBuilder::new().use_temp_cwd().build();