yggdrasil/src/device/display/console.rs

//! Console device interfaces
use core::mem::size_of;

use abi::{error::Error, primitive_enum};
use alloc::vec::Vec;
use bitflags::bitflags;

use crate::{
    debug::DebugSink,
    mem::{
        phys::{self, PageUsage},
        ConvertAddress,
    },
    sync::IrqSafeSpinlock,
    task::tasklet::TaskFlow,
    util::StaticVector,
};

const CONSOLE_ROW_LEN: usize = 128;

const DEFAULT_FG_COLOR: ColorAttribute = ColorAttribute::White;
const DEFAULT_BG_COLOR: ColorAttribute = ColorAttribute::Blue;

primitive_enum! {
    #[doc = "Color attribute of a console character"]
    pub enum ColorAttribute: u8 {
        #[doc = "..."]
        Black = 0,
        #[doc = "..."]
        Red = 1,
        #[doc = "..."]
        Green = 2,
        #[doc = "..."]
        Yellow = 3,
        #[doc = "..."]
        Blue = 4,
        #[doc = "..."]
        Magenta = 5,
        #[doc = "..."]
        Cyan = 6,
        #[doc = "..."]
        White = 7,
    }
}

bitflags! {
    #[derive(Clone, Copy)]
    pub struct Attributes: u8 {
        const BOLD = 1 << 0;
    }
}

impl ColorAttribute {
    pub fn from_vt100(val: u8) -> Self {
        match val {
            0..=7 => Self::try_from(val).unwrap(),
            _ => todo!(),
        }
    }

    /// Converts the attribute to RGBA representation
    pub fn as_rgba(&self, bold: bool) -> u32 {
        let color = match self {
            Self::Black => 0x000000,
            Self::Red => 0x7F0000,
            Self::Green => 0x007F00,
            Self::Yellow => 0x7F7F00,
            Self::Blue => 0x00007F,
            Self::Magenta => 0x7F007F,
            Self::Cyan => 0x007F7F,
            Self::White => 0x7F7F7F,
        };
        if bold {
            color * 2
        } else {
            color
        }
    }
}

/// Represents a single character with its attributes
#[derive(Clone, Copy)]
#[repr(C)]
pub struct ConsoleChar {
    attributes: u16,
    char: u8,
    _pad: u8,
}

/// Represents a single line in the console buffer
#[derive(Clone, Copy)]
pub struct ConsoleRow {
    dirty: u8,
    chars: [ConsoleChar; CONSOLE_ROW_LEN],
}

/// Buffer that contains text rows of the console with their attributes + tracks dirty rows which
/// need to be flushed to the display
pub struct ConsoleBuffer {
    rows: &'static mut [ConsoleRow],
    width: u32,
    height: u32,
}

pub enum EscapeState {
    Normal,
    Escape,
    Csi,
}

/// Common state for console output devices
pub struct ConsoleState {
    /// Current cursor row
    pub cursor_row: u32,
    /// Current cursor column
    pub cursor_col: u32,
    /// Current foreground color
    pub fg_color: ColorAttribute,
    pub bg_color: ColorAttribute,
    pub attributes: Attributes,

    esc_args: StaticVector<u32, 4>,
    esc_state: EscapeState,

    /// Row buffer
    pub buffer: ConsoleBuffer,
}

/// Helper type to iterate over dirty rows in the buffer
pub struct RowIter<'a> {
    buffer: &'a mut ConsoleBuffer,
    index: u32,
}

/// Interface to implement buffered console semantics on an abstract console output device
pub trait DisplayConsole {
    /// Returns the state lock
    fn state(&self) -> &IrqSafeSpinlock<ConsoleState>;
    /// Allocates a new buffer for the console according to the display's size
    fn reallocate_buffer(&self) -> Result<(), Error>;

    /// Flushes the data from console buffer to the display
    fn flush(&self, state: &mut ConsoleState);

    /// Writes characters to the backing buffer + handles special control sequences
    fn write_char(&self, c: u8) {
        let mut state = self.state().lock();
        if state.putc(c) {
            self.flush(&mut state);
        }
    }
}

impl ConsoleChar {
    /// Empty character
    pub const BLANK: Self = Self {
        attributes: 0,
        char: 0,
        _pad: 0,
    };

