virtual_desktop_manager/

quick_switch.rs

use core::ops::Range;
use std::collections::BTreeMap;

#[derive(Default)]
pub struct QuickSwitchMenu {
    /// Menu items that map to specific virtual desktops.
    tray_virtual_desktops_quick: Vec<(nwg::MenuItem, u32)>,

    tray_separators: Vec<nwg::MenuSeparator>,

    /// [`nwg::Menu`] won't be automatically removed from their parent when
    /// dropped so make sure to use [`crate::nwg_ext::remove_menu`] to manually
    /// remove them.
    tray_submenus: Vec<(nwg::Menu, u32)>,

    /// Support shortcut access keys to specific virtual desktop indexes. This
    /// keys should not be numbers since those access keys are already in use by
    /// the quick switch menu.
    pub shortcuts: BTreeMap<String, u32>,
    /// Only show the shortcut menu items in the root of the quick switch menu.
    pub shortcuts_only_in_root: bool,

    /// Used to limit recursion in order to prevent stack overflow, should be
    /// empty when not inside a [`Self::create_quick_switch_menu`] call.
    ///
    /// We could have used a `VecDeque` here to more faithfully remember the
    /// call order but its actually more performant to not do that since the
    /// normal call order is breadth first which means we would store more items
    /// in this queue.
    call_queue: Vec<(nwg::ControlHandle, u32, u32, u32)>,
    /// `true` if we are in a recursive call of
    /// [`Self::create_quick_switch_menu`].
    is_recursive_call: bool,
}
impl QuickSwitchMenu {
    /// Get the virtual desktop index that should be selected when a specific
    /// context menu item is pressed.
    pub fn get_clicked_desktop_index(&self, handle: nwg::ControlHandle) -> Option<usize> {
        if !matches!(handle, nwg::ControlHandle::MenuItem { .. }) {
            return None;
        }
        self.tray_virtual_desktops_quick
            .iter()
            .find(|(d, _)| d.handle == handle)
            .map(|(_, ix)| *ix as usize)
    }
    /// Get the number that would be generated by the keyboard shortcuts needed
    /// to open the specified submenu.
    pub fn get_desktop_index_so_far(&self, submenu_handle: isize) -> Option<usize> {
        self.tray_submenus
            .iter()
            .find(|(d, _)| matches!(d.handle.hmenu(), Some((_, handle)) if handle as isize == submenu_handle))
            .map(|(_, ix)| *ix as usize)
    }
    /// Find first item/menu inside a submenu.
    pub fn first_item_in_submenu(&self, submenu_handle: isize) -> Option<nwg::ControlHandle> {
        // Submenus come before menu items:
        self.tray_submenus
            .iter()
            .find_map(|(d, _)| {
                let (parent, _) = d.handle.hmenu()?;
                (parent as isize == submenu_handle).then_some(d.handle)
            })
            .or_else(|| {
                self.tray_virtual_desktops_quick
                    .iter()
                    .find_map(|(item, _)| {
                        let (parent, _) = item.handle.hmenu_item()?;
                        (parent as isize == submenu_handle).then_some(item.handle)
                    })
            })
    }

    pub fn clear(&mut self) {
        self.tray_virtual_desktops_quick.clear();
        self.tray_separators.clear();
        self.tray_submenus
            .drain(..)
            .rev() // Remove deepest nested submenus first
            .for_each(|(menu, _)| crate::nwg_ext::menu_remove(&menu));
    }

    #[allow(dead_code)]
    pub fn has_submenu(&self, menu_handle: isize) -> bool {
        self.tray_submenus.iter().any(|(menu, _)| {
            menu.handle
                .hmenu()
                .map_or(false, |(_, id)| id as isize == menu_handle)
        })
    }

