refactor(lsp): determine file referrer for each document (#23867)

1 files changed, 47 insertions, 18 deletions

diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 5d5f5228c..9215be074 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -4,6 +4,7 @@ use super::analysis::CodeActionData;
 use super::code_lens;
 use super::config;
 use super::documents::AssetOrDocument;
+use super::documents::Document;
 use super::documents::DocumentsFilter;
 use super::language_server;
 use super::language_server::StateSnapshot;
@@ -393,8 +394,8 @@ impl TsServer {
     let _join_handle = thread::spawn(move || {
       run_tsc_thread(
         receiver,
-        performance.clone(),
-        specifier_map.clone(),
+        performance,
+        specifier_map,
         maybe_inspector_server,
       )
     });
@@ -3990,6 +3991,8 @@ struct State {
   response_tx: Option<oneshot::Sender<Result<String, AnyError>>>,
   state_snapshot: Arc<StateSnapshot>,
   specifier_map: Arc<TscSpecifierMap>,
+  root_referrers: HashMap<ModuleSpecifier, ModuleSpecifier>,
+  last_referrer: Option<ModuleSpecifier>,
   token: CancellationToken,
   pending_requests: Option<UnboundedReceiver<Request>>,
   mark: Option<PerformanceMark>,
@@ -4008,12 +4011,30 @@ impl State {
       response_tx: None,
       state_snapshot,
       specifier_map,
+      root_referrers: Default::default(),
+      last_referrer: None,
       token: Default::default(),
       mark: None,
       pending_requests: Some(pending_requests),
     }
   }
 
+  fn get_document(&self, specifier: &ModuleSpecifier) -> Option<Arc<Document>> {
+    if let Some(referrer) = self.root_referrers.get(specifier) {
+      self
+        .state_snapshot
+        .documents
+        .get_or_load(specifier, referrer)
+    } else if let Some(referrer) = &self.last_referrer {
+      self
+        .state_snapshot
+        .documents
+        .get_or_load(specifier, referrer)
+    } else {
+      self.state_snapshot.documents.get(specifier)
+    }
+  }
+
   fn get_asset_or_document(
     &self,
     specifier: &ModuleSpecifier,
@@ -4022,10 +4043,8 @@ impl State {
     if specifier.scheme() == "asset" {
       snapshot.assets.get(specifier).map(AssetOrDocument::Asset)
     } else {
-      snapshot
-        .documents
-        .get(specifier)
-        .map(AssetOrDocument::Document)
+      let document = self.get_document(specifier);
+      document.map(AssetOrDocument::Document)
     }
   }
 
@@ -4037,11 +4056,8 @@ impl State {
         None
       }
     } else {
-      self
-        .state_snapshot
-        .documents
-        .get(specifier)
-        .map(|d| d.script_version())
+      let document = self.get_document(specifier);
+      document.map(|d| d.script_version())
     }
   }
 }
@@ -4222,7 +4238,7 @@ fn op_resolve_inner(
       })
     })
     .collect();
-
+  state.last_referrer = Some(referrer);
   state.performance.measure(mark);
   Ok(specifiers)
 }
@@ -4235,6 +4251,7 @@ fn op_respond(
 ) {
   let state = state.borrow_mut::<State>();
   state.performance.measure(state.mark.take().unwrap());
+  state.last_referrer = None;
   let response = if !error.is_empty() {
     Err(anyhow!("tsc error: {error}"))
   } else {
@@ -4269,9 +4286,16 @@ fn op_script_names(state: &mut OpState) -> Vec<String> {
   }
 
   // inject these next because they're global
-  for specifier in state.state_snapshot.resolver.graph_import_specifiers() {
-    if seen.insert(Cow::Borrowed(specifier.as_str())) {
-      result.push(specifier.to_string());
+  for (referrer, specifiers) in
+    state.state_snapshot.resolver.graph_imports_by_referrer()
+  {
+    for specifier in specifiers {
+      if seen.insert(Cow::Borrowed(specifier.as_str())) {
+        result.push(specifier.to_string());
+      }
+      state
+        .root_referrers
+        .insert(specifier.clone(), referrer.clone());
     }
   }
 
@@ -4289,8 +4313,12 @@ fn op_script_names(state: &mut OpState) -> Vec<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)?;
+        let (types, _) = documents.resolve_dependency(
+          types,
+          specifier,
+          doc.file_referrer(),
+        )?;
+        let types_doc = documents.get_or_load(&types, specifier)?;
         Some(types_doc.specifier().clone())
       })();
       // If there is a types dep, use that as the root instead. But if the doc
@@ -5169,11 +5197,12 @@ mod tests {
         *version,
         *language_id,
         (*source).into(),
+        None,
       );
     }
     let snapshot = Arc::new(StateSnapshot {
       project_version: 0,
-      documents,
+      documents: Arc::new(documents),
       assets: Default::default(),
       config: Arc::new(config),
       resolver,