fix: diff search hit highlight across ranges

src/screen/dashboard/pull_request_diff_view.rs

@@ -14,7 +14,7 @@ use crate::{
         text_input::{self, TextInput, text_input},
     },
     query::{self, QueryStatus, read_query, use_query, watch_query},
-    util::{self, diff::DiffLineIndex},
+    util::{self, diff::DiffLineIndex, syntax_highlight},
 };
 
 pub(crate) struct PullRequestDiffView {
@@ -42,6 +42,7 @@ struct DiffSearchResult {
     cursor: DiffSearchResultCursor,
     old_side: Vec<DiffSearchHit>,
     new_side: Vec<DiffSearchHit>,
+    prev_diff_line_highlights: (usize, Vec<syntax_highlight::HighlightedRange>),
 }
 
 struct DiffSearchResultCursor {
@@ -308,6 +309,7 @@ impl PullRequestDiffView {
                 cursor,
                 old_side: old_search_result,
                 new_side: new_search_result,
+                prev_diff_line_highlights: (0, Vec::new()),
             });
         }
 
@@ -387,23 +389,31 @@ impl PullRequestDiffView {
                 if next_highlight_line == diff_line_i
                     && matches!(search_result.cursor.side, util::diff::DiffSide::Old)
                 {
-                    // there is a hit on new side on same diff line as the next hit on old side
+                    // next search hit candidates found on next diff line on both old side and new side
-                    // go to old side first
+                    // since we are on new side currently, we should go back to old side first.
                     search_result.cursor.diff_line_i = diff_line_i;
                     search_result.cursor.index += 1;
+                    let (range, style) = &next.unwrap().highlight_range;
+                    Some((
+                        range.clone(),
+                        gpui::HighlightStyle {
+                            background_color: Some(gpui::red()),
+                            ..style.clone()
+                        },
+                    ))
                 } else {
                     // next cursor should be on other side with the found index
                     search_result.cursor.side = search_result.cursor.side.flipped();
                     search_result.cursor.diff_line_i = diff_line_i;
                     search_result.cursor.index = other_side_i;
+                    Some((
+                        r.clone(),
+                        gpui::HighlightStyle {
+                            background_color: Some(gpui::red()),
+                            ..symbol_highlight_style(theme)
+                        },
+                    ))
                 }
-                Some((
-                    r.clone(),
-                    gpui::HighlightStyle {
-                        background_color: Some(gpui::red()),
-                        ..symbol_highlight_style(theme)
-                    },
-                ))
             } else if let Some(next) = next {
                 // stay on old side, point to next old side highlight
                 search_result.cursor.index = search_result.cursor.index + 1;
@@ -426,8 +436,14 @@ impl PullRequestDiffView {
                 | util::diff::DiffSide::Old => diff_view_state.old_side_highlights_mut(),
                 | util::diff::DiffSide::New => diff_view_state.new_side_highlights_mut(),
                 };
-                if let Some(mut content) = prev_highlighted_content {
+                if let Some(mut content) = prev_highlighted_content
-                    content.replace_highlight_range(&current_search_hit.highlight_range);
+                    && !search_result.prev_diff_line_highlights.1.is_empty()
+                {
+                    // restore highlights for the line where the current higlighted search hit is in
+                    content.set_highlights_for_line(
+                        search_result.prev_diff_line_highlights.0,
+                        search_result.prev_diff_line_highlights.1.drain(..),
+                    );
                 }
             }
 
@@ -435,8 +451,16 @@ impl PullRequestDiffView {
             | util::diff::DiffSide::Old => diff_view_state.old_side_highlights_mut(),
             | util::diff::DiffSide::New => diff_view_state.new_side_highlights_mut(),
             };
-            if let Some(mut content) = content {
+
-                content.replace_highlight_range(&highlight_range);
+            if let Some(mut content) = content
+                && let Some(line_i) = content.line_index_of_range(&highlight_range.0)
+            {
+                search_result.prev_diff_line_highlights.0 = line_i;
+                search_result
+                    .prev_diff_line_highlights
+                    .1
+                    .splice(.., content.highlights_at_line(line_i).iter().cloned());
+                content.add_highlight_range(highlight_range);
             }
         }
 

src/util/diff.rs

@@ -1,4 +1,4 @@
-use std::{ops::Deref, sync::Arc};
+use std::sync::Arc;
 
 use similar::DiffableStr;
 

src/util/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) mod diff;
 pub(crate) mod file;
