From 2f14c71be27fe9d776dadd47842b97c152028d09 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mano=20S=C3=A9gransan?= <mano.segransan@protonmail.com>
Date: Tue, 13 Aug 2024 14:03:50 +0200
Subject: [PATCH] Fix off by 1 highlight when message payload contains newlines

---
 src/handlers/data.rs | 68 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 2 deletions(-)

diff --git a/src/handlers/data.rs b/src/handlers/data.rs
index 5df06dd..6809625 100644
--- a/src/handlers/data.rs
+++ b/src/handlers/data.rs
@@ -1,7 +1,7 @@
 use chrono::{offset::Local, DateTime};
 use fuzzy_matcher::{skim::SkimMatcherV2, FuzzyMatcher};
 use log::{error, warn};
-use memchr::memmem;
+use memchr::{memchr_iter, memmem};
 use once_cell::sync::Lazy;
 use std::{borrow::Cow, mem::swap, string::ToString};
 use tui::{
@@ -562,6 +562,11 @@ impl MessageData {
         let mut first_line = lines.next().unwrap();
         let first_line_msg = split_cow_in_place(&mut first_line, prefix_len);
 
+        // Find all newlines in `self.payload`, to increment `next_index` when a newline is encountered.
+        // Wrapping the message above removes newlines, which makes the highlight indices off by 1
+        // when a newline is encountered.
+        let mut newlines_iter = memchr_iter(b'\n', self.payload.as_bytes());
+
         let mut emotes = &self.emotes[..];
 
         first_row.extend(Self::build_line(
@@ -573,9 +578,16 @@ impl MessageData {
             &mut emotes,
         ));
 
+        let mut nl_idx = newlines_iter.next();
+
         let mut rows = vec![Line::from(first_row)];
 
         rows.extend(lines.map(|line| {
+            if nl_idx == Some(next_index) {
+                next_index += 1;
+                nl_idx = newlines_iter.next();
+            }
+
             Line::from(Self::build_line(
                 line,
                 &mut next_index,
@@ -642,6 +654,7 @@ impl<'conf> DataBuilder<'conf> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::collections::BTreeMap;
 
     #[test]
     fn test_username_hash() {
@@ -806,7 +819,7 @@ mod tests {
         let emote_w_3 = UnicodePlaceholder::new(3).string();
 
         let line =
-        format!("foo{ZERO_WIDTH_SPACE}{emote_w_3}{ZERO_WIDTH_SPACE}{emote_w_1}{ZERO_WIDTH_SPACE}bar baz{ZERO_WIDTH_SPACE}{emote_w_2}");
+            format!("foo{ZERO_WIDTH_SPACE}{emote_w_3}{ZERO_WIDTH_SPACE}{emote_w_1}{ZERO_WIDTH_SPACE}bar baz{ZERO_WIDTH_SPACE}{emote_w_2}");
 
         let (spans, _, emotes) = assert_build_line(
             &line,
@@ -920,4 +933,55 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn build_vec_with_newlines_and_highlights() {
+        let raw_message = RawMessageData::new(
+            "foo".to_string(),
+            None,
+            false,
+            "foo\nbar baz".to_string(),
+            BTreeMap::new(),
+            None,
+            false,
+        );
+
+        let data = MessageData::from_twitch_message(raw_message, &SharedEmotes::default(), false);
+
+        let frontendconfig = FrontendConfig {
+            show_datetimes: false,
+            ..FrontendConfig::default()
+        };
+
+        let lines = data.to_vec(&frontendconfig, 80, Some("bar"), None);
+
+        assert_eq!(
+            lines,
+            vec![
+                Line::from(vec![
+                    Span::styled("foo", data.hash_username(&Palette::Pastel)),
+                    Span::raw(": "),
+                    Span::raw("foo")
+                ]),
+                Line::from(vec![
+                    Span::styled(
+                        "b",
+                        Style::default().fg(Color::Red).add_modifier(Modifier::BOLD)
+                    ),
+                    Span::styled(
+                        "a",
+                        Style::default().fg(Color::Red).add_modifier(Modifier::BOLD)
+                    ),
+                    Span::styled(
+                        "r",
+                        Style::default().fg(Color::Red).add_modifier(Modifier::BOLD)
+                    ),
+                    Span::raw(" "),
+                    Span::raw("b"),
+                    Span::raw("a"),
+                    Span::raw("z")
+                ])
+            ]
+        );
+    }
 }