refactor: optimize render (#202)

7 files changed, 75 insertions, 30 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 1b00832..8589a8b 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -139,6 +139,8 @@ impl Config {
         config.setup_highlight();
         config.setup_light_theme()?;
 
+        setup_logger()?;
+
         Ok(config)
     }
 
@@ -822,3 +824,19 @@ fn set_bool(target: &mut bool, value: &str) {
         _ => {}
     }
 }
+
+#[cfg(debug_assertions)]
+fn setup_logger() -> Result<()> {
+    use simplelog::WriteLogger;
+    let file = std::fs::File::create(Config::local_path("debug.log")?)?;
+    let config = simplelog::ConfigBuilder::new()
+        .add_filter_allow_str("aichat")
+        .build();
+    WriteLogger::init(log::LevelFilter::Debug, config, file)?;
+    Ok(())
+}
+
+#[cfg(not(debug_assertions))]
+fn setup_logger() -> Result<()> {
+    Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
index eae9a4a..04f873f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,9 @@ mod client;
 mod config;
 mod render;
 mod repl;
+
+#[macro_use]
+extern crate log;
 #[macro_use]
 mod utils;
 
diff --git a/src/render/cmd.rs b/src/render/cmd.rs
index 75768b9..e942025 100644
--- a/src/render/cmd.rs
+++ b/src/render/cmd.rs
@@ -1,6 +1,6 @@
 use super::{MarkdownRender, ReplyEvent};
 
-use crate::utils::{spaces, split_line_sematic, split_line_tail, AbortSignal};
+use crate::utils::{split_line_sematic, split_line_tail, AbortSignal};
 
 use anyhow::Result;
 use crossbeam::channel::Receiver;
@@ -12,7 +12,7 @@ pub fn cmd_render_stream(
     abort: &AbortSignal,
 ) -> Result<()> {
     let mut buffer = String::new();
-    let mut col = 0;
+    let mut indent = 0;
     loop {
         if abort.aborted() {
             return Ok(());
@@ -24,10 +24,9 @@ pub fn cmd_render_stream(
                         let text = format!("{buffer}{text}");
                         let (head, tail) = split_line_tail(&text);
                         buffer = tail.to_string();
-                        let input = format!("{}{head}", spaces(col));
-                        let output = render.render(&input);
-                        println!("{}", &output[col..]);
-                        col = 0;
+                        let output = render.render_with_indent(head, indent);
+                        println!("{}", output);
+                        indent = 0;
                     } else {
                         buffer = format!("{buffer}{text}");
                         if !(render.is_code()
@@ -38,16 +37,15 @@ pub fn cmd_render_stream(
                         {
                             if let Some((head, remain)) = split_line_sematic(&buffer) {
                                 buffer = remain;
-                                let input = format!("{}{head}", spaces(col));
-                                let output = render.render(&input);
-                                let output = &output[col..];
-                                let (_, tail) = split_line_tail(output);
-                                if render.wrap_width().is_some() {
+                                let output = render.render_with_indent(&head, indent);
+                                let (_, tail) = split_line_tail(&output);
+                                if let Some(width) = render.wrap_width() {
                                     if output.contains('\n') {
-                                        col = display_width(tail);
+                                        indent = display_width(tail);
                                     } else {
-                                        col += display_width(output);
+                                        indent += display_width(&output);
                                     }
+                                    indent %= width as usize;
                                 }
                                 print!("{}", output);
                             }
@@ -55,9 +53,7 @@ pub fn cmd_render_stream(
                     }
                 }
                 ReplyEvent::Done => {
-                    let input = format!("{}{buffer}", spaces(col));
-                    let output = render.render(&input);
-                    let output = &output[col..];
+                    let output = render.render_with_indent(&buffer, indent);
                     println!("{}", output);
                     break;
                 }
diff --git a/src/render/markdown.rs b/src/render/markdown.rs
index 8ae0a09..834c47b 100644
--- a/src/render/markdown.rs
+++ b/src/render/markdown.rs
@@ -82,6 +82,16 @@ impl MarkdownRender {
             .join("\n")
     }
 
+    pub fn render_with_indent(&mut self, text: &str, padding: usize) -> String {
+        let text = format!("{}{}", " ".repeat(padding), text);
+        let output = self.render(&text);
+        if output.starts_with('\n') {
+            output
+        } else {
+            output.chars().skip(padding).collect()
+        }
+    }
+
     pub fn render_line(&self, line: &str) -> String {
         let (_, code_syntax, is_code) = self.check_line(line);
         if is_code {
@@ -359,4 +369,17 @@ std::error::Error>> {
         let output = render.render(TEXT);
         assert_eq!(TEXT_WRAP_ALL, output);
     }
+
+    #[test]
+    fn wrap_with_indent() {
+        let options = RenderOptions::default();
+        let mut render = MarkdownRender::init(options).unwrap();
+        render.wrap_width = Some(80);
+
+        let input = "To unzip a file in Rust, you can use the `zip` crate. Here's an example code";
+        let output = render.render_with_indent(input, 40);
+        let expect =
+            "To unzip a file in Rust, you can use the\n`zip` crate. Here's an example code";
+        assert_eq!(output, expect);
+    }
 }
diff --git a/src/render/mod.rs b/src/render/mod.rs
index 554353d..93eb30e 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -77,6 +77,7 @@ impl ReplyHandler {
     }
 
     pub fn text(&mut self, text: &str) -> Result<()> {
+        debug!("ReplyText: {}", text);
         if self.buffer.is_empty() && text == "\n\n" {
             return Ok(());
         }
@@ -90,6 +91,7 @@ impl ReplyHandler {
     }
 
     pub fn done(&mut self) -> Result<()> {
+        debug!("ReplyDone");
         let ret = self
             .sender
             .send(ReplyEvent::Done)
diff --git a/src/render/repl.rs b/src/render/repl.rs
index f554359..31c597e 100644
--- a/src/render/repl.rs
+++ b/src/render/repl.rs
@@ -39,10 +39,12 @@ fn repl_render_stream_inner(
 ) -> Result<()> {
     let mut last_tick = Instant::now();
     let tick_rate = Duration::from_millis(50);
+
     let mut buffer = String::new();
+    let mut buffer_rows = 1;
+
     let columns = terminal::size()?.0;
 
-    let mut clear_rows = 0;
     loop {
         if abort.aborted() {
             return Ok(());
@@ -53,21 +55,22 @@ fn repl_render_stream_inner(
                 ReplyEvent::Text(text) => {
                     let (col, mut row) = cursor::position()?;
 
-                    // fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
+                    // Fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
                     if col == 0 && row > 0 && display_width(&buffer) == columns as usize {
                         row -= 1;
                     }
 
-                    if row + 1 >= clear_rows {
-                        queue!(writer, cursor::MoveTo(0, row.saturating_sub(clear_rows)))?;
+                    if row + 1 >= buffer_rows {
+                        queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows),)?;
                     } else {
-                        let scroll_rows = clear_rows - row - 1;
+                        let scroll_rows = buffer_rows - row - 1;
                         queue!(
                             writer,
                             terminal::ScrollUp(scroll_rows),
                             cursor::MoveTo(0, 0),
                         )?;
                     }
+                    queue!(writer, terminal::Clear(terminal::ClearType::UntilNewLine))?;
 
                     if text.contains('\n') {
                         let text = format!("{buffer}{text}");
@@ -76,19 +79,18 @@ fn repl_render_stream_inner(
                         let output = render.render(head);
                         print_block(writer, &output, columns)?;
                         queue!(writer, style::Print(&buffer),)?;
-                        clear_rows = 0;
+                        buffer_rows = need_rows(&buffer, columns);
                     } else {
                         buffer = format!("{buffer}{text}");
                         let output = render.render_line(&buffer);
                         if output.contains('\n') {
                             let (head, tail) = split_line_tail(&output);
-                            clear_rows = print_block(writer, head, columns)?;
+                            buffer_rows = print_block(writer, head, columns)?;
                             queue!(writer, style::Print(&tail),)?;
+                            buffer_rows += need_rows(tail, columns);
                         } else {
                             queue!(writer, style::Print(&output))?;
-                            let buffer_width = display_width(&output) as u16;
-                            let need_rows = (buffer_width + columns - 1) / columns;
-                            clear_rows = need_rows.saturating_sub(1);
+                            buffer_rows = need_rows(&output, columns);
                         }
                     }
 
@@ -147,3 +149,8 @@ fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result<u16> {
     }
     Ok(num)
 }
+
+fn need_rows(text: &str, columns: u16) -> u16 {
+    let buffer_width = display_width(text) as u16;
+    (buffer_width + columns - 1) / columns
+}
diff --git a/src/utils/split_line.rs b/src/utils/split_line.rs
index 105c5f3..da10ad1 100644
--- a/src/utils/split_line.rs
+++ b/src/utils/split_line.rs
@@ -35,10 +35,6 @@ pub fn split_line_tail(text: &str) -> (&str, &str) {
     }
 }
 
-pub fn spaces(n: usize) -> String {
-    " ".repeat(n)
-}
-
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 enum Kind {
     ParentheseStart,