From c05eab904474599c48eae96f62c178e2f0e2bdaf Mon Sep 17 00:00:00 2001
From: alex-ds13 <145657253+alex-ds13@users.noreply.github.com>
Date: Wed, 29 Jan 2025 15:23:50 +0000
Subject: [PATCH] feat(wm): add monitor_usr_idx_map to wm

This commit creates a new field on the `WindowManager` called
`monitor_usr_idx_map` which contains a map of user intended index for
monitors to their actual monitors' index.

It will be rebuilt on `load_monitor_information` by taking into account
the `display_index_preferences`.
---
 komorebi/src/static_config.rs  |  1 +
 komorebi/src/window_manager.rs |  6 ++++-
 komorebi/src/windows_api.rs    | 40 +++++++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/komorebi/src/static_config.rs b/komorebi/src/static_config.rs
index 4bc766d1..0c5aefcf 100644
--- a/komorebi/src/static_config.rs
+++ b/komorebi/src/static_config.rs
@@ -1178,6 +1178,7 @@ impl StaticConfig {
 
         let mut wm = WindowManager {
             monitors: Ring::default(),
+            monitor_usr_idx_map: HashMap::new(),
             incoming_events: incoming,
             command_listener: listener,
             is_paused: false,
diff --git a/komorebi/src/window_manager.rs b/komorebi/src/window_manager.rs
index d6923aeb..944b1673 100644
--- a/komorebi/src/window_manager.rs
+++ b/komorebi/src/window_manager.rs
@@ -99,6 +99,7 @@ use crate::WORKSPACE_MATCHING_RULES;
 #[derive(Debug)]
 pub struct WindowManager {
     pub monitors: Ring<Monitor>,
+    pub monitor_usr_idx_map: HashMap<usize, usize>,
     pub incoming_events: Receiver<WindowManagerEvent>,
     pub command_listener: UnixListener,
     pub is_paused: bool,
@@ -122,6 +123,7 @@ pub struct WindowManager {
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 pub struct State {
     pub monitors: Ring<Monitor>,
+    pub monitor_usr_idx_map: HashMap<usize, usize>,
     pub is_paused: bool,
     pub resize_delta: i32,
     pub new_window_behaviour: WindowContainerBehaviour,
@@ -283,6 +285,7 @@ impl From<&WindowManager> for State {
     fn from(wm: &WindowManager) -> Self {
         Self {
             monitors: wm.monitors.clone(),
+            monitor_usr_idx_map: wm.monitor_usr_idx_map.clone(),
             is_paused: wm.is_paused,
             work_area_offset: wm.work_area_offset,
             resize_delta: wm.resize_delta,
@@ -343,6 +346,7 @@ impl WindowManager {
 
         Ok(Self {
             monitors: Ring::default(),
+            monitor_usr_idx_map: HashMap::new(),
             incoming_events: incoming,
             command_listener: listener,
             is_paused: false,
@@ -366,7 +370,7 @@ impl WindowManager {
     #[tracing::instrument(skip(self))]
     pub fn init(&mut self) -> Result<()> {
         tracing::info!("initialising");
-        WindowsApi::load_monitor_information(&mut self.monitors)?;
+        WindowsApi::load_monitor_information(self)?;
         WindowsApi::load_workspace_information(&mut self.monitors)
     }
 
diff --git a/komorebi/src/windows_api.rs b/komorebi/src/windows_api.rs
index 56e81215..360a16e5 100644
--- a/komorebi/src/windows_api.rs
+++ b/komorebi/src/windows_api.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::collections::VecDeque;
 use std::convert::TryFrom;
 use std::ffi::c_void;
@@ -143,6 +144,7 @@ use crate::ring::Ring;
 use crate::set_window_position::SetWindowPosition;
 use crate::windows_callbacks;
 use crate::Window;
+use crate::WindowManager;
 use crate::DISPLAY_INDEX_PREFERENCES;
 use crate::MONITOR_INDEX_PREFERENCES;
 
@@ -252,7 +254,10 @@ impl WindowsApi {
             .collect::<Vec<_>>())
     }
 
-    pub fn load_monitor_information(monitors: &mut Ring<Monitor>) -> Result<()> {
+    pub fn load_monitor_information(wm: &mut WindowManager) -> Result<()> {
+        let monitors = &mut wm.monitors;
+        let monitor_usr_idx_map = &mut wm.monitor_usr_idx_map;
+
         'read: for display in win32_display_data::connected_displays_all().flatten() {
             let path = display.device_path.clone();
 
@@ -326,6 +331,39 @@ impl WindowsApi {
             .elements_mut()
             .retain(|m| m.name().ne("PLACEHOLDER"));
 
+        // Rebuild monitor index map
+        *monitor_usr_idx_map = HashMap::new();
+        let mut added_monitor_idxs = Vec::new();
+        for (index, id) in &*DISPLAY_INDEX_PREFERENCES.lock() {
+            if let Some(m_idx) = monitors.elements().iter().position(|m| {
+                m.serial_number_id().as_ref().is_some_and(|sn| sn == id) || m.device_id() == id
+            }) {
+                monitor_usr_idx_map.insert(*index, m_idx);
+                added_monitor_idxs.push(m_idx);
+            }
+        }
+
+        let max_usr_idx = monitors
+            .elements()
+            .len()
+            .max(monitor_usr_idx_map.keys().max().map_or(0, |v| *v));
+
+        let available_usr_idxs = (0..max_usr_idx)
+            .filter(|i| !monitor_usr_idx_map.contains_key(i))
+            .collect::<Vec<_>>();
+
+        let not_added_monitor_idxs = (0..monitors.elements().len())
+            .filter(|i| !added_monitor_idxs.contains(i))
+            .collect::<Vec<_>>();
+
+        for i in not_added_monitor_idxs {
+            if let Some(next_usr_idx) = available_usr_idxs.first() {
+                monitor_usr_idx_map.insert(*next_usr_idx, i);
+            } else if let Some(idx) = monitor_usr_idx_map.keys().max() {
+                monitor_usr_idx_map.insert(*idx, i);
+            }
+        }
+
         Ok(())
     }