yggdrasil/userspace/term/src/state.rs

use std::collections::VecDeque;

use crate::attr::{CellAttributes, Color};

#[derive(Clone, Copy, Debug)]
pub struct GridCell {
    pub char: char,
    pub attrs: CellAttributes,
}

#[derive(Clone)]
pub struct GridRow {
    cols: Vec<GridCell>,
    dirty: bool,
}

pub struct Buffer {
    scrollback: VecDeque<GridRow>,
    rows: Vec<GridRow>,
    width: usize,
    height: usize,
    scrollback_limit: usize,
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Cursor {
    pub row: usize,
    pub col: usize,
}

struct CsiState {
    byte: u8,
}

enum EscapeState {
    Normal,
    Escape,
    Csi(CsiState),
}

#[derive(Default)]
struct Utf8Decoder {
    buffer: [u8; 4],
    len: usize,
}

pub struct State {
    pub buffer: Buffer,

    utf8_decode: Utf8Decoder,
    esc_state: EscapeState,
    esc_args: Vec<u32>,

    scroll: usize,
    pub cursor: Cursor,
    pub cursor_visible: bool,
    pub alternate: bool,
    #[allow(unused)]
    saved_cursor: Option<Cursor>,

    pub default_attributes: CellAttributes,
    pub attributes: CellAttributes,
}

impl Utf8Decoder {
    pub fn push(&mut self, byte: u8) -> Option<char> {
        self.buffer[self.len] = byte;
        self.len += 1;
        if let Ok(str) = std::str::from_utf8(&self.buffer[..self.len]) {
            let ch = str.chars().next().unwrap();
            self.len = 0;
            Some(ch)
        } else {
            if self.len == 4 {
                // Got 4 bytes and could not decode a single character
                self.len = 0;
                Some('\u{25A1}')
            } else {
                None
            }
        }
    }
}

impl GridCell {
    pub fn new(char: char, attrs: CellAttributes) -> Self {
        Self { char, attrs }
    }

    pub fn empty(bg: Color) -> Self {
        Self {
            char: '\0',
            attrs: CellAttributes {
                fg: Color::Black,
                bg,
                bright: false,
            },
        }
    }
}

impl GridRow {
    pub fn new(width: usize, bg: Color) -> Self {
        Self {
            cols: vec![GridCell::empty(bg); width],
            dirty: true,
        }
    }

    pub fn set_cell(&mut self, col: usize, char: char, attrs: CellAttributes) {
        self.cols[col] = GridCell::new(char, attrs);
        self.dirty = true;
    }

    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    pub fn clear_dirty(&mut self) {
        self.dirty = false;
    }

    pub fn cells(&self) -> impl Iterator<Item = &GridCell> {
        self.cols.iter()
    }

    fn clear(&mut self, bg: Color) {
        self.cols.fill(GridCell::empty(bg));
        self.dirty = true;
    }

    fn erase_to_right(&mut self, start: usize, bg: Color) {
        self.cols[start..].fill(GridCell::empty(bg));
        self.dirty = true;
    }

    fn resize(&mut self, width: usize, bg: Color) {
        self.cols.resize(width, GridCell::empty(bg));
        self.dirty = true;
    }
}

impl Buffer {
    pub fn new(width: usize, height: usize, bg: Color) -> Self {
        Self {
            rows: vec![GridRow::new(width, bg); height],
            scrollback: VecDeque::new(),
            scrollback_limit: 1024,
            width,
            height,
        }
    }

    pub fn clear(&mut self, bg: Color) {
        self.rows.fill(GridRow::new(self.width, bg));
    }

    pub fn iter_rows_mut<F: FnMut(usize, &mut GridRow)>(&mut self, scroll: usize, mut handler: F) {
        let scroll = scroll.min(self.scrollback.len());
        let non_scroll = self.height.saturating_sub(scroll);
        let scroll_end = self.height - non_scroll;

        for y in 0..scroll_end {
            let i = scroll - 1 - y;
            let row = &mut self.scrollback[i];
            handler(y, row);
        }
        for i in 0..non_scroll {
            let y = scroll + i;
            let row = &mut self.rows[i];
            handler(y, row);
        }
    }

    pub fn visible_rows_mut<F: FnMut(usize, &mut GridRow)>(
        &mut self,
        scroll: usize,
        mut handler: F,
    ) {
        self.iter_rows_mut(scroll, |y, row| {
            if row.dirty {
                row.dirty = false;
                handler(y, row);
            }
        });
    }

    pub fn resize(&mut self, width: usize, height: usize, bg: Color) {
        self.rows.resize(height, GridRow::new(width, bg));
        for row in self.rows.iter_mut() {
            row.resize(width, bg);
        }

        self.width = width;
        self.height = height;
    }

