From 3144a8cb9e03b99fa5f14d278a1fa1fbc9605941 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Tue, 29 Oct 2024 17:33:41 +0900
Subject: [PATCH] annotate: add line_ranges() and compact_line_ranges()
 iterator

They allow callers to test range overlaps with e.g. diff hunks. "jj absorb"
will leverage compact_line_ranges().
---
 lib/src/annotate.rs | 122 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 122 insertions(+)

diff --git a/lib/src/annotate.rs b/lib/src/annotate.rs
index c8e14830ec..a238d38581 100644
--- a/lib/src/annotate.rs
+++ b/lib/src/annotate.rs
@@ -20,6 +20,8 @@
 
 use std::collections::hash_map;
 use std::collections::HashMap;
+use std::iter;
+use std::ops::Range;
 use std::rc::Rc;
 
 use bstr::BStr;
@@ -64,6 +66,43 @@ impl FileAnnotation {
             .map(|(commit_id, line)| (commit_id.as_ref(), line.as_ref()))
     }
 
+    /// Returns iterator over `(commit_id, line_range)`s.
+    ///
+    /// For each line, the `commit_id` points to the originator commit of the
+    /// line. The `line_range` is a slice range in the file `text`. Consecutive
+    /// ranges having the same `commit_id` are not compacted.
+    pub fn line_ranges(&self) -> impl Iterator<Item = (Option<&CommitId>, Range<usize>)> {
+        let ranges = self
+            .text
+            .split_inclusive(|b| *b == b'\n')
+            .scan(0, |total, line| {
+                let start = *total;
+                *total += line.len();
+                Some(start..*total)
+            });
+        itertools::zip_eq(&self.line_map, ranges)
+            .map(|(commit_id, range)| (commit_id.as_ref(), range))
+    }
+
+    /// Returns iterator over compacted `(commit_id, line_range)`s.
+    ///
+    /// Consecutive ranges having the same `commit_id` are merged into one.
+    pub fn compact_line_ranges(&self) -> impl Iterator<Item = (Option<&CommitId>, Range<usize>)> {
+        let mut ranges = self.line_ranges();
+        let mut acc = ranges.next();
+        iter::from_fn(move || {
+            let (acc_commit_id, acc_range) = acc.as_mut()?;
+            for (cur_commit_id, cur_range) in ranges.by_ref() {
+                if *acc_commit_id == cur_commit_id {
+                    acc_range.end = cur_range.end;
+                } else {
+                    return acc.replace((cur_commit_id, cur_range));
+                }
+            }
+            acc.take()
+        })
+    }
+
     /// File content at the starting commit.
     pub fn text(&self) -> &BStr {
         self.text.as_ref()
@@ -333,3 +372,86 @@ fn get_file_contents(
         _ => Ok(Vec::new()),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_lines_iterator_empty() {
+        let annotation = FileAnnotation {
+            line_map: vec![],
+            text: "".into(),
+        };
+        assert_eq!(annotation.lines().collect_vec(), vec![]);
+        assert_eq!(annotation.line_ranges().collect_vec(), vec![]);
+        assert_eq!(annotation.compact_line_ranges().collect_vec(), vec![]);
+    }
+
+    #[test]
+    fn test_lines_iterator_with_content() {
+        let commit_id1 = CommitId::from_hex("111111");
+        let commit_id2 = CommitId::from_hex("222222");
+        let commit_id3 = CommitId::from_hex("333333");
+        let annotation = FileAnnotation {
+            line_map: vec![
+                Some(commit_id1.clone()),
+                Some(commit_id2.clone()),
+                Some(commit_id3.clone()),
+            ],
+            text: "foo\n\nbar\n".into(),
+        };
+        assert_eq!(
+            annotation.lines().collect_vec(),
+            vec![
+                (Some(&commit_id1), "foo\n".as_ref()),
+                (Some(&commit_id2), "\n".as_ref()),
+                (Some(&commit_id3), "bar\n".as_ref()),
+            ]
+        );
+        assert_eq!(
+            annotation.line_ranges().collect_vec(),
+            vec![
+                (Some(&commit_id1), 0..4),
+                (Some(&commit_id2), 4..5),
+                (Some(&commit_id3), 5..9),
+            ]
+        );
+        assert_eq!(
+            annotation.compact_line_ranges().collect_vec(),
+            vec![
+                (Some(&commit_id1), 0..4),
+                (Some(&commit_id2), 4..5),
+                (Some(&commit_id3), 5..9),
+            ]
+        );
+    }
+
+    #[test]
+    fn test_lines_iterator_compaction() {
+        let commit_id1 = CommitId::from_hex("111111");
+        let commit_id2 = CommitId::from_hex("222222");
+        let commit_id3 = CommitId::from_hex("333333");
+        let annotation = FileAnnotation {
+            line_map: vec![
+                Some(commit_id1.clone()),
+                Some(commit_id1.clone()),
+                Some(commit_id2.clone()),
+                Some(commit_id1.clone()),
+                Some(commit_id3.clone()),
+                Some(commit_id3.clone()),
+                Some(commit_id3.clone()),
+            ],
+            text: "\n".repeat(7).into(),
+        };
+        assert_eq!(
+            annotation.compact_line_ranges().collect_vec(),
+            vec![
+                (Some(&commit_id1), 0..2),
+                (Some(&commit_id2), 2..3),
+                (Some(&commit_id1), 3..4),
+                (Some(&commit_id3), 4..7),
+            ]
+        );
+    }
+}