acid/firmware2/src/keymap.rs

use core::alloc::Layout;
use core::convert::TryFrom;
use core::fmt::Debug;
use core::slice;

use alloc::boxed::Box;
use alloc::collections::btree_map::BTreeMap;
use alloc::string::String;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::channel::Channel;
use embassy_time::Instant;
use esp_alloc::MemoryCapability;
use log::{debug, info, warn};
use rmk::descriptor::KeyboardReport;
use rmk::futures::FutureExt;
use rmk::hid::Report;
use rmk::types::action::{Action, KeyAction};
use rmk::{a, k, layer};
use slint::platform::Key;
use xkbcommon::xkb::{self, FeedResult, KeyDirection, Keysym, Status};

use crate::matrix::{MATRIX_COLS, MATRIX_ROWS};
use crate::vial::CustomKeycodes;
use crate::{KEYBOARD_REPORT_PROXY, PSRAM_ALLOCATOR};

pub const NUM_LAYER: usize = 1;

pub static KEY_MESSAGE_CHANNEL: Channel<CriticalSectionRawMutex, KeyMessage, 16> = Channel::new();

pub struct KeyMessage {
    pub keysym: Keysym,
    pub string: Option<String>,
    pub direction: KeyDirection,
}

impl Debug for KeyMessage {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("KeyMessage")
            .field("keysym", &self.keysym)
            .field("string", &self.string)
            .field_with("direction", |f| match self.direction {
                KeyDirection::Down => f.write_str("Down"),
                KeyDirection::Up => f.write_str("Up"),
            })
            .finish()
    }
}

#[rustfmt::skip]
pub const fn get_default_keymap() -> [[[KeyAction; MATRIX_COLS]; MATRIX_ROWS]; NUM_LAYER] {
    [
        layer!([
            [k!(Escape), k!(Kc1), k!(Kc2), k!(Kc3), k!(Kc4), k!(Kc5), k!(Kc6), k!(Kc7), k!(Kc8), k!(Kc9), k!(Kc0), k!(Backspace)],
            [k!(Tab), k!(Q), k!(W), k!(E), k!(R), k!(T), k!(Z), k!(U), k!(I), k!(O), k!(P), k!(Delete)],
            [k!(LCtrl), k!(A), k!(S), k!(D), k!(F), k!(G), k!(H), k!(J), k!(K), k!(L), k!(Comma), k!(Enter)],
            [k!(LShift), k!(Y), k!(X), k!(C), k!(V), k!(B), k!(N), k!(M), a!(No), a!(No), k!(Up), KeyAction::Single(Action::User(CustomKeycodes::FOCUS_LCD as u8))],
            [a!(No), a!(No), k!(LGui), k!(LAlt), KeyAction::Single(Action::TriLayerLower), k!(Space), k!(Space), KeyAction::Single(Action::TriLayerLower), k!(RAlt), k!(Left), k!(Down), k!(Right)]
            // [a!(No), a!(No), k!(LGui), k!(LAlt), k!(TriLayerLower), k!(Space), k!(Space), k!(TriLayerLower), k!(RAlt), k!(Left), k!(Down), k!(Right)]
        ])
    ]
}

