use set_keymap and get_keymap to proper handle keymap

Cargo.toml

@@ -74,6 +74,7 @@ aliasable = { version = "0.1.3", optional = true }
 atomic_float = "1.1.0"
 sha2 = "0.10.9"
 tracy-client = { version = "0.18.4", default-features = false, optional = true }
+memmap2 = "0.9.9"
 
 [dev-dependencies]
 clap = { version = "4", features = ["derive"] }

src/input/keyboard/mod.rs

@@ -203,6 +203,35 @@ impl fmt::Debug for Xkb {
 // same thread
 unsafe impl Send for Xkb {}
 
+#[derive(Clone)]
+/// A keymap that is bound to a specific keyboard handler
+pub struct Keymap {
+    /// XKB context used to create this keymap
+    context: xkb::Context,
+    /// The XKB keymap
+    inner: xkb::Keymap,
+}
+
+impl Keymap {
+    /// Get current keymap
+    pub fn keymap(&self) -> &xkb::Keymap {
+        &self.inner
+    }
+    /// Get current context
+    pub fn context(&self) -> &xkb::Context {
+        &self.context
+    }
+}
+
+impl fmt::Debug for Keymap {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Keymap")
+            .field("inner", &self.inner.get_raw_ptr())
+            .field("context", &self.context.get_raw_ptr())
+            .finish()
+    }
+}
+
 pub(crate) struct KbdInternal<D: SeatHandler> {
     pub(crate) focus: Option<(<D as SeatHandler>::KeyboardFocus, Serial)>,
     pending_focus: Option<<D as SeatHandler>::KeyboardFocus>,
@@ -767,15 +796,22 @@ impl<D: SeatHandler + 'static> KeyboardHandle<D> {
         true
     }
 
-    fn update_xkb_state(&self, data: &mut D, keymap: xkb::Keymap) {
+    fn update_context(&self, context: &xkb::Context) {
+        let internal = self.arc.internal.lock().unwrap();
+        let mut xkb = internal.xkb.lock().unwrap();
+        xkb.context = context.clone();
+        drop(xkb);
+    }
+
+    fn update_xkb_state(&self, data: &mut D, keymap: &xkb::Keymap) {
         let mut internal = self.arc.internal.lock().unwrap();
 
-        let mut state = xkb::State::new(&keymap);
+        let mut state = xkb::State::new(keymap);
         for key in &internal.pressed_keys {
             state.update_key(*key, xkb::KeyDirection::Down);
         }
 
-        let led_mapping = LedMapping::from_keymap(&keymap);
+        let led_mapping = LedMapping::from_keymap(keymap);
         internal.led_mapping = led_mapping;
         internal.mods_state.update_with(&state);
         let leds_changed = internal.led_state.update_with(&state, &led_mapping);
@@ -794,7 +830,7 @@ impl<D: SeatHandler + 'static> KeyboardHandle<D> {
         };
 
         #[cfg(feature = "wayland_frontend")]
-        self.change_keymap(data, &focus, &keymap, mods);
+        self.change_keymap(data, &focus, keymap, mods);
 
         if leds_changed {
             let led_state = internal.led_state;
@@ -820,7 +856,7 @@ impl<D: SeatHandler + 'static> KeyboardHandle<D> {
             debug!("Loading keymap from string failed");
             Error::BadKeymap
         })?;
-        self.update_xkb_state(data, keymap);
+        self.update_xkb_state(data, &keymap);
         Ok(())
     }
 
@@ -832,7 +868,26 @@ impl<D: SeatHandler + 'static> KeyboardHandle<D> {
                 debug!("Loading keymap from XkbConfig failed");
                 Error::BadKeymap
             })?;
-        self.update_xkb_state(data, keymap);
+        self.update_xkb_state(data, &keymap);
+        Ok(())
+    }
+
+    /// Get the current [`Keymap`] used by the keyboard.
+    pub fn get_keymap(&self) -> Keymap {
+        let internal = self.arc.internal.lock().unwrap();
+        let xkb = internal.xkb.lock().unwrap();
+        Keymap {
+            inner: xkb.keymap.clone(),
+            context: xkb.context.clone(),
+        }
+    }
+
+    /// Change the [`Keymap`] used by the keyboard.
+    ///
+    /// The keymap must have been created from this keyboard handle.
+    pub fn set_keymap(&self, data: &mut D, keymap: &Keymap) -> Result<(), Error> {
+        self.update_xkb_state(data, &keymap.inner);
+        self.update_context(&keymap.context);
         Ok(())
     }
 