    /// Constructs a console character from a char and its attributes
    #[inline]
    pub fn from_parts(
        char: u8,
        fg: ColorAttribute,
        bg: ColorAttribute,
        attributes: Attributes,
    ) -> Self {
        let attributes = ((attributes.bits() as u16) << 8)
            | ((u8::from(bg) as u16) << 4)
            | (u8::from(fg) as u16);
        Self {
            attributes,
            char,
            _pad: 0,
        }
    }

    #[inline]
    pub fn attributes(self) -> (ColorAttribute, ColorAttribute, Attributes) {
        let fg = ColorAttribute::try_from((self.attributes & 0xF) as u8).unwrap();
        let bg = ColorAttribute::try_from(((self.attributes >> 4) & 0xF) as u8).unwrap();
        let attributes = Attributes::from_bits((self.attributes >> 8) as u8).unwrap();

        (fg, bg, attributes)
    }

    /// Returns the character data of this [ConsoleChar]
    #[inline]
    pub const fn character(self) -> u8 {
        self.char
    }
}

impl<'a> RowIter<'a> {
    /// Returns the next dirty row
    pub fn next_dirty(&mut self) -> Option<(u32, &[ConsoleChar])> {
        loop {
            if self.index == self.buffer.height {
                return None;
            }

            if !self.buffer.rows[self.index as usize].clear_dirty() {
                self.index += 1;
                continue;
            }

            let row_index = self.index;
            let row = &self.buffer.rows[self.index as usize];

            self.index += 1;

            return Some((row_index, &row.chars));
        }
    }
}

impl ConsoleRow {
    /// Constructs a row filled with blank characters
    pub const fn zeroed() -> Self {
        Self {
            dirty: 1,
            chars: [ConsoleChar {
                attributes: ((DEFAULT_BG_COLOR as u8) as u16) << 4,
                char: b' ',
                _pad: 0,
            }; CONSOLE_ROW_LEN],
        }
    }

    /// Returns `true` if the row's dirty flag is set
    #[inline]
    pub const fn is_dirty(&self) -> bool {
        self.dirty != 0
    }

    /// Clears "dirty" flag for the row
    #[inline]
    pub fn clear_dirty(&mut self) -> bool {
        let old = self.dirty;
        self.dirty = 0;
        old == 1
    }

    /// Clears the console row with blank characters
    pub fn clear(&mut self, bg: ColorAttribute) {
        self.dirty = 1;
        self.chars
            .fill(ConsoleChar::from_parts(b' ', bg, bg, Attributes::empty()));
    }
}

impl ConsoleBuffer {
    /// Constructs a fixed-size console buffer
    pub const fn fixed(rows: &'static mut [ConsoleRow], width: u32, height: u32) -> Self {
        Self {
            rows,
            width,
            height,
        }
    }

    /// Reallocates the internal buffer with a new size
    pub fn reallocate(&mut self, new_height: u32) -> Result<(), Error> {
        // TODO suppress debugging output here
        if new_height <= self.height {
            // Keep using the old buffer
            return Ok(());
        }

        let size = size_of::<ConsoleRow>() * (new_height as usize);
        let page_count = (size + 0xFFF) / 0x1000;
        let pages = phys::alloc_pages_contiguous(page_count, PageUsage::Used)?;

        let data = unsafe {
            core::slice::from_raw_parts_mut(
                pages.virtualize() as *mut ConsoleRow,
                new_height as usize,
            )
        };

        // Copy rows from the old buffer
        data[0..self.height as usize].copy_from_slice(self.rows);
        data[self.height as usize..].fill(ConsoleRow::zeroed());

        self.rows = data;
        self.height = new_height;

        Ok(())
    }

    fn set_char(&mut self, row: u32, col: u32, c: ConsoleChar) {
        self.rows[row as usize].dirty = 1;
        self.rows[row as usize].chars[col as usize] = c;
    }

    /// Returns an iterator over dirty rows, while clearing dirty flag for them
    pub fn flush_rows(&mut self) -> RowIter {
        RowIter {
            buffer: self,
            index: 0,
        }
    }

    fn clear(&mut self, bg: ColorAttribute) {
        for row in self.rows.iter_mut() {
            row.clear(bg);
        }
    }

    fn scroll_once(&mut self, bg: ColorAttribute) {
        self.rows.copy_within(1.., 0);
        self.rows[(self.height - 1) as usize].clear(bg);

        // Mark everything dirty
        self.rows.iter_mut().for_each(|row| {
            row.dirty = 1;
        });
    }
}

