Better ScrollView for tui-rs

scrollview.rs

use tui::{
    buffer::Buffer,
    layout::Rect,
    style::Style,
    text::{Spans, Text},
    widgets::{Block, StatefulWidget, Widget},
};
use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

#[cfg(test)]
mod test {
    use tui::{buffer::Buffer, layout::Rect, widgets::Block};

    use super::*;

    #[test]
    fn test_foo() {
        let width = 80;
        let area = Rect::new(0, 0, width, 4);
        let mut buffer = Buffer::empty(area);
        let view = ScrollView::new(Block::default());
        let mut state = ScrollViewState::default();
        state.add_line("123456".to_owned());
        view.render(area, &mut buffer, &mut state);
        let result = Buffer::with_lines(vec!["1234", "56  "]);
        assert!(buffer.diff(&result).is_empty());
    }
}

type LineIndex = usize;
type WrappedLineIndex = usize;
type GraphemeIndex = usize;

#[derive(Debug)]
struct WrappedLine {
    line_index: LineIndex,
    start_offset: GraphemeIndex,
    length: GraphemeIndex,
}

#[derive(Default)]
pub struct ScrollViewState {
    pub text: Text<'static>,
    // Indexes refer to owned `text`, and is always kept in sync with it
    wrapped_lines: Vec<WrappedLine>,
    scroll: WrappedLineIndex,
    // The area the scroll view was last rendered into
    area: Rect,
}

impl ScrollViewState {
    pub fn is_empty(&self) -> bool {
        self.text.height() == 0
    }

    pub fn content_height(&self) -> usize {
        self.wrapped_lines.len()
    }

    pub fn scroll_to_end(&mut self) {
        self.scroll_by(self.content_height() as isize);
    }

    pub fn scroll_by(&mut self, delta: isize) {
        self.scroll = self.bound_scroll(self.scroll.saturating_add_signed(delta));
    }

    pub fn add_line(&mut self, line: String) {
        let line_index = self.text.lines.len();
        let spans = line.into();
        self.wrap_line(&spans, line_index);
        self.text.lines.push(spans);
    }

    fn bound_scroll(&self, scroll: usize) -> usize {
        scroll.min(
            self.content_height()
                .saturating_sub(self.area.height.into()),
        )
    }

    fn wrap_line(&mut self, line: &Spans<'static>, line_index: LineIndex) {
        Self::wrap_line_(self.area.width, &mut self.wrapped_lines, line, line_index)
    }

    pub fn rewrap_lines(&mut self) {
        self.wrapped_lines.clear();
        for (line_index, line) in self.text.lines.iter().enumerate() {
            Self::wrap_line_(self.area.width, &mut self.wrapped_lines, line, line_index);
        }
        self.scroll = self.bound_scroll(self.scroll);
    }

    fn wrap_line_(
        area_width: u16,
        wrapped_lines: &mut Vec<WrappedLine>,
        line: &Spans<'static>,
        line_index: LineIndex,
    ) {
        let mut iter = line.0.iter().flat_map(|span| span.content.graphemes(true));
        let mut wrap = WordWrapper::new(&mut iter, area_width);
        while let Some((start_offset, length)) = wrap.next_line() {
            wrapped_lines.push(WrappedLine {
                line_index,
                start_offset,
                length,
            });
        }
    }
}

#[derive(Clone, Default)]
pub struct ScrollView<'a> {
    block: Block<'a>,
}

impl<'a> ScrollView<'a> {
    pub fn new(block: Block<'a>) -> ScrollView<'a> {
        Self { block }
    }

    fn maybe_relayout(&self, state: &mut ScrollViewState, area: Rect) {
        let text_area = self.block.inner(area);
        let did_resize = !state.area.eq(&text_area);
        state.area = text_area;
        if !did_resize {
            return;
        }
        state.rewrap_lines();
    }
}

impl<'a> StatefulWidget for ScrollView<'a> {
    type State = ScrollViewState;

