yggdrasil/userspace/colors/src/main.rs

#![feature(yggdrasil_os, rustc_private)]

// TODO rewrite and split this into meaningful components

use std::{
    collections::{BTreeMap, HashMap}, env, os::{
        fd::{AsRawFd, RawFd},
        yggdrasil::io::{
            mapping::FileMapping, poll::PollChannel,
            shared_memory::SharedMemory,
        },
    }, process::{Command, ExitCode}
};

use display::Display;
use error::Error;
use input::{InputState, KeyboardInput};
use libcolors::{
    event::{Event, KeyModifiers, WindowEvent, WindowInfo},
    message::{ClientMessage, ServerMessage},
};
use serde_ipc::{PeerAddr, Receiver, Sender};
use yggdrasil_abi::io::{KeyboardKey, KeyboardKeyEvent};

pub mod display;
pub mod error;
pub mod input;

pub struct Window<'a> {
    window_id: u32,
    client_id: PeerAddr,

    surface_mapping: FileMapping<'a>,
    surface_data: &'a [u32],
}

pub struct Frame {
    x: u32,
    y: u32,
    w: u32,
    h: u32,

    dirty: bool,

    window: Option<u32>,
}

pub struct Row {
    frames: Vec<Frame>,
    x: u32,
    y: u32,
    width: u32,
    height: u32,
}

pub struct ServerSender(Sender<ServerMessage>);

pub struct Server<'a, 'd> {
    display: Display<'d>,

    input_state: InputState,

    last_client_id: u32,
    client_map: HashMap<u32, PeerAddr>,

    // Window management
    windows: BTreeMap<u32, Window<'a>>,
    rows: Vec<Row>,
    last_window_id: u32,
    focused_frame: Option<(usize, usize)>,
    // Outer frame
    padding: usize,
    background: u32,

    // Event generators
    poll: PollChannel,
    receiver: Receiver<ClientMessage>,
    input: KeyboardInput,

    // Comms
    sender: ServerSender,
}

impl Row {
    pub fn new(x: u32, y: u32, w: u32, h: u32) -> Self {
        Self {
            frames: vec![],
            x,
            y,
            width: w,
            height: h,
        }
    }

    pub fn balance_frames(&mut self) {
        if self.frames.is_empty() {
            return;
        }

        let spacing = 4;
        let wc = self.frames.len() as u32;

        let w = (self.width - spacing * (wc - 1)) / wc;
        let h = self.height;
        let mut x = self.x;
        let y = self.y;

        for frame in self.frames.iter_mut() {
            frame.dirty = true;
            frame.x = x;
            frame.y = y;
            frame.w = w;
            frame.h = h;

            x += w + spacing;
        }
    }

    pub fn place_frame(&mut self) -> &mut Frame {
        self.frames.push(Frame {
            x: 0,
            y: 0,
            w: 0,
            h: 0,

            dirty: true,
            window: None,
        });

        self.balance_frames();

        self.frames.last_mut().unwrap()
    }

    pub fn remove_frame(&mut self, col: usize) {
        self.frames.remove(col);

        self.balance_frames();
    }
}

