refactor: split long paragraphs for smoother stream output (#33)

2 files changed, 267 insertions, 34 deletions

diff --git a/src/render/cmd.rs b/src/render/cmd.rs
index f29fe7d..d349c79 100644
--- a/src/render/cmd.rs
+++ b/src/render/cmd.rs
@@ -23,6 +23,17 @@ pub fn cmd_render_stream(rx: Receiver<ReplyStreamEvent>, abort: SharedAbortSigna
                         dump(markdown_render.render(&output), 1);
                     } else {
                         buffer = format!("{buffer}{text}");
+                        if !(markdown_render.is_code_block()
+                            || buffer.len() < 60
+                            || buffer.starts_with('#')
+                            || buffer.starts_with('>')
+                            || buffer.starts_with('|'))
+                        {
+                            if let Some((output, remain)) = split_line(&buffer) {
+                                dump(markdown_render.render_line_stateless(&output), 0);
+                                buffer = remain
+                            }
+                        }
                     }
                 }
                 ReplyStreamEvent::Done => {
@@ -35,3 +46,215 @@ pub fn cmd_render_stream(rx: Receiver<ReplyStreamEvent>, abort: SharedAbortSigna
     }
     Ok(())
 }
+
+fn split_line(line: &str) -> Option<(String, String)> {
+    let mut balance: Vec<Kind> = Vec::new();
+    let chars: Vec<char> = line.chars().collect();
+    let mut index = 0;
+    let len = chars.len();
+    while index < len - 1 {
+        let ch = chars[index];
+        if balance.is_empty()
+            && ((matches!(ch, ',' | '.' | ';') && chars[index + 1].is_whitespace())
+                || matches!(ch, '，' | '。' | '；'))
+        {
+            let (output, remain) = chars.split_at(index + 1);
+            return Some((output.iter().collect(), remain.iter().collect()));
+        }
+        if index + 2 < len && do_balance(&mut balance, &chars[index..=index + 2]) {
+            index += 3;
+            continue;
+        }
+        if do_balance(&mut balance, &chars[index..=index + 1]) {
+            index += 2;
+            continue;
+        }
+        do_balance(&mut balance, &chars[index..index + 1]);
+        index += 1
+    }
+
+    None
+}
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+enum Kind {
+    ParentheseStart,
+    ParentheseEnd,
+    BracketStart,
+    BracketEnd,
+    Asterisk,
+    Asterisk2,
+    SingleQuota,
+    DoubleQuota,
+    Tilde,
+    Tilde2,
+    Backtick,
+    Backtick3,
+}
+
+impl Kind {
+    fn from_chars(chars: &[char]) -> Option<Self> {
+        let kind = match chars.len() {
+            1 => match chars[0] {
+                '(' => Kind::ParentheseStart,
+                ')' => Kind::ParentheseEnd,
+                '[' => Kind::BracketStart,
+                ']' => Kind::BracketEnd,
+                '*' => Kind::Asterisk,
+                '\'' => Kind::SingleQuota,
+                '"' => Kind::DoubleQuota,
+                '~' => Kind::Tilde,
+                '`' => Kind::Backtick,
+                _ => return None,
+            },
+            2 if chars[0] == chars[1] => match chars[0] {
+                '*' => Kind::Asterisk2,
+                '~' => Kind::Tilde2,
+                _ => return None,
+            },
+            3 => {
+                if chars == ['`', '`', '`'] {
+                    Kind::Backtick3
+                } else {
+                    return None;
+                }
+            }
+            _ => return None,
+        };
+        Some(kind)
+    }
+}
+
+fn do_balance(balance: &mut Vec<Kind>, chars: &[char]) -> bool {
+    if let Some(kind) = Kind::from_chars(chars) {
+        let last = balance.last();
+        match (kind, last) {
+            (Kind::ParentheseStart | Kind::BracketStart, _) => {
+                balance.push(kind);
+                true
+            }
+            (Kind::ParentheseEnd, Some(&Kind::ParentheseStart)) => {
+                balance.pop();
+                true
+            }
+            (Kind::BracketEnd, Some(&Kind::BracketStart)) => {
+                balance.pop();
+                true
+            }
+            (Kind::Asterisk, Some(&Kind::Asterisk))
+            | (Kind::Asterisk2, Some(&Kind::Asterisk2))
+            | (Kind::SingleQuota, Some(&Kind::SingleQuota))
+            | (Kind::DoubleQuota, Some(&Kind::DoubleQuota))
+            | (Kind::Tilde, Some(&Kind::Tilde))
+            | (Kind::Tilde2, Some(&Kind::Tilde2))
+            | (Kind::Backtick, Some(&Kind::Backtick))
+            | (Kind::Backtick3, Some(&Kind::Backtick3)) => {
+                balance.pop();
+                true
+            }
+            (Kind::Asterisk, _)
+            | (Kind::Asterisk2, _)
+            | (Kind::SingleQuota, _)
+            | (Kind::DoubleQuota, _)
+            | (Kind::Tilde, _)
+            | (Kind::Tilde2, _)
+            | (Kind::Backtick, _)
+            | (Kind::Backtick3, _) => {
+                balance.push(kind);
+                true
+            }
+            _ => false,
+        }
+    } else {
+        false
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! assert_split_line {
+        ($a:literal, $b:literal, true) => {
+            assert_eq!(
+                split_line(&format!("{}{}", $a, $b)),
+                Some(($a.into(), $b.into()))
+            );
+        };
+        ($a:literal, $b:literal, false) => {
+            assert_eq!(split_line(&format!("{}{}", $a, $b)), None);
+        };
+    }
+
+    #[test]
+    fn test_split_line() {
+        assert_split_line!(
+            "Lorem ipsum dolor sit amet,",
+            " consectetur adipiscing elit.",
+            true
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit amet.",
+            " consectetur adipiscing elit.",
+            true
+        );
+        assert_split_line!("黃更室幼許刀知，", "波食小午足田世根候法。", true);
+        assert_split_line!("黃更室幼許刀知。", "波食小午足田世根候法。", true);
+        assert_split_line!("黃更室幼許刀知；", "波食小午足田世根候法。", true);
+        assert_split_line!(
+            "Lorem ipsum (dolor sit amet).",
+            " consectetur adipiscing elit.",
+            true
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit `amet,",
+            " consectetur` adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit ```amet,",
+            " consectetur``` adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit *amet,",
+            " consectetur* adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit **amet,",
+            " consectetur** adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit ~amet,",
+            " consectetur~ adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit ~~amet,",
+            " consectetur~~ adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit ``amet,",
+            " consectetur`` adipiscing elit.",
+            true
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit \"amet,",
+            " consectetur\" adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit 'amet,",
+            " consectetur' adipiscing elit.",
+            false
+        );
+        assert_split_line!(
+            "Lorem ipsum dolor sit amet.",
+            "consectetur adipiscing elit.",
+            false
+        );
+    }
+}
diff --git a/src/render/markdown.rs b/src/render/markdown.rs
index f0cb64a..23d3d59 100644
--- a/src/render/markdown.rs
+++ b/src/render/markdown.rs
@@ -1,4 +1,3 @@
-// use colored::{Color, Colorize};
 use crossterm::style::{Color, Stylize};
 use syntect::highlighting::{Color as SyntectColor, FontStyle, Style, Theme};
 use syntect::parsing::SyntaxSet;
@@ -15,7 +14,7 @@ pub struct MarkdownRender {
     code_color: Color,
     md_syntax: SyntaxReference,
     code_syntax: Option<SyntaxReference>,
-    line_type: LineType,
+    prev_line_type: LineType,
 }
 
 impl MarkdownRender {
@@ -32,7 +31,7 @@ impl MarkdownRender {
             code_color,
             md_syntax,
             code_syntax: None,
-            line_type,
+            prev_line_type: line_type,
         }
     }
 
@@ -43,11 +42,27 @@ impl MarkdownRender {
             .join("\n")
     }
 
-    pub fn render_line(&mut self, line: &str) -> Option<String> {
+    pub fn render_line_stateless(&self, line: &str) -> String {
+        let output = if self.is_code_block() && detect_code_block(line).is_none() {
+            self.render_code_line(line)
+        } else {
+            self.render_line_inner(line, &self.md_syntax)
+        };
+        output.unwrap_or_else(|| line.to_string())
+    }
+
+    pub fn is_code_block(&self) -> bool {
+        matches!(
+            self.prev_line_type,
+            LineType::CodeBegin | LineType::CodeInner
+        )
+    }
+
+    fn render_line(&mut self, line: &str) -> Option<String> {
         if let Some(lang) = detect_code_block(line) {
-            match self.line_type {
+            match self.prev_line_type {
                 LineType::Normal | LineType::CodeEnd => {
-                    self.line_type = LineType::CodeBegin;
+                    self.prev_line_type = LineType::CodeBegin;
                     self.code_syntax = if lang.is_empty() {
                         None
                     } else {
@@ -55,20 +70,20 @@ impl MarkdownRender {
                     };
                 }
                 LineType::CodeBegin | LineType::CodeInner => {
-                    self.line_type = LineType::CodeEnd;
+                    self.prev_line_type = LineType::CodeEnd;
                     self.code_syntax = None;
                 }
             }
             self.render_line_inner(line, &self.md_syntax)
         } else {
-            match self.line_type {
+            match self.prev_line_type {
                 LineType::Normal => self.render_line_inner(line, &self.md_syntax),
                 LineType::CodeEnd => {
-                    self.line_type = LineType::Normal;
+                    self.prev_line_type = LineType::Normal;
                     self.render_line_inner(line, &self.md_syntax)
                 }
                 LineType::CodeBegin => {
-                    self.line_type = LineType::CodeInner;
+                    self.prev_line_type = LineType::CodeInner;
                     self.render_code_line(line)
                 }
                 LineType::CodeInner => self.render_code_line(line),
@@ -76,25 +91,14 @@ impl MarkdownRender {
         }
     }
 
-    pub fn render_line_stateless(&self, line: &str) -> String {
-        let output = if detect_code_block(line).is_some() {
-            self.render_line_inner(line, &self.md_syntax)
-        } else {
-            match self.line_type {
-                LineType::Normal | LineType::CodeEnd => {
-                    self.render_line_inner(line, &self.md_syntax)
-                }
-                _ => self.render_code_line(line),
-            }
-        };
-
-        output.unwrap_or_else(|| line.to_string())
-    }
-
     fn render_line_inner(&self, line: &str, syntax: &SyntaxReference) -> Option<String> {
+        let ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
+        let trimed_line = &line[ws.len()..];
         let mut highlighter = HighlightLines::new(syntax, &self.md_theme);
-        let ranges = highlighter.highlight_line(line, &self.syntax_set).ok()?;
-        Some(as_terminal_escaped(&ranges))
+        let ranges = highlighter
+            .highlight_line(trimed_line, &self.syntax_set)
+            .ok()?;
+        Some(format!("{ws}{}", as_terminal_escaped(&ranges)))
     }
 
     fn render_code_line(&self, line: &str) -> Option<String> {
@@ -112,7 +116,7 @@ impl MarkdownRender {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum LineType {
+pub enum LineType {
     Normal,
     CodeBegin,
     CodeInner,
@@ -184,10 +188,16 @@ fn get_code_color(theme: &Theme) -> Color {
         .unwrap_or_else(|| Color::Yellow)
 }
 
-#[test]
-fn test_assets() {
-    let syntax_set: SyntaxSet = bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes.bin");
-    assert!(syntax_set.find_syntax_by_extension("md").is_some());
-    let md_theme: Theme = bincode::deserialize_from(MD_THEME).expect("invalid md_theme binary");
-    assert_eq!(md_theme.name, Some("Monokai Extended".into()));
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_assets() {
+        let syntax_set: SyntaxSet =
+            bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes.bin");
+        assert!(syntax_set.find_syntax_by_extension("md").is_some());
+        let md_theme: Theme = bincode::deserialize_from(MD_THEME).expect("invalid md_theme binary");
+        assert_eq!(md_theme.name, Some("Monokai Extended".into()));
+    }
 }