    fn render(self, area: Rect, screen_buffer: &mut Buffer, state: &mut Self::State) {
        self.maybe_relayout(state, area);
        let text_area = self.block.inner(area);
        self.block.render(area, screen_buffer);
        let scroll = state.scroll;

        let wrapped_lines_in_view = state
            .wrapped_lines
            .iter()
            .skip(scroll as usize)
            .take(text_area.height as usize);
        let first_line_index = wrapped_lines_in_view
            .clone()
            .next()
            .map_or(0, |first_line| first_line.line_index);
        let mut lines_in_view = state.text.lines.iter().skip(first_line_index);

        let mut graphemes = None;

        for (y, wrapped_line) in wrapped_lines_in_view.enumerate() {
            let mut offset = 0;

            if graphemes.is_none() || wrapped_line.start_offset == 0 {
                graphemes = Some(
                    lines_in_view
                        .next()
                        .unwrap()
                        .0
                        .iter()
                        .flat_map(|span| span.styled_graphemes(Style::default()))
                        .enumerate(),
                );
            }
            while let Some((grapheme_index, grapheme)) = graphemes.as_mut().unwrap().next() {
                if grapheme_index >= wrapped_line.start_offset {
                    screen_buffer.set_stringn(
                        text_area.x + offset as u16,
                        text_area.y + y as u16,
                        grapheme.symbol,
                        grapheme.symbol.width() as usize,
                        grapheme.style,
                    );
                    offset += grapheme.symbol.width();
                }
                if grapheme_index >= wrapped_line.start_offset + wrapped_line.length - 1 {
                    break;
                }
            }
        }
    }
}

/// A state machine that wraps lines on word boundaries.
pub struct WordWrapper<'a, 'b> {
    line_symbols: &'b mut dyn Iterator<Item = &'a str>,
    max_line_width: u16,
    current_line: Vec<&'a str>,
    next_line: Vec<&'a str>,
    returned: bool,
    current_line_start: usize,
    next_line_start: usize,
}

impl<'a, 'b> WordWrapper<'a, 'b> {
    pub fn new(
        line_symbols: &'b mut dyn Iterator<Item = &'a str>,
        max_line_width: u16,
    ) -> WordWrapper<'a, 'b> {
        WordWrapper {
            line_symbols,
            max_line_width,
            current_line: vec![],
            next_line: vec![],
            returned: false,
            current_line_start: 0,
            next_line_start: 0,
        }
    }

    fn next_line(&mut self) -> Option<(usize, usize)> {
        if self.max_line_width == 0 {
            return None;
        }
        std::mem::swap(&mut self.current_line, &mut self.next_line);
        self.current_line_start = self.next_line_start;
        self.next_line.clear();

        let mut current_line_width: u16 = self
            .current_line
            .iter()
            .map(|symbol| symbol.width() as u16)
            .sum();

        let mut symbols_to_last_word_end: usize = 0;
        let mut prev_whitespace = false;
        for symbol in &mut self.line_symbols {
            // Ignore characters wider that the total max width.
            if symbol.width() as u16 > self.max_line_width {
                continue;
            }

            const NBSP: &str = "\u{00a0}";
            let symbol_whitespace = symbol.chars().all(&char::is_whitespace) && symbol != NBSP;

            // Mark the previous symbol as word end.
            if symbol_whitespace && !prev_whitespace {
                symbols_to_last_word_end = self.current_line.len();
            }

            self.current_line.push(symbol);
            current_line_width += symbol.width() as u16;

            if current_line_width > self.max_line_width {
                // If there was no word break in the text, wrap at the end of the line.
                let (truncate_at,) = if symbols_to_last_word_end != 0 {
                    (symbols_to_last_word_end,)
                } else {
                    (self.current_line.len() - 1,)
                };

                // Push the remainder to the next line but strip leading whitespace:
                {
                    let remainder = &self.current_line[truncate_at..];
                    if let Some(remainder_nonwhite) = remainder
                        .iter()
                        .position(|symbol| !symbol.chars().all(&char::is_whitespace))
                    {
                        self.next_line_start += remainder_nonwhite;
                        self.next_line
                            .extend_from_slice(&remainder[remainder_nonwhite..]);
                    }
                }
                self.current_line.truncate(truncate_at);
                // current_line_width = truncated_width;
                self.next_line_start += truncate_at;
                break;
            }

            prev_whitespace = symbol_whitespace;
        }

        // Even if the iterator is exhausted, pass the previous remainder.
        if self.current_line.is_empty() && self.returned {
            None
        } else {
            self.returned = true;
            Some((self.current_line_start, self.current_line.len()))
        }
    }
}