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);
     }
 }