From b058cca46d4c8a7e46734b2eb86e6993dd5bee8f Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 4 Mar 2023 15:00:11 +0800 Subject: refactor: use tui for better stream render (#5) * feat: show question in the begining of stream render * refactor: use tui for better stream render * trivial updates * fine-tune stream text merging --- src/render.rs | 327 +++++++++++++++++++++++++++++++++++++--------------------- src/repl.rs | 30 ++++-- 2 files changed, 234 insertions(+), 123 deletions(-) (limited to 'src') diff --git a/src/render.rs b/src/render.rs index 6a052e2..ff71b27 100644 --- a/src/render.rs +++ b/src/render.rs @@ -1,13 +1,10 @@ use anyhow::Result; -use crossbeam::sync::WaitGroup; +use crossbeam::channel::Receiver; use crossterm::{ - cursor, - event::{self, Event, KeyCode, KeyEvent, KeyModifiers}, - execute, queue, style, - terminal::{ - self, disable_raw_mode, enable_raw_mode, size, ClearType, EnterAlternateScreen, - LeaveAlternateScreen, - }, + event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyModifiers}, + execute, + style::{self, Color}, + terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, }; use mdcat::{ push_tty, @@ -16,143 +13,243 @@ use mdcat::{ }; use pulldown_cmark::Parser; use std::{ - io::{self, Write}, + io, sync::{ atomic::{AtomicBool, Ordering}, - mpsc::Receiver, Arc, }, - thread, - time::Duration, + time::{Duration, Instant}, }; use syntect::parsing::SyntaxSet; -use crate::repl::{dump, ReplyEvent}; +use crate::repl::{dump, RenderStreamEvent}; pub fn render_stream( - rx: Receiver, + rx: Receiver, ctrlc: Arc, markdown_render: Arc, ) -> Result<()> { - let wg = WaitGroup::new(); - let ctrlc_clone = ctrlc.clone(); - let stream_done = Arc::new(AtomicBool::new(false)); - let stream_done_clone = stream_done.clone(); - let wg_clone = wg.clone(); - thread::spawn(move || { - let _ = detect_ctrlc(ctrlc_clone, stream_done_clone); - drop(wg_clone); - }); - let ret = render_stream_inner(rx, ctrlc, markdown_render); - stream_done.store(true, Ordering::SeqCst); - wg.wait(); - ret + // setup terminal + enable_raw_mode()?; + let mut stdout = io::stdout(); + execute!(stdout, EnterAlternateScreen, EnableMouseCapture)?; + let backend = tui::backend::CrosstermBackend::new(stdout); + let mut terminal = tui::Terminal::new(backend)?; + + // create app and run it + let app = render_stream_tui::App::new(ctrlc, markdown_render); + let res = render_stream_tui::run(&mut terminal, app, rx); + + // restore terminal + disable_raw_mode()?; + execute!( + terminal.backend_mut(), + LeaveAlternateScreen, + DisableMouseCapture + )?; + terminal.show_cursor()?; + + res } -fn detect_ctrlc(ctrlc: Arc, stream_done: Arc) -> Result<()> { - loop { - if ctrlc.load(Ordering::SeqCst) || stream_done.load(Ordering::SeqCst) { - return Ok(()); - } - if event::poll(Duration::from_millis(100))? { - if let Event::Key(KeyEvent { - code: KeyCode::Char('c'), - modifiers: KeyModifiers::CONTROL, - .. - }) = event::read()? - { - ctrlc.store(true, Ordering::SeqCst); - break; +mod render_stream_tui { + use super::*; + use ansi_to_tui::IntoText; + use crossterm::event::MouseEventKind; + use tui::{ + backend::Backend, + layout::{Constraint, Direction, Layout}, + style::Style, + widgets::{List, ListItem, ListState}, + Frame, Terminal, + }; + + pub fn run( + terminal: &mut Terminal, + mut app: App, + rx: Receiver, + ) -> Result<()> { + let mut last_tick = Instant::now(); + let tick_rate = Duration::from_millis(250); + let mut count_evt = 0; + let mut count_done = 0; + + loop { + if app.ctrlc.load(Ordering::SeqCst) { + return Ok(()); } - } - } - Ok(()) -} -fn render_stream_inner( - rx: Receiver, - ctrlc: Arc, - markdown_render: Arc, -) -> Result<()> { - // setup terminal - enable_raw_mode()?; - let mut output = String::new(); - let mut stdout = io::stdout(); - execute!(stdout, EnterAlternateScreen)?; - - fn clear(stdout: &mut impl Write) -> io::Result<()> { - queue!( - stdout, - style::ResetColor, - terminal::Clear(ClearType::All), - cursor::Hide, - cursor::MoveTo(0, 0) - ) - } + if let Ok(evt) = rx.try_recv() { + count_evt += 1; + app.handle(evt)?; + if count_evt <= 16 { + continue; + } else { + count_evt = 0; + } + } - clear(&mut stdout)?; + terminal.draw(|f| ui(f, &mut app))?; - while let Ok(ev) = rx.recv() { - if ctrlc.load(Ordering::SeqCst) { - break; - } - match ev { - ReplyEvent::Text(text) => { - output.push_str(&text); - let rows = size()?.1 as usize; - let lines: Vec<&str> = output.split('\n').collect(); - let len = lines.len(); - let skip = if len > rows { len - rows } else { 0 }; - let mut selected_lines = vec![]; - let mut count_begin_code = 0; - let mut code = None; - for (index, line) in lines.iter().enumerate() { - if index < skip { - if line.starts_with("```") { - count_begin_code += 1; - code = Some(*line); - } - } else { - selected_lines.push(*line); - } + if app.no_interrupt && app.done { + count_done += 1; + if count_done >= 5 { + return Ok(()); } - if count_begin_code % 2 == 1 { - if let Some(code) = code { - selected_lines[0] = code + } + + let timeout = tick_rate + .checked_sub(last_tick.elapsed()) + .unwrap_or_else(|| Duration::from_secs(0)); + if crossterm::event::poll(timeout)? { + match event::read()? { + Event::Key(key) => { + app.no_interrupt = false; + match key.code { + KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => { + app.quit(); + return Ok(()); + } + KeyCode::Down => app.next(), + KeyCode::Up => app.previous(), + _ => {} + } } - }; - let content = selected_lines.join("\n"); - let markdown = markdown_render.render(&content)?; - if text.contains('\n') { - clear(&mut stdout)?; - for line in markdown.split('\n') { - queue!(stdout, style::Print(line), cursor::MoveToNextLine(1))?; + Event::Mouse(ev) => { + app.no_interrupt = false; + match ev.kind { + MouseEventKind::ScrollDown => app.next(), + MouseEventKind::ScrollUp => app.previous(), + _ => {} + } } - } else if let Some(line) = markdown.split('\n').last() { - queue!( - stdout, - style::ResetColor, - terminal::Clear(ClearType::CurrentLine), - cursor::MoveToColumn(0), - style::Print(line) - )?; + _ => {} } + } + if last_tick.elapsed() >= tick_rate { + app.on_tick(); + last_tick = Instant::now(); + } + } + } + + pub struct App { + buffer: String, + items: Vec, + list_state: ListState, + done: bool, + ctrlc: Arc, + markdown_render: Arc, + no_interrupt: bool, + num_rows: usize, + entry_index: usize, + } - stdout.flush()?; + impl App { + pub fn new(ctrlc: Arc, markdown_render: Arc) -> Self { + Self { + buffer: String::new(), + ctrlc, + done: false, + markdown_render, + num_rows: 0, + no_interrupt: true, + entry_index: 0, + items: vec![], + list_state: ListState::default(), } - ReplyEvent::Done => { - break; + } + + pub fn handle(&mut self, evt: RenderStreamEvent) -> Result<()> { + match evt { + RenderStreamEvent::Start(question) => { + let mut buf = Vec::with_capacity(8); + execute!( + buf, + style::SetForegroundColor(Color::Cyan), + style::Print("〉"), + style::ResetColor + )?; + let indicator = String::from_utf8_lossy(&buf); + self.buffer.push_str(&format!("{indicator}{question}\n")); + } + RenderStreamEvent::Text(text) => { + self.buffer.push_str(&text); + } + RenderStreamEvent::Done => { + self.done = true; + } + } + + let markdown = self.markdown_render.render(&self.buffer)?; + self.items = markdown.split('\n').map(|v| v.to_string()).collect(); + if self.no_interrupt { + self.end(); } + + Ok(()) + } + + pub fn quit(&mut self) { + self.ctrlc.store(true, Ordering::SeqCst); + } + + pub fn set_rows(&mut self, rows: u16) { + self.num_rows = rows as usize; + } + + pub fn next(&mut self) { + let index = if self.entry_index < self.num_rows { + self.num_rows.min(self.items.len() - 1) + } else { + self.entry_index + 1 + }; + self.entry_index = index; + self.list_state.select(Some(index)); } + + pub fn previous(&mut self) { + let index = self + .entry_index + .saturating_sub(1) + .min(self.items.len().saturating_sub(self.num_rows + 1)); + self.entry_index = index; + self.list_state.select(Some(index)); + } + + pub fn end(&mut self) { + let len = self.items.len(); + self.entry_index = if len < self.num_rows { + 0 + } else { + len - self.num_rows + }; + self.list_state.select(Some(len - 1)) + } + + pub fn on_tick(&mut self) {} } - execute!(stdout, style::ResetColor, cursor::Show)?; + fn ui(f: &mut Frame, app: &mut App) { + // Create two chunks with equal horizontal screen space + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Percentage(100)].as_ref()) + .split(f.size()); - // restore terminal - disable_raw_mode()?; - execute!(stdout, LeaveAlternateScreen)?; + let items: Vec = app + .items + .iter() + .map(|line| { + let text = line.into_text().unwrap_or_default(); + ListItem::new(text) + }) + .collect(); - Ok(()) + app.set_rows(chunks[0].height); + let items = List::new(items).highlight_style(Style::default()); + f.render_stateful_widget(items, chunks[0], &mut app.list_state); + } } pub struct MarkdownRender { diff --git a/src/repl.rs b/src/repl.rs index bb33bc1..e1cf4cf 100644 --- a/src/repl.rs +++ b/src/repl.rs @@ -2,6 +2,7 @@ use crate::client::ChatGptClient; use crate::config::{Config, Role}; use crate::render::{self, MarkdownRender}; use anyhow::{anyhow, Result}; +use crossbeam::channel::{unbounded, Sender}; use crossbeam::sync::WaitGroup; use reedline::{ default_emacs_keybindings, ColumnarMenu, DefaultCompleter, DefaultPrompt, DefaultPromptSegment, @@ -12,8 +13,6 @@ use std::cell::RefCell; use std::fs::File; use std::io::{stdout, Write}; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::mpsc::channel; -use std::sync::mpsc::Sender; use std::sync::Arc; use std::thread::spawn; @@ -232,7 +231,7 @@ impl ReplCmdHandler { }; let wg = WaitGroup::new(); let mut receiver = if let Some(markdown_render) = self.render.clone() { - let (tx, rx) = channel(); + let (tx, rx) = unbounded(); let ctrlc = self.ctrlc.clone(); let wg = wg.clone(); spawn(move || { @@ -243,6 +242,7 @@ impl ReplCmdHandler { } else { ReplyReceiver::new(None) }; + receiver.start(&input); self.client .acquire_stream(&input, prompt, &mut receiver, self.ctrlc.clone())?; Config::save_message( @@ -281,20 +281,32 @@ impl ReplCmdHandler { pub struct ReplyReceiver { output: String, - sender: Option>, + sender: Option>, } impl ReplyReceiver { - pub fn new(sender: Option>) -> Self { + pub fn new(sender: Option>) -> Self { Self { output: String::new(), sender, } } + + fn start(&self, input: &str) { + match self.sender.as_ref() { + Some(tx) => { + let _ = tx.send(RenderStreamEvent::Start(input.to_string())); + } + None => { + dump("", 2); + } + } + } + pub fn text(&mut self, text: &str) { match self.sender.as_ref() { Some(tx) => { - let _ = tx.send(ReplyEvent::Text(text.to_string())); + let _ = tx.send(RenderStreamEvent::Text(text.to_string())); } None => { dump(text, 0); @@ -302,10 +314,11 @@ impl ReplyReceiver { } self.output.push_str(text); } + pub fn done(&mut self) { match self.sender.as_ref() { Some(tx) => { - let _ = tx.send(ReplyEvent::Done); + let _ = tx.send(RenderStreamEvent::Done); } None => { dump("", 2); @@ -314,7 +327,8 @@ impl ReplyReceiver { } } -pub enum ReplyEvent { +pub enum RenderStreamEvent { + Start(String), Text(String), Done, } -- cgit v1.2.3