+pub(crate) mod range;
 pub(crate) mod str;
 pub(crate) mod syntax_highlight;
 pub(crate) mod timeout;

src/util/range.rs

@@ -0,0 +1,17 @@
+use std::ops::Sub;
+
+pub(crate) trait RangeExt<T> {
+    fn relative_to(&self, other: &std::ops::Range<T>) -> std::ops::Range<T::Output>
+    where
+        T: Sub + Ord + PartialOrd + Clone + Copy;
+}
+
+impl<T> RangeExt<T> for std::ops::Range<T>
+where
+    T: Sub + Ord + PartialOrd + Clone + Copy,
+{
+    fn relative_to(&self, other: &std::ops::Range<T>) -> std::ops::Range<T::Output> {
+        debug_assert!(self.start >= other.start);
+        (self.start - other.start)..(self.end - other.start)
+    }
+}

src/util/syntax_highlight.rs

@@ -1,6 +1,10 @@
 use std::sync::LazyLock;
 
-use crate::{component::code_view::symbol_highlight_style, theme, util};
+use crate::{
+    component::{code_view::symbol_highlight_style, file_tree::Item},
+    theme,
+    util::{self, range::RangeExt, syntax_highlight},
+};
 
 pub(crate) type HighlightedRange = (std::ops::Range<usize>, gpui::HighlightStyle);
 
@@ -13,6 +17,68 @@ pub(crate) struct HighlightedLine {
     highlights: Vec<(std::ops::Range<usize>, gpui::HighlightStyle)>,
 }
 