src/wayland/virtual_keyboard/virtual_keyboard_handle.rs

@@ -1,21 +1,21 @@
-use std::os::unix::io::OwnedFd;
 use std::{
     fmt,
+    fs::File,
+    os::unix::io::OwnedFd,
     sync::{Arc, Mutex},
 };
 
+use memmap2::MmapOptions;
 use tracing::debug;
 use wayland_protocols_misc::zwp_virtual_keyboard_v1::server::zwp_virtual_keyboard_v1::Error::NoKeymap;
 use wayland_protocols_misc::zwp_virtual_keyboard_v1::server::zwp_virtual_keyboard_v1::{
     self, ZwpVirtualKeyboardV1,
 };
-use wayland_server::{
-    backend::ClientId, protocol::wl_keyboard::KeymapFormat, Client, DataInit, Dispatch, DisplayHandle,
-    Resource,
-};
+use wayland_server::protocol::wl_keyboard::KeymapFormat;
+use wayland_server::{backend::ClientId, Client, DataInit, Dispatch, DisplayHandle, Resource};
 use xkbcommon::xkb;
 
-use crate::input::keyboard::{KeyboardTarget, KeymapFile, ModifiersState};
+use crate::input::keyboard::{Keymap, ModifiersState};
 use crate::wayland::input_method::InputMethodSeat;
 use crate::{
     input::{Seat, SeatHandler},
@@ -33,15 +33,17 @@ pub(crate) struct VirtualKeyboard {
 }
 
 struct VirtualKeyboardState {
-    keymap: KeymapFile,
+    keymap: Keymap,
     mods: ModifiersState,
     state: xkb::State,
+    pressed_keys: Vec<u32>,
+    pressed_keys_internal: Vec<u32>,
 }
 
 impl fmt::Debug for VirtualKeyboardState {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("VirtualKeyboardState")
-            .field("keymap", &self.keymap.id())
+            .field("keymap", &self.keymap)
             .field("mods", &self.mods)
             .field("state", &self.state.get_raw_ptr())
             .finish()
@@ -96,7 +98,7 @@ where
         };
         match request {
             zwp_virtual_keyboard_v1::Request::Keymap { format, fd, size } => {
-                update_keymap(data, format, fd, size as usize);
+                update_keymap(user_data, data, format, fd, size as usize);
             }
 
             zwp_virtual_keyboard_v1::Request::Key { time, key, state } => {
@@ -111,40 +113,50 @@ where
                 };
                 // Ensure virtual keyboard's keymap is active.
                 let keyboard_handle = data.seat.get_keyboard().unwrap();
-
                 if ime_keyboard_grabbed
                     .map(|grab| grab.client().as_ref() == Some(client))
                     .unwrap_or(false)
                 {
+                    let old_modifiers = keyboard_handle.modifier_state();
+                    let old_keymap = keyboard_handle.get_keymap();
+                    let _ = keyboard_handle.set_keymap(user_data, &vk_state.keymap);
                     use wayland_server::protocol::wl_keyboard::KeyState;
                     let mut internal = keyboard_handle.arc.internal.lock().unwrap();
                     let focus = internal.focus.as_mut().map(|(focus, _)| focus);
-                    keyboard_handle.send_keymap(user_data, &focus, &vk_state.keymap, vk_state.mods);
                     if let Some(wl_surface) = focus.and_then(|f| f.wl_surface()) {
                         for_each_focused_kbds(&data.seat, &wl_surface, |kbd| {
                             // This should be wl_keyboard::KeyState, but the protocol does not state
                             // the parameter is an enum.
                             let key_state = if state == 1 {
+                                vk_state.pressed_keys_internal.push(key);
                                 KeyState::Pressed
                             } else {
+                                vk_state.pressed_keys_internal.retain(|&x| x == key);
                                 KeyState::Released
                             };
 
                             kbd.key(SERIAL_COUNTER.next_serial().0, time, key, key_state);
                         });
                     }
+                    drop(internal);
+                    let _ = keyboard_handle.set_keymap(user_data, &old_keymap);
+                    let _ = keyboard_handle.set_modifier_state(old_modifiers);
+                    keyboard_handle.advertise_modifier_state(user_data);
                 } else {
-                    {
-                        let mut internal = keyboard_handle.arc.internal.lock().unwrap();
-                        let focus = internal.focus.as_mut().map(|(focus, _)| focus);
-                        keyboard_handle.send_keymap(user_data, &focus, &vk_state.keymap, vk_state.mods);
-                    }
+                    let old_modifiers = keyboard_handle.modifier_state();
+                    let old_keymap = keyboard_handle.get_keymap();
+                    let _ = keyboard_handle.set_keymap(user_data, &vk_state.keymap);
                     let key_state = if state == 1 {
+                        vk_state.pressed_keys.push(key);
                         KeyState::Pressed
                     } else {
+                        vk_state.pressed_keys.retain(|&x| x == key);
                         KeyState::Released
                     };
-                    user_data.on_keyboard_event((key + 8).into(), key_state, time, keyboard_handle);
+                    user_data.on_keyboard_event((key + 8).into(), key_state, time, keyboard_handle.clone());
+                    let _ = keyboard_handle.set_keymap(user_data, &old_keymap);
+                    let _ = keyboard_handle.set_modifier_state(old_modifiers);
+                    keyboard_handle.advertise_modifier_state(user_data);
                 }
             }
             zwp_virtual_keyboard_v1::Request::Modifiers {
@@ -171,27 +183,25 @@ where
 
                 // Ensure virtual keyboard's keymap is active.
                 let keyboard_handle = data.seat.get_keyboard().unwrap();
-                {
-                    let mut internal = keyboard_handle.arc.internal.lock().unwrap();
-                    let focus = internal.focus.as_mut().map(|(focus, _)| focus);
-                    let keymap_changed =
-                        keyboard_handle.send_keymap(user_data, &focus, &state.keymap, state.mods);
-                    if !keymap_changed {
-                        if let Some(focus) = focus {
-                            focus.modifiers(&data.seat, user_data, state.mods, SERIAL_COUNTER.next_serial());
-                        }
-                    }
-                }
-                if ime_keyboard_grabbed.is_none()
-                    || (ime_keyboard_grabbed.unwrap().client().unwrap() != *client)
+                let old_modifiers = keyboard_handle.modifier_state();
+                let old_keymap = keyboard_handle.get_keymap();
+                let _ = keyboard_handle.set_keymap(user_data, &state.keymap);
+                let _ = keyboard_handle.set_modifier_state(state.mods);
+                keyboard_handle.advertise_modifier_state(user_data);
+                if ime_keyboard_grabbed
+                    .map(|grab| grab.client().as_ref() == Some(client))
+                    .unwrap_or(false)
                 {
                     user_data.on_keyboard_modifiers(
                         mods_depressed,
                         mods_latched,
                         mods_locked,
-                        keyboard_handle,
+                        keyboard_handle.clone(),
                     );
                 }
+                let _ = keyboard_handle.set_keymap(user_data, &old_keymap);
+                let _ = keyboard_handle.set_modifier_state(old_modifiers);
+                keyboard_handle.advertise_modifier_state(user_data);
             }
             zwp_virtual_keyboard_v1::Request::Destroy => {
                 // Nothing to do
@@ -201,55 +211,92 @@ where
     }
 
     fn destroyed(
-        _state: &mut D,
+        state: &mut D,
         _client: ClientId,
         _virtual_keyboard: &ZwpVirtualKeyboardV1,
-        _data: &VirtualKeyboardUserData<D>,
+        data: &VirtualKeyboardUserData<D>,
     ) {
+        release_key(state, data);
+    }
+}
+
+fn release_key<D>(user_data: &mut D, data: &VirtualKeyboardUserData<D>)
+where
+    D: SeatHandler + 'static + VirtualKeyboardHandler,
+    <D as SeatHandler>::KeyboardFocus: WaylandFocus,
+{
+    let mut virtual_data = data.handle.inner.lock().unwrap();
+    let vk_state = match virtual_data.state.as_mut() {
+        Some(vk_state) => vk_state,
+        None => {
+            return;
+        }
+    };
+    let pressed_keys = &mut vk_state.pressed_keys;
+    let keyboard_handle = data.seat.get_keyboard().unwrap();
+    let old_keymap = keyboard_handle.get_keymap();
+    let _ = keyboard_handle.set_keymap(user_data, &vk_state.keymap);
+    for i in pressed_keys.drain(..) {
+        user_data.on_keyboard_event((i + 8).into(), KeyState::Released, 0, keyboard_handle.clone());
+    }
+    let pressed_keys_internal = &mut vk_state.pressed_keys_internal;
+    let mut internal = keyboard_handle.arc.internal.lock().unwrap();
+    let focus = internal.focus.as_mut().map(|(focus, _)| focus);
+    for i in pressed_keys_internal.drain(..) {
+        if let Some(wl_surface) = focus.as_ref().and_then(|f| f.wl_surface()) {
+            for_each_focused_kbds(&data.seat, &wl_surface, |kbd| {
+                kbd.key(SERIAL_COUNTER.next_serial().0, 0, i, KeyState::Released.into());
+            });
+        }
     }
+    drop(internal);
+    let _ = keyboard_handle.set_keymap(user_data, &old_keymap);
+    keyboard_handle.advertise_modifier_state(user_data);
 }
 
 /// Handle the zwp_virtual_keyboard_v1::keymap request.
 ///
 /// The `true` returns when keymap was properly loaded.
