fix(lsp): `textDocument/references` should respect `includeDeclaration` (#18496)

3 files changed, 122 insertions, 73 deletions

diff --git a/cli/lsp/code_lens.rs b/cli/lsp/code_lens.rs
index 1253fe5aa..650e5e241 100644
--- a/cli/lsp/code_lens.rs
+++ b/cli/lsp/code_lens.rs
@@ -297,70 +297,72 @@ async fn resolve_references_code_lens(
   data: CodeLensData,
   language_server: &language_server::Inner,
 ) -> Result<lsp::CodeLens, AnyError> {
-  let asset_or_document =
-    language_server.get_asset_or_document(&data.specifier)?;
-  let line_index = asset_or_document.line_index();
-  let req = tsc::RequestMethod::GetReferences((
-    data.specifier.clone(),
-    line_index.offset_tsc(code_lens.range.start)?,
-  ));
-  let snapshot = language_server.snapshot();
-  let maybe_references: Option<Vec<tsc::ReferenceEntry>> =
-    language_server.ts_server.request(snapshot, req).await?;
-  if let Some(references) = maybe_references {
+  fn get_locations(
+    maybe_referenced_symbols: Option<Vec<tsc::ReferencedSymbol>>,
+    language_server: &language_server::Inner,
+  ) -> Result<Vec<lsp::Location>, AnyError> {
+    let symbols = match maybe_referenced_symbols {
+      Some(symbols) => symbols,
+      None => return Ok(Vec::new()),
+    };
     let mut locations = Vec::new();
-    for reference in references {
+    for reference in symbols.iter().flat_map(|s| &s.references) {
       if reference.is_definition {
         continue;
       }
       let reference_specifier =
-        resolve_url(&reference.document_span.file_name)?;
+        resolve_url(&reference.entry.document_span.file_name)?;
       let asset_or_doc =
         language_server.get_asset_or_document(&reference_specifier)?;
       locations.push(
         reference
+          .entry
           .to_location(asset_or_doc.line_index(), &language_server.url_map),
       );
     }
-    let command = if !locations.is_empty() {
-      let title = if locations.len() > 1 {
-        format!("{} references", locations.len())
-      } else {
-        "1 reference".to_string()
-      };
-      lsp::Command {
-        title,
-        command: "deno.showReferences".to_string(),
-        arguments: Some(vec![
-          json!(data.specifier),
-          json!(code_lens.range.start),
-          json!(locations),
-        ]),
-      }
-    } else {
-      lsp::Command {
-        title: "0 references".to_string(),
-        command: "".to_string(),
-        arguments: None,
-      }
-    };
-    Ok(lsp::CodeLens {
-      range: code_lens.range,
-      command: Some(command),
-      data: None,
-    })
+    Ok(locations)
+  }
+
+  let asset_or_document =
+    language_server.get_asset_or_document(&data.specifier)?;
+  let line_index = asset_or_document.line_index();
+  let snapshot = language_server.snapshot();
+  let maybe_referenced_symbols = language_server
+    .ts_server
+    .find_references(
+      snapshot,
+      &data.specifier,
+      line_index.offset_tsc(code_lens.range.start)?,
+    )
+    .await?;
+  let locations = get_locations(maybe_referenced_symbols, language_server)?;
+  let title = if locations.len() == 1 {
+    "1 reference".to_string()
   } else {
-    let command = lsp::Command {
-      title: "0 references".to_string(),
-      command: "".to_string(),
+    format!("{} references", locations.len())
+  };
+  let command = if locations.is_empty() {
+    lsp::Command {
+      title,
+      command: String::new(),
       arguments: None,
-    };
-    Ok(lsp::CodeLens {
-      range: code_lens.range,
-      command: Some(command),
-      data: None,
-    })
-  }
+    }
+  } else {
+    lsp::Command {
+      title,
+      command: "deno.showReferences".to_string(),
+      arguments: Some(vec![
+        json!(data.specifier),
+        json!(code_lens.range.start),
+        json!(locations),
+      ]),
+    }
+  };
+  Ok(lsp::CodeLens {
+    range: code_lens.range,
+    command: Some(command),
+    data: None,
+  })
 }
 
 pub async fn resolve_code_lens(
diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index 2d0bbd140..164c9734f 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -1953,27 +1953,23 @@ impl Inner {
     let mark = self.performance.mark("references", Some(&params));
     let asset_or_doc = self.get_asset_or_document(&specifier)?;
     let line_index = asset_or_doc.line_index();
-    let req = tsc::RequestMethod::GetReferences((
-      specifier.clone(),
-      line_index.offset_tsc(params.text_document_position.position)?,
-    ));
-    let maybe_references: Option<Vec<tsc::ReferenceEntry>> = self
+    let maybe_referenced_symbols = self
       .ts_server
-      .request(self.snapshot(), req)
-      .await
-      .map_err(|err| {
-        error!("Unable to get references from TypeScript: {}", err);
-        LspError::internal_error()
-      })?;
+      .find_references(
+        self.snapshot(),
+        &specifier,
+        line_index.offset_tsc(params.text_document_position.position)?,
+      )
+      .await?;
 
-    if let Some(references) = maybe_references {
+    if let Some(symbols) = maybe_referenced_symbols {
       let mut results = Vec::new();
-      for reference in references {
+      for reference in symbols.iter().flat_map(|s| &s.references) {
         if !params.context.include_declaration && reference.is_definition {
           continue;
         }
         let reference_specifier =
-          resolve_url(&reference.document_span.file_name).unwrap();
+          resolve_url(&reference.entry.document_span.file_name).unwrap();
         let reference_line_index = if reference_specifier == specifier {
           line_index.clone()
         } else {
@@ -1981,8 +1977,11 @@ impl Inner {
             self.get_asset_or_document(&reference_specifier)?;
           asset_or_doc.line_index()
         };
-        results
-          .push(reference.to_location(reference_line_index, &self.url_map));
+        results.push(
+          reference
+            .entry
+            .to_location(reference_line_index, &self.url_map),
+        );
       }
 
       self.performance.measure(mark);
diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 316436988..e846cc496 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -149,7 +149,28 @@ impl TsServer {
     if self.0.send((req, snapshot, tx, token)).is_err() {
       return Err(anyhow!("failed to send request to tsc thread"));
     }
-    rx.await?.map(|v| serde_json::from_value::<R>(v).unwrap())
+    let value = rx.await??;
+    Ok(serde_json::from_value::<R>(value)?)
+  }
+
+  // todo(dsherret): refactor the rest of the request methods to have
+  // methods to call on this struct, then make `RequestMethod` and
+  // friends internal
+
+  pub async fn find_references(
+    &self,
+    snapshot: Arc<StateSnapshot>,
+    specifier: &ModuleSpecifier,
+    position: u32,
+  ) -> Result<Option<Vec<ReferencedSymbol>>, LspError> {
+    let req = RequestMethod::FindReferences {
+      specifier: specifier.clone(),
+      position,
+    };
+    self.request(snapshot, req).await.map_err(|err| {
+      log::error!("Unable to get references from TypeScript: {}", err);
+      LspError::internal_error()
+    })
   }
 }
 
@@ -1688,10 +1709,31 @@ pub struct CombinedCodeActions {
 
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
-pub struct ReferenceEntry {
-  // is_write_access: bool,
+pub struct ReferencedSymbol {
+  pub definition: ReferencedSymbolDefinitionInfo,
+  pub references: Vec<ReferencedSymbolEntry>,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ReferencedSymbolDefinitionInfo {
+  #[serde(flatten)]
+  pub definition_info: DefinitionInfo,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ReferencedSymbolEntry {
   #[serde(default)]
   pub is_definition: bool,
+  #[serde(flatten)]
+  pub entry: ReferenceEntry,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ReferenceEntry {
+  // is_write_access: bool,
   // is_in_string: Option<bool>,
   #[serde(flatten)]
   pub document_span: DocumentSpan,
@@ -3178,8 +3220,11 @@ pub enum RequestMethod {
   GetOutliningSpans(ModuleSpecifier),
   /// Return quick info at position (hover information).
   GetQuickInfo((ModuleSpecifier, u32)),
-  /// Get document references for a specific position.
-  GetReferences((ModuleSpecifier, u32)),
+  /// Finds the document references for a specific position.
+  FindReferences {
+    specifier: ModuleSpecifier,
+    position: u32,
+  },
   /// Get signature help items for a specific position.
   GetSignatureHelpItems((ModuleSpecifier, u32, SignatureHelpItemsOptions)),
   /// Get a selection range for a specific position.
@@ -3349,9 +3394,12 @@ impl RequestMethod {
         "specifier": state.denormalize_specifier(specifier),
         "position": position,
       }),
-      RequestMethod::GetReferences((specifier, position)) => json!({
+      RequestMethod::FindReferences {
+        specifier,
+        position,
+      } => json!({
         "id": id,
-        "method": "getReferences",
+        "method": "findReferences",
         "specifier": state.denormalize_specifier(specifier),
         "position": position,
       }),