asdf_overlay/
backend.rs

1//! Manage window states for rendering overlays.
2//! You can access states for specific window using [`Backends::with_backend`].
3//! This allows you to interact with the overlay state of a window, including its layout and rendering data.
4
5#[doc(hidden)]
6pub mod render;
7
8pub mod window;
9
10use core::{mem, num::NonZeroU32};
11use std::collections::VecDeque;
12
13use anyhow::Context;
14use asdf_overlay_common::cursor::Cursor;
15use asdf_overlay_event::{GpuLuid, OverlayEvent, WindowEvent};
16use dashmap::mapref::multiple::RefMulti;
17use once_cell::sync::Lazy;
18use parking_lot::Mutex;
19use tracing::trace;
20use window::proc::hooked_wnd_proc;
21use windows::Win32::{
22    Foundation::{HWND, LPARAM, RECT, WPARAM},
23    Graphics::Dxgi::IDXGIAdapter,
24    UI::{
25        Input::{
26            Ime::{HIMC, ImmAssociateContext, ImmCreateContext, ImmDestroyContext},
27            KeyboardAndMouse::{GetCapture, ReleaseCapture, SetFocus},
28        },
29        WindowsAndMessaging::{
30            self as msg, ClipCursor, DefWindowProcA, GWLP_WNDPROC, GetClipCursor, GetSystemMetrics,
31            PostMessageA, SM_CXVIRTUALSCREEN, SM_CYVIRTUALSCREEN, SetCursor, SetWindowLongPtrA,
32            ShowCursor, WNDPROC,
33        },
34    },
35};
36
37use crate::{
38    backend::{
39        render::RenderData,
40        window::{InputBlockData, ListenInputFlags, WindowProcData, cursor::load_cursor},
41    },
42    event_sink::OverlayEventSink,
43    interop::DxInterop,
44    layout::OverlayLayout,
45    types::IntDashMap,
46    util::get_client_size,
47};
48
49static BACKENDS: Lazy<Backends> = Lazy::new(|| Backends {
50    map: IntDashMap::default(),
51});
52
53/// Global store for window backends.
54pub struct Backends {
55    map: IntDashMap<u32, WindowBackend>,
56}
57
58impl Backends {
59    /// Iterate over all window backends.
60    pub fn iter<'a>() -> impl Iterator<Item = RefMulti<'a, u32, WindowBackend>> {
61        BACKENDS.map.iter()
62    }
63
64    #[must_use]
65    /// Run closure with the specified backend, if it exists.
66    pub fn with_backend<R>(id: u32, f: impl FnOnce(&WindowBackend) -> R) -> Option<R> {
67        Some(f(&*BACKENDS.map.get(&id)?))
68    }
69
70    #[doc(hidden)]
71    pub fn with_or_init_backend<R>(
72        id: u32,
73        adapter_fn: impl FnOnce() -> Option<IDXGIAdapter>,
74        f: impl FnOnce(&WindowBackend) -> R,
75    ) -> anyhow::Result<R> {
76        if let Some(backend) = BACKENDS.map.get(&id) {
77            return Ok(f(&backend));
78        }
79
80        let backend = BACKENDS
81            .map
82            .entry(id)
83            .or_try_insert_with(|| {
84                let original_proc: WNDPROC = unsafe {
85                    mem::transmute::<isize, WNDPROC>(SetWindowLongPtrA(
86                        HWND(id as _),
87                        GWLP_WNDPROC,
88                        hooked_wnd_proc as usize as _,
89                    ) as _)
90                };
91
92                let interop = DxInterop::create(adapter_fn().as_ref())
93                    .context("failed to create backend interop dxdevice")?;
94
95                let window_size = get_client_size(HWND(id as _))?;
96
97                OverlayEventSink::emit(OverlayEvent::Window {
98                    id,
99                    event: WindowEvent::Added {
100                        width: window_size.0,
101                        height: window_size.1,
102                        gpu_id: interop.gpu_id(),
103                    },
104                });
105
106                Ok::<_, anyhow::Error>(WindowBackend {
107                    id,
108                    original_proc,
109                    layout: Mutex::new(OverlayLayout::new()),
110                    proc: Mutex::new(WindowProcData::new()),
111                    render: Mutex::new(RenderData::new(interop, window_size)),
112                    proc_queue: Mutex::new(VecDeque::new()),
113                })
114            })?
115            .downgrade();
116
117        Ok(f(&backend))
118    }
119
120    fn remove_backend(hwnd: HWND) {
121        let key = hwnd.0 as u32;
122        BACKENDS.map.remove(&key);
123
124        OverlayEventSink::emit(OverlayEvent::Window {
125            id: key,
126            event: WindowEvent::Destroyed,
127        });
128    }
129
130    /// Reset backend states for all windows.
131    pub fn cleanup_backends() {
132        for backend in BACKENDS.map.iter() {
133            backend.reset();
134        }
135    }
136}
137
138pub type ProcDispatchFn = Box<dyn FnOnce(&WindowBackend) + Send>;
139
140/// Data associated to a specific window for overlay rendering.
141pub struct WindowBackend {
142    /// Unique identifier for the window.
143    pub id: u32,
144    pub(crate) original_proc: WNDPROC,
145    pub(crate) layout: Mutex<OverlayLayout>,
146    pub(crate) proc: Mutex<WindowProcData>,
147    #[doc(hidden)]
148    pub render: Mutex<RenderData>,
149    pub(crate) proc_queue: Mutex<VecDeque<ProcDispatchFn>>,
150}
151
152impl WindowBackend {
153    #[tracing::instrument(skip(self))]
154    /// Reset the backend state.
155    /// This reset all set user settable state.
156    pub fn reset(&self) {
157        trace!("backend id: {:?} reset", self.id);
158        *self.layout.lock() = OverlayLayout::new();
159        self.render.lock().reset();
160        self.proc.lock().reset();
161        self.block_input(false);
162    }
163
164    /// Get the locally unique identifier for the GPU.
165    /// This is the GPU adapter used by the window to present to surface.
166    /// Overlay surface texture must be created with this GPU.
167    /// Otherwise, surface cannot be rendered.
168    pub fn gpu_luid(&self) -> GpuLuid {
169        self.render.lock().interop.gpu_id()
170    }
171
172    /// Update overlay surface using the given shared handle.
173    pub fn update_surface(&self, handle: Option<NonZeroU32>) -> anyhow::Result<()> {
174        self.render.lock().update_surface(handle)?;
175        self.invalidate_layout();
176        Ok(())
177    }
178
179    /// Get overlay layout.
180    pub fn layout(&self) -> OverlayLayout {
181        OverlayLayout::clone(&self.layout.lock())
182    }
183
184    /// Update overlay layout.
185    pub fn update_layout(&self, f: impl FnOnce(&mut OverlayLayout)) {
186        f(&mut self.layout.lock());
187        self.invalidate_layout();
188    }
189
190    /// Invalidate layout and recompute position.
191    pub fn invalidate_layout(&self) {
192        let mut render = self.render.lock();
193        let position = self.layout.lock().calc(
194            render
195                .surface
196                .get()
197                .map(|surface| surface.size())
198                .unwrap_or((0, 0)),
199            render.window_size,
200        );
201
202        self.proc.lock().position = position;
203        render.position = position;
204    }
205
206    /// Set which input events are being listened to.
207    pub fn listen_input(&self, flags: ListenInputFlags) {
208        self.proc.lock().listen_input = flags;
209    }
210
211    /// Sets the cursor to be displayed while input is blocked.
212    pub fn set_blocking_cursor(&self, cursor: Option<Cursor>) {
213        self.proc.lock().blocking_cursor = cursor;
214    }
215
216    /// Blocks or unblocks input for the window.
217    pub fn block_input(&self, block: bool) {
218        if block == self.proc.lock().blocking_state.is_some() {
219            return;
220        }
221
222        if block {
223            self.execute_gui(|backend| unsafe {
224                if backend.proc.lock().blocking_state.is_some() {
225                    return;
226                }
227
228                ShowCursor(true);
229                SetCursor(backend.proc.lock().blocking_cursor.and_then(load_cursor));
230                let clip_cursor = {
231                    let mut rect = RECT::default();
232                    _ = GetClipCursor(&mut rect);
233                    let screen = RECT {
234                        left: 0,
235                        top: 0,
236                        right: GetSystemMetrics(SM_CXVIRTUALSCREEN),
237                        bottom: GetSystemMetrics(SM_CYVIRTUALSCREEN),
238                    };
239                    _ = ClipCursor(None);
240
241                    if rect != screen { Some(rect) } else { None }
242                };
243
244                let old_ime_cx =
245                    ImmAssociateContext(HWND(backend.id as _), ImmCreateContext()).0 as usize;
246
247                // give focus to target window
248                _ = SetFocus(Some(HWND(backend.id as _)));
249
250                // In case of ime is already enabled, hide composition windows
251                DefWindowProcA(
252                    HWND(backend.id as _),
253                    msg::WM_IME_SETCONTEXT,
254                    WPARAM(1),
255                    LPARAM(0),
256                );
257                backend.proc.lock().blocking_state = Some(InputBlockData {
258                    clip_cursor,
259                    old_ime_cx,
260                });
261            });
262        } else {
263            self.execute_gui(|backend| unsafe {
264                ShowCursor(false);
265                if GetCapture().0 as u32 == backend.id {
266                    _ = ReleaseCapture();
267                }
268
269                let Some(data) = backend.proc.lock().blocking_state.take() else {
270                    return;
271                };
272                _ = ClipCursor(data.clip_cursor.as_ref().map(|r| r as _));
273                let ime_cx = ImmAssociateContext(HWND(backend.id as _), HIMC(data.old_ime_cx as _));
274                _ = ImmDestroyContext(ime_cx);
275
276                OverlayEventSink::emit(OverlayEvent::Window {
277                    id: backend.id,
278                    event: WindowEvent::InputBlockingEnded,
279                });
280            });
281        }
282    }
283
284    /// Execute function on the GUI thread.
285    /// Calling `execute_gui` inside the closure will deadlock.
286    pub fn execute_gui(&self, f: impl FnOnce(&WindowBackend) + Send + 'static) {
287        let mut proc_queue = self.proc_queue.lock();
288        proc_queue.push_back(Box::new(f));
289        unsafe {
290            _ = PostMessageA(Some(HWND(self.id as _)), msg::WM_NULL, WPARAM(0), LPARAM(0));
291        }
292    }
293}