fix: preserve diff semantic highlighting across pane focus

src/editor.rs

@@ -1477,6 +1477,27 @@ impl Editor {
             .map(|_| focus)
     }
 
+    pub(crate) fn diff_semantic_tokens(&self, focus: Focus) -> &[SemanticToken] {
+        let Some(diff) = self.diff.as_ref() else {
+            return &[];
+        };
+        let buffer = if focus == Focus::After && diff.after_editable() {
+            diff.live_buffer
+                .and_then(|id| self.buffer(id))
+                .filter(|buffer| Some(buffer.text_revision) == diff.live_revision)
+        } else {
+            let Some(comparison) = diff.comparison_id else {
+                return &[];
+            };
+            std::iter::once(&self.active)
+                .chain(self.parked.values())
+                .find(|buffer| buffer.diff_snapshot == Some((comparison, diff.file, focus)))
+        };
+        buffer
+            .and_then(|buffer| buffer.semantic_tokens.as_deref())
+            .unwrap_or_default()
+    }
+
     pub(crate) fn mutation_blocked(&self) -> bool {
         self.active.diff_snapshot.is_some() || (self.diff.is_some() && !self.diff_editing())
     }

src/tui.rs

@@ -613,6 +613,7 @@ fn draw(
                 1,
             );
             paint_syntax_line(frame, line, syntax.line(row), area, *left);
+            paint_semantic_line(frame, line, editor.semantic_tokens(), row, area, *left);
             paint_editor_line(frame, editor, row, area, *left, &carets, &selections);
         }
     }
@@ -840,19 +841,14 @@ fn paint_grapheme_range(
     }
 }
 
