yggdrasil/userspace/term/src/state.rs

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 {
    rows: Vec<GridRow>,
    width: usize,
    height: usize,
}

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

enum EscapeState {
    Normal,
    Escape,
    Csi,
}

pub struct State {
    pub buffer: Buffer,

    esc_state: EscapeState,
    esc_args: Vec<u32>,

    pub cursor: Cursor,
    #[allow(unused)]
    saved_cursor: Option<Cursor>,

    pub default_attributes: CellAttributes,
    pub attributes: CellAttributes,
}

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],
            width,
            height,
        }
    }

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

    pub fn dirty_rows(&mut self) -> impl Iterator<Item = (usize, &mut GridRow)> {
        self.rows.iter_mut().enumerate().filter_map(|(i, row)| {
            if row.dirty {
                row.dirty = false;
                Some((i, row))
            } else {
                None
            }
        })
    }

    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) {
        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) {
        self.rows[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),

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

            cursor: Cursor { row: 0, col: 0 },
            saved_cursor: None,

            default_attributes,
            attributes: default_attributes,
        }
    }

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

        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: u8) -> 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
            }
            b'\x1B' => {
                self.esc_state = EscapeState::Escape;
                self.esc_args.clear();
                self.esc_args.push(0);
                return false;
            }
            b'\r' => {
                self.buffer.rows[self.cursor.row].dirty = true;
                self.cursor.col = 0;
                true
            }
            b'\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.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, c: u8) -> bool {
        let redraw = match c {
            // Move back one character
            b'D' => {
                if self.cursor.col > 0 {
                    self.buffer.set_row_dirty(self.cursor.row);
                    self.cursor.col -= 1;
                    true
                } else {
                    false
                }
            }
            // Character attributes
            b'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
            b'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
            }
            // Clear rows/columns/screen
            b'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,
            },
            b'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, c: u8) -> bool {
        match c {
            b'0'..=b'9' => {
                let arg = self.esc_args.last_mut().unwrap();
                *arg *= 10;
                *arg += (c - b'0') as u32;
                false
            }
            b';' => {
                self.esc_args.push(0);
                false
            }
            _ => self.handle_ctlseq(c),
        }
    }

    pub fn handle_shell_output(&mut self, ch: u8) -> bool {
        match self.esc_state {
            EscapeState::Normal => self.putc_normal(ch),
            EscapeState::Escape => match ch {
                b'[' => {
                    self.esc_state = EscapeState::Csi;
                    false
                }
                _ => {
                    self.esc_state = EscapeState::Normal;
                    false
                }
            },
            EscapeState::Csi => self.handle_ctlseq_byte(ch),
        }
    }
}