path: root/src/keys.rs

use std::ops::RangeInclusive;

use crate::{config::KEYBINDS, error::*, WM};
use xcb::{
    x::{
        GetKeyboardMapping, GetKeyboardMappingReply, GetModifierMapping, GrabKey, GrabMode,
        KeyPressEvent, Mapping, MappingNotifyEvent, ModMask, UngrabKey, GRAB_ANY,
    },
    Connection,
};
use xkeysym::{KeyCode, Keysym, RawKeyCode};

impl WM<'_> {
    /// Dispatch the given keypress event according to [`self::KEYBINDS`]
    pub fn handle_key_press(&mut self, e: KeyPressEvent) -> Result<()> {
        let Some(sym) = self.keyboard_state.keycode_to_keysym(e.detail().into(), 0) else {
            return Ok(()); // probably not bound
        };

        KEYBINDS.dispatch(
            self,
            sym,
            ModMask::from_bits_truncate(e.state().bits())
                .difference(self.keyboard_state.numlock_mask() | ModMask::LOCK),
        );

        Ok(())
    }

    /// Update our keyboard info when the mapping changes.
    pub fn handle_mapping_notify(&mut self, e: MappingNotifyEvent) -> Result<()> {
        if e.request() == Mapping::Keyboard {
            self.grab_keys()?;
        }

        Ok(())
    }

    /// Refresh our keyboard info, and ensure that we get events for bound keys.
    pub fn grab_keys(&mut self) -> Result<()> {
        // Refresh keyboard state
        self.keyboard_state = KeyboardInfo::new_with(self.conn)?;

        // Ungrab all keys
        self.conn.send_request(&UngrabKey {
            key: GRAB_ANY,
            grab_window: self.root,
            modifiers: ModMask::ANY,
        });

        // Bind all of the keycodes which have keysyms we see in our binds.
        for (keycode, keysym) in self.keyboard_state.iter_keycodes_keysyms() {
            for bind in KEYBINDS.binds() {
                if bind.key == keysym {
                    // grab key with any combination of modifiers
                    for modmask in [
                        ModMask::empty(),
                        ModMask::LOCK,
                        self.keyboard_state.numlock_mask(),
                        self.keyboard_state.numlock_mask() | ModMask::LOCK,
                    ] {
                        self.conn.send_request(&GrabKey {
                            grab_window: self.root,
                            key: keycode.raw() as u8,
                            modifiers: bind.modifiers | modmask,
                            owner_events: true,
                            pointer_mode: GrabMode::Async,
                            keyboard_mode: GrabMode::Async,
                        });
                    }
                }
            }
        }

        // Ensure all requests succeeded
        self.conn.flush()?;

        Ok(())
    }
}

/// Cached information about our keyboard layout.
pub struct KeyboardInfo {
    /// The range of keycodes used
    keycodes: RangeInclusive<RawKeyCode>,

    /// The ModMask corresponding to NumLock.
    /// This varies sometimes, and we need to know to ignore it.
    numlock_mask: ModMask,

    /// The mapping from keycodes to (multiple) key symbols
    mapping: GetKeyboardMappingReply,
}

impl KeyboardInfo {
    /// Query information about the keyboard layout from the given connection.
    pub fn new_with(conn: &Connection) -> Result<Self> {
        let min_keycode = conn.get_setup().min_keycode();
        let max_keycode = conn.get_setup().max_keycode();

        let mapping = conn.wait_for_reply(conn.send_request(&GetKeyboardMapping {
            first_keycode: min_keycode,
            count: max_keycode - min_keycode + 1,
        }))?;

        let mut this = Self {
            keycodes: min_keycode as RawKeyCode..=max_keycode as RawKeyCode,
            numlock_mask: ModMask::empty(),
            mapping,
        };

        let Some(numlock_keycode) = this.keysym_to_keycode(Keysym::Num_Lock) else {
            // No numlock button, so no modmask for numlock
            return Ok(this);
        };
        let mod_map = conn.wait_for_reply(conn.send_request(&GetModifierMapping {}))?;
        let keypermod = mod_map.keycodes().len() / 8;
        for i in 0..8 {
            for j in 0..keypermod {
                if mod_map.keycodes()[i * keypermod + j] as u32 == numlock_keycode.raw() {
                    this.numlock_mask =
                        ModMask::from_bits(1 << i).expect("x11 returned unrecognised modifier");
                }
            }
        }

        Ok(this)
    }

    /// Get the modifier mask being used for numlock
    pub fn numlock_mask(&self) -> ModMask {
        self.numlock_mask
    }

    /// Iterate over all keycodes and their bound keysyms.
    /// This is likely to contain duplicate pairs.
    pub fn iter_keycodes_keysyms(&self) -> impl Iterator<Item = (KeyCode, Keysym)> + '_ {
        (0..self.mapping.keysyms_per_keycode())
            .flat_map(|shift| self.keycodes.clone().map(move |keycode| (shift, keycode)))
            .flat_map(|(shift, keycode)| -> Option<_> {
                Some((
                    keycode.into(),
                    self.keycode_to_keysym(keycode.into(), shift)?,
                ))
            })
    }

    /// Lookup the first keycode which has the given keysym in any column
    pub(crate) fn keysym_to_keycode(&self, keysym: Keysym) -> Option<KeyCode> {
        for shift in 0..self.mapping.keysyms_per_keycode() {
            for keycode in self.keycodes.clone() {
                if self.mapping.keysyms()[(keycode as usize - *self.keycodes.start() as usize)
                    * self.mapping.keysyms_per_keycode() as usize
                    + shift as usize]
                    == keysym.raw()
                {
                    return Some(keycode.into());
                }
            }
        }

        None
    }

    /// Lookup the keysym in the given column for the given keycode
    pub fn keycode_to_keysym(&self, keycode: KeyCode, col: u8) -> Option<Keysym> {
        xkeysym::keysym(
            keycode,
            col,
            (*self.keycodes.start()).into(),
            self.mapping.keysyms_per_keycode(),
            self.mapping.keysyms(),
        )
    }
}

/// A key bound to some action
pub struct Keybind {
    pub modifiers: ModMask,
    pub key: Keysym,
    pub action: &'static dyn Fn(&mut WM<'_>),
}

/// A set of keybinds. Currently, there is only one instance of this defined statically: [`crate::config::KEYBINDS`].
pub struct Keybinds(pub &'static [Keybind]);

impl Keybinds {
    /// Get an iterator over all bound keys
    pub fn binds(&self) -> impl Iterator<Item = &Keybind> {
        self.0.iter()
    }

    /// Attempt to run the action for the matching keybind, if it exists.
    pub fn dispatch(&self, wm: &mut WM<'_>, key: Keysym, modifiers: ModMask) {
        if let Some(bind) = self
            .binds()
            .find(|b| b.key == key && modifiers.contains(b.modifiers))
        {
            (bind.action)(wm);
        }
    }
}

impl std::fmt::Debug for Keybind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Keybind")
            .field("modifiers", &self.modifiers)
            .field("key", &self.key)
            .finish()
    }
}