wry/keyboard/src/lib.rs

use {
    jay_input_types::{
        LED_CAPS_LOCK, LED_COMPOSE, LED_KANA, LED_NUM_LOCK, LED_SCROLL_LOCK, Leds,
    },
    jay_utils::{
        event_listener::EventSource,
        numcell::NumCell,
        oserror::{OsError, OsErrorExt, OsErrorExt2},
        syncqueue::SyncQueue,
        vecset::VecSet,
    },
    kbvm::{
        Components,
        lookup::LookupTable,
        state_machine::{self, Event, StateMachine},
        xkb::{
            self, Keymap,
            diagnostic::{Diagnostic, WriteToLog},
            keymap::{Indicator, IndicatorMatcher},
            rmlvo::Group,
        },
    },
    std::{
        cell::{Ref, RefCell},
        io::Write,
        rc::Rc,
    },
    thiserror::Error,
    uapi::{OwnedFd, c},
};

#[derive(Debug, Error)]
pub enum KeyboardError {
    #[error("Could not create a keymap memfd")]
    KeymapMemfd(#[source] OsError),
    #[error("Could not copy the keymap")]
    KeymapCopy(#[source] OsError),
}

#[derive(Debug)]
pub struct KeyboardStateIds {
    next: NumCell<u64>,
}

impl Default for KeyboardStateIds {
    fn default() -> Self {
        Self {
            next: NumCell::new(1),
        }
    }
}

impl KeyboardStateIds {
    pub fn next(&self) -> KeyboardStateId {
        KeyboardStateId(self.next.fetch_add(1))
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct KeyboardStateId(u64);

impl KeyboardStateId {
    pub fn raw(&self) -> u64 {
        self.0
    }

    pub fn from_raw(id: u64) -> Self {
        Self(id)
    }
}

impl std::fmt::Display for KeyboardStateId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.0, f)
    }
}

#[derive(Debug, Error)]
pub enum KbvmError {
    #[error("could not parse the keymap")]
    CouldNotParseKeymap(#[source] Diagnostic),
    #[error("Could not create a keymap memfd")]
    KeymapMemfd(#[source] OsError),
}

pub struct KbvmContext {
    pub ctx: xkb::Context,
}

impl Default for KbvmContext {
    fn default() -> Self {
        let mut ctx = xkb::Context::builder();
        ctx.enable_environment(true);
        Self { ctx: ctx.build() }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct KbvmMapId([u8; 32]);

pub struct KbvmMap {
    pub id: KbvmMapId,
    pub state_machine: StateMachine,
    pub lookup_table: LookupTable,
    pub map: KeymapFd,
    pub xwayland_map: KeymapFd,
    pub has_indicators: bool,
    pub num_lock: Option<IndicatorMatcher>,
    pub caps_lock: Option<IndicatorMatcher>,
    pub scroll_lock: Option<IndicatorMatcher>,
    pub compose: Option<IndicatorMatcher>,
    pub kana: Option<IndicatorMatcher>,
}

pub struct KbvmState {
    pub map: Rc<KbvmMap>,
    pub state: state_machine::State,
    pub kb_state: KeyboardState,
}

pub struct KeyboardState {
    pub id: KeyboardStateId,
    pub map: Rc<KbvmMap>,
    pub pressed_keys: VecSet<u32>,
    pub mods: Components,
    pub leds: Leds,
    pub leds_changed: EventSource<dyn LedsListener>,
}

pub trait LedsListener {
    fn leds(&self, leds: Leds);
}

pub trait DynKeyboardState {
    fn borrow(&self) -> Ref<'_, KeyboardState>;
}

impl DynKeyboardState for RefCell<KeyboardState> {
    fn borrow(&self) -> Ref<'_, KeyboardState> {
        self.borrow()
    }
}

impl DynKeyboardState for RefCell<KbvmState> {
    fn borrow(&self) -> Ref<'_, KeyboardState> {
        Ref::map(self.borrow(), |v| &v.kb_state)
    }
}

impl KeyboardState {
    pub fn apply_event(&mut self, event: Event) -> bool {
        let changed = self.mods.apply_event(event);
        if changed && self.map.has_indicators {
            self.update_leds();
        }
        changed
    }

    pub fn update_leds(&mut self) {
        if !self.map.has_indicators {
            return;
        }
        let mut new = Leds::none();
        macro_rules! map_led {
            ($field:ident, $led:ident) => {
                if let Some(m) = &self.map.$field
                    && m.matches(&self.mods)
                {
                    new |= $led;
                }
            };
        }
        map_led!(num_lock, LED_NUM_LOCK);
        map_led!(caps_lock, LED_CAPS_LOCK);
        map_led!(scroll_lock, LED_SCROLL_LOCK);
        map_led!(compose, LED_COMPOSE);
        map_led!(kana, LED_KANA);
        if new != self.leds {
            self.leds = new;
            for listener in self.leds_changed.iter() {
                listener.leds(new);
            }
        }
    }
}

#[derive(Clone)]
pub struct KeymapFd {
    pub map: Rc<OwnedFd>,
    pub len: usize,
}

impl KeymapFd {
    pub fn create_unprotected_fd(&self) -> Result<Self, KeyboardError> {
        let fd = uapi::memfd_create("shared-keymap", c::MFD_CLOEXEC)
            .map_os_err(KeyboardError::KeymapMemfd)?;
        let target = self.len as c::off_t;
        let mut pos = 0;
        while pos < target {
            let rem = target - pos;
            let res = uapi::sendfile(fd.raw(), self.map.raw(), Some(&mut pos), rem as usize)
                .to_os_error();
            match res {
                Ok(_) | Err(OsError(c::EINTR)) => {}
                Err(e) => return Err(KeyboardError::KeymapCopy(e)),
            }
        }
        Ok(Self {
            map: Rc::new(fd),
            len: self.len,
        })
    }
}

impl KbvmContext {
    pub fn parse_keymap(&self, keymap: &[u8]) -> Result<Rc<KbvmMap>, KbvmError> {
        let map = self
            .ctx
            .keymap_from_bytes(WriteToLog, None, keymap)
            .map_err(KbvmError::CouldNotParseKeymap)?;
        let id = KbvmMapId(*blake3::hash(keymap).as_bytes());
        self.create_keymap(id, map)
    }

    pub fn keymap_from_rmlvo(
        &self,
        rules: Option<&str>,
        model: Option<&str>,
        layout: Option<&str>,
        variant: Option<&str>,
        options: Option<&str>,
    ) -> Result<Rc<KbvmMap>, KbvmError> {
        let mut groups = None::<Vec<_>>;
        if layout.is_some() || variant.is_some() {
            groups = Some(
                Group::from_layouts_and_variants(
                    layout.unwrap_or_default(),
                    variant.unwrap_or_default(),
                )
                .collect(),
            );
        }
        let mut options_vec = None::<Vec<_>>;
        if let Some(options) = options {
            options_vec = Some(options.split(",").collect());
        }
        self.keymap_from_names(rules, model, groups.as_deref(), options_vec.as_deref())
    }

    pub fn keymap_from_names(
        &self,
        rules: Option<&str>,
        model: Option<&str>,
        groups: Option<&[Group<'_>]>,
        options: Option<&[&str]>,
    ) -> Result<Rc<KbvmMap>, KbvmError> {
        let map = self
            .ctx
            .keymap_from_names(WriteToLog, rules, model, groups, options);
        let id = KbvmMapId(*blake3::hash(map.format().to_string().as_bytes()).as_bytes());
        self.create_keymap(id, map)
    }

    fn create_keymap(&self, id: KbvmMapId, map: Keymap) -> Result<Rc<KbvmMap>, KbvmError> {
        let mut has_indicators = false;
        let mut num_lock = None;
        let mut caps_lock = None;
        let mut scroll_lock = None;
        let mut compose = None;
        let mut kana = None;
        for indicator in map.indicators() {
            match indicator.name() {
                Indicator::NUM_LOCK => num_lock = Some(indicator.matcher()),
                Indicator::CAPS_LOCK => caps_lock = Some(indicator.matcher()),
                Indicator::SCROLL_LOCK => scroll_lock = Some(indicator.matcher()),
                Indicator::COMPOSE => compose = Some(indicator.matcher()),
                Indicator::KANA => kana = Some(indicator.matcher()),
                _ => continue,
            }
            has_indicators = true;
        }
        let builder = map.to_builder();
        let (_, xwayland_map) = create_keymap_memfd(&map, true).map_err(KbvmError::KeymapMemfd)?;
        let (_, map) = create_keymap_memfd(&map, false).map_err(KbvmError::KeymapMemfd)?;
        Ok(Rc::new(KbvmMap {
            id,
            state_machine: builder.build_state_machine(),
            map,
            xwayland_map,
            lookup_table: builder.build_lookup_table(),
            has_indicators,
            num_lock,
            caps_lock,
            scroll_lock,
            compose,
            kana,
        }))
    }
}

fn create_keymap_memfd(map: &Keymap, xwayland: bool) -> Result<(String, KeymapFd), OsError> {
    let mut format = map.format();
    if xwayland {
        format = format.lookup_only(true).rename_long_keys(true);
    }
    let str = format!("{}\n", format);
    let mut memfd =
        uapi::memfd_create("keymap", c::MFD_CLOEXEC | c::MFD_ALLOW_SEALING).to_os_error()?;
    memfd.write_all(str.as_bytes())?;
    memfd.write_all(&[0])?;
    uapi::lseek(memfd.raw(), 0, c::SEEK_SET).to_os_error()?;
    uapi::fcntl_add_seals(
        memfd.raw(),
        c::F_SEAL_SEAL | c::F_SEAL_GROW | c::F_SEAL_SHRINK | c::F_SEAL_WRITE,
    )
    .to_os_error()?;
    let fd = KeymapFd {
        map: Rc::new(memfd),
        len: str.len() + 1,
    };
    Ok((str, fd))
}

impl KbvmMap {
    pub fn state(self: &Rc<Self>, id: KeyboardStateId) -> KbvmState {
        KbvmState {
            map: self.clone(),
            state: self.state_machine.create_state(),
            kb_state: KeyboardState {
                id,
                map: self.clone(),
                pressed_keys: Default::default(),
                mods: Default::default(),
                leds: Default::default(),
                leds_changed: Default::default(),
            },
        }
    }
}

impl KbvmState {
    pub fn apply_events(&mut self, events: &SyncQueue<Event>) {
        let state = &mut self.kb_state;
        while let Some(event) = events.pop() {
            state.apply_event(event);
            match event {
                Event::KeyDown(kc) => {
                    state.pressed_keys.insert(kc.to_evdev());
                }
                Event::KeyUp(kc) => {
                    state.pressed_keys.remove(&kc.to_evdev());
                }
                _ => {}
            }
        }
    }
}