pub fn create_hid_report_interceptor() -> impl Future<Output = ()> {
    const KEYMAP_STRING: &str = include_str!("../keymaps/cz.xkb");
    const COMPOSE_MAP_STRING: &str = include_str!("../compose/cs_CZ Compose.txt");
    const COMPOSE_MAP_LOCALE: &str = "cs_CZ.UTF-8";

    let keymap_string_buffer = unsafe {
        let allocation = PSRAM_ALLOCATOR.alloc_caps(
            MemoryCapability::External.into(),
            Layout::from_size_align(KEYMAP_STRING.len(), 32).unwrap(),
        );
        let slice = str::from_utf8_unchecked_mut(slice::from_raw_parts_mut(
            allocation,
            KEYMAP_STRING.len(),
        ));

        slice
            .as_bytes_mut()
            .copy_from_slice(KEYMAP_STRING.as_bytes());

        Box::from_raw_in(slice as *mut str, &PSRAM_ALLOCATOR)
    };
    let context = xkb::Context::new(xkb::CONTEXT_NO_FLAGS);
    info!("Loading XKB keymap...");
    let instant_start = Instant::now();
    let keymap = xkb::Keymap::new_from_string(
        &context,
        keymap_string_buffer,
        xkb::KEYMAP_FORMAT_TEXT_V1,
        xkb::KEYMAP_COMPILE_NO_FLAGS,
    )
    .unwrap();
    let duration = Instant::now().duration_since(instant_start);
    info!(
        "XKB keymap loaded successfully! Took {seconds}.{millis:03} seconds.",
        seconds = duration.as_secs(),
        millis = duration.as_millis() % 1_000
    );
    info!("Loading XKB compose map...");
    let instant_start = Instant::now();
    let compose_table = xkb::compose::Table::new_from_buffer(
        &context,
        COMPOSE_MAP_STRING,
        COMPOSE_MAP_LOCALE,
        xkb::compose::FORMAT_TEXT_V1,
        xkb::compose::COMPILE_NO_FLAGS,
    )
    .unwrap();
    let duration = Instant::now().duration_since(instant_start);
    info!(
        "XKB compose map loaded successfully! Took {seconds}.{millis:03} seconds.",
        seconds = duration.as_secs(),
        millis = duration.as_millis() % 1_000
    );
    let mut state = xkb::State::new(&keymap);
    let mut previous_state = KeyboardReport::default();
    let mut compose_state = xkb::compose::State::new(&compose_table, xkb::compose::STATE_NO_FLAGS);
    // TODO: Use a stack-allocated map instead
    // This is a map from the basic keysyms (not a composed ones) to the string that should be produced.
    let mut pressed_keys_to_strings = BTreeMap::<Keysym, Option<String>>::new();

    async move {
        loop {
            let report = KEYBOARD_REPORT_PROXY.receive().await;

            if let Report::KeyboardReport(report) = &report {
                const KEYCODES_LEN_MODIFIERS: usize = 8;
                const KEYCODES_LEN_REGULAR: usize = 6;
                const KEYCODES_LEN: usize = KEYCODES_LEN_MODIFIERS + KEYCODES_LEN_REGULAR;

                let mut pressed_keys = rmk::heapless::Vec::<u8, KEYCODES_LEN>::new();
                let mut released_keys = rmk::heapless::Vec::<u8, KEYCODES_LEN>::new();

                let pressed_mods_bits = !previous_state.modifier & report.modifier;
                let released_mods_bits = previous_state.modifier & !report.modifier;

                for index in 0..KEYCODES_LEN_MODIFIERS {
                    const USB_HID_LEFT_CTRL: u8 = 0xE0;
                    let mod_bit = 1_u8 << index;
                    let mod_keycode = USB_HID_LEFT_CTRL + index as u8;

                    if pressed_mods_bits & mod_bit != 0 {
                        pressed_keys.push(mod_keycode).unwrap();
                    }

                    if released_mods_bits & mod_bit != 0 {
                        released_keys.push(mod_keycode).unwrap();
                    }
                }

                // TODO: This currently depends on pressed keys not changing position in the array.
                // Should be made independent of that.
                for (&keycode_old, &keycode_new) in
                    core::iter::zip(&previous_state.keycodes, &report.keycodes)
                {
                    if keycode_old == keycode_new {
                        continue;
                    }

                    if keycode_old != 0 {
                        released_keys.push(keycode_old).unwrap();
                    }

                    if keycode_new != 0 {
                        pressed_keys.push(keycode_new).unwrap();
                    }
                }

                previous_state = report.clone();

                for keycode in released_keys {
                    debug!("Release: 0x{:02x} ({})", keycode, keycode);
                    let keycode_xkb = into_xkb_keycode(keycode);
                    state.update_key(keycode_xkb, KeyDirection::Up);
                    let keysym = state.key_get_one_sym(keycode_xkb);
                    let string = pressed_keys_to_strings.remove(&keysym).unwrap_or_else(|| {
                        warn!("Could not determine the string of a released key: keysym={keysym:?} pressed_keys_to_strings={pressed_keys_to_strings:?}");
                        None
                    });
                    KEY_MESSAGE_CHANNEL
                        .send(KeyMessage {
                            keysym,
                            string,
                            direction: KeyDirection::Up,
                        })
                        .await;
                }

                for keycode in pressed_keys {
                    debug!("Pressed: 0x{:02x} ({})", keycode, keycode);
                    let keycode_xkb = into_xkb_keycode(keycode);
                    let sym = state.key_get_one_sym(keycode_xkb);

                    let result: Option<(Keysym, Keysym, Option<String>)> = match compose_state
                        .feed(sym)
                    {
                        FeedResult::Ignored => {
                            let string = state.key_get_utf8(keycode_xkb);
                            Some((sym, sym, Some(string)))
                        }
                        FeedResult::Accepted => {
                            let status = compose_state.status();
                            debug!("Compose status: {status:?}");
                            match status {
                                Status::Nothing => {
                                    let string = state.key_get_utf8(keycode_xkb);
                                    Some((sym, sym, Some(string)))
                                }
                                Status::Composing => None,
                                Status::Composed => {
                                    let composed_sym = compose_state.keysym().unwrap_or_default();
                                    let string = compose_state.utf8().unwrap_or_default();
                                    Some((sym, composed_sym, Some(string)))
                                }
                                Status::Cancelled => None,
                            }
                        }
                    };

                    if let Some((basic_keysym, composed_keysym, string)) = result {
                        // Change `Some("")` into `None`.
                        let string = string.filter(|string| !string.is_empty());

                        info!(
                            "Basic keysym: {basic_keysym:?}, composed keysym: {composed_keysym:?}, string: {:?}",
                            string.as_ref()
                        );

                        pressed_keys_to_strings.insert(basic_keysym, string.clone());
                        KEY_MESSAGE_CHANNEL
                            .send(KeyMessage {
                                keysym: composed_keysym,
                                string,
                                direction: KeyDirection::Down,
                            })
                            .await;
                        state.update_key(keycode_xkb, KeyDirection::Down);
                    }
                }
            }
        }
    }
}

