From 9f79389d9ca5c382b7b8ae5f360993164a8286ab Mon Sep 17 00:00:00 2001
From: Jonathan <jhrldev@gmail.com>
Date: Sat, 6 Apr 2024 00:18:37 -0300
Subject: [PATCH] fix: index out of bound in various instances with trailing
 new line

Fixes #688, fixes #694, fixes #682 and fixes #681.

This works around the fact that `str::lines` does not include last
newline character.
This means that `"a\n".lines().count() == "a".lines().count()`, which
breaks the 1:1 assumption made by display implementations.
---
 src/display/inline.rs       |  5 +++--
 src/display/side_by_side.rs |  9 +++++----
 src/display/style.rs        |  3 ++-
 src/lines.rs                | 38 ++++++++++++++++++++++++++++++++++++-
 4 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src/display/inline.rs b/src/display/inline.rs
index dc34248d5e..5ef6338f73 100644
--- a/src/display/inline.rs
+++ b/src/display/inline.rs
@@ -1,5 +1,6 @@
 //! Inline, or "unified" diff display.
 
+use crate::lines::lines_raw;
 use crate::{
     constants::Side,
     display::context::{calculate_after_context, calculate_before_context, opposite_positions},
@@ -43,8 +44,8 @@ pub(crate) fn print(
         )
     } else {
         (
-            lhs_src.lines().map(|s| format!("{}\n", s)).collect(),
-            rhs_src.lines().map(|s| format!("{}\n", s)).collect(),
+            lines_raw(lhs_src).map(|s| format!("{}\n", s)).collect(),
+            lines_raw(rhs_src).map(|s| format!("{}\n", s)).collect(),
         )
     };
 
diff --git a/src/display/side_by_side.rs b/src/display/side_by_side.rs
index 0b32e1fa6d..ce4a8513ec 100644
--- a/src/display/side_by_side.rs
+++ b/src/display/side_by_side.rs
@@ -9,6 +9,7 @@ use line_numbers::LineNumber;
 use line_numbers::SingleLineSpan;
 use owo_colors::{OwoColorize, Style};
 
+use crate::lines::lines_raw;
 use crate::{
     constants::Side,
     display::context::all_matched_lines_filled,
@@ -334,8 +335,8 @@ pub(crate) fn print(
         )
     } else {
         (
-            lhs_src.lines().map(|s| format!("{}\n", s)).collect(),
-            rhs_src.lines().map(|s| format!("{}\n", s)).collect(),
+            lines_raw(&lhs_src).map(|s| format!("{}\n", s)).collect(),
+            lines_raw(&rhs_src).map(|s| format!("{}\n", s)).collect(),
         )
     };
 
@@ -397,8 +398,8 @@ pub(crate) fn print(
     let mut prev_lhs_line_num = None;
     let mut prev_rhs_line_num = None;
 
-    let lhs_lines = lhs_src.lines().collect::<Vec<_>>();
-    let rhs_lines = rhs_src.lines().collect::<Vec<_>>();
+    let lhs_lines = lines_raw(lhs_src).collect::<Vec<_>>();
+    let rhs_lines = lines_raw(rhs_src).collect::<Vec<_>>();
     let matched_lines = all_matched_lines_filled(lhs_mps, rhs_mps, &lhs_lines, &rhs_lines);
     let mut matched_lines_to_print = &matched_lines[..];
 
diff --git a/src/display/style.rs b/src/display/style.rs
index e402af14d7..645b34e37e 100644
--- a/src/display/style.rs
+++ b/src/display/style.rs
@@ -7,6 +7,7 @@ use line_numbers::SingleLineSpan;
 use owo_colors::{OwoColorize, Style};
 use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
 
+use crate::lines::lines_raw;
 use crate::parse::syntax::StringKind;
 use crate::{
     constants::Side,
@@ -401,7 +402,7 @@ pub(crate) fn apply_colors(
     positions: &[MatchedPos],
 ) -> Vec<String> {
     let styles = color_positions(side, background, syntax_highlight, file_format, positions);
-    let lines = s.lines().collect::<Vec<_>>();
+    let lines = lines_raw(s).collect::<Vec<_>>();
     style_lines(&lines, &styles)
 }
 
diff --git a/src/lines.rs b/src/lines.rs
index a47182cd6e..756e26ae91 100644
--- a/src/lines.rs
+++ b/src/lines.rs
@@ -1,6 +1,9 @@
 //! Manipulate lines of text and groups of lines.
 
+use std::iter;
+use std::iter::{Chain, Once};
 use std::ops::Sub;
+use std::str::Lines;
 
 use line_numbers::LineNumber;
 
@@ -44,6 +47,30 @@ pub(crate) fn is_all_whitespace(s: &str) -> bool {
     s.chars().all(|c| c.is_whitespace())
 }
 
+pub(crate) fn lines_raw(s: &str) -> impl Iterator<Item = &str> {
+    if s.ends_with('\n') {
+        LinesIter::Chained(s.lines().chain(iter::once("")))
+    } else {
+        LinesIter::Lines(s.lines())
+    }
+}
+
+enum LinesIter<'a> {
+    Chained(Chain<Lines<'a>, Once<&'a str>>),
+    Lines(Lines<'a>),
+}
+
+impl<'a> Iterator for LinesIter<'a> {
+    type Item = &'a str;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            Self::Chained(iter) => iter.next(),
+            Self::Lines(iter) => iter.next(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use pretty_assertions::assert_eq;
@@ -75,7 +102,16 @@ mod tests {
     }
 
     #[test]
-    fn test_is_all_whiteapce() {
+    fn test_is_all_whitespace() {
         assert!(is_all_whitespace(" \n\t"));
     }
+
+    #[test]
+    fn test_last_new_line_preserved() {
+        assert_eq!(lines_raw("").count(), 0);
+        assert_eq!(lines_raw("a").count(), 1);
+        assert_eq!(lines_raw("\n").count(), 2);
+        assert_eq!(lines_raw("a\n").count(), 2);
+        assert_eq!(lines_raw("a\n\n").count(), 3);
+    }
 }