impl<'a> Server<'a, '_> {
    pub fn new(framebuffer: &str) -> Result<Self, Error> {
        let mut poll = PollChannel::new()?;

        let mut display = Display::open(framebuffer)?;
        let input = KeyboardInput::open()?;

        let (sender, receiver) = serde_ipc::listen(libcolors::CHANNEL_NAME)?;
        let sender = ServerSender(sender);

        poll.add(input.as_poll_fd())?;
        poll.add(receiver.as_raw_fd())?;

        let background = 0xFFCCCCCC;
        display.fill(background);
        display.flush();

        Ok(Self {
            display,

            input_state: InputState::default(),

            poll,
            receiver,
            input,
            sender,

            padding: 4,
            background,

            last_client_id: 0,
            client_map: HashMap::new(),

            windows: BTreeMap::new(),
            rows: vec![],
            last_window_id: 1,
            focused_frame: None,
        })
    }

    fn create_window(
        &mut self,
        client_id: &PeerAddr,
    ) -> Result<(WindowInfo, RawFd), Error> {
        if self.rows.is_empty() {
            self.rows.push(Row::new(
                self.padding as _,
                self.padding as _,
                (self.display.width() - self.padding * 2) as _,
                (self.display.height() - self.padding * 2) as _,
            ));
        }

        // Create a frame
        let row = self.rows.last_mut().unwrap();
        let frame = row.place_frame();

        // Create the actual window
        let window_id = self.last_window_id;
        self.last_window_id += 1;

        let mapping_size = self.display.width() * self.display.height() * 4;
        let surface_shm = SharedMemory::new(mapping_size).unwrap();
        let fd = surface_shm.as_raw_fd();
        let mut surface_mapping = surface_shm.into_mapping().unwrap();

        let surface_data = unsafe {
            std::slice::from_raw_parts_mut(
                surface_mapping.as_mut_ptr() as *mut u32,
                (frame.w * frame.h) as usize,
            )
        };

        frame.window = Some(window_id);

        let window = Window {
            window_id,
            client_id: client_id.clone(),
            surface_mapping,
            surface_data,
        };

        self.windows.insert(window_id, window);

        let info = WindowInfo {
            window_id,
            surface_stride: self.display.width() * 4,
            surface_mapping_size: mapping_size,
            width: frame.w,
            height: frame.h,
        };

        self.display.fill(self.background);
        self.set_focused_window(window_id)?;
        self.flush_dirty_frames();

        Ok((info, fd))
    }

    fn remove_window(&mut self, window_id: u32) {
        // Find the window
        if !self.windows.contains_key(&window_id) {
            return;
        }

        // TODO this is ugly
        let mut res = None;
        for (i, row) in self.rows.iter().enumerate() {
            let j = row
                .frames
                .iter()
                .position(|f| f.window.map(|w| w == window_id).unwrap_or(false));

            if let Some(j) = j {
                res = Some((i, j));
            }
        }

        // Remove the frame
        if let Some((row, col)) = res {
            self.rows[row].remove_frame(col);
            self.display.fill(self.background);
            self.flush_dirty_frames();
        }

        self.windows.remove(&window_id);

        if self.focused_frame == res {
            self.focused_frame = None;

            let new_focus = if let Some((row, col)) = res {
                // Focus some other frame in the same row
                if let Some(f_row) = self.rows.get(row) {
                    let row_len = f_row.frames.len();

                    if col == 0 && row_len != 0 {
                        Some((row, 1))
                    } else if col > 0 {
                        Some((row, col - 1))
                    } else {
                        // Empty row
                        None
                    }
                } else {
                    // No row exists
                    None
                }
            } else {
                // No frames?
                None
            };

            self.set_focused_frame(new_focus);
        }
    }

    fn handle_keyboard_event(&mut self, event: KeyboardKeyEvent) -> Result<(), Error> {
        let (key, state) = event.split();

        self.input_state.update(key, state);

        if state {
            let input = self.input_state.make_input(key);

            // Non-window keys
            #[allow(clippy::single_match)]
            match (input.modifiers, input.key) {
                (KeyModifiers::ALT, KeyboardKey::Enter) => {
                    // TODO do something with spawned child
                    Command::new("/bin/term").spawn().ok();
                    return Ok(());
                }
                _ => (),
            }

            // Window keys
            if let Some((row, col)) = self.focused_frame {
                let row_len = self.rows[row].frames.len();

                match (input.modifiers, input.key) {
                    (KeyModifiers::ALT, KeyboardKey::Char(b'l')) => {
                        if col + 1 < row_len {
                            self.set_focused_frame(Some((row, col + 1)));
                        } else {
                            self.set_focused_frame(Some((row, 0)));
                        }

                        return Ok(());
                    }
                    (KeyModifiers::ALT, KeyboardKey::Char(b'h')) => {
                        if col > 0 {
                            self.set_focused_frame(Some((row, col - 1)));
                        } else if row_len != 0 {
                            self.set_focused_frame(Some((row, row_len - 1)));
                        }

                        return Ok(());
                    }
                    _ => (),
                }
            }

            if let Some((_, window)) = self.get_focused_window() {
                // Deliver event to the window
                self.sender
                    .send_event(
                        Event::WindowEvent(window.window_id, WindowEvent::KeyInput(input)),
                        &window.client_id,
                    )
                    .ok();
            } else {
                self.focused_frame = None;
            }
        }

        Ok(())
    }

    fn get_window(&self, window_id: u32) -> Option<(&Frame, &Window<'a>)> {
        let window = self.windows.get(&window_id)?;
        for row in self.rows.iter() {
            if let Some(f) = row
                .frames
                .iter()
                .find(|f| f.window.map(|w| w == window_id).unwrap_or(false))
            {
                return Some((f, window));
            }
        }
        // TODO Orphaned frame/window?
        None
    }

    fn get_focused_window(&self) -> Option<(&Frame, &Window<'a>)> {
        let (row, col) = self.focused_frame?;

        let frame = &self.rows[row].frames[col];
        let window = frame.window.and_then(|w| self.windows.get(&w))?;

        Some((frame, window))
    }

    fn set_focused_frame(&mut self, focus: Option<(usize, usize)>) {
        if self.focused_frame == focus {
            return;
        }

        if let Some((_, old_window)) = self.get_focused_window() {
            self.sender
                .send_event(
                    Event::WindowEvent(old_window.window_id, WindowEvent::FocusChanged(false)),
                    &old_window.client_id,
                )
                .ok();
        }

        self.focused_frame = focus;

        if let Some((row, col)) = focus {
            let Some(f_row) = self.rows.get(row) else {
                return;
            };
            let Some(frame) = f_row.frames.get(col) else {
                return;
            };
            let Some(window) = frame.window.and_then(|w| self.windows.get(&w)) else {
                return;
            };

            self.sender
                .send_event(
                    Event::WindowEvent(window.window_id, WindowEvent::FocusChanged(true)),
                    &window.client_id,
                )
                .ok();
        }
    }

    fn set_focused_window(&mut self, window_id: u32) -> Result<(), Error> {
        // TODO this is ugly
        let mut res = None;
        for (i, row) in self.rows.iter().enumerate() {
            let j = row
                .frames
                .iter()
                .position(|f| f.window.map(|w| w == window_id).unwrap_or(false));

            if let Some(j) = j {
                res = Some((i, j));
            }
        }

        self.set_focused_frame(res);

        Ok(())
    }

    fn flush_dirty_frames(&mut self) {
        for row in self.rows.iter() {
            for frame in row.frames.iter() {
                if !frame.dirty {
                    continue;
                }

                let Some(window) = frame.window.and_then(|w| self.windows.get_mut(&w)) else {
                    // TODO handle orphaned frame
                    continue;
                };

                let new_surface_data = unsafe {
                    std::slice::from_raw_parts_mut(
                        window.surface_mapping.as_mut_ptr() as *mut u32,
                        (frame.w * frame.h) as usize,
                    )
                };

                window.surface_data = new_surface_data;

                self.sender
                    .send_event(
                        Event::WindowEvent(
                            window.window_id,
                            WindowEvent::Resized {
                                width: frame.w,
                                height: frame.h,
                            },
                        ),
                        &window.client_id,
                    )
                    .ok();
            }
        }
    }

    fn handle_client_message(
        &mut self,
        client_id: PeerAddr,
        message: ClientMessage,
    ) -> Result<(), Error> {
        match message {
            ClientMessage::ClientHello => {
                debug_trace!("{:?}: ClientHello", client_id);
                // Echo the ID back
                self.last_client_id += 1;
                let id = self.last_client_id;
                self.client_map.insert(id, client_id.clone());
                self.sender
                    .send_event(Event::ServerHello(id), &client_id)
            }
            ClientMessage::CreateWindow => {
                debug_trace!("{:?}: CreateWindow", client_id);
                let (info, shm_fd) = self.create_window(&client_id)?;
                let window_id = info.window_id;

                self.sender
                    .send_event(Event::NewWindowInfo(info), &client_id)?;
                self.sender.send_fd(&shm_fd, &client_id)?;
                self.sender.send_event(
                    Event::WindowEvent(window_id, WindowEvent::RedrawRequested),
                    &client_id,
                )?;

                Ok(())
            }
            ClientMessage::BlitWindow {
                window_id,
                x,
                y,
                w,
                h,
            } => {
                if let Some((frame, window)) = self.get_window(window_id) {
                    let x = x.min(frame.w);
                    let y = y.min(frame.h);
                    let w = w.min(frame.w - x);
                    let h = h.min(frame.h - y);

                    if w == 0 || h == 0 {
                        // Invalid rectangle, skip it
                        return Ok(());
                    }

                    self.display.blit_buffer(
                        window.surface_data,
                        display::Point(frame.x as _, frame.y as _),
                        display::Point(x as _, y as _),
                        w as _,
                        h as _,
                        frame.w as usize,
                    );
                }
                Ok(())
            }
            ClientMessage::DestroyWindow(window_id) => {
                debug_trace!("{:?}: DestroyWindow {}", client_id, window_id);
                self.remove_window(window_id);
                Ok(())
            }
        }
    }

    fn run_inner(mut self) -> Result<(), Error> {
        loop {
            match self.poll.wait(None, true)? {
                Some((fd, Ok(_))) if fd == self.input.as_poll_fd() => {
                    let event = self.input.read_event()?;
                    self.handle_keyboard_event(event)?;
                }
                Some((fd, Ok(_))) if fd == self.receiver.as_raw_fd() => {
                    let (client_id, message) = self.receiver.receive_message()?;
                    self.handle_client_message(client_id, message)?;
                }
                Some((_, Ok(_))) => {
                    todo!()
                }
                Some((_, Err(error))) => {
                    return Err(Error::from(error));
                }
                None => (),
            }
        }
    }

    pub fn run(self) -> ExitCode {
        match self.run_inner() {
            Ok(_) => ExitCode::SUCCESS,
            Err(error) => {
                debug_trace!("colors server finished with an error: {}", error);
                ExitCode::FAILURE
            }
        }
    }
}

impl ServerSender {
    pub fn send_event(&self, event: Event, client_id: &PeerAddr) -> Result<(), Error> {
        self.0
            .send_to(&ServerMessage::Event(event), client_id)
            .map_err(Error::from)
    }

    pub fn send_fd<F: AsRawFd>(&self, file: &F, client_id: &PeerAddr) -> Result<(), Error> {
        self.0.send_file_to(file, client_id).map_err(Error::from)
    }
}

fn main() -> ExitCode {
    let args = env::args().skip(1).collect::<Vec<_>>();
    let framebuffer = args.get(0).map_or("/dev/fb0", |s| s.as_str());

    let server = Server::new(framebuffer).unwrap();

    server.run()
}