Use a closure to allocate structs onto the `Arena`

Antonio Scandurra created

This is a trick borrowed from Bumpalo that helps LLVM understand
it should instantiate the object directly on the heap, as opposed to
doing so on the stack and then moving it.

Change summary

crates/gpui2/src/arena.rs   | 36 ++++++++++++++++--------
crates/gpui2/src/element.rs |  9 ++---
crates/gpui2/src/gpui2.rs   |  1 
crates/gpui2/src/window.rs  | 56 ++++++++++++++++++--------------------
4 files changed, 55 insertions(+), 47 deletions(-)

Detailed changes

crates/gpui2/src/arena.rs 🔗

@@ -12,6 +12,7 @@ struct ArenaElement {
 }
 
 impl Drop for ArenaElement {
+    #[inline(always)]
     fn drop(&mut self) {
         unsafe {
             (self.drop)(self.value);
@@ -48,15 +49,23 @@ impl Arena {
     }
 
     #[inline(always)]
-    pub fn alloc<T>(&mut self, value: T) -> ArenaRef<T> {
+    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::for_value(&value).pad_to_align();
+            let layout = alloc::Layout::new::<T>().pad_to_align();
             let ptr = NonNull::new_unchecked(self.start.as_ptr().add(self.offset).cast::<T>());
-            ptr::write(ptr.as_ptr(), value);
+            inner_writer(ptr.as_ptr(), f);
 
             self.elements.push(ArenaElement {
                 value: ptr.cast(),
@@ -92,6 +101,7 @@ impl<T: ?Sized> Clone for ArenaRef<T> {
 }
 
 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)) },
@@ -110,6 +120,7 @@ impl<T: ?Sized> ArenaRef<T> {
 impl<T: ?Sized> Deref for ArenaRef<T> {
     type Target = T;
 
+    #[inline(always)]
     fn deref(&self) -> &Self::Target {
         self.validate();
         unsafe { self.ptr.as_ref() }
@@ -117,6 +128,7 @@ impl<T: ?Sized> Deref for ArenaRef<T> {
 }
 
 impl<T: ?Sized> DerefMut for ArenaRef<T> {
+    #[inline(always)]
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.validate();
         unsafe { self.ptr.as_mut() }
@@ -132,20 +144,20 @@ mod tests {
     #[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);
+        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);
+        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);
@@ -159,7 +171,7 @@ mod tests {
                 self.0.set(true);
             }
         }
-        arena.alloc(DropGuard(dropped.clone()));
+        arena.alloc(|| DropGuard(dropped.clone()));
         arena.clear();
         assert!(dropped.get());
     }

crates/gpui2/src/element.rs 🔗

@@ -1,6 +1,6 @@
 use crate::{
-    arena::ArenaRef, AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels, Point,
-    Size, ViewContext, WindowContext, FRAME_ARENA,
+    frame_alloc, ArenaRef, AvailableSpace, BorrowWindow, Bounds, ElementId, LayoutId, Pixels,
+    Point, Size, ViewContext, WindowContext,
 };
 use derive_more::{Deref, DerefMut};
 pub(crate) use smallvec::SmallVec;
@@ -413,9 +413,8 @@ impl AnyElement {
         E: 'static + Element,
         E::State: Any,
     {
-        let element =
-            FRAME_ARENA.with_borrow_mut(|arena| arena.alloc(Some(DrawableElement::new(element))));
-        let element = element.map(|element| element as &mut dyn ElementObject);
+        let element = frame_alloc(|| Some(DrawableElement::new(element)))
+            .map(|element| element as &mut dyn ElementObject);
         AnyElement(element)
     }
 

crates/gpui2/src/gpui2.rs 🔗

@@ -39,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/window.rs 🔗

@@ -102,6 +102,11 @@ 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 {
@@ -829,15 +834,12 @@ 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 = FRAME_ARENA
-            .with_borrow_mut(|arena| {
-                arena.alloc(
-                    move |event: &dyn Any, phase: DispatchPhase, cx: &mut WindowContext<'_>| {
-                        handler(event.downcast_ref().unwrap(), phase, cx)
-                    },
-                )
-            })
-            .map(|handler| handler as _);
+        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
@@ -856,15 +858,14 @@ impl<'a> WindowContext<'a> {
         &mut self,
         listener: impl Fn(&Event, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        let listener = FRAME_ARENA
-            .with_borrow_mut(|arena| {
-                arena.alloc(move |event: &dyn Any, phase, cx: &mut WindowContext<'_>| {
-                    if let Some(event) = event.downcast_ref::<Event>() {
-                        listener(event, phase, cx)
-                    }
-                })
-            })
-            .map(|handler| handler as _);
+        let listener = frame_alloc(|| {
+            move |event: &dyn Any, phase, cx: &mut WindowContext<'_>| {
+                if let Some(event) = event.downcast_ref::<Event>() {
+                    listener(event, phase, cx)
+                }
+            }
+        })
+        .map(|handler| handler as _);
         self.window.next_frame.dispatch_tree.on_key_event(listener);
     }
 
@@ -879,9 +880,7 @@ impl<'a> WindowContext<'a> {
         action_type: TypeId,
         listener: impl Fn(&dyn Any, DispatchPhase, &mut WindowContext) + 'static,
     ) {
-        let listener = FRAME_ARENA
-            .with_borrow_mut(|arena| arena.alloc(listener))
-            .map(|handler| handler as _);
+        let listener = frame_alloc(|| listener).map(|handler| handler as _);
         self.window
             .next_frame
             .dispatch_tree
@@ -1286,15 +1285,12 @@ impl<'a> WindowContext<'a> {
             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 = FRAME_ARENA
-                            .with_borrow_mut(|arena| {
-                                arena.alloc(
-                                    move |action: &dyn Any, phase, cx: &mut WindowContext<'_>| {
-                                        action_listener(action, phase, cx)
-                                    },
-                                )
-                            })
-                            .map(|listener| listener as _);
+                        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