Avoid losing focus when block decorations go offscreen (#14815)

Release Notes:

- Fixed a bug that caused focus to be lost when renames and inline
assists were scrolled offscreen.

---------

Co-authored-by: Nathan <nathan@zed.dev>

crates/gpui/src/element.rs

@@ -32,8 +32,8 @@
 //! your own custom layout algorithm or rendering a code editor.
 
 use crate::{
-    util::FluentBuilder, ArenaBox, AvailableSpace, Bounds, DispatchNodeId, ElementId, LayoutId,
-    Pixels, Point, Size, Style, ViewContext, WindowContext, ELEMENT_ARENA,
+    util::FluentBuilder, ArenaBox, AvailableSpace, Bounds, DispatchNodeId, ElementId, FocusHandle,
+    LayoutId, Pixels, Point, Size, Style, ViewContext, WindowContext, ELEMENT_ARENA,
 };
 use derive_more::{Deref, DerefMut};
 pub(crate) use smallvec::SmallVec;
@@ -209,7 +209,7 @@ impl<C: RenderOnce> Element for Component<C> {
         _: &mut Self::PrepaintState,
         cx: &mut WindowContext,
     ) {
-        element.paint(cx)
+        element.paint(cx);
     }
 }
 
@@ -493,13 +493,23 @@ impl AnyElement {
 
     /// Prepares the element to be painted by storing its bounds, giving it a chance to draw hitboxes and
     /// request autoscroll before the final paint pass is confirmed.
-    pub fn prepaint(&mut self, cx: &mut WindowContext) {
-        self.0.prepaint(cx)
+    pub fn prepaint(&mut self, cx: &mut WindowContext) -> Option<FocusHandle> {
+        let focus_assigned = cx.window.next_frame.focus.is_some();
+
+        self.0.prepaint(cx);
+
+        if !focus_assigned {
+            if let Some(focus_id) = cx.window.next_frame.focus {
+                return FocusHandle::for_id(focus_id, &cx.window.focus_handles);
+            }
+        }
+
+        None
     }
 
     /// Paints the element stored in this `AnyElement`.
     pub fn paint(&mut self, cx: &mut WindowContext) {
-        self.0.paint(cx)
+        self.0.paint(cx);
     }
 
     /// Performs layout for this element within the given available space and returns its size.
@@ -512,19 +522,25 @@ impl AnyElement {
     }
 
     /// Prepaints this element at the given absolute origin.
-    pub fn prepaint_at(&mut self, origin: Point<Pixels>, cx: &mut WindowContext) {
-        cx.with_absolute_element_offset(origin, |cx| self.0.prepaint(cx));
+    /// If any element in the subtree beneath this element is focused, its FocusHandle is returned.
+    pub fn prepaint_at(
+        &mut self,
+        origin: Point<Pixels>,
+        cx: &mut WindowContext,
+    ) -> Option<FocusHandle> {
+        cx.with_absolute_element_offset(origin, |cx| self.prepaint(cx))
     }
 
     /// Performs layout on this element in the available space, then prepaints it at the given absolute origin.
+    /// If any element in the subtree beneath this element is focused, its FocusHandle is returned.
     pub fn prepaint_as_root(
         &mut self,
         origin: Point<Pixels>,
         available_space: Size<AvailableSpace>,
         cx: &mut WindowContext,
-    ) {
+    ) -> Option<FocusHandle> {
         self.layout_as_root(available_space, cx);
-        cx.with_absolute_element_offset(origin, |cx| self.0.prepaint(cx));
+        cx.with_absolute_element_offset(origin, |cx| self.prepaint(cx))
     }
 }
 
@@ -552,7 +568,7 @@ impl Element for AnyElement {
         _: &mut Self::RequestLayoutState,
         cx: &mut WindowContext,
     ) {
-        self.prepaint(cx)
+        self.prepaint(cx);
     }
 
     fn paint(
@@ -563,7 +579,7 @@ impl Element for AnyElement {
         _: &mut Self::PrepaintState,
         cx: &mut WindowContext,
     ) {
-        self.paint(cx)
+        self.paint(cx);
     }
 }
 

crates/gpui/src/elements/div.rs

@@ -1359,6 +1359,9 @@ impl Interactivity {
         f: impl FnOnce(&Style, Point<Pixels>, Option<Hitbox>, &mut WindowContext) -> R,
     ) -> R {
         self.content_size = content_size;
+        if let Some(focus_handle) = self.tracked_focus_handle.as_ref() {
+            cx.set_focus_handle(&focus_handle);
+        }
         cx.with_optional_element_state::<InteractiveElementState, _>(
             global_id,
             |element_state, cx| {
@@ -1998,9 +2001,6 @@ impl Interactivity {
         if let Some(context) = self.key_context.clone() {
             cx.set_key_context(context);
         }
-        if let Some(focus_handle) = self.tracked_focus_handle.as_ref() {
-            cx.set_focus_handle(focus_handle);
-        }
 
         for listener in key_down_listeners {
             cx.on_key_event(move |event: &KeyDownEvent, phase, cx| {

crates/gpui/src/key_dispatch.rs

@@ -92,6 +92,7 @@ pub(crate) struct DispatchNode {
 pub(crate) struct ReusedSubtree {
     old_range: Range<usize>,
     new_range: Range<usize>,
+    contains_focus: bool,
 }
 
 impl ReusedSubtree {
@@ -104,6 +105,10 @@ impl ReusedSubtree {
         );
         DispatchNodeId((node_id.0 - self.old_range.start) + self.new_range.start)
     }
+
+    pub fn contains_focus(&self) -> bool {
+        self.contains_focus
+    }
 }
 
 type KeyListener = Rc<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>;
@@ -246,9 +251,15 @@ impl DispatchTree {
         target.modifiers_changed_listeners = mem::take(&mut source.modifiers_changed_listeners);
     }
 
-    pub fn reuse_subtree(&mut self, old_range: Range<usize>, source: &mut Self) -> ReusedSubtree {
+    pub fn reuse_subtree(
+        &mut self,
+        old_range: Range<usize>,
+        source: &mut Self,
+        focus: Option<FocusId>,
+    ) -> ReusedSubtree {
         let new_range = self.nodes.len()..self.nodes.len() + old_range.len();
 
+        let mut contains_focus = false;
         let mut source_stack = vec![];
         for (source_node_id, source_node) in source
             .nodes
@@ -268,6 +279,9 @@ impl DispatchTree {
             }
 
             source_stack.push(source_node_id);
+            if source_node.focus_id.is_some() && source_node.focus_id == focus {
+                contains_focus = true;
+            }
             self.move_node(source_node);
         }
 
@@ -279,6 +293,7 @@ impl DispatchTree {
         ReusedSubtree {
             old_range,
             new_range,
+            contains_focus,
         }
     }
 

crates/gpui/src/window.rs

@@ -464,6 +464,7 @@ impl Frame {
         self.cursor_styles.clear();
         self.hitboxes.clear();
         self.deferred_draws.clear();
+        self.focus = None;
     }
 
     pub(crate) fn hit_test(&self, position: Point<Pixels>) -> HitTest {
@@ -1460,7 +1461,6 @@ impl<'a> WindowContext<'a> {
                 &mut self.window.rendered_frame.dispatch_tree,
                 self.window.focus,
             );
-        self.window.next_frame.focus = self.window.focus;
         self.window.next_frame.window_active = self.window.active.get();
 
         // Register requested input handler with the platform window.
@@ -1574,7 +1574,7 @@ impl<'a> WindowContext<'a> {
         self.paint_deferred_draws(&sorted_deferred_draws);
 
         if let Some(mut prompt_element) = prompt_element {
-            prompt_element.paint(self)
+            prompt_element.paint(self);
         } else if let Some(mut drag_element) = active_drag_element {
             drag_element.paint(self);
         } else if let Some(mut tooltip_element) = tooltip_element {
@@ -1730,7 +1730,13 @@ impl<'a> WindowContext<'a> {
         let reused_subtree = window.next_frame.dispatch_tree.reuse_subtree(
             range.start.dispatch_tree_index..range.end.dispatch_tree_index,
             &mut window.rendered_frame.dispatch_tree,
+            window.focus,
         );
+
+        if reused_subtree.contains_focus() {
+            window.next_frame.focus = window.focus;
+        }
+
         window.next_frame.deferred_draws.extend(
             window.rendered_frame.deferred_draws
                 [range.start.deferred_draws_index..range.end.deferred_draws_index]
@@ -2845,13 +2851,16 @@ impl<'a> WindowContext<'a> {
     /// Sets the focus handle for the current element. This handle will be used to manage focus state
     /// and keyboard event dispatch for the element.
     ///
-    /// This method should only be called as part of the paint phase of element drawing.
+    /// This method should only be called as part of the prepaint phase of element drawing.
     pub fn set_focus_handle(&mut self, focus_handle: &FocusHandle) {
         debug_assert_eq!(
             self.window.draw_phase,
-            DrawPhase::Paint,
-            "this method can only be called during paint"
+            DrawPhase::Prepaint,
+            "this method can only be called during prepaint"
         );
+        if focus_handle.is_focused(self) {
+            self.window.next_frame.focus = Some(focus_handle.id);
+        }
         self.window
             .next_frame
             .dispatch_tree