Move Arena to a thread-local and use it to allocate AnyElement

crates/gpui2/src/arena.rs

@@ -76,6 +76,10 @@ impl<T: ?Sized> ArenaRef<T> {
         ArenaRef(NonNull::new_unchecked(u))
     }
 
+    pub unsafe fn get(&self) -> &T {
+        self.0.as_ref()
+    }
+
     pub unsafe fn get_mut(&mut self) -> &mut T {
         self.0.as_mut()
     }

crates/gpui2/src/element.rs

@@ -1,6 +1,6 @@
 use crate::{
-    AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels, Point, Size, ViewContext,
+    arena::ArenaRef, AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels, Point,
-    WindowContext,
+    Size, ViewContext, WindowContext, FRAME_ARENA,
 };
 use derive_more::{Deref, DerefMut};
 pub(crate) use smallvec::SmallVec;
@@ -405,7 +405,7 @@ where
     }
 }
 
-pub struct AnyElement(Box<dyn ElementObject>);
+pub struct AnyElement(ArenaRef<dyn ElementObject>);
 
 impl AnyElement {
     pub fn new<E>(element: E) -> Self
@@ -413,15 +413,18 @@ impl AnyElement {
         E: 'static + Element,
         E::State: Any,
     {
-        AnyElement(Box::new(Some(DrawableElement::new(element))) as Box<dyn ElementObject>)
+        let element =
+            FRAME_ARENA.with_borrow_mut(|arena| arena.alloc(Some(DrawableElement::new(element))));
+        let element = unsafe { element.map(|element| element as &mut dyn ElementObject) };
+        AnyElement(element)
     }
 
     pub fn layout(&mut self, cx: &mut WindowContext) -> LayoutId {
-        self.0.layout(cx)
+        unsafe { self.0.get_mut() }.layout(cx)
     }
 
     pub fn paint(&mut self, cx: &mut WindowContext) {
-        self.0.paint(cx)
+        unsafe { self.0.get_mut() }.paint(cx)
     }
 
     /// Initializes this element and performs layout within the given available space to determine its size.
@@ -430,7 +433,7 @@ impl AnyElement {
         available_space: Size<AvailableSpace>,
         cx: &mut WindowContext,
     ) -> Size<Pixels> {
-        self.0.measure(available_space, cx)
+        unsafe { self.0.get_mut() }.measure(available_space, cx)
     }
 
     /// Initializes this element and performs layout in the available space, then paints it at the given origin.
@@ -440,16 +443,11 @@ impl AnyElement {
         available_space: Size<AvailableSpace>,
         cx: &mut WindowContext,
     ) {
-        self.0.draw(origin, available_space, cx)
+        unsafe { self.0.get_mut() }.draw(origin, available_space, cx)
-    }
-
-    /// Converts this `AnyElement` into a trait object that can be stored and manipulated.
-    pub fn into_any(self) -> AnyElement {
-        AnyElement::new(self)
     }
 
     pub fn inner_id(&self) -> Option<ElementId> {
-        self.0.element_id()
+        unsafe { self.0.get() }.element_id()
     }
 }
 
@@ -480,6 +478,10 @@ impl IntoElement for AnyElement {
     fn into_element(self) -> Self::Element {
         self
     }
+
+    fn into_any_element(self) -> AnyElement {
+        self
+    }
 }
 
 /// The empty element, which renders nothing.

crates/gpui2/src/window.rs

@@ -27,6 +27,7 @@ use smallvec::SmallVec;
 use std::{
     any::{Any, TypeId},
     borrow::{Borrow, BorrowMut, Cow},
+    cell::RefCell,
     fmt::Debug,
     future::Future,
     hash::{Hash, Hasher},
@@ -97,6 +98,10 @@ struct FocusEvent {
 
 slotmap::new_key_type! { pub struct FocusId; }
 
+thread_local! {
+    pub static FRAME_ARENA: RefCell<Arena> = RefCell::new(Arena::default());
+}
+
 impl FocusId {
     /// Obtains whether the element associated with this handle is currently focused.
     pub fn is_focused(&self, cx: &WindowContext) -> bool {
@@ -272,7 +277,6 @@ pub(crate) struct ElementStateBox {
 
 pub(crate) struct Frame {
     focus: Option<FocusId>,
-    arena: Arena,
     pub(crate) element_states: FxHashMap<GlobalElementId, ElementStateBox>,
     mouse_listeners: FxHashMap<TypeId, Vec<(StackingOrder, AnyMouseListener)>>,
     pub(crate) dispatch_tree: DispatchTree,
@@ -287,7 +291,6 @@ impl Frame {
     fn new(dispatch_tree: DispatchTree) -> Self {
         Frame {
             focus: None,
-            arena: Arena::default(),
             element_states: FxHashMap::default(),
             mouse_listeners: FxHashMap::default(),
             dispatch_tree,
@@ -302,7 +305,6 @@ impl Frame {
     fn clear(&mut self) {
         self.element_states.clear();
         self.mouse_listeners.values_mut().for_each(Vec::clear);
-        self.arena.clear();
         self.dispatch_tree.clear();
         self.depth_map.clear();
     }
@@ -827,11 +829,13 @@ impl<'a> WindowContext<'a> {
         mut handler: impl FnMut(&Event, DispatchPhase, &mut WindowContext) + 'static,
     ) {
         let order = self.window.next_frame.z_index_stack.clone();
-        let handler = self.window.next_frame.arena.alloc(
+        let handler = FRAME_ARENA.with_borrow_mut(|arena| {
-            move |event: &dyn Any, phase: DispatchPhase, cx: &mut WindowContext<'_>| {
+            arena.alloc(
-                handler(event.downcast_ref().unwrap(), phase, cx)
+                move |event: &dyn Any, phase: DispatchPhase, cx: &mut WindowContext<'_>| {
-            },
+                    handler(event.downcast_ref().unwrap(), phase, cx)
-        );
+                },
+            )
+        });
         let handler = unsafe { handler.map(|handler| handler as _) };
         self.window
             .next_frame
@@ -851,13 +855,13 @@ impl<'a> WindowContext<'a> {
         &mut self,
         listener: impl Fn(&Event, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        let listener = self.window.next_frame.arena.alloc(
+        let listener = FRAME_ARENA.with_borrow_mut(|arena| {
-            move |event: &dyn Any, phase, cx: &mut WindowContext<'_>| {
+            arena.alloc(move |event: &dyn Any, phase, cx: &mut WindowContext<'_>| {
                 if let Some(event) = event.downcast_ref::<Event>() {
                     listener(event, phase, cx)
                 }
-            },
+            })
-        );
+        });
         let listener = unsafe { listener.map(|handler| handler as _) };
         self.window.next_frame.dispatch_tree.on_key_event(listener);
     }
@@ -873,7 +877,7 @@ impl<'a> WindowContext<'a> {
         action_type: TypeId,
         listener: impl Fn(&dyn Any, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        let listener = self.window.next_frame.arena.alloc(listener);
+        let listener = FRAME_ARENA.with_borrow_mut(|arena| arena.alloc(listener));
         let listener = unsafe { listener.map(|handler| handler as _) };
         self.window
             .next_frame
@@ -1273,17 +1277,20 @@ impl<'a> WindowContext<'a> {
         self.window.platform_window.clear_input_handler();
         self.window.layout_engine.as_mut().unwrap().clear();
         self.window.next_frame.clear();
+        FRAME_ARENA.with_borrow_mut(|arena| arena.clear());
         let root_view = self.window.root_view.take().unwrap();
 
         self.with_z_index(0, |cx| {
             cx.with_key_dispatch(Some(KeyContext::default()), None, |_, cx| {
                 for (action_type, action_listeners) in &cx.app.global_action_listeners {
                     for action_listener in action_listeners.iter().cloned() {
-                        let listener = cx.window.next_frame.arena.alloc(
+                        let listener = FRAME_ARENA.with_borrow_mut(|arena| {
-                            move |action: &dyn Any, phase, cx: &mut WindowContext<'_>| {
+                            arena.alloc(
-                                action_listener(action, phase, cx)
+                                move |action: &dyn Any, phase, cx: &mut WindowContext<'_>| {
-                            },
+                                    action_listener(action, phase, cx)
-                        );
+                                },
+                            )
+                        });
                         let listener = unsafe { listener.map(|listener| listener as _) };
                         cx.window
                             .next_frame