-fn paint_editor_line(
+fn paint_semantic_line(
     frame: &mut Frame,
-    editor: &Editor,
+    line: &str,
+    semantic: &[SemanticToken],
     row: usize,
     area: Rect,
     left: usize,
-    carets: &[Cursor],
-    selections: &[Selection],
 ) {
-    let Some(line) = editor.lines().get(row) else {
-        return;
-    };
-    let semantic = editor.semantic_tokens();
     let mut token = semantic.partition_point(|token| token.start.row < row);
     let mut token_style = None;
     let mut column = 0;
@@ -884,6 +880,20 @@ fn paint_editor_line(
         }
         column = end;
     }
+}
+
+fn paint_editor_line(
+    frame: &mut Frame,
+    editor: &Editor,
+    row: usize,
+    area: Rect,
+    left: usize,
+    carets: &[Cursor],
+    selections: &[Selection],
+) {
+    let Some(line) = editor.lines().get(row) else {
+        return;
+    };
     for severity in [
         DiagnosticSeverity::Hint,
         DiagnosticSeverity::Information,
@@ -1439,6 +1449,7 @@ fn draw_diff(
                     before_area,
                     true,
                     &syntax.before,
+                    editor.map_or(&[], |editor| editor.diff_semantic_tokens(Focus::Before)),
                     (
                         top,
                         &mut before_left,
@@ -1468,6 +1479,7 @@ fn draw_diff(
                     after_area,
                     false,
                     &syntax.after,
+                    editor.map_or(&[], |editor| editor.diff_semantic_tokens(Focus::After)),
                     (
                         top,
                         &mut after_left,
@@ -1547,6 +1559,7 @@ fn diff_status(status: FileStatus) -> &'static str {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 fn draw_diff_pane(
     frame: &mut Frame,
     diff: &DiffState,
@@ -1554,6 +1567,7 @@ fn draw_diff_pane(
     area: Rect,
     before: bool,
     syntax: &SyntaxCache,
+    semantic: &[SemanticToken],
     viewport: (usize, &mut usize, Option<u16>),
 ) -> Rect {
     if area.width == 0 || area.height == 0 {
@@ -1699,6 +1713,14 @@ fn draw_diff_pane(
             Rect::new(content.x + gutter as u16, row_area.y, width as u16, 1),
             left,
         );
+        paint_semantic_line(
+            frame,
+            source_text,
+            semantic,
+            index,
+            Rect::new(content.x + gutter as u16, row_area.y, width as u16, 1),
+            left,
+        );
     }
     Rect::new(
         content.x + gutter as u16,
@@ -3925,6 +3947,111 @@ mod tests {
         }
     }
 
+    #[test]
+    fn diff_semantic_colors_stay_buffer_local_across_focus_changes() {
+        for writable in [false, true] {
+            let directory = Directory::new();
+            let path = directory.file("live.rs", "let value = 2;");
+            let mut editor = Editor::open(Some(path), vec![]).unwrap();
+            editor.execute("diff");
+            let (id, _) = editor.take_diff_request().unwrap();
+            editor.apply_diff_result(id, Ok(crate::diff::Outcome::Catalog(diff_catalog())));
+            editor.handle(Input::Down);
+            editor.handle(Input::Enter);
+            editor.handle(Input::Enter);
+            let (id, _) = editor.take_diff_request().unwrap();
+            editor.apply_diff_result(
+                id,
+                Ok(crate::diff::Outcome::Comparison(crate::diff::Comparison {
+                    root: directory.0.clone(),
+                    editable_paths: if writable {
+                        [std::path::PathBuf::from("live.rs")].into()
+                    } else {
+                        Default::default()
+                    },
+                    before: "baseline".into(),
+                    after: "Working checkout".into(),
+                    files: vec![syntax_diff_file(
+                        "live.rs",
+                        "let value = 1;",
+                        "let value = 2;",
+                    )],
+                })),
+            );
+            editor.handle(Input::Char('\t'));
+            assert_eq!(editor.diff_buffer_focus(), Some(Focus::Before));
+            editor
+                .replace_semantic_tokens(
+                    editor.buffer_id(),
+                    editor.text_revision(),
+                    Some(Arc::from([semantic_token(
+                        0,
+                        4,
+                        9,
+                        SemanticTokenKind::Parameter,
+                        0,
+                    )])),
+                )
+                .unwrap();
+            editor.handle(Input::Char('\t'));
+            editor
+                .replace_semantic_tokens(
+                    editor.buffer_id(),
+                    editor.text_revision(),
+                    Some(Arc::from([semantic_token(
+                        0,
+                        4,
+                        9,
+                        SemanticTokenKind::Function,
+                        SemanticTokenModifier::Declaration as u16,
+                    )])),
+                )
+                .unwrap();
+            let mut terminal = Terminal::new(ratatui::backend::TestBackend::new(120, 14)).unwrap();
+            let mut views = BufferViewports::default();
+            for _ in 0..4 {
+                terminal.draw(|frame| views.draw(frame, &editor)).unwrap();
+                let regions = views.panels.diff.as_ref().unwrap().regions;
+                for (index, color, bold) in
+                    [(1, Color::LightCyan, false), (2, Color::LightBlue, true)]
+                {
+                    let pane = regions[index].0;
+                    let cell = &terminal.backend().buffer()[(pane.x + 8, pane.y + 1)];
+                    assert_eq!(cell.symbol(), "v");
+                    assert_eq!(cell.fg, color);
+                    assert_eq!(cell.modifier.contains(ratatui::style::Modifier::BOLD), bold);
+                    assert_eq!(
+                        cell.bg,
+                        if index == 1 {
+                            Color::Rgb(64, 24, 24)
+                        } else {
+                            Color::Rgb(24, 48, 32)
+                        }
+                    );
+                }
+                editor.handle(Input::Char('\t'));
+            }
+            editor
+                .replace_semantic_tokens(editor.buffer_id(), editor.text_revision(), None)
+                .unwrap();
+            terminal.draw(|frame| views.draw(frame, &editor)).unwrap();
+            let regions = views.panels.diff.as_ref().unwrap().regions;
+            let before = regions[1].0;
+            let after = regions[2].0;
+            assert_eq!(
+                terminal.backend().buffer()[(before.x + 8, before.y + 1)].fg,
+                Color::LightCyan
+            );
+            let value = &terminal.backend().buffer()[(after.x + 8, after.y + 1)];
+            assert_eq!(value.fg, Color::Reset);
+            assert!(!value.modifier.contains(ratatui::style::Modifier::BOLD));
+            assert_eq!(
+                terminal.backend().buffer()[(after.x + 4, after.y + 1)].fg,
+                Color::Magenta
+            );
+        }
+    }
+
     #[test]
     fn editable_after_renders_live_syntax_caret_visual_selection_and_command_prompt() {
         let directory = Directory::new();