fn into_xkb_keycode(rmk_keycode: u8) -> xkb::Keycode {
    // https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/drivers/hid/hid-input.c?id=refs/tags/v6.18#n27
    const UNK: u8 = 240;
    #[rustfmt::skip]
    const HID_KEYBOARD: [u8; 256] = [
          0,   0,   0,   0,  30,  48,  46,  32,  18,  33,  34,  35,  23,  36,  37,  38,
         50,  49,  24,  25,  16,  19,  31,  20,  22,  47,  17,  45,  21,  44,   2,   3,
          4,   5,   6,   7,   8,   9,  10,  11,  28,   1,  14,  15,  57,  12,  13,  26,
         27,  43,  43,  39,  40,  41,  51,  52,  53,  58,  59,  60,  61,  62,  63,  64,
         65,  66,  67,  68,  87,  88,  99,  70, 119, 110, 102, 104, 111, 107, 109, 106,
        105, 108, 103,  69,  98,  55,  74,  78,  96,  79,  80,  81,  75,  76,  77,  71,
         72,  73,  82,  83,  86, 127, 116, 117, 183, 184, 185, 186, 187, 188, 189, 190,
        191, 192, 193, 194, 134, 138, 130, 132, 128, 129, 131, 137, 133, 135, 136, 113,
        115, 114, UNK, UNK, UNK, 121, UNK,  89,  93, 124,  92,  94,  95, UNK, UNK, UNK,
        122, 123,  90,  91,  85, UNK, UNK, UNK, UNK, UNK, UNK, UNK, 111, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, 179, 180, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
        UNK, UNK, UNK, UNK, UNK, UNK, UNK, UNK, 111, UNK, UNK, UNK, UNK, UNK, UNK, UNK,
         29,  42,  56, 125,  97,  54, 100, 126, 164, 166, 165, 163, 161, 115, 114, 113,
        150, 158, 159, 128, 136, 177, 178, 176, 142, 152, 173, 140, UNK, UNK, UNK, UNK,
    ];
    // https://cgit.freedesktop.org/xorg/driver/xf86-input-evdev/tree/src/evdev.c#n73
    const MIN_KEYCODE: u8 = 8;

    // TODO: The combination of these two operations should be precomputed
    // in a const expr into a single look-up table.
    xkb::Keycode::new((HID_KEYBOARD[rmk_keycode as usize] + MIN_KEYCODE) as u32)
}

pub trait TryFromKeysym: Sized {
    fn try_from_keysym(k: Keysym) -> Option<Self>;
}

macro_rules! declare_consts_for_special_keys {
    ($($_char:literal # $name:ident # $($_qt:ident)|* # $($_winit:ident $(($_pos:ident))?)|*    # $($xkb:ident)|*;)*) => {
        impl TryFromKeysym for Key {
            fn try_from_keysym(k: Keysym) -> Option<Self> {
                match k {
                    $(
                        $(Keysym::$xkb => Some(Key::$name),)*
                    )*
                    _ => None
                }
            }
        }
    };
}

i_slint_common::for_each_special_keys!(declare_consts_for_special_keys);