Improve diff rendering for notebooks (#20036)

## Summary

As noted in a code TODO, our `Diff` rendering code previously didn't
have any
special handling for notebooks. This was particularly obvious when the
diffs
were rendered right next to the corresponding diagnostic because the
diagnostic
used cell-based line numbers, while the diff was still using line
numbers from
the concatenated source. This PR updates the diff rendering to handle
notebooks
too.

The main improvements shown in the example below are:

- Line numbers are now remapped to be relative to their cell
- Context lines from other cells are suppressed

```
error[unused-import][*]: `math` imported but unused                             
 --> notebook.ipynb:cell 2:2:8                                                  
  |                                                                             
1 | # cell 2                                                                    
2 | import math                                                                 
  |        ^^^^                                                                 
3 |                                                                             
4 | print('hello world')                                                        
  |                                                                             
help: Remove unused import: `math`                                              
                                                                                
ℹ Safe fix                                                                      
1 1 | # cell 2                                                                  
2   |-import math                                                               
3 2 |                                                                           
4 3 | print('hello world')                                                      
```

I tried a few different approaches here before finally just splitting
the notebook into separate text ranges by cell and diffing each one
separately. It seems to work and passes all of our tests, but I don't
know if it's actually enforced anywhere that a single edit doesn't span
cells. Such an edit would silently be dropped right now since it would
fail the `contains_range` check. I also feel like I may have overlooked
an existing way to partition a file into cells like this.

## Test Plan

Existing notebook tests, plus a new one in `ruff_db`

Author: ntBre

crates/ruff_db/src/diagnostic/mod.rs

@@ -325,6 +325,11 @@ impl Diagnostic {
         self.inner.fix.as_ref()
     }
 
+    #[cfg(test)]
+    pub(crate) fn fix_mut(&mut self) -> Option<&mut Fix> {
+        Arc::make_mut(&mut self.inner).fix.as_mut()
+    }
+
     /// Set the fix for this diagnostic.
     pub fn set_fix(&mut self, fix: Fix) {
         debug_assert!(

crates/ruff_db/src/diagnostic/render.rs

@@ -2622,6 +2622,13 @@ watermelon
             self.config = config;
         }
 
+        /// Show a diff for the fix when rendering.
+        pub(super) fn show_fix_diff(&mut self, yes: bool) {
+            let mut config = std::mem::take(&mut self.config);
+            config = config.show_fix_diff(yes);
+            self.config = config;
+        }
+
         /// The lowest fix applicability to show when rendering.
         pub(super) fn fix_applicability(&mut self, applicability: Applicability) {
             let mut config = std::mem::take(&mut self.config);

crates/ruff_db/src/diagnostic/render/full.rs

@@ -2,12 +2,13 @@ use std::borrow::Cow;
 use std::num::NonZeroUsize;
 
 use anstyle::Style;
+use ruff_notebook::NotebookIndex;
 use similar::{ChangeTag, TextDiff};
 
 use ruff_annotate_snippets::Renderer as AnnotateRenderer;
 use ruff_diagnostics::{Applicability, Fix};
 use ruff_source_file::OneIndexed;
-use ruff_text_size::{Ranged, TextRange, TextSize};
+use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::diagnostic::render::{FileResolver, Resolved};
 use crate::diagnostic::stylesheet::{DiagnosticStylesheet, fmt_styled};
@@ -81,6 +82,7 @@ impl<'a> FullRenderer<'a> {
 struct Diff<'a> {
     fix: &'a Fix,
     diagnostic_source: DiagnosticSource,
+    notebook_index: Option<NotebookIndex>,
     stylesheet: &'a DiagnosticStylesheet,
 }
 
@@ -90,12 +92,11 @@ impl<'a> Diff<'a> {
         stylesheet: &'a DiagnosticStylesheet,
         resolver: &'a dyn FileResolver,
     ) -> Option<Diff<'a>> {
+        let file = &diagnostic.primary_span_ref()?.file;
         Some(Diff {
             fix: diagnostic.fix()?,
-            diagnostic_source: diagnostic
-                .primary_span_ref()?
-                .file
-                .diagnostic_source(resolver),
+            diagnostic_source: file.diagnostic_source(resolver),
+            notebook_index: resolver.notebook_index(file),
             stylesheet,
         })
     }
@@ -106,19 +107,24 @@ impl std::fmt::Display for Diff<'_> {
         let source_code = self.diagnostic_source.as_source_code();
         let source_text = source_code.text();
 
-        // TODO(dhruvmanila): Add support for Notebook cells once it's user-facing
-        let mut output = String::with_capacity(source_text.len());
-        let mut last_end = TextSize::default();
-
-        for edit in self.fix.edits() {
-            output.push_str(source_code.slice(TextRange::new(last_end, edit.start())));
-            output.push_str(edit.content().unwrap_or_default());
-            last_end = edit.end();
-        }
-
-        output.push_str(&source_text[usize::from(last_end)..]);
-
-        let diff = TextDiff::from_lines(source_text, &output);
+        // Partition the source code into end offsets for each cell. If `self.notebook_index` is
+        // `None`, indicating a regular script file, all the lines will be in one "cell" under the
+        // `None` key.
+        let cells = if let Some(notebook_index) = &self.notebook_index {
+            let mut last_cell = OneIndexed::MIN;
+            let mut cells: Vec<(Option<OneIndexed>, TextSize)> = Vec::new();
+            for (row, cell) in notebook_index.iter() {
+                if cell != last_cell {
+                    let offset = source_code.line_start(row);
+                    cells.push((Some(last_cell), offset));
+                    last_cell = cell;
+                }
+            }
+            cells.push((Some(last_cell), source_text.text_len()));
+            cells
+        } else {
+            vec![(None, source_text.text_len())]
+        };
 
         let message = match self.fix.applicability() {
             // TODO(zanieb): Adjust this messaging once it's user-facing
@@ -133,59 +139,97 @@ impl std::fmt::Display for Diff<'_> {
         // tests, which is the only place these are currently used.
         writeln!(f, "ℹ {}", fmt_styled(message, self.stylesheet.separator))?;
 
-        let (largest_old, largest_new) = diff
-            .ops()
-            .last()
-            .map(|op| (op.old_range().start, op.new_range().start))
-            .unwrap_or_default();
+        let mut last_end = TextSize::ZERO;
+        for (cell, offset) in cells {
+            let range = TextRange::new(last_end, offset);
+            last_end = offset;
+            let input = source_code.slice(range);
+
+            let mut output = String::with_capacity(input.len());
+            let mut last_end = range.start();
+
+            let mut applied = 0;
+            for edit in self.fix.edits() {
+                if range.contains_range(edit.range()) {
+                    output.push_str(source_code.slice(TextRange::new(last_end, edit.start())));
+                    output.push_str(edit.content().unwrap_or_default());
+                    last_end = edit.end();
+                    applied += 1;
+                }
+            }
 
-        let digit_with = OneIndexed::from_zero_indexed(largest_new.max(largest_old)).digits();
+            // No edits were applied, so there's no need to diff.
+            if applied == 0 {
+                continue;
+            }
+
+            output.push_str(&source_text[usize::from(last_end)..usize::from(range.end())]);
+
+            let diff = TextDiff::from_lines(input, &output);
+
+            let (largest_old, largest_new) = diff
+                .ops()
+                .last()
+                .map(|op| (op.old_range().start, op.new_range().start))
+                .unwrap_or_default();
+
+            let digit_with = OneIndexed::from_zero_indexed(largest_new.max(largest_old)).digits();
 
-        for (idx, group) in diff.grouped_ops(3).iter().enumerate() {
-            if idx > 0 {
-                writeln!(f, "{:-^1$}", "-", 80)?;
+            if let Some(cell) = cell {
+                // Room for 2 digits, 2 x 1 space before each digit, 1 space, and 1 `|`. This
+                // centers the three colons on the pipe.
+                writeln!(f, "{:>1$} cell {cell}", ":::", 2 * digit_with.get() + 4)?;
             }
-            for op in group {
-                for change in diff.iter_inline_changes(op) {
-                    let sign = match change.tag() {
-                        ChangeTag::Delete => "-",
-                        ChangeTag::Insert => "+",
-                        ChangeTag::Equal => " ",
-                    };
-
-                    let line_style = LineStyle::from(change.tag(), self.stylesheet);
-
-                    let old_index = change.old_index().map(OneIndexed::from_zero_indexed);
-                    let new_index = change.new_index().map(OneIndexed::from_zero_indexed);
-
-                    write!(
-                        f,
-                        "{} {} |{}",
-                        Line {
-                            index: old_index,
-                            width: digit_with
-                        },
-                        Line {
-                            index: new_index,
-                            width: digit_with
-                        },
-                        fmt_styled(line_style.apply_to(sign), self.stylesheet.emphasis),
-                    )?;
-
-                    for (emphasized, value) in change.iter_strings_lossy() {
-                        let value = show_nonprinting(&value);
-                        if emphasized {
-                            write!(
-                                f,
-                                "{}",
-                                fmt_styled(line_style.apply_to(&value), self.stylesheet.underline)
-                            )?;
-                        } else {
-                            write!(f, "{}", line_style.apply_to(&value))?;
+
+            for (idx, group) in diff.grouped_ops(3).iter().enumerate() {
+                if idx > 0 {
+                    writeln!(f, "{:-^1$}", "-", 80)?;
+                }
+                for op in group {
+                    for change in diff.iter_inline_changes(op) {
+                        let sign = match change.tag() {
+                            ChangeTag::Delete => "-",
+                            ChangeTag::Insert => "+",
+                            ChangeTag::Equal => " ",
+                        };
+
+                        let line_style = LineStyle::from(change.tag(), self.stylesheet);
+
+                        let old_index = change.old_index().map(OneIndexed::from_zero_indexed);
+                        let new_index = change.new_index().map(OneIndexed::from_zero_indexed);
+
+                        write!(
+                            f,
+                            "{} {} |{}",
+                            Line {
+                                index: old_index,
+                                width: digit_with,
+                            },
+                            Line {
+                                index: new_index,
+                                width: digit_with,
+                            },
+                            fmt_styled(line_style.apply_to(sign), self.stylesheet.emphasis),
+                        )?;
+
+                        for (emphasized, value) in change.iter_strings_lossy() {
+                            let value = show_nonprinting(&value);
+                            if emphasized {
+                                write!(
+                                    f,
+                                    "{}",
+                                    fmt_styled(
+                                        line_style.apply_to(&value),
+                                        self.stylesheet.underline
+                                    )
+                                )?;
+                            } else {
+                                write!(f, "{}", line_style.apply_to(&value))?;
+                            }
+                        }
+                        if change.missing_newline() {
+                            writeln!(f)?;
                         }
-                    }
-                    if change.missing_newline() {
-                        writeln!(f)?;
                     }
                 }
             }
@@ -253,7 +297,7 @@ fn show_nonprinting(s: &str) -> Cow<'_, str> {
 
 #[cfg(test)]
 mod tests {
-    use ruff_diagnostics::Applicability;
+    use ruff_diagnostics::{Applicability, Fix};
     use ruff_text_size::{TextLen, TextRange, TextSize};
 
     use crate::diagnostic::{
@@ -654,6 +698,107 @@ print()
         ");
     }
 
+    /// Test that we remap notebook cell line numbers in the diff as well as the main diagnostic.
+    #[test]
+    fn notebook_output_with_diff() {
+        let (mut env, diagnostics) = create_notebook_diagnostics(DiagnosticFormat::Full);
+        env.show_fix_diff(true);
+        insta::assert_snapshot!(env.render_diagnostics(&diagnostics), @r"
+        error[unused-import][*]: `os` imported but unused
+         --> notebook.ipynb:cell 1:2:8
+          |
+        1 | # cell 1
+        2 | import os
+          |        ^^
+          |
+        help: Remove unused import: `os`
+
+        ℹ Safe fix
+           ::: cell 1
+        1 1 | # cell 1
+        2   |-import os
+
+        error[unused-import][*]: `math` imported but unused
+         --> notebook.ipynb:cell 2:2:8
+          |
+        1 | # cell 2
+        2 | import math
+          |        ^^^^
+        3 |
+        4 | print('hello world')
+          |
+        help: Remove unused import: `math`
+
+        ℹ Safe fix
+           ::: cell 2
+        1 1 | # cell 2
+        2   |-import math
+        3 2 | 
+        4 3 | print('hello world')
+
+        error[unused-variable]: Local variable `x` is assigned to but never used
+         --> notebook.ipynb:cell 3:4:5
+          |
+        2 | def foo():
+        3 |     print()
+        4 |     x = 1
+          |     ^
+          |
+        help: Remove assignment to unused variable `x`
+
+        ℹ Unsafe fix
+           ::: cell 3
+        1 1 | # cell 3
+        2 2 | def foo():
+        3 3 |     print()
+        4   |-    x = 1
+        5 4 |
+        ");
+    }
+
+    #[test]
+    fn notebook_output_with_diff_spanning_cells() {
+        let (mut env, mut diagnostics) = create_notebook_diagnostics(DiagnosticFormat::Full);
+        env.show_fix_diff(true);
+
+        // Move all of the edits from the later diagnostics to the first diagnostic to simulate a
+        // single diagnostic with edits in different cells.
+        let mut diagnostic = diagnostics.swap_remove(0);
+        let fix = diagnostic.fix_mut().unwrap();
+        let mut edits = fix.edits().to_vec();
+        for diag in diagnostics {
+            edits.extend_from_slice(diag.fix().unwrap().edits());
+        }
+        *fix = Fix::unsafe_edits(edits.remove(0), edits);
+
+        insta::assert_snapshot!(env.render(&diagnostic), @r"
+        error[unused-import]: `os` imported but unused
+         --> notebook.ipynb:cell 1:2:8
+          |
+        1 | # cell 1
+        2 | import os
+          |        ^^
+          |
+        help: Remove unused import: `os`
+
+        ℹ Unsafe fix
+           ::: cell 1
+        1 1 | # cell 1
+        2   |-import os
+           ::: cell 2
+        1 1 | # cell 2
+        2   |-import math
+        3 2 | 
+        4 3 | print('hello world')
+           ::: cell 3
+        1 1 | # cell 3
+        2 2 | def foo():
+        3 3 |     print()
+        4   |-    x = 1
+        5 4 |
+        ");
+    }
+
     /// Carriage return (`\r`) is a valid line-ending in Python, so we should normalize this to a
     /// line feed (`\n`) for rendering. Otherwise we report a single long line for this case.
     #[test]