    pub fn set_cell(&mut self, cur: Cursor, cell: GridCell) {
        self.rows[cur.row].cols[cur.col] = cell;
        self.rows[cur.row].dirty = true;
    }

    pub fn scroll_once(&mut self, bg: Color) {
        self.scrollback.push_front(self.rows[0].clone());
        if self.scrollback.len() >= self.scrollback_limit {
            self.scrollback.pop_back();
        }

        for i in 1..self.height {
            self.rows[i - 1] = self.rows[i].clone();
            self.rows[i - 1].dirty = true;
        }
        self.rows[self.height - 1] = GridRow::new(self.width, bg);
    }

    pub fn erase_row(&mut self, row: usize, bg: Color) {
        self.rows[row].clear(bg);
    }

    pub fn set_row_dirty(&mut self, row: usize) {
        if row >= self.rows.len() {
            return;
        }
        self.rows[row].dirty = true;
    }

    pub fn invalidate_rows(&mut self, scroll: usize) {
        self.iter_rows_mut(scroll, |_, row| row.dirty = true);
    }
}

impl State {
    pub fn new(width: usize, height: usize) -> Self {
        let default_attributes = CellAttributes {
            fg: Color::White,
            bg: Color::Black,
            bright: false,
        };

        Self {
            buffer: Buffer::new(width, height, default_attributes.bg),
            utf8_decode: Utf8Decoder::default(),

            esc_args: Vec::new(),
            esc_state: EscapeState::Normal,

            cursor: Cursor { row: 0, col: 0 },
            cursor_visible: true,
            alternate: false,
            scroll: 0,
            saved_cursor: None,

            default_attributes,
            attributes: default_attributes,
        }
    }

    pub fn clear(&mut self) {
        self.buffer.clear(self.attributes.bg);
    }

    pub fn resize(&mut self, width: usize, height: usize) {
        self.buffer
            .resize(width, height, self.default_attributes.bg);
        self.scroll = 0;

        if self.cursor.row >= height {
            self.cursor.row = height - 1;
        }

        if self.cursor.col >= width {
            self.cursor.col = width - 1;
        }
    }

    fn putc_normal(&mut self, ch: char) -> bool {
        let mut redraw = match ch {
            // c if c >= 127 => {
            //     let attr = CellAttributes {
            //         fg: Color::Black,
            //         bg: Color::Red,
            //         bright: false,
            //     };
            //     self.buffer.set_cell(self.cursor, GridCell::new('?', attr));
            //     self.cursor.col += 1;
            //     true
            // }
            '\x1B' => {
                self.esc_state = EscapeState::Escape;
                self.esc_args.clear();
                self.esc_args.push(0);
                return false;
            }
            '\r' => {
                self.buffer.rows[self.cursor.row].dirty = true;
                self.cursor.col = 0;
                true
            }
            '\n' => {
                self.buffer.rows[self.cursor.row].dirty = true;
                self.cursor.row += 1;
                self.cursor.col = 0;
                true
            }
            _ => {
                self.buffer
                    .set_cell(self.cursor, GridCell::new(ch as char, self.attributes));
                self.cursor.col += 1;
                true
            }
        };

        if self.alternate {
            self.cursor.col = self.cursor.col.min(self.buffer.width - 1);
            self.cursor.row = self.cursor.row.min(self.buffer.height - 1);
        }

        if self.cursor.col >= self.buffer.width {
            if self.cursor.row < self.buffer.height {
                self.buffer.rows[self.cursor.row].dirty = true;
            }

            self.cursor.row += 1;
            self.cursor.col = 0;
            redraw = true;
        }

        while self.cursor.row >= self.buffer.height {
            self.buffer.scroll_once(self.default_attributes.bg);
            self.cursor.row -= 1;
            redraw = true;
        }

        redraw
    }