-fn update_keymap<D>(data: &VirtualKeyboardUserData<D>, format: u32, fd: OwnedFd, size: usize)
-where
-    D: SeatHandler + 'static,
+fn update_keymap<D>(
+    user_data: &mut D,
+    data: &VirtualKeyboardUserData<D>,
+    format: u32,
+    fd: OwnedFd,
+    size: usize,
+) where
+    D: SeatHandler + 'static + VirtualKeyboardHandler,
+    <D as SeatHandler>::KeyboardFocus: WaylandFocus,
 {
+    release_key(user_data, data);
     // Only libxkbcommon compatible keymaps are supported.
     if format != KeymapFormat::XkbV1 as u32 {
         debug!("Unsupported keymap format: {format:?}");
         return;
     }
 
-    let context = xkb::Context::new(xkb::CONTEXT_NO_FLAGS);
-    // SAFETY: we can map the keymap into the memory.
-    let new_keymap = match unsafe {
-        xkb::Keymap::new_from_fd(
-            &context,
-            fd,
-            size,
-            xkb::KEYMAP_FORMAT_TEXT_V1,
-            xkb::KEYMAP_COMPILE_NO_FLAGS,
-        )
-    } {
-        Ok(Some(new_keymap)) => new_keymap,
-        Ok(None) => {
-            debug!("Invalid libxkbcommon keymap");
-            return;
-        }
-        Err(err) => {
-            debug!("Could not map the keymap: {err:?}");
-            return;
-        }
+    let map = unsafe {
+        MmapOptions::new()
+            .len(size)
+            .map_copy_read_only(&File::from(fd))
+            .unwrap()
     };
-
+    let keymap_string = String::from_utf8_lossy(&map[..]).to_string();
     // Store active virtual keyboard map.
     let mut inner = data.handle.inner.lock().unwrap();
     let mods = inner.state.take().map(|state| state.mods).unwrap_or_default();
+    let keyboard_handle = data.seat.get_keyboard().unwrap();
+    let old_keymap = keyboard_handle.get_keymap().clone();
+    let keymap = {
+        let _ = keyboard_handle.set_keymap_from_string(user_data, keymap_string);
+        keyboard_handle.get_keymap().clone()
+    };
+    let state = xkb::State::new(keymap.keymap());
     inner.state = Some(VirtualKeyboardState {
         mods,
-        keymap: KeymapFile::new(&new_keymap),
-        state: xkb::State::new(&new_keymap),
+        keymap: keymap.clone(),
+        state,
+        pressed_keys: Vec::<u32>::new(),
+        pressed_keys_internal: Vec::<u32>::new(),
     });
+    let _ = keyboard_handle.set_keymap(user_data, &old_keymap);
 }