Re-enable view caching

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

crates/assistant/src/assistant_panel.rs

@@ -1279,7 +1279,7 @@ impl Render for AssistantPanel {
                             let conversation_count = self.saved_conversations.len();
                             canvas(move |bounds, cx| {
                                 uniform_list(
-                                    view,
+                                    view.clone(),
                                     "saved_conversations",
                                     conversation_count,
                                     |this, range, cx| {
@@ -1288,7 +1288,7 @@ impl Render for AssistantPanel {
                                             .collect()
                                     },
                                 )
-                                .track_scroll(scroll_handle)
+                                .track_scroll(scroll_handle.clone())
                                 .into_any_element()
                                 .draw(
                                     bounds.origin,

crates/diagnostics/src/diagnostics.rs

@@ -1592,7 +1592,7 @@ mod tests {
                 .enumerate()
                 .filter_map(|(ix, (row, block))| {
                     let name: SharedString = match block {
-                        TransformBlock::Custom(block) => cx.with_element_context({
+                        TransformBlock::Custom(block) => cx.with_element_context(false, {
                             let editor_view = editor_view.clone();
                             |cx| -> Option<SharedString> {
                                 block

crates/editor/src/element.rs

@@ -375,7 +375,7 @@ impl EditorElement {
     ) {
         let mouse_position = cx.mouse_position();
         if !text_bounds.contains(&mouse_position)
-            || !cx.was_top_layer(&mouse_position, stacking_order)
+            || !cx.is_top_layer(&mouse_position, stacking_order)
         {
             return;
         }
@@ -407,7 +407,7 @@ impl EditorElement {
         } else if !text_bounds.contains(&event.position) {
             return;
         }
-        if !cx.was_top_layer(&event.position, stacking_order) {
+        if !cx.is_top_layer(&event.position, stacking_order) {
             return;
         }
 
@@ -483,7 +483,7 @@ impl EditorElement {
             && !pending_nonempty_selections
             && event.modifiers.command
             && text_bounds.contains(&event.position)
-            && cx.was_top_layer(&event.position, stacking_order)
+            && cx.is_top_layer(&event.position, stacking_order)
         {
             let point = position_map.point_for_position(text_bounds, event.position);
             editor.handle_click_hovered_link(point, event.modifiers, cx);
@@ -551,7 +551,7 @@ impl EditorElement {
         let modifiers = event.modifiers;
         let text_hovered = text_bounds.contains(&event.position);
         let gutter_hovered = gutter_bounds.contains(&event.position);
-        let was_top = cx.was_top_layer(&event.position, stacking_order);
+        let was_top = cx.is_top_layer(&event.position, stacking_order);
 
         editor.set_gutter_hovered(gutter_hovered, cx);
 
@@ -1243,7 +1243,7 @@ impl EditorElement {
                         popover_origin.x = popover_origin.x + x_out_of_bounds;
                     }
 
-                    if cx.was_top_layer(&popover_origin, cx.stacking_order()) {
+                    if cx.is_top_layer(&popover_origin, cx.stacking_order()) {
                         cx.break_content_mask(|cx| {
                             hover_popover.draw(popover_origin, available_space, cx)
                         });
@@ -2203,7 +2203,7 @@ impl EditorElement {
             let scroll_width = longest_line_width.max(max_visible_line_width) + overscroll.width;
 
             let editor_view = cx.view().clone();
-            let (scroll_width, blocks) = cx.with_element_context(|cx| {
+            let (scroll_width, blocks) = cx.with_element_context(false, |cx| {
              cx.with_element_id(Some("editor_blocks"), |cx| {
                 self.layout_blocks(
                     start_row..end_row,
@@ -2303,7 +2303,7 @@ impl EditorElement {
             };
 
             let editor_view = cx.view().clone();
-            let fold_indicators = cx.with_element_context(|cx| {
+            let fold_indicators = cx.with_element_context(false, |cx| {
 
                 cx.with_element_id(Some("gutter_fold_indicators"), |_cx| {
                 editor.render_fold_indicators(
@@ -3022,13 +3022,14 @@ impl Element for EditorElement {
                         let mut style = Style::default();
                         style.size.width = relative(1.).into();
                         style.size.height = self.style.text.line_height_in_pixels(rem_size).into();
-                        cx.with_element_context(|cx| cx.request_layout(&style, None))
+                        cx.with_element_context(false, |cx| cx.request_layout(&style, None))
                     }
+
                     EditorMode::AutoHeight { max_lines } => {
                         let editor_handle = cx.view().clone();
                         let max_line_number_width =
                             self.max_line_number_width(&editor.snapshot(cx), cx);
-                        cx.with_element_context(|cx| {
+                        cx.with_element_context(false, |cx| {
                             cx.request_measured_layout(
                                 Style::default(),
                                 move |known_dimensions, _, cx| {
@@ -3047,11 +3048,12 @@ impl Element for EditorElement {
                             )
                         })
                     }
+
                     EditorMode::Full => {
                         let mut style = Style::default();
                         style.size.width = relative(1.).into();
                         style.size.height = relative(1.).into();
-                        cx.with_element_context(|cx| cx.request_layout(&style, None))
+                        cx.with_element_context(false, |cx| cx.request_layout(&style, None))
                     }
                 };
 
@@ -3661,7 +3663,7 @@ mod tests {
             .unwrap();
         let state = cx
             .update_window(window.into(), |view, cx| {
-                cx.with_element_context(|cx| {
+                cx.with_element_context(false, |cx| {
                     cx.with_view_id(view.entity_id(), |cx| {
                         element.compute_layout(
                             Bounds {
@@ -3757,7 +3759,7 @@ mod tests {
 
         let state = cx
             .update_window(window.into(), |view, cx| {
-                cx.with_element_context(|cx| {
+                cx.with_element_context(false, |cx| {
                     cx.with_view_id(view.entity_id(), |cx| {
                         element.compute_layout(
                             Bounds {
@@ -3823,7 +3825,7 @@ mod tests {
         let mut element = EditorElement::new(&editor, style);
         let state = cx
             .update_window(window.into(), |view, cx| {
-                cx.with_element_context(|cx| {
+                cx.with_element_context(false, |cx| {
                     cx.with_view_id(view.entity_id(), |cx| {
                         element.compute_layout(
                             Bounds {
@@ -3851,7 +3853,7 @@ mod tests {
         // Don't panic.
         let bounds = Bounds::<Pixels>::new(Default::default(), size);
         cx.update_window(window.into(), |_, cx| {
-            cx.with_element_context(|cx| element.paint(bounds, &mut (), cx))
+            cx.with_element_context(false, |cx| element.paint(bounds, &mut (), cx))
         })
         .unwrap()
     }
@@ -4028,7 +4030,7 @@ mod tests {
             .unwrap();
         let layout_state = cx
             .update_window(window.into(), |_, cx| {
-                cx.with_element_context(|cx| {
+                cx.with_element_context(false, |cx| {
                     element.compute_layout(
                         Bounds {
                             origin: point(px(500.), px(500.)),

crates/extensions_ui/src/extensions_ui.rs

@@ -414,13 +414,13 @@ impl Render for ExtensionsPage {
                         let item_count = self.extensions_entries.len();
                         move |bounds, cx| {
                             uniform_list::<_, Div, _>(
-                                view,
+                                view.clone(),
                                 "entries",
                                 item_count,
                                 Self::render_extensions,
                             )
                             .size_full()
-                            .track_scroll(scroll_handle)
+                            .track_scroll(scroll_handle.clone())
                             .into_any_element()
                             .draw(
                                 bounds.origin,

crates/gpui/src/app/test_context.rs

@@ -686,7 +686,7 @@ impl<'a> VisualTestContext {
                 .expect("Can't draw to this window without a root view")
                 .entity_id();
 
-            cx.with_element_context(|cx| {
+            cx.with_element_context(false, |cx| {
                 cx.with_view_id(entity_id, |cx| {
                     f(cx).draw(origin, space, cx);
                 })

crates/gpui/src/element.rs

@@ -103,13 +103,13 @@ pub trait IntoElement: Sized {
     {
         let element = self.into_element();
         let element_id = element.element_id();
-        let element = DrawableElement {
+        let mut element = DrawableElement {
             element: Some(element),
             phase: ElementDrawPhase::Start,
         };
 
         let frame_state =
-            DrawableElement::draw(element, origin, available_space.map(Into::into), cx);
+            DrawableElement::draw(&mut element, origin, available_space.map(Into::into), cx);
 
         if let Some(mut frame_state) = frame_state {
             f(&mut frame_state, cx)
@@ -301,8 +301,8 @@ impl<E: Element> DrawableElement<E> {
         layout_id
     }
 
-    fn paint(mut self, cx: &mut ElementContext) -> Option<E::State> {
-        match self.phase {
+    fn paint(&mut self, cx: &mut ElementContext) -> Option<&mut E::State> {
+        match &mut self.phase {
             ElementDrawPhase::LayoutRequested {
                 layout_id,
                 frame_state,
@@ -312,13 +312,13 @@ impl<E: Element> DrawableElement<E> {
                 frame_state,
                 ..
             } => {
-                let bounds = cx.layout_bounds(layout_id);
+                let bounds = cx.layout_bounds(*layout_id);
 
-                if let Some(mut frame_state) = frame_state {
+                if let Some(frame_state) = frame_state {
                     self.element
-                        .take()
+                        .as_mut()
                         .unwrap()
-                        .paint(bounds, &mut frame_state, cx);
+                        .paint(bounds, frame_state, cx);
                     Some(frame_state)
                 } else {
                     let element_id = self
@@ -330,7 +330,7 @@ impl<E: Element> DrawableElement<E> {
                     cx.with_element_state(element_id, |element_state, cx| {
                         let mut element_state = element_state.unwrap();
                         self.element
-                            .take()
+                            .as_mut()
                             .unwrap()
                             .paint(bounds, &mut element_state, cx);
                         ((), element_state)
@@ -349,7 +349,7 @@ impl<E: Element> DrawableElement<E> {
         cx: &mut ElementContext,
     ) -> Size<Pixels> {
         if matches!(&self.phase, ElementDrawPhase::Start) {
-            self.request_layout(cx);
+            cx.with_z_index_reset(|cx| self.request_layout(cx));
         }
 
         let layout_id = match &mut self.phase {
@@ -384,11 +384,11 @@ impl<E: Element> DrawableElement<E> {
     }
 
     fn draw(
-        mut self,
+        &mut self,
         origin: Point<Pixels>,
         available_space: Size<AvailableSpace>,
         cx: &mut ElementContext,
-    ) -> Option<E::State> {
+    ) -> Option<&mut E::State> {
         self.measure(available_space, cx);
         cx.with_absolute_element_offset(origin, |cx| self.paint(cx))
     }
@@ -408,7 +408,7 @@ where
     }
 
     fn paint(&mut self, cx: &mut ElementContext) {
-        DrawableElement::paint(self.take().unwrap(), cx);
+        DrawableElement::paint(self.as_mut().unwrap(), cx);
     }
 
     fn measure(
@@ -425,7 +425,7 @@ where
         available_space: Size<AvailableSpace>,
         cx: &mut ElementContext,
     ) {
-        DrawableElement::draw(self.take().unwrap(), origin, available_space, cx);
+        DrawableElement::draw(self.as_mut().unwrap(), origin, available_space, cx);
     }
 }
 

crates/gpui/src/elements/canvas.rs

@@ -4,7 +4,7 @@ use crate::{Bounds, Element, ElementContext, IntoElement, Pixels, Style, StyleRe
 
 /// Construct a canvas element with the given paint callback.
 /// Useful for adding short term custom drawing to a view.
-pub fn canvas(callback: impl 'static + FnOnce(&Bounds<Pixels>, &mut ElementContext)) -> Canvas {
+pub fn canvas(callback: impl 'static + FnMut(&Bounds<Pixels>, &mut ElementContext)) -> Canvas {
     Canvas {
         paint_callback: Some(Box::new(callback)),
         style: StyleRefinement::default(),
@@ -14,7 +14,7 @@ pub fn canvas(callback: impl 'static + FnOnce(&Bounds<Pixels>, &mut ElementConte
 /// A canvas element, meant for accessing the low level paint API without defining a whole
 /// custom element
 pub struct Canvas {
-    paint_callback: Option<Box<dyn FnOnce(&Bounds<Pixels>, &mut ElementContext)>>,
+    paint_callback: Option<Box<dyn FnMut(&Bounds<Pixels>, &mut ElementContext)>>,
     style: StyleRefinement,
 }
 
@@ -46,7 +46,7 @@ impl Element for Canvas {
 
     fn paint(&mut self, bounds: Bounds<Pixels>, style: &mut Style, cx: &mut ElementContext) {
         style.paint(bounds, cx, |cx| {
-            (self.paint_callback.take().unwrap())(&bounds, cx)
+            (self.paint_callback.as_mut().unwrap())(&bounds, cx)
         });
     }
 }

crates/gpui/src/elements/list.rs

@@ -638,7 +638,7 @@ impl Element for List {
         cx.on_mouse_event(move |event: &ScrollWheelEvent, phase, cx| {
             if phase == DispatchPhase::Bubble
                 && bounds.contains(&event.position)
-                && cx.was_top_layer(&event.position, cx.stacking_order())
+                && cx.is_top_layer(&event.position, cx.stacking_order())
             {
                 list_state.0.borrow_mut().scroll(
                     &layout_response.scroll_top,

crates/gpui/src/elements/text.rs

@@ -9,7 +9,6 @@ use parking_lot::{Mutex, MutexGuard};
 use smallvec::SmallVec;
 use std::{
     cell::{Cell, RefCell},
-    mem,
     ops::Range,
     rc::Rc,
     sync::Arc,
@@ -312,10 +311,10 @@ pub struct InteractiveText {
     element_id: ElementId,
     text: StyledText,
     click_listener:
-        Option<Box<dyn Fn(&[Range<usize>], InteractiveTextClickEvent, &mut WindowContext<'_>)>>,
-    hover_listener: Option<Box<dyn Fn(Option<usize>, MouseMoveEvent, &mut WindowContext<'_>)>>,
+        Option<Rc<dyn Fn(&[Range<usize>], InteractiveTextClickEvent, &mut WindowContext<'_>)>>,
+    hover_listener: Option<Rc<dyn Fn(Option<usize>, MouseMoveEvent, &mut WindowContext<'_>)>>,
     tooltip_builder: Option<Rc<dyn Fn(usize, &mut WindowContext<'_>) -> Option<AnyView>>>,
-    clickable_ranges: Vec<Range<usize>>,
+    clickable_ranges: Rc<[Range<usize>]>,
 }
 
 struct InteractiveTextClickEvent {
@@ -341,7 +340,7 @@ impl InteractiveText {
             click_listener: None,
             hover_listener: None,
             tooltip_builder: None,
-            clickable_ranges: Vec::new(),
+            clickable_ranges: Rc::new([]),
         }
     }
 
@@ -352,7 +351,7 @@ impl InteractiveText {
         ranges: Vec<Range<usize>>,
         listener: impl Fn(usize, &mut WindowContext<'_>) + 'static,
     ) -> Self {
-        self.click_listener = Some(Box::new(move |ranges, event, cx| {
+        self.click_listener = Some(Rc::new(move |ranges, event, cx| {
             for (range_ix, range) in ranges.iter().enumerate() {
                 if range.contains(&event.mouse_down_index) && range.contains(&event.mouse_up_index)
                 {
@@ -360,7 +359,7 @@ impl InteractiveText {
                 }
             }
         }));
-        self.clickable_ranges = ranges;
+        self.clickable_ranges = Rc::from(ranges.into_boxed_slice());
         self
     }
 
@@ -370,7 +369,7 @@ impl InteractiveText {
         mut self,
         listener: impl Fn(Option<usize>, MouseMoveEvent, &mut WindowContext<'_>) + 'static,
     ) -> Self {
-        self.hover_listener = Some(Box::new(listener));
+        self.hover_listener = Some(Rc::new(listener));
         self
     }
 
@@ -420,14 +419,14 @@ impl Element for InteractiveText {
     }
 
     fn paint(&mut self, bounds: Bounds<Pixels>, state: &mut Self::State, cx: &mut ElementContext) {
-        if let Some(click_listener) = self.click_listener.take() {
+        if let Some(click_listener) = self.click_listener.clone() {
             let mouse_position = cx.mouse_position();
             if let Some(ix) = state.text_state.index_for_position(bounds, mouse_position) {
                 if self
                     .clickable_ranges
                     .iter()
                     .any(|range| range.contains(&ix))
-                    && cx.was_top_layer(&mouse_position, cx.stacking_order())
+                    && cx.is_top_layer(&mouse_position, cx.stacking_order())
                 {
                     cx.set_cursor_style(crate::CursorStyle::PointingHand)
                 }
@@ -436,7 +435,7 @@ impl Element for InteractiveText {
             let text_state = state.text_state.clone();
             let mouse_down = state.mouse_down_index.clone();
             if let Some(mouse_down_index) = mouse_down.get() {
-                let clickable_ranges = mem::take(&mut self.clickable_ranges);
+                let clickable_ranges = self.clickable_ranges.clone();
                 cx.on_mouse_event(move |event: &MouseUpEvent, phase, cx| {
                     if phase == DispatchPhase::Bubble {
                         if let Some(mouse_up_index) =
@@ -469,7 +468,7 @@ impl Element for InteractiveText {
                 });
             }
         }
-        if let Some(hover_listener) = self.hover_listener.take() {
+        if let Some(hover_listener) = self.hover_listener.clone() {
             let text_state = state.text_state.clone();
             let hovered_index = state.hovered_index.clone();
             cx.on_mouse_event(move |event: &MouseMoveEvent, phase, cx| {

crates/gpui/src/style.rs

@@ -411,8 +411,8 @@ impl Style {
         });
 
         let background_color = self.background.as_ref().and_then(Fill::color);
-        if background_color.map_or(false, |color| !color.is_transparent()) {
-            cx.with_z_index(1, |cx| {
+        cx.with_z_index(1, |cx| {
+            if background_color.map_or(false, |color| !color.is_transparent()) {
                 let mut border_color = background_color.unwrap_or_default();
                 border_color.a = 0.;
                 cx.paint_quad(quad(
@@ -422,8 +422,8 @@ impl Style {
                     Edges::default(),
                     border_color,
                 ));
-            });
-        }
+            }
+        });
 
         cx.with_z_index(2, |cx| {
             continuation(cx);

crates/gpui/src/view.rs

@@ -105,7 +105,7 @@ impl<V: Render> Element for View<V> {
     }
 
     fn paint(&mut self, _: Bounds<Pixels>, element: &mut Self::State, cx: &mut ElementContext) {
-        cx.paint_view(self.entity_id(), |cx| element.take().unwrap().paint(cx));
+        cx.paint_view(self.entity_id(), |cx| element.as_mut().unwrap().paint(cx));
     }
 }
 
@@ -247,22 +247,71 @@ impl AnyView {
         self.model.entity_id()
     }
 
-    pub(crate) fn draw(
+    pub(crate) fn layout(
         &self,
         origin: Point<Pixels>,
         available_space: Size<AvailableSpace>,
         cx: &mut ElementContext,
-    ) {
-        cx.paint_view(self.entity_id(), |cx| {
+    ) -> AnyElement {
+        let original_disallow = cx.disallow_adding_opaque_layers;
+        cx.disallow_adding_opaque_layers = true;
+
+        let element = cx.paint_view(self.entity_id(), |cx| {
             cx.with_absolute_element_offset(origin, |cx| {
                 let (layout_id, mut rendered_element) = (self.request_layout)(self, cx);
                 cx.compute_layout(layout_id, available_space);
-                rendered_element.paint(cx)
+                rendered_element
+            })
+        });
+
+        cx.disallow_adding_opaque_layers = original_disallow;
+
+        element
+    }
+
+    pub(crate) fn paint(
+        &self,
+        origin: Point<Pixels>,
+        rendered_element: &mut AnyElement,
+        cx: &mut ElementContext,
+    ) {
+        cx.paint_view(self.entity_id(), |cx| {
+            cx.with_absolute_element_offset(origin, |cx| {
+                rendered_element.paint(cx);
             });
-        })
+        });
     }
 }
 
+//     pub(crate) fn draw(
+//         &self,
+//         origin: Point<Pixels>,
+//         available_space: Size<AvailableSpace>,
+//         cx: &mut ElementContext,
+//     ) {
+//         cx.paint_view(self.entity_id(), |cx| {
+//             cx.with_absolute_element_offset(origin, |cx| {
+//                 cx.in_layout(true);
+//                 let (layout_id, mut rendered_element) = (self.request_layout)(self, cx);
+//                 assert_eq!(
+//                     *cx.window.next_frame.next_stacking_order_ids.last().unwrap(),
+//                     0
+//                 );
+//                 // let prev_last_stacking_order_id = cx.window.next_frame.next_stacking_order_ids.last();
+//                 cx.compute_layout(layout_id, available_space);
+//                 cx.in_layout(false);
+//                 // assert_eq!(cx.window.next_frame.z_index_stack.len(), 0);
+//                 *cx.window
+//                     .next_frame
+//                     .next_stacking_order_ids
+//                     .last_mut()
+//                     .unwrap() = 0;
+//                 rendered_element.paint(cx)
+//             });
+//         })
+//     }
+// }
+
 impl<V: Render> From<View<V>> for AnyView {
     fn from(value: View<V>) -> Self {
         AnyView {
@@ -307,7 +356,7 @@ impl Element for AnyView {
     fn paint(&mut self, bounds: Bounds<Pixels>, state: &mut Self::State, cx: &mut ElementContext) {
         cx.paint_view(self.entity_id(), |cx| {
             if !self.cache {
-                state.element.take().unwrap().paint(cx);
+                state.element.as_mut().unwrap().paint(cx);
                 return;
             }
 
@@ -322,7 +371,7 @@ impl Element for AnyView {
                 }
             }
 
-            if let Some(mut element) = state.element.take() {
+            if let Some(mut element) = state.element.as_mut() {
                 element.paint(cx);
             } else {
                 let mut element = (self.request_layout)(self, cx).1;
@@ -336,12 +385,15 @@ impl Element for AnyView {
                 .last()
                 .copied()
                 .unwrap();
-            state.cache_key = Some(ViewCacheKey {
-                bounds,
-                stacking_order: cx.stacking_order().clone(),
-                content_mask: cx.content_mask(),
-                text_style: cx.text_style(),
-            });
+
+            if !cx.painting_dry_run {
+                state.cache_key = Some(ViewCacheKey {
+                    bounds,
+                    stacking_order: cx.stacking_order().clone(),
+                    content_mask: cx.content_mask(),
+                    text_style: cx.text_style(),
+                });
+            }
         })
     }
 }

crates/gpui/src/window.rs

@@ -862,10 +862,10 @@ impl<'a> WindowContext<'a> {
     /// Returns true if there is no opaque layer containing the given point
     /// on top of the given level. Layers who are extensions of the queried layer
     /// are not considered to be on top of queried layer.
-    pub fn was_top_layer(&self, point: &Point<Pixels>, layer: &StackingOrder) -> bool {
+    pub fn is_top_layer(&self, point: &Point<Pixels>, layer: &StackingOrder) -> bool {
         // Precondition: the depth map is ordered from topmost to bottomost.
 
-        for (opaque_layer, _, bounds) in self.window.rendered_frame.depth_map.iter() {
+        for (opaque_layer, _, bounds) in self.window.next_frame.depth_map.iter() {
             if layer >= opaque_layer {
                 // The queried layer is either above or is the same as the this opaque layer.
                 // Anything after this point is guaranteed to be below the queried layer.
@@ -902,7 +902,7 @@ impl<'a> WindowContext<'a> {
     ) -> bool {
         // Precondition: the depth map is ordered from topmost to bottomost.
 
-        for (opaque_layer, _, bounds) in self.window.rendered_frame.depth_map.iter() {
+        for (opaque_layer, _, bounds) in self.window.next_frame.depth_map.iter() {
             if layer >= opaque_layer {
                 // The queried layer is either above or is the same as the this opaque layer.
                 // Anything after this point is guaranteed to be below the queried layer.
@@ -960,8 +960,83 @@ impl<'a> WindowContext<'a> {
             requested_handler.handler = input_handler;
         }
 
+        // Layout root views
+
         let root_view = self.window.root_view.take().unwrap();
-        self.with_element_context(|cx| {
+        let mut root_element = self.with_element_context(true, |cx| {
+            cx.with_z_index(0, |cx| {
+                let available_space = cx.window.viewport_size.map(Into::into);
+                let mut root_element = root_view.layout(Point::default(), available_space, cx);
+                root_element
+            })
+        });
+
+        let mut drag_element = None;
+        let mut tooltip_element = None;
+
+        self.with_element_context(true, |cx| {
+            cx.with_z_index(1, |cx| {
+                if let Some(active_drag) = cx.app.active_drag.take() {
+                    let offset = cx.mouse_position() - active_drag.cursor_offset;
+                    let available_space =
+                        size(AvailableSpace::MinContent, AvailableSpace::MinContent);
+                    let mut element = active_drag.view.layout(offset, available_space, cx);
+                    drag_element = Some(element);
+                    cx.active_drag = Some(active_drag);
+                } else if let Some(tooltip_request) = cx.window.next_frame.tooltip_request.take() {
+                    let available_space =
+                        size(AvailableSpace::MinContent, AvailableSpace::MinContent);
+                    let mut element = tooltip_request.tooltip.view.layout(
+                        tooltip_request.tooltip.cursor_offset,
+                        available_space,
+                        cx,
+                    );
+                    tooltip_element = Some(element);
+                    cx.window.next_frame.tooltip_request = Some(tooltip_request);
+                }
+            })
+        });
+
+        assert_eq!(self.window.next_frame.z_index_stack.len(), 0);
+        self.window.next_frame.next_stacking_order_ids = vec![0];
+
+        // Painting dry run
+
+        self.with_element_context(true, |cx| {
+            cx.with_z_index(0, |cx| {
+                cx.with_key_dispatch(Some(KeyContext::default()), None, |_, cx| {
+                    root_view.paint(Point::default(), &mut root_element, cx);
+                })
+            })
+        });
+
+        self.with_element_context(false, |cx| {
+            cx.with_z_index(1, |cx| {
+                if let Some(drag_element) = &mut drag_element {
+                    if let Some(active_drag) = cx.app.active_drag.take() {
+                        let offset = cx.mouse_position() - active_drag.cursor_offset;
+                        active_drag.view.paint(offset, drag_element, cx);
+                    }
+                }
+
+                if let Some(tooltip_element) = &mut tooltip_element {
+                    if let Some(tooltip_request) = cx.window.next_frame.tooltip_request.take() {
+                        tooltip_request.tooltip.view.paint(
+                            tooltip_request.tooltip.cursor_offset,
+                            tooltip_element,
+                            cx,
+                        );
+                    }
+                }
+            })
+        });
+
+        assert_eq!(self.window.next_frame.z_index_stack.len(), 0);
+        self.window.next_frame.next_stacking_order_ids = vec![0];
+
+        // Actual painting run
+
+        self.with_element_context(false, |cx| {
             cx.with_z_index(0, |cx| {
                 cx.with_key_dispatch(Some(KeyContext::default()), None, |_, cx| {
                     // We need to use cx.cx here so we can utilize borrow splitting
@@ -978,36 +1053,33 @@ impl<'a> WindowContext<'a> {
                         }
                     }
 
-                    let available_space = cx.window.viewport_size.map(Into::into);
-                    root_view.draw(Point::default(), available_space, cx);
+                    root_view.paint(Point::default(), &mut root_element, cx);
                 })
             })
         });
 
-        if let Some(active_drag) = self.app.active_drag.take() {
-            self.with_element_context(|cx| {
-                cx.with_z_index(ACTIVE_DRAG_Z_INDEX, |cx| {
-                    let offset = cx.mouse_position() - active_drag.cursor_offset;
-                    let available_space =
-                        size(AvailableSpace::MinContent, AvailableSpace::MinContent);
-                    active_drag.view.draw(offset, available_space, cx);
-                })
-            });
-            self.active_drag = Some(active_drag);
-        } else if let Some(tooltip_request) = self.window.next_frame.tooltip_request.take() {
-            self.with_element_context(|cx| {
-                cx.with_z_index(1, |cx| {
-                    let available_space =
-                        size(AvailableSpace::MinContent, AvailableSpace::MinContent);
-                    tooltip_request.tooltip.view.draw(
-                        tooltip_request.tooltip.cursor_offset,
-                        available_space,
-                        cx,
-                    );
-                })
-            });
-            self.window.next_frame.tooltip_request = Some(tooltip_request);
-        }
+        self.with_element_context(false, |cx| {
+            cx.with_z_index(1, |cx| {
+                if let Some(mut drag_element) = drag_element {
+                    if let Some(active_drag) = cx.app.active_drag.take() {
+                        let offset = cx.mouse_position() - active_drag.cursor_offset;
+                        active_drag.view.paint(offset, &mut drag_element, cx);
+                        cx.active_drag = Some(active_drag);
+                    }
+                }
+
+                if let Some(mut tooltip_element) = tooltip_element {
+                    if let Some(tooltip_request) = cx.window.next_frame.tooltip_request.take() {
+                        tooltip_request.tooltip.view.paint(
+                            tooltip_request.tooltip.cursor_offset,
+                            &mut tooltip_element,
+                            cx,
+                        );
+                    }
+                }
+            })
+        });
+
         self.window.dirty_views.clear();
 
         self.window
@@ -1206,7 +1278,7 @@ impl<'a> WindowContext<'a> {
             // Capture phase, events bubble from back to front. Handlers for this phase are used for
             // special purposes, such as detecting events outside of a given Bounds.
             for (_, _, handler) in &mut handlers {
-                self.with_element_context(|cx| {
+                self.with_element_context(false, |cx| {
                     handler(event, DispatchPhase::Capture, cx);
                 });
                 if !self.app.propagate_event {
@@ -1217,7 +1289,7 @@ impl<'a> WindowContext<'a> {
             // Bubble phase, where most normal handlers do their work.
             if self.app.propagate_event {
                 for (_, _, handler) in handlers.iter_mut().rev() {
-                    self.with_element_context(|cx| {
+                    self.with_element_context(false, |cx| {
                         handler(event, DispatchPhase::Bubble, cx);
                     });
                     if !self.app.propagate_event {
@@ -1345,7 +1417,7 @@ impl<'a> WindowContext<'a> {
             let node = self.window.rendered_frame.dispatch_tree.node(*node_id);
 
             for key_listener in node.key_listeners.clone() {
-                self.with_element_context(|cx| {
+                self.with_element_context(false, |cx| {
                     key_listener(event, DispatchPhase::Capture, cx);
                 });
                 if !self.propagate_event {
@@ -1359,7 +1431,7 @@ impl<'a> WindowContext<'a> {
             // Handle low level key events
             let node = self.window.rendered_frame.dispatch_tree.node(*node_id);
             for key_listener in node.key_listeners.clone() {
-                self.with_element_context(|cx| {
+                self.with_element_context(false, |cx| {
                     key_listener(event, DispatchPhase::Bubble, cx);
                 });
                 if !self.propagate_event {
@@ -1446,7 +1518,7 @@ impl<'a> WindowContext<'a> {
             {
                 let any_action = action.as_any();
                 if action_type == any_action.type_id() {
-                    self.with_element_context(|cx| {
+                    self.with_element_context(false, |cx| {
                         listener(any_action, DispatchPhase::Capture, cx);
                     });
 
@@ -1468,7 +1540,7 @@ impl<'a> WindowContext<'a> {
                 if action_type == any_action.type_id() {
                     self.propagate_event = false; // Actions stop propagation by default during the bubble phase
 
-                    self.with_element_context(|cx| {
+                    self.with_element_context(false, |cx| {
                         listener(any_action, DispatchPhase::Bubble, cx);
                     });
 

crates/gpui/src/window/element_cx.rs

@@ -164,16 +164,26 @@ impl Frame {
 /// This context is used for assisting in the implementation of the element trait
 #[derive(Deref, DerefMut)]
 pub struct ElementContext<'a> {
+    #[deref]
+    #[deref_mut]
     pub(crate) cx: WindowContext<'a>,
+    pub(crate) painting_dry_run: bool,
+    pub(crate) disallow_adding_opaque_layers: bool,
 }
 
 impl<'a> WindowContext<'a> {
     /// Convert this window context into an ElementContext in this callback.
     /// If you need to use this method, you're probably intermixing the imperative
     /// and declarative APIs, which is not recommended.
-    pub fn with_element_context<R>(&mut self, f: impl FnOnce(&mut ElementContext) -> R) -> R {
+    pub fn with_element_context<R>(
+        &mut self,
+        painting_dry_run: bool,
+        f: impl FnOnce(&mut ElementContext) -> R,
+    ) -> R {
         f(&mut ElementContext {
             cx: WindowContext::new(self.app, self.window),
+            painting_dry_run,
+            disallow_adding_opaque_layers: false,
         })
     }
 }
@@ -310,50 +320,53 @@ impl<'a> VisualContext for ElementContext<'a> {
 }
 
 impl<'a> ElementContext<'a> {
-    pub(crate) fn reuse_view(&mut self, next_stacking_order_id: u16) {
-        let view_id = self.parent_view_id();
-        let grafted_view_ids = self
-            .cx
-            .window
-            .next_frame
-            .dispatch_tree
-            .reuse_view(view_id, &mut self.cx.window.rendered_frame.dispatch_tree);
-        for view_id in grafted_view_ids {
-            assert!(self.window.next_frame.reused_views.insert(view_id));
-
-            // Reuse the previous input handler requested during painting of the reused view.
-            if self
+    pub(crate) fn reuse_view(&mut self, stored_next_stacking_order_id: u16) {
+        if !self.painting_dry_run {
+            let view_id = self.parent_view_id();
+            let grafted_view_ids = self
+                .cx
                 .window
-                .rendered_frame
-                .requested_input_handler
-                .as_ref()
-                .map_or(false, |requested| requested.view_id == view_id)
-            {
-                self.window.next_frame.requested_input_handler =
-                    self.window.rendered_frame.requested_input_handler.take();
-            }
+                .next_frame
+                .dispatch_tree
+                .reuse_view(view_id, &mut self.cx.window.rendered_frame.dispatch_tree);
 
-            // Reuse the tooltip previously requested during painting of the reused view.
-            if self
-                .window
-                .rendered_frame
-                .tooltip_request
-                .as_ref()
-                .map_or(false, |requested| requested.view_id == view_id)
-            {
-                self.window.next_frame.tooltip_request =
-                    self.window.rendered_frame.tooltip_request.take();
-            }
+            for view_id in grafted_view_ids {
+                assert!(self.window.next_frame.reused_views.insert(view_id));
 
-            // Reuse the cursor styles previously requested during painting of the reused view.
-            if let Some(style) = self.window.rendered_frame.cursor_styles.remove(&view_id) {
-                self.window.next_frame.cursor_styles.insert(view_id, style);
-                self.window.next_frame.requested_cursor_style = Some(style);
+                // Reuse the previous input handler requested during painting of the reused view.
+                if self
+                    .window
+                    .rendered_frame
+                    .requested_input_handler
+                    .as_ref()
+                    .map_or(false, |requested| requested.view_id == view_id)
+                {
+                    self.window.next_frame.requested_input_handler =
+                        self.window.rendered_frame.requested_input_handler.take();
+                }
+
+                // Reuse the tooltip previously requested during painting of the reused view.
+                if self
+                    .window
+                    .rendered_frame
+                    .tooltip_request
+                    .as_ref()
+                    .map_or(false, |requested| requested.view_id == view_id)
+                {
+                    self.window.next_frame.tooltip_request =
+                        self.window.rendered_frame.tooltip_request.take();
+                }
+
+                // Reuse the cursor styles previously requested during painting of the reused view.
+                if let Some(style) = self.window.rendered_frame.cursor_styles.remove(&view_id) {
+                    self.window.next_frame.cursor_styles.insert(view_id, style);
+                    self.window.next_frame.requested_cursor_style = Some(style);
+                }
             }
         }
 
         debug_assert!(
-            next_stacking_order_id
+            stored_next_stacking_order_id
                 >= self
                     .window
                     .next_frame
@@ -367,7 +380,7 @@ impl<'a> ElementContext<'a> {
             .next_frame
             .next_stacking_order_ids
             .last_mut()
-            .unwrap() = next_stacking_order_id;
+            .unwrap() = stored_next_stacking_order_id;
     }
 
     /// Push a text style onto the stack, and call a function with that style active.
@@ -497,6 +510,27 @@ impl<'a> ElementContext<'a> {
         result
     }
 
+    pub(crate) fn with_z_index_reset<R>(&mut self, continuation: impl FnOnce(&mut Self) -> R) -> R {
+        let frame = &self.window_mut().next_frame;
+        let original_next_stacking_order_id = *frame.next_stacking_order_ids.last().unwrap();
+        let original_stacking_order_ids_len = frame.next_stacking_order_ids.len();
+
+        let original_disallow = self.disallow_adding_opaque_layers;
+        self.disallow_adding_opaque_layers = true;
+
+        let result = continuation(self);
+
+        self.disallow_adding_opaque_layers = original_disallow;
+
+        let frame = &mut self.window_mut().next_frame;
+        frame
+            .next_stacking_order_ids
+            .truncate(original_stacking_order_ids_len);
+        *frame.next_stacking_order_ids.last_mut().unwrap() = original_next_stacking_order_id;
+
+        result
+    }
+
     /// Updates the global element offset relative to the current offset. This is used to implement
     /// scrolling.
     pub fn with_element_offset<R>(
@@ -572,84 +606,84 @@ impl<'a> ElementContext<'a> {
         S: 'static,
     {
         self.with_element_id(Some(id), |cx| {
-                let global_id = cx.window().element_id_stack.clone();
+            let global_id = cx.window().element_id_stack.clone();
 
-                if let Some(any) = cx
-                    .window_mut()
+            if let Some(any) = cx
+                .window_mut()
+                .next_frame
+                .element_states
+                .remove(&global_id)
+                .or_else(|| {
+                    cx.window_mut()
+                        .rendered_frame
+                        .element_states
+                        .remove(&global_id)
+                })
+            {
+                let ElementStateBox {
+                    inner,
+                    parent_view_id,
+                    #[cfg(debug_assertions)]
+                    type_name
+                } = any;
+                // Using the extra inner option to avoid needing to reallocate a new box.
+                let mut state_box = inner
+                    .downcast::<Option<S>>()
+                    .map_err(|_| {
+                        #[cfg(debug_assertions)]
+                        {
+                            anyhow::anyhow!(
+                                "invalid element state type for id, requested_type {:?}, actual type: {:?}",
+                                std::any::type_name::<S>(),
+                                type_name
+                            )
+                        }
+
+                        #[cfg(not(debug_assertions))]
+                        {
+                            anyhow::anyhow!(
+                                "invalid element state type for id, requested_type {:?}",
+                                std::any::type_name::<S>(),
+                            )
+                        }
+                    })
+                    .unwrap();
+
+                // Actual: Option<AnyElement> <- View
+                // Requested: () <- AnyElement
+                let state = state_box
+                    .take()
+                    .expect("element state is already on the stack");
+                let (result, state) = f(Some(state), cx);
+                state_box.replace(state);
+                cx.window_mut()
                     .next_frame
                     .element_states
-                    .remove(&global_id)
-                    .or_else(|| {
-                        cx.window_mut()
-                            .rendered_frame
-                            .element_states
-                            .remove(&global_id)
-                    })
-                {
-                    let ElementStateBox {
-                        inner,
+                    .insert(global_id, ElementStateBox {
+                        inner: state_box,
                         parent_view_id,
                         #[cfg(debug_assertions)]
                         type_name
-                    } = any;
-                    // Using the extra inner option to avoid needing to reallocate a new box.
-                    let mut state_box = inner
-                        .downcast::<Option<S>>()
-                        .map_err(|_| {
-                            #[cfg(debug_assertions)]
-                            {
-                                anyhow::anyhow!(
-                                    "invalid element state type for id, requested_type {:?}, actual type: {:?}",
-                                    std::any::type_name::<S>(),
-                                    type_name
-                                )
-                            }
-
-                            #[cfg(not(debug_assertions))]
-                            {
-                                anyhow::anyhow!(
-                                    "invalid element state type for id, requested_type {:?}",
-                                    std::any::type_name::<S>(),
-                                )
-                            }
-                        })
-                        .unwrap();
-
-                    // Actual: Option<AnyElement> <- View
-                    // Requested: () <- AnyElement
-                    let state = state_box
-                        .take()
-                        .expect("element state is already on the stack");
-                    let (result, state) = f(Some(state), cx);
-                    state_box.replace(state);
-                    cx.window_mut()
-                        .next_frame
-                        .element_states
-                        .insert(global_id, ElementStateBox {
-                            inner: state_box,
+                    });
+                result
+            } else {
+                let (result, state) = f(None, cx);
+                let parent_view_id = cx.parent_view_id();
+                cx.window_mut()
+                    .next_frame
+                    .element_states
+                    .insert(global_id,
+                        ElementStateBox {
+                            inner: Box::new(Some(state)),
                             parent_view_id,
                             #[cfg(debug_assertions)]
-                            type_name
-                        });
-                    result
-                } else {
-                    let (result, state) = f(None, cx);
-                    let parent_view_id = cx.parent_view_id();
-                    cx.window_mut()
-                        .next_frame
-                        .element_states
-                        .insert(global_id,
-                            ElementStateBox {
-                                inner: Box::new(Some(state)),
-                                parent_view_id,
-                                #[cfg(debug_assertions)]
-                                type_name: std::any::type_name::<S>()
-                            }
+                            type_name: std::any::type_name::<S>()
+                        }
 
-                        );
-                    result
-                }
-            })
+                    );
+                result
+            }
+        })
     }
     /// Paint one or more drop shadows into the scene for the next frame at the current z-index.
     pub fn paint_shadows(
@@ -658,6 +692,10 @@ impl<'a> ElementContext<'a> {
         corner_radii: Corners<Pixels>,
         shadows: &[BoxShadow],
     ) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let scale_factor = self.scale_factor();
         let content_mask = self.content_mask();
         let view_id = self.parent_view_id();
@@ -687,6 +725,10 @@ impl<'a> ElementContext<'a> {
     /// Quads are colored rectangular regions with an optional background, border, and corner radius.
     /// see [`fill`](crate::fill), [`outline`](crate::outline), and [`quad`](crate::quad) to construct this type.
     pub fn paint_quad(&mut self, quad: PaintQuad) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let scale_factor = self.scale_factor();
         let content_mask = self.content_mask();
         let view_id = self.parent_view_id();
@@ -710,6 +752,10 @@ impl<'a> ElementContext<'a> {
 
     /// Paint the given `Path` into the scene for the next frame at the current z-index.
     pub fn paint_path(&mut self, mut path: Path<Pixels>, color: impl Into<Hsla>) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let scale_factor = self.scale_factor();
         let content_mask = self.content_mask();
         let view_id = self.parent_view_id();
@@ -731,6 +777,10 @@ impl<'a> ElementContext<'a> {
         width: Pixels,
         style: &UnderlineStyle,
     ) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let scale_factor = self.scale_factor();
         let height = if style.wavy {
             style.thickness * 3.
@@ -767,6 +817,10 @@ impl<'a> ElementContext<'a> {
         width: Pixels,
         style: &StrikethroughStyle,
     ) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let scale_factor = self.scale_factor();
         let height = style.thickness;
         let bounds = Bounds {
@@ -806,6 +860,10 @@ impl<'a> ElementContext<'a> {
         font_size: Pixels,
         color: Hsla,
     ) -> Result<()> {
+        if self.painting_dry_run {
+            return Ok(());
+        }
+
         let scale_factor = self.scale_factor();
         let glyph_origin = origin.scale(scale_factor);
         let subpixel_variant = Point {
@@ -866,6 +924,10 @@ impl<'a> ElementContext<'a> {
         glyph_id: GlyphId,
         font_size: Pixels,
     ) -> Result<()> {
+        if self.painting_dry_run {
+            return Ok(());
+        }
+
         let scale_factor = self.scale_factor();
         let glyph_origin = origin.scale(scale_factor);
         let params = RenderGlyphParams {
@@ -920,6 +982,10 @@ impl<'a> ElementContext<'a> {
         path: SharedString,
         color: Hsla,
     ) -> Result<()> {
+        if self.painting_dry_run {
+            return Ok(());
+        }
+
         let scale_factor = self.scale_factor();
         let bounds = bounds.scale(scale_factor);
         // Render the SVG at twice the size to get a higher quality result.
@@ -965,6 +1031,10 @@ impl<'a> ElementContext<'a> {
         data: Arc<ImageData>,
         grayscale: bool,
     ) -> Result<()> {
+        if self.painting_dry_run {
+            return Ok(());
+        }
+
         let scale_factor = self.scale_factor();
         let bounds = bounds.scale(scale_factor);
         let params = RenderImageParams { image_id: data.id };
@@ -1094,6 +1164,10 @@ impl<'a> ElementContext<'a> {
 
     /// Called during painting to track which z-index is on top at each pixel position
     pub fn add_opaque_layer(&mut self, bounds: Bounds<Pixels>) {
+        if self.disallow_adding_opaque_layers || !self.painting_dry_run {
+            return;
+        }
+
         let stacking_order = self.window.next_frame.z_index_stack.clone();
         let view_id = self.parent_view_id();
         let depth_map = &mut self.window.next_frame.depth_map;
@@ -1110,16 +1184,20 @@ impl<'a> ElementContext<'a> {
         focus_handle: Option<FocusHandle>,
         f: impl FnOnce(Option<FocusHandle>, &mut Self) -> R,
     ) -> R {
-        let window = &mut self.window;
-        let focus_id = focus_handle.as_ref().map(|handle| handle.id);
-        window
-            .next_frame
-            .dispatch_tree
-            .push_node(context.clone(), focus_id, None);
+        if !self.painting_dry_run {
+            let window = &mut self.window;
+            let focus_id = focus_handle.as_ref().map(|handle| handle.id);
+            window
+                .next_frame
+                .dispatch_tree
+                .push_node(context.clone(), focus_id, None);
+        }
 
         let result = f(focus_handle, self);
 
-        self.window.next_frame.dispatch_tree.pop_node();
+        if !self.painting_dry_run {
+            self.window.next_frame.dispatch_tree.pop_node();
+        }
 
         result
     }
@@ -1168,6 +1246,10 @@ impl<'a> ElementContext<'a> {
     ///
     /// [element_input_handler]: crate::ElementInputHandler
     pub fn handle_input(&mut self, focus_handle: &FocusHandle, input_handler: impl InputHandler) {
+        if self.painting_dry_run {
+            return;
+        }
+
         if focus_handle.is_focused(self) {
             let view_id = self.parent_view_id();
             self.window.next_frame.requested_input_handler = Some(RequestedInputHandler {
@@ -1187,6 +1269,10 @@ impl<'a> ElementContext<'a> {
         &mut self,
         mut handler: impl FnMut(&Event, DispatchPhase, &mut ElementContext) + 'static,
     ) {
+        if self.painting_dry_run {
+            return;
+        }
+
         let view_id = self.parent_view_id();
         let order = self.window.next_frame.z_index_stack.clone();
         self.window
@@ -1215,6 +1301,10 @@ impl<'a> ElementContext<'a> {
         &mut self,
         listener: impl Fn(&Event, DispatchPhase, &mut ElementContext) + 'static,
     ) {
+        if self.painting_dry_run {
+            return;
+        }
+
         self.window.next_frame.dispatch_tree.on_key_event(Rc::new(
             move |event: &dyn Any, phase, cx: &mut ElementContext<'_>| {
                 if let Some(event) = event.downcast_ref::<Event>() {

crates/language_tools/src/syntax_tree_view.rs

@@ -358,13 +358,11 @@ impl Render for SyntaxTreeView {
             .track_scroll(self.list_scroll_handle.clone())
             .text_bg(cx.theme().colors().background);
 
+            let mut list_element = list.into_any_element();
+
             rendered = rendered.child(
                 canvas(move |bounds, cx| {
-                    list.into_any_element().draw(
-                        bounds.origin,
-                        bounds.size.map(AvailableSpace::Definite),
-                        cx,
-                    )
+                    list_element.draw(bounds.origin, bounds.size.map(AvailableSpace::Definite), cx)
                 })
                 .size_full(),
             );

crates/terminal_view/src/terminal_element.rs

@@ -854,7 +854,7 @@ impl Element for TerminalElement {
                     });
                 }
 
-                if let Some(mut element) = layout.hyperlink_tooltip.take() {
+                if let Some(element) = layout.hyperlink_tooltip.as_mut() {
                     element.draw(origin, bounds.size.map(AvailableSpace::Definite), cx)
                 }
             });

crates/ui/src/components/popover_menu.rs

@@ -191,15 +191,15 @@ impl<M: ManagedView> Element for PopoverMenu<M> {
         element_state: &mut Self::State,
         cx: &mut ElementContext,
     ) {
-        if let Some(mut child) = element_state.child_element.take() {
+        if let Some(child) = element_state.child_element.as_mut() {
             child.paint(cx);
         }
 
-        if let Some(child_layout_id) = element_state.child_layout_id.take() {
-            element_state.child_bounds = Some(cx.layout_bounds(child_layout_id));
+        if let Some(child_layout_id) = &element_state.child_layout_id {
+            element_state.child_bounds = Some(cx.layout_bounds(*child_layout_id));
         }
 
-        if let Some(mut menu) = element_state.menu_element.take() {
+        if let Some(menu) = element_state.menu_element.as_mut() {
             menu.paint(cx);
 
             if let Some(child_bounds) = element_state.child_bounds {

crates/ui/src/components/right_click_menu.rs

@@ -118,16 +118,16 @@ impl<M: ManagedView> Element for RightClickMenu<M> {
         element_state: &mut Self::State,
         cx: &mut ElementContext,
     ) {
-        if let Some(mut child) = element_state.child_element.take() {
+        if let Some(child) = element_state.child_element.as_mut() {
             child.paint(cx);
         }
 
-        if let Some(mut menu) = element_state.menu_element.take() {
+        if let Some(menu) = element_state.menu_element.as_mut() {
             menu.paint(cx);
             return;
         }
 
-        let Some(builder) = self.menu_builder.take() else {
+        let Some(builder) = self.menu_builder.clone() else {
             return;
         };
         let menu = element_state.menu.clone();