From b84b42b21c44cb8ffba481f812473feb752c18cc Mon Sep 17 00:00:00 2001
From: Kirill Chibisov <contact@kchibisov.com>
Date: Thu, 2 Nov 2023 00:10:22 +0400
Subject: [PATCH] Add 'input.keyboard.track_layout' config option

Track layout for each window individually when user requested `Window`
tracking mode, keeping the default as global.
---
 resources/default-config.kdl |  7 ++++++
 src/config.rs                | 13 +++++++++++
 src/niri.rs                  | 43 ++++++++++++++++++++++++++++++++----
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/resources/default-config.kdl b/resources/default-config.kdl
index 312036feb..e3cb345e2 100644
--- a/resources/default-config.kdl
+++ b/resources/default-config.kdl
@@ -16,6 +16,13 @@ input {
         // Delay is in milliseconds before the repeat starts. Rate is in characters per second.
         // repeat-delay 600
         // repeat-rate 25
+
+        // How compositor should track the layout.
+        //
+        // - "global" - layout change is global for all windows.
+        // - "window" - layout is tracked for each window individually.
+        //
+        //track-layout "global"
     }
 
     // Next sections include libinput settings.
diff --git a/src/config.rs b/src/config.rs
index ac1917bc9..b7510d8dd 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -62,6 +62,8 @@ pub struct Keyboard {
     pub repeat_delay: u16,
     #[knuffel(child, unwrap(argument), default = 25)]
     pub repeat_rate: u8,
+    #[knuffel(child, unwrap(argument), default)]
+    pub track_layout: TrackLayout,
 }
 
 #[derive(knuffel::Decode, Debug, Default, PartialEq, Eq)]
@@ -78,6 +80,15 @@ pub struct Xkb {
     pub options: Option<String>,
 }
 
+#[derive(knuffel::DecodeScalar, Debug, Default, PartialEq, Eq)]
+pub enum TrackLayout {
+    /// The layout change is global.
+    #[default]
+    Global,
+    /// The layout change is window local.
+    Window,
+}
+
 // FIXME: Add the rest of the settings.
 #[derive(knuffel::Decode, Debug, Default, PartialEq)]
 pub struct Touchpad {
@@ -514,6 +525,7 @@ mod tests {
                 keyboard {
                     repeat-delay 600
                     repeat-rate 25
+                    track-layout "window"
                     xkb {
                         layout "us,ru"
                         options "grp:win_space_toggle"
@@ -587,6 +599,7 @@ mod tests {
                         },
                         repeat_delay: 600,
                         repeat_rate: 25,
+                        track_layout: TrackLayout::Window,
                     },
                     touchpad: Touchpad {
                         tap: true,
diff --git a/src/niri.rs b/src/niri.rs
index eb1cd243b..275358b83 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1,4 +1,4 @@
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::collections::{HashMap, HashSet};
 use std::ffi::OsString;
 use std::path::PathBuf;
@@ -29,7 +29,7 @@ use smithay::desktop::utils::{
     under_from_surface_tree, update_surface_primary_scanout_output, OutputPresentationFeedback,
 };
 use smithay::desktop::{layer_map_for_output, PopupManager, Space, Window, WindowSurfaceType};
-use smithay::input::keyboard::XkbConfig;
+use smithay::input::keyboard::{Layout as KeyboardLayout, XkbConfig, XkbContextHandler};
 use smithay::input::pointer::{CursorIcon, CursorImageAttributes, CursorImageStatus, MotionEvent};
 use smithay::input::{Seat, SeatState};
 use smithay::output::Output;
@@ -76,7 +76,7 @@ use smithay::wayland::text_input::TextInputManagerState;
 use smithay::wayland::virtual_keyboard::VirtualKeyboardManagerState;
 
 use crate::backend::{Backend, RenderResult, Tty, Winit};
-use crate::config::Config;
+use crate::config::{Config, TrackLayout};
 use crate::cursor::{CursorManager, CursorTextureCache, RenderCursor, XCursor};
 #[cfg(feature = "dbus")]
 use crate::dbus::gnome_shell_screenshot::{NiriToScreenshot, ScreenshotToNiri};
@@ -382,8 +382,43 @@ impl State {
         };
 
         let keyboard = self.niri.seat.get_keyboard().unwrap();
-        if keyboard.current_focus() != focus {
+        let current_focus = keyboard.current_focus();
+        if current_focus != focus {
+            if self.niri.config.borrow().input.keyboard.track_layout == TrackLayout::Window {
+                let current_layout =
+                    keyboard.with_kkb_state(self, |context| context.active_layout());
+
+                let mut new_layout = current_layout;
+                // Store the currently active layout for the surface.
+                if let Some(current_focus) = current_focus.as_ref() {
+                    with_states(current_focus, |data| {
+                        let cell = data
+                            .data_map
+                            .get_or_insert::<Cell<KeyboardLayout>, _>(Cell::default);
+                        cell.set(current_layout);
+                    });
+                }
+
+                if let Some(focus) = focus.as_ref() {
+                    new_layout = with_states(focus, |data| {
+                        let cell = data.data_map.get_or_insert::<Cell<KeyboardLayout>, _>(|| {
+                            // The default layout is effectively the first layout in the
+                            // keymap, so use it for new windows.
+                            Cell::new(KeyboardLayout::default())
+                        });
+                        cell.get()
+                    });
+                }
+                if new_layout != current_layout && focus.is_some() {
+                    keyboard.set_focus(self, None, SERIAL_COUNTER.next_serial());
+                    keyboard.with_kkb_state(self, |mut context| {
+                        context.set_layout(new_layout);
+                    });
+                }
+            }
+
             keyboard.set_focus(self, focus, SERIAL_COUNTER.next_serial());
+
             // FIXME: can be more granular.
             self.niri.queue_redraw_all();
         }