Reduce steady-state allocations in ui_stack_system

crates/bevy_ui/src/stack.rs

@@ -15,34 +15,55 @@ pub struct UiStack {
     pub uinodes: Vec<Entity>,
 }
 
+/// Caches stacking context buffers for use in [`ui_stack_system`].
+#[derive(Default)]
+pub(crate) struct StackingContextCache {
+    inner: Vec<StackingContext>,
+}
+
+impl StackingContextCache {
+    fn pop(&mut self) -> StackingContext {
+        self.inner.pop().unwrap_or_default()
+    }
+
+    fn push(&mut self, mut context: StackingContext) {
+        for entry in context.entries.drain(..) {
+            self.push(entry.stack);
+        }
+        self.inner.push(context);
+    }
+}
+
 #[derive(Default)]
 struct StackingContext {
-    pub entries: Vec<StackingContextEntry>,
+    entries: Vec<StackingContextEntry>,
 }
 
 struct StackingContextEntry {
-    pub z_index: i32,
-    pub entity: Entity,
-    pub stack: StackingContext,
+    z_index: i32,
+    entity: Entity,
+    stack: StackingContext,
 }
 
 /// Generates the render stack for UI nodes.
 ///
 /// First generate a UI node tree (`StackingContext`) based on z-index.
 /// Then flatten that tree into back-to-front ordered `UiStack`.
-pub fn ui_stack_system(
+pub(crate) fn ui_stack_system(
+    mut cache: Local<StackingContextCache>,
     mut ui_stack: ResMut<UiStack>,
     root_node_query: Query<Entity, (With<Node>, Without<Parent>)>,
     zindex_query: Query<&ZIndex, With<Node>>,
     children_query: Query<&Children>,
     mut update_query: Query<&mut Node>,
 ) {
     // Generate `StackingContext` tree
-    let mut global_context = StackingContext::default();
+    let mut global_context = cache.pop();
     let mut total_entry_count: usize = 0;
 
     for entity in &root_node_query {
         insert_context_hierarchy(
+            &mut cache,
             &zindex_query,
             &children_query,
             entity,
@@ -55,7 +76,8 @@ pub fn ui_stack_system(
     // Flatten `StackingContext` into `UiStack`
     ui_stack.uinodes.clear();
     ui_stack.uinodes.reserve(total_entry_count);
-    fill_stack_recursively(&mut ui_stack.uinodes, &mut global_context);
+    fill_stack_recursively(&mut cache, &mut ui_stack.uinodes, &mut global_context);
+    cache.push(global_context);
 
     for (i, entity) in ui_stack.uinodes.iter().enumerate() {
         if let Ok(mut node) = update_query.get_mut(*entity) {
@@ -66,14 +88,15 @@ pub fn ui_stack_system(
 
 /// Generate z-index based UI node tree
 fn insert_context_hierarchy(
+    cache: &mut StackingContextCache,
     zindex_query: &Query<&ZIndex, With<Node>>,
     children_query: &Query<&Children>,
     entity: Entity,
     global_context: &mut StackingContext,
     parent_context: Option<&mut StackingContext>,
     total_entry_count: &mut usize,
 ) {
-    let mut new_context = StackingContext::default();
+    let mut new_context = cache.pop();
 
     if let Ok(children) = children_query.get(entity) {
         // Reserve space for all children. In practice, some may not get pushed since
@@ -82,6 +105,7 @@ fn insert_context_hierarchy(
 
         for entity in children {
             insert_context_hierarchy(
+                cache,
                 zindex_query,
                 children_query,
                 *entity,
@@ -108,16 +132,21 @@ fn insert_context_hierarchy(
 }
 
 /// Flatten `StackingContext` (z-index based UI node tree) into back-to-front entities list
-fn fill_stack_recursively(result: &mut Vec<Entity>, stack: &mut StackingContext) {
+fn fill_stack_recursively(
+    cache: &mut StackingContextCache,
+    result: &mut Vec<Entity>,
+    stack: &mut StackingContext,
+) {
     // Sort entries by ascending z_index, while ensuring that siblings
     // with the same local z_index will keep their ordering. This results
     // in `back-to-front` ordering, low z_index = back; high z_index = front.
     stack.entries.sort_by_key(|e| e.z_index);
 
-    for entry in &mut stack.entries {
+    for mut entry in stack.entries.drain(..) {
         // Parent node renders before/behind child nodes
         result.push(entry.entity);
-        fill_stack_recursively(result, &mut entry.stack);
+        fill_stack_recursively(cache, result, &mut entry.stack);
+        cache.push(entry.stack);
     }
 }