Primitives · Focus: an overlay dropped without Close hands focus back to what it took it from

src/primitives/focus.rs

@@ -1,6 +1,7 @@
 use gpui::{
-    AnyElement, App, Div, ElementId, Entity, FocusHandle, InteractiveElement, IntoElement,
-    ParentElement, RenderOnce, StyleRefinement, Styled, Subscription, Window, actions, div,
+    AnyElement, AnyWindowHandle, App, Div, ElementId, Entity, FocusHandle, Global,
+    InteractiveElement, IntoElement, ParentElement, RenderOnce, StyleRefinement, Styled,
+    Subscription, WeakEntity, Window, actions, div,
 };
 
 use crate::theme::ActiveTheme;
@@ -71,8 +72,16 @@ impl RenderOnce for FocusScope {
             let handle = self.handle.clone();
             window.use_keyed_state("focus-root", cx, |window, cx| Root {
                 _lost: cx.on_focus_lost(window, move |_: &mut Root, window, cx| {
-                    log::info!("focus: its element left the tree; the root takes it");
-                    window.focus(&handle, cx);
+                    match dropped_overlay(window, cx) {
+                        Some(previous) => {
+                            log::info!("focus: an overlay left unclosed; focus goes back");
+                            window.focus(&previous, cx);
+                        }
+                        None => {
+                            log::info!("focus: its element left the tree; the root takes it");
+                            window.focus(&handle, cx);
+                        }
+                    }
                 }),
             });
         }
@@ -117,6 +126,26 @@ fn step(scope: &FocusHandle, trap: bool, forward: bool, window: &mut Window, cx:
     }
 }
 
+/// Open overlays holding focus, oldest first, with what they took it from.
+#[derive(Default)]
+struct Holders(Vec<(WeakEntity<Takeover>, AnyWindowHandle, Option<FocusHandle>)>);
+
+impl Global for Holders {}
+
+/// Forgets overlays gone without closing; the newest one's prior focus.
+fn dropped_overlay(window: &Window, cx: &mut App) -> Option<FocusHandle> {
+    let here = window.window_handle();
+    let mut back = None;
+    cx.default_global::<Holders>().0.retain(|(holder, at, previous)| {
+        let gone = holder.upgrade().is_none();
+        if gone && *at == here {
+            back = previous.clone();
+        }
+        !gone
+    });
+    back
+}
+
 /// Focus an overlay holds while open, and what it held before.
 pub(crate) struct Takeover {
     pub focus: FocusHandle,
@@ -147,6 +176,8 @@ pub(crate) fn take_focus(
     if !state.read(cx).taken {
         let previous = window.focused(cx);
         window.focus(&state.read(cx).focus.clone(), cx);
+        let holder = (state.downgrade(), window.window_handle(), previous.clone());
+        cx.default_global::<Holders>().0.push(holder);
         state.update(cx, |takeover, _| {
             takeover.previous = previous;
             takeover.taken = true;
@@ -157,6 +188,8 @@ pub(crate) fn take_focus(
 
 /// Returns focus to whatever held it before the overlay opened.
 pub(crate) fn give_back(state: &Entity<Takeover>, window: &mut Window, cx: &mut App) {
+    let closed = state.downgrade();
+    cx.default_global::<Holders>().0.retain(|(holder, ..)| *holder != closed);
     let previous = state.update(cx, |takeover, _| {
         takeover.returned = true;
         takeover.previous.clone()