refactor: optimize render (#202)

Cargo.lock

@@ -49,6 +49,7 @@ dependencies = [
  "inquire",
  "is-terminal",
  "lazy_static",
+ "log",
  "nu-ansi-term",
  "parking_lot",
  "reedline",
@@ -58,6 +59,7 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_yaml",
+ "simplelog",
  "syntect",
  "textwrap",
  "tokio",
@@ -1125,6 +1127,15 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "num_threads"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2819ce041d2ee131036f4fc9d6ae7ae125a3a40e97ba64d04fe799ad9dabbb44"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "objc"
 version = "0.2.7"
@@ -1694,6 +1705,17 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "simplelog"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "acee08041c5de3d5048c8b3f6f13fafb3026b24ba43c6a695a0c76179b844369"
+dependencies = [
+ "log",
+ "termcolor",
+ "time",
+]
+
 [[package]]
 name = "slab"
 version = "0.4.9"
@@ -1844,6 +1866,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "termcolor"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755"
+dependencies = [
+ "winapi-util",
+]
+
 [[package]]
 name = "textwrap"
 version = "0.16.0"
@@ -1883,6 +1914,8 @@ checksum = "c4a34ab300f2dee6e562c10a046fc05e358b29f9bf92277f30c3c8d82275f6f5"
 dependencies = [
  "deranged",
  "itoa",
+ "libc",
+ "num_threads",
  "powerfmt",
  "serde",
  "time-core",

Cargo.toml

@@ -40,6 +40,8 @@ async-trait = "0.1.74"
 textwrap = "0.16.0"
 ansi_colours = "1.2.2"
 reqwest-eventsource = "0.5.0"
+simplelog = "0.12.1"
+log = "0.4.20"
 
 [dependencies.reqwest]
 version = "0.11.14"

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(())
+}

src/main.rs

@@ -3,6 +3,9 @@ mod client;
 mod config;
 mod render;
 mod repl;
+
+#[macro_use]
+extern crate log;
 #[macro_use]
 mod utils;
 

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;
                 }

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);
+    }
 }

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)

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
+}

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,