Combine OPEN_WINDOWS and HOOK into HOOK_STATE RwLock

Author: estroBiologist

src/win/hook.rs

@@ -2,7 +2,7 @@ use std::{
     collections::HashSet,
     ffi::c_int,
     ptr,
-    sync::{LazyLock, Mutex, RwLock},
+    sync::{LazyLock, RwLock},
 };
 
 use winapi::{
@@ -15,33 +15,34 @@ use winapi::{
         processthreadsapi::GetCurrentThreadId,
         winuser::{
             CallNextHookEx, SetWindowsHookExW, UnhookWindowsHookEx, HC_ACTION, MSG, PM_REMOVE,
-            WH_GETMESSAGE, WM_CHAR, WM_KEYDOWN, WM_KEYUP, WM_SYSCHAR, WM_SYSKEYDOWN, WM_SYSKEYUP,
-            WM_USER,
+            WH_GETMESSAGE, WM_CHAR, WM_KEYDOWN, WM_KEYUP, WM_SYSCHAR, WM_SYSKEYDOWN, WM_SYSKEYUP, WM_USER,
         },
     },
 };
 
 use crate::win::wnd_proc;
 
-static HOOK: Mutex<Option<KeyboardHook>> = Mutex::new(None);
-
 // track all windows opened by this instance of baseview
 // we use an RwLock here since the vast majority of uses (event interceptions)
 // will only need to read from the HashSet
-static OPEN_WINDOWS: LazyLock<RwLock<HashSet<HWNDWrapper>>> = LazyLock::new(|| RwLock::default());
+static HOOK_STATE: LazyLock<RwLock<KeyboardHookState>> = LazyLock::new(|| RwLock::default());
 
 pub(crate) struct KeyboardHookHandle(HWNDWrapper);
 
-struct KeyboardHook(HHOOK);
+#[derive(Default)]
+struct KeyboardHookState {
+    hook: Option<HHOOK>,
+    open_windows: HashSet<HWNDWrapper>,
+}
 
 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
 struct HWNDWrapper(HWND);
 
-// SAFETY: it's a pointer behind a mutex. we'll live
-unsafe impl Send for KeyboardHook {}
-unsafe impl Sync for KeyboardHook {}
+// SAFETY: it's a pointer behind an RwLock. we'll live
+unsafe impl Send for KeyboardHookState {}
+unsafe impl Sync for KeyboardHookState {}
 
-// SAFETY: ditto
+// SAFETY: we never access the underlying HWND ourselves, just use it as a HashSet 
 unsafe impl Send for HWNDWrapper {}
 unsafe impl Sync for HWNDWrapper {}
 
@@ -55,45 +56,43 @@ impl Drop for KeyboardHookHandle {
 // some DAWs (particularly Ableton) intercept incoming keyboard messages,
 // but we're naughty so we intercept them right back
 pub(crate) fn init_keyboard_hook(hwnd: HWND) -> KeyboardHookHandle {
-    // register hwnd to global window set
-    OPEN_WINDOWS.write().unwrap().insert(HWNDWrapper(hwnd));
+    let state = &mut *HOOK_STATE.write().unwrap();
 
-    let hook = &mut *HOOK.lock().unwrap();
+    // register hwnd to global window set
+    state.open_windows.insert(HWNDWrapper(hwnd));
 
-    if hook.is_some() {
+    if state.hook.is_some() {
         // keyboard hook already exists, just return handle
         KeyboardHookHandle(HWNDWrapper(hwnd))
     } else {
         // keyboard hook doesn't exist (no windows open before this), create it
-        let new_hook = KeyboardHook(unsafe {
+        let new_hook = unsafe {
             SetWindowsHookExW(
                 WH_GETMESSAGE,
                 Some(keyboard_hook_callback),
                 GetModuleHandleW(ptr::null()),
                 GetCurrentThreadId(),
             )
-        });
+        };
 
-        *hook = Some(new_hook);
+        state.hook = Some(new_hook);
 
         KeyboardHookHandle(HWNDWrapper(hwnd))
     }
 }
 
 fn deinit_keyboard_hook(hwnd: HWNDWrapper) {
-    let windows = &mut *OPEN_WINDOWS.write().unwrap();
-
-    windows.remove(&hwnd);
+    let state = &mut *HOOK_STATE.write().unwrap();
 
-    if windows.is_empty() {
-        if let Ok(mut hook) = HOOK.lock() {
-            if let Some(KeyboardHook(hhook)) = &mut *hook {
-                unsafe {
-                    UnhookWindowsHookEx(*hhook);
-                }
+    state.open_windows.remove(&hwnd);
 
-                *hook = None;
+    if state.open_windows.is_empty() {
+        if let Some(hhook) = state.hook {
+            unsafe {
+                UnhookWindowsHookEx(hhook);
             }
+
+            state.hook = None;
         }
     }
 }
@@ -119,8 +118,8 @@ unsafe extern "system" fn keyboard_hook_callback(
     }
 }
 
-// check if `msg` is a keyboard message addressed
-// to a window in OPEN_WINDOWS, and intercept it if so
+// check if `msg` is a keyboard message addressed to a window
+// in KeyboardHookState::open_windows, and intercept it if so
 unsafe fn offer_message_to_baseview(msg: *mut MSG) -> bool {
     let msg = &*msg;
 
@@ -132,9 +131,9 @@ unsafe fn offer_message_to_baseview(msg: *mut MSG) -> bool {
     }
 
     // check if this is one of our windows. if so, intercept it
-    let Ok(windows) = OPEN_WINDOWS.read() else { return false };
+    let state = HOOK_STATE.read().unwrap();
 
-    if windows.contains(&HWNDWrapper(msg.hwnd)) {
+    if state.open_windows.contains(&HWNDWrapper(msg.hwnd)) {
         let _ = wnd_proc(msg.hwnd, msg.message, msg.wParam, msg.lParam);
 
         return true;