diff --git a/crates/gpui2/src/element.rs b/crates/gpui2/src/element.rs
index b18ffb8ca6..3c8f678b89 100644
--- a/crates/gpui2/src/element.rs
+++ b/crates/gpui2/src/element.rs
@@ -111,7 +111,7 @@ pub struct Component<C> {
 
 pub struct CompositeElementState<C: RenderOnce> {
     rendered_element: Option<<C::Rendered as IntoElement>::Element>,
-    rendered_element_state: <<C::Rendered as IntoElement>::Element as Element>::State,
+    rendered_element_state: Option<<<C::Rendered as IntoElement>::Element as Element>::State>,
 }
 
 impl<C> Component<C> {
@@ -131,20 +131,40 @@ impl<C: RenderOnce> Element for Component<C> {
         cx: &mut WindowContext,
     ) -> (LayoutId, Self::State) {
         let mut element = self.component.take().unwrap().render(cx).into_element();
-        let (layout_id, state) = element.layout(state.map(|s| s.rendered_element_state), cx);
-        let state = CompositeElementState {
-            rendered_element: Some(element),
-            rendered_element_state: state,
-        };
-        (layout_id, state)
+        if let Some(element_id) = element.element_id() {
+            let layout_id =
+                cx.with_element_state(element_id, |state, cx| element.layout(state, cx));
+            let state = CompositeElementState {
+                rendered_element: Some(element),
+                rendered_element_state: None,
+            };
+            (layout_id, state)
+        } else {
+            let (layout_id, state) =
+                element.layout(state.and_then(|s| s.rendered_element_state), cx);
+            let state = CompositeElementState {
+                rendered_element: Some(element),
+                rendered_element_state: Some(state),
+            };
+            (layout_id, state)
+        }
     }
 
     fn paint(self, bounds: Bounds<Pixels>, state: &mut Self::State, cx: &mut WindowContext) {
-        state
-            .rendered_element
-            .take()
-            .unwrap()
-            .paint(bounds, &mut state.rendered_element_state, cx);
+        let element = state.rendered_element.take().unwrap();
+        if let Some(element_id) = element.element_id() {
+            cx.with_element_state(element_id, |element_state, cx| {
+                let mut element_state = element_state.unwrap();
+                element.paint(bounds, &mut element_state, cx);
+                ((), element_state)
+            });
+        } else {
+            element.paint(
+                bounds,
+                &mut state.rendered_element_state.as_mut().unwrap(),
+                cx,
+            );
+        }
     }
 }
 
diff --git a/crates/gpui2/src/window.rs b/crates/gpui2/src/window.rs
index 6f342f7065..b02005aa78 100644
--- a/crates/gpui2/src/window.rs
+++ b/crates/gpui2/src/window.rs
@@ -1939,23 +1939,6 @@ pub trait BorrowWindow: BorrowMut<Window> + BorrowMut<AppContext> {
         })
     }
 
-    /// Like `with_element_state`, but for situations where the element_id is optional. If the
-    /// id is `None`, no state will be retrieved or stored.
-    fn with_optional_element_state<S, R>(
-        &mut self,
-        element_id: Option<ElementId>,
-        f: impl FnOnce(Option<S>, &mut Self) -> (R, S),
-    ) -> R
-    where
-        S: 'static,
-    {
-        if let Some(element_id) = element_id {
-            self.with_element_state(element_id, f)
-        } else {
-            f(None, self).0
-        }
-    }
-
     /// Obtain the current content mask.
     fn content_mask(&self) -> ContentMask<Pixels> {
         self.window()
diff --git a/crates/ui2/src/components/icon_button.rs b/crates/ui2/src/components/icon_button.rs
index 104eb00ec8..2ccda3ea0e 100644
--- a/crates/ui2/src/components/icon_button.rs
+++ b/crates/ui2/src/components/icon_button.rs
@@ -70,8 +70,7 @@ impl RenderOnce for IconButton {
             }
         }
 
-        // HACK: Add an additional identified element wrapper to fix tooltips not showing up.
-        div().id(self.id.clone()).child(button)
+        button
     }
 }