    /// Extra menu items that don't use numbers as access keys. These allow the
    /// user to give specific keys as shortcuts for frequently visited virtual
    /// desktops.
    fn create_shortcut_items(&mut self, parent: nwg::ControlHandle) {
        // Start with a separator:
        let mut separator = Default::default();
        nwg::MenuSeparator::builder()
            .parent(parent)
            .build(&mut separator)
            .expect("Failed to build separator for quick switch menu's shortcut section");
        self.tray_separators.push(separator);

        for (shortcut, &virtual_desktop_ix) in &self.shortcuts {
            let mut item = Default::default();
            nwg::MenuItem::builder()
                .text(&format!(
                    "Virtual desktop {}\t&{shortcut}",
                    virtual_desktop_ix.saturating_add(1)
                ))
                .parent(parent)
                .build(&mut item)
                .expect("Failed to build \"shortcut\" menu item for quick switch menu");
            self.tray_virtual_desktops_quick
                .push((item, virtual_desktop_ix));
        }
    }
    /// If there are 10 items or less in a submenu then each one could get a
    /// unique access key. This function will only create context menu items if
    /// that is the case.
    fn try_create_leaf_items(&mut self, parent: nwg::ControlHandle, range: Range<u32>) -> bool {
        if range.len() <= 10 {
            // Check if each item can have its own access key:
            let mut last_digit = [0_u32; 10];
            for ix in range.clone() {
                last_digit[(ix % 10) as usize] += 1;
            }

            if last_digit.iter().all(|&x| x <= 1) {
                // All last digits are unique, so can make a leaf menu
                for ix in range {
                    let tens = ix / 10;
                    let ones = ix % 10;
                    let mut item = Default::default();
                    nwg::MenuItem::builder()
                        .text(&format!(
                            "Virtual desktop {}&{ones}",
                            if tens == 0 {
                                String::new()
                            } else {
                                tens.to_string()
                            }
                        ))
                        .parent(parent)
                        .build(&mut item)
                        .expect("Failed to build \"leaf\" menu item for quick switch menu");

                    // Note: we store 0-based index but show 1 based
                    // index in UI
                    self.tray_virtual_desktops_quick.push((item, ix - 1));
                }
                return true;
            }
        }
        false
    }
    /// If there are too many items to assign each a unique access key then we
    /// create up to 10 submenus that divide the range into smaller parts.
    ///
    /// Returns part of the range for the first created submenu, but changes its
    /// start value so that it doesn't overlap with the submenus' access keys.
    fn create_submenus(
        &mut self,
        parent: nwg::ControlHandle,
        desktop_count: u32,
        prefix: u32,
        remaining_digits: u32,
        start_digit: u32,
    ) -> usize {
        if remaining_digits == 0 {
            return 0;
        }
        // remaining_digits: 2, unit: 100
        let unit = 10_u32.pow(remaining_digits - 1);
        let start = ((prefix * 10 + start_digit) * unit).max(1);
        let end = ((prefix + 1) * unit * 10).min(desktop_count);
        let end_digit = (end - 1) / unit % 10;
        let max_digits = if prefix == 0 { 0 } else { prefix.ilog10() + 1 } + remaining_digits;

        let get_range_for = |digit: u32| {
            let start = if digit == start_digit {
                start
            } else {
                // digit: 2, Start digit: 200
                (prefix * 10 * unit) + digit * unit
            };
            let end = if digit == end_digit {
                end
            } else {
                // digit: 2, Exclusive End digit: 300
                (prefix * 10 * unit) + (digit + 1) * unit
            };
            (start, end)
        };
        for digit in start_digit..end_digit + 1 {
            let (start, end) = get_range_for(digit);
            let mut menu = Default::default();
            let start_prefix = start / (unit * 10);
            let start_suffix = start % (unit * 10);

            let prefix_width = max_digits.saturating_sub(remaining_digits) as usize;
            nwg::Menu::builder()
                .text(&format!(
                    "{}&{start_suffix:0suffix_width$} - {:0width$}",
                    if prefix_width == 0 {
                        debug_assert_eq!(start_prefix, 0);
                        String::new()
                    } else {
                        format!("{start_prefix:0prefix_width$}")
                    },
                    end - 1,
                    suffix_width = remaining_digits as usize,
                    width = max_digits as usize
                ))
                .parent(parent)
                .build(&mut menu)
                .expect("Failed to build submenu for quick switch menu");

            let handle = menu.handle;
            // Note: we want to be taken to the number that the user has entered
            // on their keypad (when this menu is selected that includes the
            // first digit in the suffix).
            self.tray_submenus.push((
                menu,
                (start_prefix * 10 + start_suffix % 10)
                    // one-based to zero-based index:
                    .saturating_sub(1),
            ));
            self.call_queue.push((
                handle,
                desktop_count,
                prefix * 10 + digit,
                remaining_digits - 1,
            ));
        }
        (start_digit..end_digit + 1).len()
    }
    /// Create a menu that is easily navigable with keyboard access keys.
    ///
    /// Here is an example of what it can look like:
    ///
    /// ```text
    /// quick -> &00 - 09 ------------> 0&1
    ///          &10 - 19 --> 1&0       0&2
    ///          &20 - 29     1&1       0&3
    ///          &30 - 39     1&2       0&4
    ///          &40 - 49     1&3       0&5
    ///          &50 - 59     1&4       0&6
    ///         ----------    1%5       0&7
    ///          &6           1&6       0&8
    ///          &7           1&7       0&9
    ///          &8           1&8
    ///          &9           1&9
    /// ```
    pub fn create_quick_switch_menu(&mut self, parent: nwg::ControlHandle, desktop_count: u32) {
        self._create_quick_switch_menu(parent, desktop_count, 0, desktop_count.ilog10() + 1);
    }
    fn _create_quick_switch_menu(
        &mut self,
        parent: nwg::ControlHandle,
        desktop_count: u32,
        prefix: u32,
        remaining_digits: u32,
    ) {
        let is_root = !self.is_recursive_call;
        if remaining_digits == 0 {
            return;
        }
        if remaining_digits == 1 {
            let start = (prefix * 10).max(1);
            let end = ((prefix + 1) * 10).min(desktop_count);
            let _did_create = self.try_create_leaf_items(parent, start..end);
            debug_assert!(
                _did_create,
                "Should have created leaf items, otherwise our range was incorrect, used the range: {:?}",
                start..end
            );
            if !self.shortcuts_only_in_root || is_root {
                self.create_shortcut_items(parent);
            }
            return;
        }

        /// Handles recursive calls.
        struct Guard<'a> {
            this: &'a mut QuickSwitchMenu,
            /// `true` if we are in a recursive call.
            is_recursive_call: bool,
        }
        impl<'a> Guard<'a> {
            fn new(this: &'a mut QuickSwitchMenu) -> Self {
                let is_recursive_call = this.is_recursive_call;
                this.is_recursive_call = true;
                Self {
                    this,
                    is_recursive_call,
                }
            }
        }
        impl Drop for Guard<'_> {
            fn drop(&mut self) {
                if !self.is_recursive_call {
                    while let Some((parent, desktop_count, prefix, remaining_digits)) =
                        self.this.call_queue.pop()
                    {
                        self.this._create_quick_switch_menu(
                            parent,
                            desktop_count,
                            prefix,
                            remaining_digits,
                        );
                    }
                }
                // Restore this after all recursive calls:
                self.this.is_recursive_call = self.is_recursive_call;
            }
        }
        let guard = Guard::new(self);
        let this = &mut *guard.this;

        let created_menus =
            this.create_submenus(parent, desktop_count, prefix, remaining_digits, 0);
        if created_menus >= 10 && prefix != 0 {
            // All access keys already in use!
            return;
        }

        // We can use any remaining digits after `end_digit` (the last used access
        // key) as shortcuts for the items inside the first `0001-0999` range (or
        // 01-09 more commonly).
        let mut separator = Default::default();
        nwg::MenuSeparator::builder()
            .parent(parent)
            .build(&mut separator)
            .expect("Failed to build separator for quick switch menu");
        this.tray_separators.push(separator);

        if remaining_digits == 2 {
            let start = (prefix * 100 + created_menus as u32).max(1);
            let end = ((prefix * 10 + 1) * 10).min(desktop_count);
            let _did_create = this.try_create_leaf_items(parent, start..end);
            debug_assert!(
                _did_create,
                "Should have created trailing leaf items, otherwise our range was incorrect, used the range: {:?}",
                start..end
            );
        } else {
            this.create_submenus(
                parent,
                desktop_count,
                prefix * 10,
                remaining_digits - 1,
                created_menus as u32,
            );
        }

        if !this.shortcuts_only_in_root || is_root {
            this.create_shortcut_items(parent);
        }

        // Handle recursive calls if we are the first call (not recursive
        // ourself):
        drop(guard);
    }
}
impl Drop for QuickSwitchMenu {
    fn drop(&mut self) {
        self.clear();
    }
}