Fix more side-by-side line number miscounts

2 files changed, 148 insertions, 32 deletions

diff --git a/src/features/line_numbers.rs b/src/features/line_numbers.rs
index 47017960..35674f99 100644
--- a/src/features/line_numbers.rs
+++ b/src/features/line_numbers.rs
@@ -779,21 +779,90 @@ pub mod tests {
 
     #[test]
     fn test_line_numbers_continue_correctly() {
-        let config = make_config_from_args(&[
+        DeltaTest::with(&["--side-by-side", "--width", "44", "--line-fill-method=ansi"])
+            .with_input(DIFF_PLUS_MINUS_WITH_1_CONTEXT_DIFF)
+            .expect(
+                r#"
+                │ 1  │abc             │ 1  │abc
+                │ 2  │a = left side   │ 2  │a = right side
+                │ 3  │xyz             │ 3  │xyz"#,
+            );
+    }
+
+    #[test]
+    fn test_line_numbers_continue_correctly_after_wrapping() {
+        DeltaTest::with(&[
             "--side-by-side",
             "--width",
-            "40",
-            "--line-fill-method=spaces",
-        ]);
-
-        let output = run_delta(HUNK_PLUS_MINUS_WITH_1_CONTEXT_DIFF, &config);
-        let mut lines = output.lines().skip(crate::config::HEADER_LEN);
+            "32",
+            "--line-fill-method=ansi",
+            "--wrap-left-symbol",
+            "@",
+            "--wrap-right-symbol",
+            "@",
+        ])
+        .with_input(DIFF_PLUS_MINUS_WITH_1_CONTEXT_DIFF)
+        .expect(
+            r#"
+            │ 1  │abc       │ 1  │abc
+            │ 2  │a = left @│ 2  │a = right@
+            │    │side      │    │ side
+            │ 3  │xyz       │ 3  │xyz"#,
+        );
 
-        // closure to help with `cargo fmt`-ing:
-        let mut next_line = || strip_ansi_codes(lines.next().unwrap());
-        assert_eq!("│ 1  │same          │ 1  │same", next_line());
-        assert_eq!("│ 2  │a = left      │ 2  │a = right     ", next_line());
-        assert_eq!("│ 3  │also same     │ 3  │also same", next_line());
+        let cfg = &[
+            "--side-by-side",
+            "--width",
+            "42",
+            "--line-fill-method=ansi",
+            "--wrap-left-symbol",
+            "@",
+            "--wrap-right-symbol",
+            "@",
+        ];
+
+        DeltaTest::with(cfg)
+            .with_input(DIFF_WITH_LONGER_MINUS_1_CONTEXT)
+            .expect(
+                r#"
+                │ 1  │abc            │ 1  │abc
+                │ 2  │a = one side   │ 2  │a = one longer@
+                │    │               │    │ side
+                │ 3  │xyz            │ 3  │xyz"#,
+            );
+
+        DeltaTest::with(cfg)
+            .with_input(DIFF_WITH_LONGER_PLUS_1_CONTEXT)
+            .expect(
+                r#"
+                │ 1  │abc            │ 1  │abc
+                │ 2  │a = one longer@│ 2  │a = one side
+                │    │ side          │    │
+                │ 3  │xyz            │ 3  │xyz"#,
+            );
+
+        DeltaTest::with(cfg)
+            .with_input(DIFF_MISMATCH_LONGER_MINUS_1_CONTEXT)
+            .expect(
+                r#"
+                │ 1  │abc            │ 1  │abc
+                │ 2  │a = left side @│    │
+                │    │which is longer│    │
+                │    │               │ 2  │a = other one
+                │ 3  │xyz            │ 3  │xyz"#,
+            );
+
+        DeltaTest::with(cfg)
+            .with_input(DIFF_MISMATCH_LONGER_PLUS_1_CONTEXT)
+            .expect(
+                r#"
+                │ 1  │abc            │ 1  │abc
+                │ 2  │a = other one  │    │
+                │    │               │ 2  │a = right side@
+                │    │               │    │ which is long@
+                │    │               │    │er
+                │ 3  │xyz            │ 3  │xyz"#,
+            );
     }
 
     pub const TWO_MINUS_LINES_DIFF: &str = "\
@@ -865,12 +934,48 @@ index 223ca50..367a6f6 100644
 +bb = 2
 ";
 
-    const HUNK_PLUS_MINUS_WITH_1_CONTEXT_DIFF: &str = "\
+    const DIFF_PLUS_MINUS_WITH_1_CONTEXT_DIFF: &str = "\
+--- a/a.py
++++ b/b.py
+@@ -1,3 +1,3 @@
+ abc
+-a = left side
++a = right side
+ xyz";
+
+    const DIFF_WITH_LONGER_MINUS_1_CONTEXT: &str = "\
+--- a/a.py
++++ b/b.py
+@@ -1,3 +1,3 @@
+ abc
+-a = one side
++a = one longer side
+ xyz";
+
+    const DIFF_WITH_LONGER_PLUS_1_CONTEXT: &str = "\
+--- a/a.py
++++ b/b.py
+@@ -1,3 +1,3 @@
+ abc
+-a = one longer side
++a = one side
+ xyz";
+
+    const DIFF_MISMATCH_LONGER_MINUS_1_CONTEXT: &str = "\
+--- a/a.py
++++ b/b.py
+@@ -1,3 +1,3 @@
+ abc
+-a = left side which is longer
++a = other one
+ xyz";
+
+    const DIFF_MISMATCH_LONGER_PLUS_1_CONTEXT: &str = "\
 --- a/a.py
 +++ b/b.py
 @@ -1,3 +1,3 @@
- same
--a = left
-+a = right
- also same";
+ abc
+-a = other one
++a = right side which is longer
+ xyz";
 }