impl ConsoleState {
    /// Constructs a new console state with given buffer
    pub fn new(buffer: ConsoleBuffer) -> Self {
        Self {
            cursor_row: 0,
            cursor_col: 0,
            fg_color: DEFAULT_FG_COLOR,
            bg_color: DEFAULT_BG_COLOR,
            attributes: Attributes::empty(),

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

            buffer,
        }
    }

    fn putc_normal(&mut self, c: u8) -> bool {
        let mut flush = false;

        match c {
            c if c >= 127 => {
                self.buffer.set_char(
                    self.cursor_row,
                    self.cursor_col,
                    ConsoleChar::from_parts(
                        b'?',
                        self.fg_color,
                        ColorAttribute::Red,
                        self.attributes,
                    ),
                );

                self.cursor_col += 1;
            }
            b'\x1b' => {
                self.esc_state = EscapeState::Escape;
                return false;
            }
            b'\n' => {
                self.cursor_row += 1;
                self.cursor_col = 0;
                flush = true;
            }
            _ => {
                self.buffer.set_char(
                    self.cursor_row,
                    self.cursor_col,
                    ConsoleChar::from_parts(c, self.fg_color, self.bg_color, self.attributes),
                );

                self.cursor_col += 1;
            }
        }

        if self.cursor_col == self.buffer.width {
            self.cursor_col = 0;
            self.cursor_row += 1;
        }

        if self.cursor_row == self.buffer.height {
            self.buffer.scroll_once(self.bg_color);
            self.cursor_row = self.buffer.height - 1;
            flush = true;
        }

        flush
    }

    fn handle_csi(&mut self, c: u8) -> bool {
        match c {
            // Move back one character
            b'D' => {
                if self.cursor_col > 0 {
                    self.cursor_col = self.cursor_col - 1;
                }
            }
            // Manipulate display attributes
            b'm' => match self.esc_args[0] {
                // Reset
                0 => {
                    self.fg_color = DEFAULT_FG_COLOR;
                    self.bg_color = DEFAULT_BG_COLOR;
                    self.attributes = Attributes::empty();
                }
                // Bold
                1 => {
                    self.attributes |= Attributes::BOLD;
                }
                // Foreground colors
                31..=37 => {
                    let vt_color = self.esc_args[0] % 10;
                    self.fg_color = ColorAttribute::from_vt100(vt_color as u8);
                }
                _ => loop {},
            },
            // Move cursor to position
            b'f' => {
                let row = self.esc_args[1].clamp(1, self.buffer.height) - 1;
                let col = self.esc_args[1].clamp(1, CONSOLE_ROW_LEN as u32) - 1;

                self.cursor_row = row;
                self.cursor_col = col;
            }
            // Clear rows/columns/screen
            b'J' => match self.esc_args[0] {
                // Erase lines down
                0 => loop {},
                // Erase lines up
                1 => loop {},
                // Erase all
                2 => {
                    self.buffer.clear(self.bg_color);
                }
                _ => (),
            },
            _ => loop {},
        }

        self.esc_state = EscapeState::Normal;
        false
    }

    fn handle_csi_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_csi(c),
        }
    }

    fn putc(&mut self, c: u8) -> bool {
        match self.esc_state {
            EscapeState::Normal => self.putc_normal(c),
            EscapeState::Escape => match c {
                b'[' => {
                    self.esc_state = EscapeState::Csi;
                    self.esc_args.clear();
                    self.esc_args.push(0);
                    false
                }
                _ => {
                    self.esc_state = EscapeState::Normal;
                    false
                }
            },
            EscapeState::Csi => self.handle_csi_byte(c),
        }
    }
}

impl DebugSink for dyn DisplayConsole {
    fn putc(&self, c: u8) -> Result<(), Error> {
        self.write_char(c);
        Ok(())
    }

    fn supports_control_sequences(&self) -> bool {
        true
    }
}

static CONSOLES: IrqSafeSpinlock<Vec<&'static dyn DisplayConsole>> =
    IrqSafeSpinlock::new(Vec::new());

/// Adds a console device to a auto-flush list
pub fn add_console_autoflush(console: &'static dyn DisplayConsole) {
    CONSOLES.lock().push(console);
}

/// Periodically flushes data from console buffers onto their displays
pub fn task_update_consoles() -> TaskFlow {
    for console in CONSOLES.lock().iter() {
        let mut state = console.state().lock();
        console.flush(&mut state);
    }
    TaskFlow::Continue
}