Introduce an Arena to allocate and deallocate frame-related structs more cheaply (#3689)

Picture worth a thousand words:


![before-after](https://github.com/zed-industries/zed/assets/482957/0aa92baf-f1cd-485e-a234-6d8f8b63a79a)

Note how having an area makes a substantial amount of frames between
0.5ms and 1ms faster (15-20% faster frames).

Release Notes:

- N/A

crates/gpui2/src/arena.rs

@@ -0,0 +1,178 @@
+use std::{
+    alloc,
+    cell::Cell,
+    ops::{Deref, DerefMut},
+    ptr::{self, NonNull},
+    rc::Rc,
+};
+
+struct ArenaElement {
+    value: NonNull<u8>,
+    drop: unsafe fn(NonNull<u8>),
+}
+
+impl Drop for ArenaElement {
+    #[inline(always)]
+    fn drop(&mut self) {
+        unsafe {
+            (self.drop)(self.value);
+        }
+    }
+}
+
+pub struct Arena {
+    start: NonNull<u8>,
+    offset: usize,
+    elements: Vec<ArenaElement>,
+    valid: Rc<Cell<bool>>,
+}
+
+impl Arena {
+    pub fn new(size_in_bytes: usize) -> Self {
+        unsafe {
+            let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap();
+            let ptr = alloc::alloc(layout);
+            Self {
+                start: NonNull::new_unchecked(ptr),
+                offset: 0,
+                elements: Vec::new(),
+                valid: Rc::new(Cell::new(true)),
+            }
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.valid.set(false);
+        self.valid = Rc::new(Cell::new(true));
+        self.elements.clear();
+        self.offset = 0;
+    }
+
+    #[inline(always)]
+    pub fn alloc<T>(&mut self, f: impl FnOnce() -> T) -> ArenaRef<T> {
+        #[inline(always)]
+        unsafe fn inner_writer<T, F>(ptr: *mut T, f: F)
+        where
+            F: FnOnce() -> T,
+        {
+            ptr::write(ptr, f());
+        }
+
+        unsafe fn drop<T>(ptr: NonNull<u8>) {
+            std::ptr::drop_in_place(ptr.cast::<T>().as_ptr());
+        }
+
+        unsafe {
+            let layout = alloc::Layout::new::<T>().pad_to_align();
+            let ptr = NonNull::new_unchecked(self.start.as_ptr().add(self.offset).cast::<T>());
+            inner_writer(ptr.as_ptr(), f);
+
+            self.elements.push(ArenaElement {
+                value: ptr.cast(),
+                drop: drop::<T>,
+            });
+            self.offset += layout.size();
+            ArenaRef {
+                ptr,
+                valid: self.valid.clone(),
+            }
+        }
+    }
+}
+
+impl Drop for Arena {
+    fn drop(&mut self) {
+        self.clear();
+    }
+}
+
+pub struct ArenaRef<T: ?Sized> {
+    ptr: NonNull<T>,
+    valid: Rc<Cell<bool>>,
+}
+
+impl<T: ?Sized> Clone for ArenaRef<T> {
+    fn clone(&self) -> Self {
+        Self {
+            ptr: self.ptr,
+            valid: self.valid.clone(),
+        }
+    }
+}
+
+impl<T: ?Sized> ArenaRef<T> {
+    #[inline(always)]
+    pub fn map<U: ?Sized>(mut self, f: impl FnOnce(&mut T) -> &mut U) -> ArenaRef<U> {
+        ArenaRef {
+            ptr: unsafe { NonNull::new_unchecked(f(&mut *self)) },
+            valid: self.valid,
+        }
+    }
+
+    fn validate(&self) {
+        assert!(
+            self.valid.get(),
+            "attempted to dereference an ArenaRef after its Arena was cleared"
+        );
+    }
+}
+
+impl<T: ?Sized> Deref for ArenaRef<T> {
+    type Target = T;
+
+    #[inline(always)]
+    fn deref(&self) -> &Self::Target {
+        self.validate();
+        unsafe { self.ptr.as_ref() }
+    }
+}
+
+impl<T: ?Sized> DerefMut for ArenaRef<T> {
+    #[inline(always)]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.validate();
+        unsafe { self.ptr.as_mut() }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{cell::Cell, rc::Rc};
+
+    use super::*;
+
+    #[test]
+    fn test_arena() {
+        let mut arena = Arena::new(1024);
+        let a = arena.alloc(|| 1u64);
+        let b = arena.alloc(|| 2u32);
+        let c = arena.alloc(|| 3u16);
+        let d = arena.alloc(|| 4u8);
+        assert_eq!(*a, 1);
+        assert_eq!(*b, 2);
+        assert_eq!(*c, 3);
+        assert_eq!(*d, 4);
+
+        arena.clear();
+        let a = arena.alloc(|| 5u64);
+        let b = arena.alloc(|| 6u32);
+        let c = arena.alloc(|| 7u16);
+        let d = arena.alloc(|| 8u8);
+        assert_eq!(*a, 5);
+        assert_eq!(*b, 6);
+        assert_eq!(*c, 7);
+        assert_eq!(*d, 8);
+
+        // Ensure drop gets called.
+        let dropped = Rc::new(Cell::new(false));
+        struct DropGuard(Rc<Cell<bool>>);
+        impl Drop for DropGuard {
+            fn drop(&mut self) {
+                self.0.set(true);
+            }
+        }
+        arena.alloc(|| DropGuard(dropped.clone()));
+        arena.clear();
+        assert!(dropped.get());
+    }
+}

crates/gpui2/src/element.rs

@@ -1,6 +1,6 @@
 use crate::{
-    AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels, Point, Size, ViewContext,
-    WindowContext,
+    frame_alloc, ArenaRef, AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels,
+    Point, Size, ViewContext, WindowContext,
 };
 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,7 +413,9 @@ impl AnyElement {
         E: 'static + Element,
         E::State: Any,
     {
-        AnyElement(Box::new(Some(DrawableElement::new(element))) as Box<dyn ElementObject>)
+        let element = frame_alloc(|| Some(DrawableElement::new(element)))
+            .map(|element| element as &mut dyn ElementObject);
+        AnyElement(element)
     }
 
     pub fn layout(&mut self, cx: &mut WindowContext) -> LayoutId {
@@ -443,11 +445,6 @@ impl AnyElement {
         self.0.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()
     }
@@ -480,6 +477,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/gpui2.rs

@@ -2,6 +2,7 @@
 mod action;
 mod app;
 
+mod arena;
 mod assets;
 mod color;
 mod element;
@@ -38,6 +39,7 @@ mod private {
 pub use action::*;
 pub use anyhow::Result;
 pub use app::*;
+pub(crate) use arena::*;
 pub use assets::*;
 pub use color::*;
 pub use ctor::ctor;

crates/gpui2/src/key_dispatch.rs

@@ -1,6 +1,6 @@
 use crate::{
-    Action, ActionRegistry, DispatchPhase, FocusId, KeyBinding, KeyContext, KeyMatch, Keymap,
-    Keystroke, KeystrokeMatcher, WindowContext,
+    arena::ArenaRef, Action, ActionRegistry, DispatchPhase, FocusId, KeyBinding, KeyContext,
+    KeyMatch, Keymap, Keystroke, KeystrokeMatcher, WindowContext,
 };
 use collections::HashMap;
 use parking_lot::Mutex;
@@ -33,12 +33,12 @@ pub(crate) struct DispatchNode {
     parent: Option<DispatchNodeId>,
 }
 
-type KeyListener = Rc<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>;
+type KeyListener = ArenaRef<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>;
 
 #[derive(Clone)]
 pub(crate) struct DispatchActionListener {
     pub(crate) action_type: TypeId,
-    pub(crate) listener: Rc<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>,
+    pub(crate) listener: ArenaRef<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>,
 }
 
 impl DispatchTree {
@@ -117,7 +117,7 @@ impl DispatchTree {
     pub fn on_action(
         &mut self,
         action_type: TypeId,
-        listener: Rc<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>,
+        listener: ArenaRef<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>,
     ) {
         self.active_node()
             .action_listeners

crates/gpui2/src/window.rs

@@ -1,15 +1,17 @@
 use crate::{
-    key_dispatch::DispatchActionListener, px, size, transparent_black, Action, AnyDrag, AnyView,
-    AppContext, AsyncWindowContext, AvailableSpace, Bounds, BoxShadow, Context, Corners,
-    CursorStyle, DevicePixels, DispatchNodeId, DispatchTree, DisplayId, Edges, Effect, Entity,
-    EntityId, EventEmitter, FileDropEvent, Flatten, FontId, GlobalElementId, GlyphId, Hsla,
-    ImageData, InputEvent, IsZero, KeyBinding, KeyContext, KeyDownEvent, KeystrokeEvent, LayoutId,
-    Model, ModelContext, Modifiers, MonochromeSprite, MouseButton, MouseMoveEvent, MouseUpEvent,
-    Path, Pixels, PlatformAtlas, PlatformDisplay, PlatformInputHandler, PlatformWindow, Point,
-    PolychromeSprite, PromptLevel, Quad, Render, RenderGlyphParams, RenderImageParams,
-    RenderSvgParams, ScaledPixels, Scene, SceneBuilder, Shadow, SharedString, Size, Style,
-    SubscriberSet, Subscription, Surface, TaffyLayoutEngine, Task, Underline, UnderlineStyle, View,
-    VisualContext, WeakView, WindowBounds, WindowOptions, SUBPIXEL_VARIANTS,
+    arena::{Arena, ArenaRef},
+    key_dispatch::DispatchActionListener,
+    px, size, transparent_black, Action, AnyDrag, AnyView, AppContext, AsyncWindowContext,
+    AvailableSpace, Bounds, BoxShadow, Context, Corners, CursorStyle, DevicePixels, DispatchNodeId,
+    DispatchTree, DisplayId, Edges, Effect, Entity, EntityId, EventEmitter, FileDropEvent, Flatten,
+    FontId, GlobalElementId, GlyphId, Hsla, ImageData, InputEvent, IsZero, KeyBinding, KeyContext,
+    KeyDownEvent, KeystrokeEvent, LayoutId, Model, ModelContext, Modifiers, MonochromeSprite,
+    MouseButton, MouseMoveEvent, MouseUpEvent, Path, Pixels, PlatformAtlas, PlatformDisplay,
+    PlatformInputHandler, PlatformWindow, Point, PolychromeSprite, PromptLevel, Quad, Render,
+    RenderGlyphParams, RenderImageParams, RenderSvgParams, ScaledPixels, Scene, SceneBuilder,
+    Shadow, SharedString, Size, Style, SubscriberSet, Subscription, Surface, TaffyLayoutEngine,
+    Task, Underline, UnderlineStyle, View, VisualContext, WeakView, WindowBounds, WindowOptions,
+    SUBPIXEL_VARIANTS,
 };
 use anyhow::{anyhow, Context as _, Result};
 use collections::FxHashMap;
@@ -25,6 +27,7 @@ use smallvec::SmallVec;
 use std::{
     any::{Any, TypeId},
     borrow::{Borrow, BorrowMut, Cow},
+    cell::RefCell,
     fmt::Debug,
     future::Future,
     hash::{Hash, Hasher},
@@ -85,7 +88,7 @@ impl DispatchPhase {
 }
 
 type AnyObserver = Box<dyn FnMut(&mut WindowContext) -> bool + 'static>;
-type AnyMouseListener = Box<dyn FnMut(&dyn Any, DispatchPhase, &mut WindowContext) + 'static>;
+type AnyMouseListener = ArenaRef<dyn FnMut(&dyn Any, DispatchPhase, &mut WindowContext) + 'static>;
 type AnyWindowFocusListener = Box<dyn FnMut(&FocusEvent, &mut WindowContext) -> bool + 'static>;
 
 struct FocusEvent {
@@ -95,6 +98,15 @@ struct FocusEvent {
 
 slotmap::new_key_type! { pub struct FocusId; }
 
+thread_local! {
+    pub static FRAME_ARENA: RefCell<Arena> = RefCell::new(Arena::new(16 * 1024 * 1024));
+}
+
+#[inline(always)]
+pub(crate) fn frame_alloc<T>(f: impl FnOnce() -> T) -> ArenaRef<T> {
+    FRAME_ARENA.with_borrow_mut(|arena| arena.alloc(f))
+}
+
 impl FocusId {
     /// Obtains whether the element associated with this handle is currently focused.
     pub fn is_focused(&self, cx: &WindowContext) -> bool {
@@ -268,7 +280,6 @@ pub(crate) struct ElementStateBox {
     type_name: &'static str,
 }
 
-// #[derive(Default)]
 pub(crate) struct Frame {
     focus: Option<FocusId>,
     pub(crate) element_states: FxHashMap<GlobalElementId, ElementStateBox>,
@@ -818,25 +829,23 @@ impl<'a> WindowContext<'a> {
     /// Register a mouse event listener on the window for the next frame. The type of event
     /// is determined by the first parameter of the given listener. When the next frame is rendered
     /// the listener will be cleared.
-    ///
-    /// This is a fairly low-level method, so prefer using event handlers on elements unless you have
-    /// a specific need to register a global listener.
     pub fn on_mouse_event<Event: 'static>(
         &mut self,
         mut handler: impl FnMut(&Event, DispatchPhase, &mut WindowContext) + 'static,
     ) {
         let order = self.window.next_frame.z_index_stack.clone();
+        let handler = frame_alloc(|| {
+            move |event: &dyn Any, phase: DispatchPhase, cx: &mut WindowContext<'_>| {
+                handler(event.downcast_ref().unwrap(), phase, cx)
+            }
+        })
+        .map(|handler| handler as _);
         self.window
             .next_frame
             .mouse_listeners
             .entry(TypeId::of::<Event>())
             .or_default()
-            .push((
-                order,
-                Box::new(move |event: &dyn Any, phase, cx| {
-                    handler(event.downcast_ref().unwrap(), phase, cx)
-                }),
-            ))
+            .push((order, handler))
     }
 
     /// Register a key event listener on the window for the next frame. The type of event
@@ -847,16 +856,17 @@ impl<'a> WindowContext<'a> {
     /// a specific need to register a global listener.
     pub fn on_key_event<Event: 'static>(
         &mut self,
-        handler: impl Fn(&Event, DispatchPhase, &mut WindowContext) + 'static,
+        listener: impl Fn(&Event, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        self.window
-            .next_frame
-            .dispatch_tree
-            .on_key_event(Rc::new(move |event, phase, cx| {
+        let listener = frame_alloc(|| {
+            move |event: &dyn Any, phase, cx: &mut WindowContext<'_>| {
                 if let Some(event) = event.downcast_ref::<Event>() {
-                    handler(event, phase, cx)
+                    listener(event, phase, cx)
                 }
-            }));
+            }
+        })
+        .map(|handler| handler as _);
+        self.window.next_frame.dispatch_tree.on_key_event(listener);
     }
 
     /// Register an action listener on the window for the next frame. The type of action
@@ -868,12 +878,13 @@ impl<'a> WindowContext<'a> {
     pub fn on_action(
         &mut self,
         action_type: TypeId,
-        handler: impl Fn(&dyn Any, DispatchPhase, &mut WindowContext) + 'static,
+        listener: impl Fn(&dyn Any, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        self.window.next_frame.dispatch_tree.on_action(
-            action_type,
-            Rc::new(move |action, phase, cx| handler(action, phase, cx)),
-        );
+        let listener = frame_alloc(|| listener).map(|handler| handler as _);
+        self.window
+            .next_frame
+            .dispatch_tree
+            .on_action(action_type, listener);
     }
 
     pub fn is_action_available(&self, action: &dyn Action) -> bool {
@@ -1267,16 +1278,23 @@ 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() {
-                        cx.window.next_frame.dispatch_tree.on_action(
-                            *action_type,
-                            Rc::new(move |action, phase, cx| action_listener(action, phase, cx)),
-                        )
+                        let listener = frame_alloc(|| {
+                            move |action: &dyn Any, phase, cx: &mut WindowContext<'_>| {
+                                action_listener(action, phase, cx)
+                            }
+                        })
+                        .map(|listener| listener as _);
+                        cx.window
+                            .next_frame
+                            .dispatch_tree
+                            .on_action(*action_type, listener)
                     }
                 }
 
@@ -2591,13 +2609,13 @@ impl<'a, V: 'static> ViewContext<'a, V> {
     pub fn on_action(
         &mut self,
         action_type: TypeId,
-        handler: impl Fn(&mut V, &dyn Any, DispatchPhase, &mut ViewContext<V>) + 'static,
+        listener: impl Fn(&mut V, &dyn Any, DispatchPhase, &mut ViewContext<V>) + 'static,
     ) {
         let handle = self.view().clone();
         self.window_cx
             .on_action(action_type, move |action, phase, cx| {
                 handle.update(cx, |view, cx| {
-                    handler(view, action, phase, cx);
+                    listener(view, action, phase, cx);
                 })
             });
     }