[red-knot] Show related information in diagnostic

Author: MichaReiser

crates/ruff_db/src/diagnostic/mod.rs

@@ -249,6 +249,25 @@ impl Diagnostic {
             diagnostic: self,
         }
     }
+
+    /// Returns all annotations, skipping the first primary annotation.
+    pub fn secondary_annotations(&self) -> impl Iterator<Item = &Annotation> {
+        let mut seen_primary = false;
+        self.inner.annotations.iter().filter(move |ann| {
+            if seen_primary {
+                true
+            } else if ann.is_primary {
+                seen_primary = true;
+                false
+            } else {
+                true
+            }
+        })
+    }
+
+    pub fn sub_diagnostics(&self) -> &[SubDiagnostic] {
+        &self.inner.subs
+    }
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -371,6 +390,57 @@ impl SubDiagnostic {
     pub fn annotate(&mut self, ann: Annotation) {
         self.inner.annotations.push(ann);
     }
+
+    pub fn annotations(&self) -> &[Annotation] {
+        &self.inner.annotations
+    }
+
+    /// Returns a shared borrow of the "primary" annotation of this diagnostic
+    /// if one exists.
+    ///
+    /// When there are multiple primary annotations, then the first one that
+    /// was added to this diagnostic is returned.
+    pub fn primary_annotation(&self) -> Option<&Annotation> {
+        self.inner.annotations.iter().find(|ann| ann.is_primary)
+    }
+
+    /// Introspects this diagnostic and returns what kind of "primary" message
+    /// it contains for concise formatting.
+    ///
+    /// When we concisely format diagnostics, we likely want to not only
+    /// include the primary diagnostic message but also the message attached
+    /// to the primary annotation. In particular, the primary annotation often
+    /// contains *essential* information or context for understanding the
+    /// diagnostic.
+    ///
+    /// The reason why we don't just always return both the main diagnostic
+    /// message and the primary annotation message is because this was written
+    /// in the midst of an incremental migration of ty over to the new
+    /// diagnostic data model. At time of writing, diagnostics were still
+    /// constructed in the old model where the main diagnostic message and the
+    /// primary annotation message were not distinguished from each other. So
+    /// for now, we carefully return what kind of messages this diagnostic
+    /// contains. In effect, if this diagnostic has a non-empty main message
+    /// *and* a non-empty primary annotation message, then the diagnostic is
+    /// 100% using the new diagnostic data model and we can format things
+    /// appropriately.
+    ///
+    /// The type returned implements the `std::fmt::Display` trait. In most
+    /// cases, just converting it to a string (or printing it) will do what
+    /// you want.
+    pub fn concise_message(&self) -> ConciseMessage {
+        let main = self.inner.message.as_str();
+        let annotation = self
+            .primary_annotation()
+            .and_then(|ann| ann.get_message())
+            .unwrap_or_default();
+        match (main.is_empty(), annotation.is_empty()) {
+            (false, true) => ConciseMessage::MainDiagnostic(main),
+            (true, false) => ConciseMessage::PrimaryAnnotation(annotation),
+            (false, false) => ConciseMessage::Both { main, annotation },
+            (true, true) => ConciseMessage::Empty,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]

crates/ruff_db/src/files.rs

@@ -11,6 +11,7 @@ use ruff_text_size::{Ranged, TextRange};
 use salsa::plumbing::AsId;
 use salsa::{Durability, Setter};
 
+use crate::diagnostic::{Span, UnifiedFile};
 use crate::file_revision::FileRevision;
 use crate::files::file_root::FileRoots;
 use crate::files::private::FileStatus;
@@ -549,6 +550,29 @@ impl Ranged for FileRange {
     }
 }
 
+impl TryFrom<&Span> for FileRange {
+    type Error = ();
+
+    fn try_from(value: &Span) -> Result<Self, Self::Error> {
+        let UnifiedFile::Ty(file) = value.file() else {
+            return Err(());
+        };
+
+        Ok(Self {
+            file: *file,
+            range: value.range().ok_or(())?,
+        })
+    }
+}
+
+impl TryFrom<Span> for FileRange {
+    type Error = ();
+
+    fn try_from(value: Span) -> Result<Self, Self::Error> {
+        Self::try_from(&value)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::file_revision::FileRevision;

crates/ty_server/src/server/api/requests/diagnostic.rs

@@ -7,11 +7,13 @@ use lsp_types::{
     NumberOrString, Range, RelatedFullDocumentDiagnosticReport,
 };
 
-use crate::document::ToRangeExt;
+use crate::PositionEncoding;
+use crate::document::{FileRangeExt, ToRangeExt};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
 use crate::server::{Result, client::Notifier};
 use crate::session::DocumentSnapshot;
-use ruff_db::diagnostic::Severity;
+use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
+use ruff_db::files::FileRange;
 use ruff_db::source::{line_index, source_text};
 use ty_project::{Db, ProjectDatabase};
 
@@ -116,6 +118,29 @@ fn to_lsp_diagnostic(
         })
         .flatten();
 
+    let mut related_information = Vec::new();
+
+    related_information.extend(
+        diagnostic
+            .secondary_annotations()
+            .chain(
+                diagnostic
+                    .sub_diagnostics()
+                    .iter()
+                    .flat_map(SubDiagnostic::annotations),
+            )
+            .filter_map(|annotation| annotation_to_related_information(db, annotation, encoding)),
+    );
+
+    related_information.extend(
+        diagnostic
+            .sub_diagnostics()
+            .iter()
+            .filter_map(|diagnostic| {
+                sub_diagnostic_to_related_information(db, diagnostic, encoding)
+            }),
+    );
+
     Diagnostic {
         range,
         severity: Some(severity),
@@ -124,7 +149,41 @@ fn to_lsp_diagnostic(
         code_description,
         source: Some("ty".into()),
         message: diagnostic.concise_message().to_string(),
-        related_information: None,
+        related_information: Some(related_information),
         data: None,
     }
 }
+
+fn annotation_to_related_information(
+    db: &dyn Db,
+    annotation: &Annotation,
+    encoding: PositionEncoding,
+) -> Option<lsp_types::DiagnosticRelatedInformation> {
+    let span = annotation.get_span();
+
+    let annotation_message = annotation.get_message()?;
+    let range = FileRange::try_from(span).ok()?;
+    let location = range.to_location(db.upcast(), encoding)?;
+
+    Some(lsp_types::DiagnosticRelatedInformation {
+        location,
+        message: annotation_message.to_string(),
+    })
+}
+
+fn sub_diagnostic_to_related_information(
+    db: &dyn Db,
+    diagnostic: &SubDiagnostic,
+    encoding: PositionEncoding,
+) -> Option<lsp_types::DiagnosticRelatedInformation> {
+    let primary_annotation = diagnostic.primary_annotation()?;
+
+    let span = primary_annotation.get_span();
+    let range = FileRange::try_from(span).ok()?;
+    let location = range.to_location(db.upcast(), encoding)?;
+
+    Some(lsp_types::DiagnosticRelatedInformation {
+        location,
+        message: diagnostic.concise_message().to_string(),
+    })
+}

crates/ty_server/src/session/capabilities.rs

@@ -38,8 +38,7 @@ impl ResolvedClientCapabilities {
         let document_changes = client_capabilities
             .workspace
             .as_ref()
-            .and_then(|workspace| workspace.workspace_edit.as_ref())
-            .and_then(|workspace_edit| workspace_edit.document_changes)
+            .and_then(|workspace| workspace.workspace_edit.as_ref()?.document_changes)
             .unwrap_or_default();
 
         let declaration_link_support = client_capabilities