diff --git a/src/features/side_by_side.rs b/src/features/side_by_side.rs
index 890bc0b3..a9dfb6a8 100644
--- a/src/features/side_by_side.rs
+++ b/src/features/side_by_side.rs
@@ -174,41 +174,52 @@ pub fn paint_minus_and_plus_lines_side_by_side(
     };
 
     for (minus_line_index, plus_line_index) in line_alignment {
+        let left_state = match minus_line_index {
+            Some(i) => &line_states[Left][i],
+            None => &State::HunkMinus(DiffType::Unified, None),
+        };
         output_buffer.push_str(&paint_left_panel_minus_line(
             minus_line_index,
             &syntax_sections[Left],
             &diff_sections[Left],
             &lines_have_homolog[Left],
-            match minus_line_index {
-                Some(i) => &line_states[Left][i],
-                None => &State::HunkMinus(DiffType::Unified, None),
-            },
+            left_state,
             &mut Some(line_numbers_data),
             bg_should_fill[Left],
             config,
         ));
 
-        // HACK: The left line number is not getting incremented in `linenumbers_and_styles()`
-        // when the alignment matches a minus with a plus line, so fix that here.
-        // This information should be passed down into `paint_line()` to set `increment` to true.
-        if minus_line_index.is_some() && plus_line_index.is_some() {
-            line_numbers_data.line_number[Left] += 1;
-        }
-
+        let right_state = match plus_line_index {
+            Some(i) => &line_states[Right][i],
+            None => &State::HunkPlus(DiffType::Unified, None),
+        };
         output_buffer.push_str(&paint_right_panel_plus_line(
             plus_line_index,
             &syntax_sections[Right],
             &diff_sections[Right],
             &lines_have_homolog[Right],
-            match plus_line_index {
-                Some(i) => &line_states[Right][i],
-                None => &State::HunkPlus(DiffType::Unified, None),
-            },
+            right_state,
             &mut Some(line_numbers_data),
             bg_should_fill[Right],
             config,
         ));
         output_buffer.push('\n');
+
+        // HACK: The left line number is not getting incremented in `linenumbers_and_styles()`
+        // when the alignment matches a minus with a plus line, so fix that here and take
+        // wrapped lines into account.
+        // Similarly an increment happens when it should not, so undo that.
+        // TODO: Pass this information down into `paint_line()` to set `increment` accordingly.
+        match (left_state, right_state, minus_line_index, plus_line_index) {
+            (State::HunkMinusWrapped, State::HunkPlus(_, _), Some(_), None) => {
+                line_numbers_data.line_number[Left] =
+                    line_numbers_data.line_number[Left].saturating_sub(1)
+            }
+            // Duplicating the logic from `linenumbers_and_styles()` a bit:
+            (State::HunkMinusWrapped | State::HunkPlusWrapped, _, _, _) => {}
+            (_, _, Some(_), Some(_)) => line_numbers_data.line_number[Left] += 1,
+            _ => {}
+        }
     }
 }