+fn highlight(a: gpui::HighlightStyle, b: gpui::HighlightStyle) -> gpui::HighlightStyle {
+    gpui::HighlightStyle {
+        color: b.color.or(a.color),
+        font_weight: b.font_weight.or(a.font_weight),
+        font_style: b.font_style.or(a.font_style),
+        background_color: b.background_color.or(a.background_color),
+        underline: b.underline.or(a.underline),
+        strikethrough: b.strikethrough.or(a.strikethrough),
+        fade_out: b.fade_out.or(a.fade_out),
+    }
+}
+
+fn combine_highlights(
+    a: impl IntoIterator<Item = HighlightedRange>,
+    b: impl IntoIterator<Item = HighlightedRange>,
+) -> impl Iterator<Item = HighlightedRange> {
+    let mut endpoints = Vec::new();
+    let mut highlights = Vec::new();
+
+    for (range, highlight) in a.into_iter().chain(b) {
+        if !range.is_empty() {
+            let highlight_id = highlights.len();
+            endpoints.push((range.start, highlight_id, true));
+            endpoints.push((range.end, highlight_id, false));
+            highlights.push(highlight);
+        }
+    }
+
+    endpoints.sort_unstable_by_key(|(position, _, _)| *position);
+
+    let mut active = vec![false; highlights.len()];
+    let mut highlighted_ranges = Vec::new();
+    let mut previous_position = endpoints
+        .first()
+        .map(|(position, _, _)| *position)
+        .unwrap_or(0);
+    let mut endpoint_i = 0;
+
+    while endpoint_i < endpoints.len() {
+        let position = endpoints[endpoint_i].0;
+        if position > previous_position && active.iter().any(|active| *active) {
+            let style = highlights
+                .iter()
+                .enumerate()
+                .filter_map(|(highlight_id, style)| active[highlight_id].then_some(*style))
+                .fold(gpui::HighlightStyle::default(), highlight);
+
+            highlighted_ranges.push((previous_position..position, style));
+        }
+
+        while endpoint_i < endpoints.len() && endpoints[endpoint_i].0 == position {
+            let (_, highlight_id, is_start) = endpoints[endpoint_i];
+            active[highlight_id] = is_start;
+            endpoint_i += 1;
+        }
+
+        previous_position = position;
+    }
+
+    highlighted_ranges.into_iter()
+}
+
 static TS_HIGHLIGHTS: LazyLock<String> = LazyLock::new(|| {
     [
         tree_sitter_javascript::HIGHLIGHT_QUERY,
@@ -209,7 +275,7 @@ impl HighlightedContent {
             {
                 // moving on to new line, flush new highlights and combine into the line highlights
                 if !new_highlights.is_empty() {
-                    self.0[current_line].highlights = gpui::combine_highlights(
+                    self.0[current_line].highlights = combine_highlights(
                         self.0[current_line].highlights.drain(..).into_iter(),
                         new_highlights.drain(..).into_iter(),
                     )
@@ -232,7 +298,7 @@ impl HighlightedContent {
         }
 
         if !new_highlights.is_empty() {
-            self.0[current_line].highlights = gpui::combine_highlights(
+            self.0[current_line].highlights = combine_highlights(
                 self.0[current_line].highlights.drain(..).into_iter(),
                 new_highlights.drain(..).into_iter(),
             )
@@ -242,6 +308,7 @@ impl HighlightedContent {
         self
     }
 
+    // given a range relative to the original source, returns the line number to which it belongs.
     pub(crate) fn line_index_of_range(&self, range: &std::ops::Range<usize>) -> Option<usize> {
         self.0
             .binary_search_by(|line| {
@@ -260,25 +327,79 @@ impl HighlightedContent {
             .ok()
     }
 
-    pub(crate) fn replace_highlight_range(&mut self, highlighted_range: &HighlightedRange) {
+    pub(crate) fn add_highlight_range(
-        let (range, _) = &highlighted_range;
+        &mut self,
+        highlighted_range: HighlightedRange,
+    ) -> Option<usize> {
+        let Some((line_i, line)) = self
+            .line_index_of_range(&highlighted_range.0)
+            .map(|line_i| (line_i, &mut self.0[line_i]))
+        else {
+            return None;
+        };
 
-        let line = self
+        line.highlights = combine_highlights(
-            .line_index_of_range(range)
+            line.highlights.drain(..).into_iter(),
-            .map(|line_i| &mut self.0[line_i]);
+            std::iter::once((
+                highlighted_range.0.relative_to(&line.line_range),
+                highlighted_range.1,
+            )),
+        )
+        .collect();
 
-        let replace_idx = line.as_ref().and_then(|l| {
+        Some(line_i)
-            let local_range = (range.start - l.line_range.start)..(range.end - l.line_range.start);
+    }
-            (l.highlights.iter().position(|(r, _)| *r == local_range)).map(|i| (i, local_range))
-        });
 
-        println!("line {:?} replace_idx {:?}", line, replace_idx);
+    pub(crate) fn set_highlights_for_line(
-
+        &mut self,
-        match (line, replace_idx) {
+        line_i: usize,
-        | (Some(line), Some((replace_idx, local_range))) => {
+        highlights: impl IntoIterator<Item = syntax_highlight::HighlightedRange>,
-            line.highlights[replace_idx] = (local_range, highlighted_range.1)
+    ) {
-        }
+        self.0[line_i].highlights.splice(.., highlights);
-        | _ => {}
+    }
-        }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn background(color: gpui::Hsla) -> gpui::HighlightStyle {
+        gpui::HighlightStyle {
+            background_color: Some(color),
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn combine_highlights_gives_precedence_to_second_iterator() {
+        let combined = combine_highlights(
+            [(0..5, background(gpui::yellow()))],
+            [(0..5, background(gpui::red()))],
+        )
+        .collect::<Vec<_>>();
+
+        assert_eq!(combined, [(0..5, background(gpui::red()))]);
+    }
+
+    #[test]
+    fn combine_highlights_splits_ranges_and_keeps_second_iterator_precedence() {
+        let combined = combine_highlights(
+            [
+                (0..3, background(gpui::yellow())),
+                (3..6, background(gpui::green())),
+            ],
+            [(1..5, background(gpui::red()))],
+        )
+        .collect::<Vec<_>>();
+
+        assert_eq!(
+            combined,
+            [
+                (0..1, background(gpui::yellow())),
+                (1..3, background(gpui::red())),
+                (3..5, background(gpui::red())),
+                (5..6, background(gpui::green())),
+            ]
+        );
     }
 }