    fn handle_ctlseq(&mut self, byte: u8, c: char) -> bool {
        let redraw = match c {
            'h' if byte == b'?' => match self.esc_args.get(0).copied().unwrap_or(0) {
                25 => {
                    // Cursor visible
                    self.cursor_visible = true;
                    true
                }
                1049 => {
                    // Enter alternate mode
                    self.alternate = true;
                    true
                }
                _ => false,
            },
            'l' if byte == b'?' => match self.esc_args.get(0).copied().unwrap_or(0) {
                25 => {
                    // Cursor not visible
                    self.cursor_visible = false;
                    true
                }
                1049 => {
                    // Leave alternate mode
                    self.alternate = false;
                    true
                }
                _ => false,
            },

            // Move back one character
            'D' => {
                if self.cursor.col > 0 {
                    self.buffer.set_row_dirty(self.cursor.row);
                    self.cursor.col -= 1;
                    true
                } else {
                    false
                }
            }
            // Character attributes
            'm' => match self.esc_args[0] {
                0 => {
                    self.attributes = self.default_attributes;
                    false
                }
                1 => {
                    self.attributes.bright = true;
                    false
                }
                30..=39 => {
                    let vt_color = self.esc_args[0] % 10;
                    if vt_color == 9 {
                        self.attributes.fg = Color::Black;
                    } else {
                        self.attributes.fg = Color::from_esc(vt_color);
                    }
                    false
                }
                40..=49 => {
                    let vt_color = self.esc_args[0] % 10;
                    if vt_color == 9 {
                        self.attributes.bg = Color::Black;
                    } else {
                        self.attributes.bg = Color::from_esc(vt_color);
                    }
                    false
                }
                _ => false,
            },
            // Move cursor to position
            'f' => {
                let row = self.esc_args[0].clamp(1, self.buffer.height as u32) - 1;
                let col = self.esc_args[1].clamp(1, self.buffer.width as u32) - 1;

                self.buffer.set_row_dirty(self.cursor.row);
                self.cursor = Cursor {
                    row: row as _,
                    col: col as _,
                };

                true
            }
            // Move cursor to home position (0; 0)
            'H' => {
                self.buffer.set_row_dirty(self.cursor.row);
                self.cursor = Cursor { row: 0, col: 0 };
                true
            }
            // Clear rows/columns/screen
            'J' => match self.esc_args[0] {
                // Erase lines down
                0 => false,
                // Erase lines up
                1 => false,
                // Erase all
                2 => {
                    self.buffer.clear(self.attributes.bg);
                    true
                }
                _ => false,
            },
            'K' => match self.esc_args[0] {
                // Erase to right
                0 => {
                    self.buffer.rows[self.cursor.row]
                        .erase_to_right(self.cursor.col, self.attributes.bg);
                    true
                }
                // Erase All
                2 => {
                    self.buffer.erase_row(self.cursor.row, self.attributes.bg);
                    true
                }
                _ => false,
            },
            _ => false,
        };

        self.esc_state = EscapeState::Normal;
        redraw
    }

    fn handle_ctlseq_byte(&mut self, byte: u8, c: char) -> bool {
        match c {
            '?' if byte == 0 => {
                self.esc_state = EscapeState::Csi(CsiState { byte: b'?' });
                false
            }
            c if let Some(digit) = c.to_digit(10) => {
                let arg = self.esc_args.last_mut().unwrap();
                *arg *= 10;
                *arg += digit;
                false
            }
            ';' => {
                self.esc_args.push(0);
                false
            }
            _ => self.handle_ctlseq(byte, c),
        }
    }

    pub fn handle_shell_output(&mut self, ch: u8) -> bool {
        if let Some(ch) = self.utf8_decode.push(ch) {
            match self.esc_state {
                EscapeState::Normal => self.putc_normal(ch),
                EscapeState::Escape => match ch {
                    '[' => {
                        self.esc_state = EscapeState::Csi(CsiState { byte: 0 });
                        false
                    }
                    _ => {
                        self.esc_state = EscapeState::Normal;
                        false
                    }
                },
                EscapeState::Csi(CsiState { byte }) => self.handle_ctlseq_byte(byte, ch),
            }
        } else {
            false
        }
    }

    pub fn invalidate_current_viewport(&mut self) {
        self.buffer.invalidate_rows(self.scroll);
    }

    pub fn adjust_scroll(&mut self) -> usize {
        if self.scroll > self.buffer.scrollback.len() {
            self.scroll = self.buffer.scrollback.len();
        }

        self.scroll
    }

    pub fn scroll_up(&mut self) -> bool {
        let max = self.buffer.scrollback.len();
        if max == 0 {
            self.scroll = 0;
            return true;
        }
        let amount = max.min(8);
        if amount != 0 {
            self.scroll += amount;
            self.invalidate_current_viewport();
            true
        } else {
            false
        }
    }

    pub fn scroll_down(&mut self) -> bool {
        let amount = self.scroll.min(8);
        if amount != 0 {
            self.scroll -= amount;
            self.invalidate_current_viewport();
            true
        } else {
            false
        }
    }

    pub fn scroll_home(&mut self) -> bool {
        if self.scroll != self.buffer.scrollback.len() {
            self.scroll = self.buffer.scrollback.len();
            self.invalidate_current_viewport();
            true
        } else {
            false
        }
    }

    pub fn scroll_end(&mut self) -> bool {
        if self.scroll != 0 {
            self.scroll = 0;
            self.invalidate_current_viewport();
            true
        } else {
            false
        }
    }
}