Refactor rendering to push-based instead of pull-based (#3113)

Our current rendering code constructs and returns complex
`-> impl Iterator<Item = SomeRenderElement>` types that are collected
into a vector at the top level Niri::render(). This causes some
problems:
- It's hard to write logic around returning iterators. Especially things
  like conditions, since the returned iterator must have a single type,
  you can't branch and return different iterators. This will be solved
  by gen fn but alas it's not here yet.
- In many cases, the returned `-> impl Iterator` will borrow from &self
  leading to complex lifetimes. In certain cases, it is also desirable
  for it to borrow the &mut NiriRenderer, which causes a lot of issues
  because it's exclusive (&mut).
- Sometimes those issues are too hard to deal with, leading to the
  escape hatch of allocating and returning a temporary
  Vec<SomeRenderElement>, like in
  Scrolling/FloatingSpace::render_elements(). These allocations are
  unfortunate because they are not really necessary.
- It's impossible to use some downstream combinators with this
  `-> impl Iterator` approach, leading to functions like Smithay's
  render_elements_from_surface_tree() returning a Vec. This is extra
  unfortunate because it results in a temporary allocation per Wayland
  toplevel/popup.
- It's hard to properly create profiling spans for the rendering
  functions since the spans are dropped when the (lazy) iterator is
  returned and not when all the code actually completes.
- The code compiles down to complex state machines in generated iterator
  types with logic located in Iterator::next(), which makes it annoying
  to follow in debuggers and profiling tools.

This refactor changes the code to push-based iteration: rendering
functions receive a push() closure that they call to push their render
elements. It solves all of the aforementioned problems:
- The logic becomes simpler. Just use conditionals and loops as normal.
- No borrowing and lifetimes since we're not returning anything.
- All temporary Vecs are removed because the problems they worked around
  no longer exist.
- The new push_elements_from_surface_tree() helper is the same as
  render_elements_from_surface_tree() but doesn't allocate a temporary
  Vec since it's not necessary; the push() closure can be passed down.
- Profiling spans work normally since the function returns when it ran
  all of the logic.
- The code compiles down to normal functions and calls as expected.

Generally, the iterator approach gives these advantages:
- You can wrap the returned items in the upstream logic. This is
  possible in exactly the same way with the push closure.
- You can decide to cut the iterator short in the upstream logic. This
  is not possible with push-based iteration, but we don't actually use
  it anywhere.

I chose the push closure type to be &mut dyn FnMut(SomeRenderElement).
It's deliberately not a generic impl FnMut() to avoid duplicating the
rendering logic when it's called from several different places. But it's
still a normal closure that can capture the outside context.

While my original idea for this refactor was to simplify the logic while
getting rid of temporary Vecs, it also appears to have brought a
consistent 2-3x speedup to the whole render list construction. On an old
Eee PC laptop I even observed a 8x speedup.

The refactor also results in smaller binary size, presumably due to
removing many iterator combinators and state tracking.

niri-visual-tests/src/cases/gradient_area.rs

@@ -89,11 +89,8 @@ impl TestCase for GradientArea {
             1.,
             1.,
         );
-        rv.extend(
-            self.border
-                .render(renderer, g_loc)
-                .map(|elem| Box::new(elem) as _),
-        );
+        self.border
+            .render(renderer, g_loc, &mut |elem| rv.push(Box::new(elem) as _));
 
         rv.extend(
             [BorderRenderElement::new(

niri-visual-tests/src/cases/layout.rs

@@ -268,12 +268,14 @@ impl TestCase for Layout {
         _size: Size<i32, Physical>,
     ) -> Vec<Box<dyn RenderElement<GlesRenderer>>> {
         self.layout.update_render_elements(Some(&self.output));
+
+        let mut rv = Vec::new();
         self.layout
             .monitor_for_output(&self.output)
             .unwrap()
-            .render_elements(renderer, RenderTarget::Output, true)
-            .flat_map(|(_, _, iter)| iter)
-            .map(|elem| Box::new(elem) as _)
-            .collect()
+            .render_workspaces(renderer, RenderTarget::Output, true, &mut |elem| {
+                rv.push(Box::new(elem) as _)
+            });
+        rv
     }
 }

niri-visual-tests/src/cases/tile.rs

@@ -119,9 +119,15 @@ impl TestCase for Tile {
             true,
             Rectangle::new(Point::from((-location.x, -location.y)), size.to_logical(1.)),
         );
-        self.tile
-            .render(renderer, location, true, RenderTarget::Output)
-            .map(|elem| Box::new(elem) as _)
-            .collect()
+
+        let mut rv = Vec::new();
+        self.tile.render(
+            renderer,
+            location,
+            true,
+            RenderTarget::Output,
+            &mut |elem| rv.push(Box::new(elem) as _),
+        );
+        rv
     }
 }

niri-visual-tests/src/cases/window.rs

@@ -52,16 +52,15 @@ impl TestCase for Window {
             .to_f64()
             .downscale(2.);
 
-        self.window
-            .render(
-                renderer,
-                location,
-                Scale::from(1.),
-                1.,
-                RenderTarget::Output,
-            )
-            .into_iter()
-            .map(|elem| Box::new(elem) as _)
-            .collect()
+        let mut rv = Vec::new();
+        self.window.render_normal(
+            renderer,
+            location,
+            Scale::from(1.),
+            1.,
+            RenderTarget::Output,
+            &mut |elem| rv.push(Box::new(elem) as _),
+        );
+        rv
     }
 }

niri-visual-tests/src/test_window.rs

@@ -9,7 +9,7 @@ use niri::layout::{
 use niri::render_helpers::offscreen::OffscreenData;
 use niri::render_helpers::renderer::NiriRenderer;
 use niri::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
-use niri::render_helpers::{RenderTarget, SplitElements};
+use niri::render_helpers::RenderTarget;
 use niri::utils::transaction::Transaction;
 use niri::window::ResolvedWindowRules;
 use smithay::backend::renderer::element::Kind;
@@ -149,36 +149,30 @@ impl LayoutElement for TestWindow {
         false
     }
 
-    fn render<R: NiriRenderer>(
+    fn render_normal<R: NiriRenderer>(
         &self,
         _renderer: &mut R,
         location: Point<f64, Logical>,
         _scale: Scale<f64>,
         alpha: f32,
         _target: RenderTarget,
-    ) -> SplitElements<LayoutElementRenderElement<R>> {
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
         let inner = self.inner.borrow();
 
-        SplitElements {
-            normal: vec![
-                SolidColorRenderElement::from_buffer(
-                    &inner.buffer,
-                    location,
-                    alpha,
-                    Kind::Unspecified,
-                )
+        push(
+            SolidColorRenderElement::from_buffer(&inner.buffer, location, alpha, Kind::Unspecified)
                 .into(),
-                SolidColorRenderElement::from_buffer(
-                    &inner.csd_shadow_buffer,
-                    location
-                        - Point::from((inner.csd_shadow_width, inner.csd_shadow_width)).to_f64(),
-                    alpha,
-                    Kind::Unspecified,
-                )
-                .into(),
-            ],
-            popups: vec![],
-        }
+        );
+        push(
+            SolidColorRenderElement::from_buffer(
+                &inner.csd_shadow_buffer,
+                location - Point::from((inner.csd_shadow_width, inner.csd_shadow_width)).to_f64(),
+                alpha,
+                Kind::Unspecified,
+            )
+            .into(),
+        );
     }
 
     fn request_size(

src/layer/mapped.rs

@@ -1,8 +1,6 @@
 use niri_config::utils::MergeWith as _;
 use niri_config::{Config, LayerRule};
-use smithay::backend::renderer::element::surface::{
-    render_elements_from_surface_tree, WaylandSurfaceRenderElement,
-};
+use smithay::backend::renderer::element::surface::WaylandSurfaceRenderElement;
 use smithay::backend::renderer::element::Kind;
 use smithay::desktop::{LayerSurface, PopupManager};
 use smithay::utils::{Logical, Point, Scale, Size};
@@ -15,7 +13,8 @@ use crate::niri_render_elements;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::shadow::ShadowRenderElement;
 use crate::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
-use crate::render_helpers::{RenderTarget, SplitElements};
+use crate::render_helpers::surface::push_elements_from_surface_tree;
+use crate::render_helpers::RenderTarget;
 use crate::utils::{baba_is_float_offset, round_logical_in_physical};
 
 #[derive(Debug)]
@@ -156,14 +155,13 @@ impl MappedLayer {
         Point::from((0., y))
     }
 
-    pub fn render<R: NiriRenderer>(
+    pub fn render_normal<R: NiriRenderer>(
         &self,
         renderer: &mut R,
         location: Point<f64, Logical>,
         target: RenderTarget,
-    ) -> SplitElements<LayerSurfaceRenderElement<R>> {
-        let mut rv = SplitElements::default();
-
+        push: &mut dyn FnMut(LayerSurfaceRenderElement<R>),
+    ) {
         let scale = Scale::from(self.scale);
         let alpha = self.rules.opacity.unwrap_or(1.).clamp(0., 1.);
         let location = location + self.bob_offset();
@@ -179,40 +177,60 @@ impl MappedLayer {
                 alpha,
                 Kind::Unspecified,
             );
-            rv.normal.push(elem.into());
+            push(elem.into());
         } else {
             // Layer surfaces don't have extra geometry like windows.
             let buf_pos = location;
 
             let surface = self.surface.wl_surface();
-            for (popup, popup_offset) in PopupManager::popups_for_surface(surface) {
-                // Layer surfaces don't have extra geometry like windows.
-                let offset = popup_offset - popup.geometry().loc;
-
-                rv.popups.extend(render_elements_from_surface_tree(
-                    renderer,
-                    popup.wl_surface(),
-                    (buf_pos + offset.to_f64()).to_physical_precise_round(scale),
-                    scale,
-                    alpha,
-                    Kind::ScanoutCandidate,
-                ));
-            }
-
-            rv.normal = render_elements_from_surface_tree(
+            push_elements_from_surface_tree(
                 renderer,
                 surface,
                 buf_pos.to_physical_precise_round(scale),
                 scale,
                 alpha,
                 Kind::ScanoutCandidate,
+                &mut |elem| push(elem.into()),
             );
         }
 
         let location = location.to_physical_precise_round(scale).to_logical(scale);
-        rv.normal
-            .extend(self.shadow.render(renderer, location).map(Into::into));
+        self.shadow
+            .render(renderer, location, &mut |elem| push(elem.into()));
+    }
 
-        rv
+    pub fn render_popups<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
+        location: Point<f64, Logical>,
+        target: RenderTarget,
+        push: &mut dyn FnMut(LayerSurfaceRenderElement<R>),
+    ) {
+        let scale = Scale::from(self.scale);
+        let alpha = self.rules.opacity.unwrap_or(1.).clamp(0., 1.);
+        let location = location + self.bob_offset();
+
+        if target.should_block_out(self.rules.block_out_from) {
+            return;
+        }
+
+        // Layer surfaces don't have extra geometry like windows.
+        let buf_pos = location;
+
+        let surface = self.surface.wl_surface();
+        for (popup, popup_offset) in PopupManager::popups_for_surface(surface) {
+            // Layer surfaces don't have extra geometry like windows.
+            let offset = popup_offset - popup.geometry().loc;
+
+            push_elements_from_surface_tree(
+                renderer,
+                popup.wl_surface(),
+                (buf_pos + offset.to_f64()).to_physical_precise_round(scale),
+                scale,
+                alpha,
+                Kind::ScanoutCandidate,
+                &mut |elem| push(elem.into()),
+            );
+        }
     }
 }

src/layout/floating.rs

@@ -1053,15 +1053,14 @@ impl<W: LayoutElement> FloatingSpace<W> {
         true
     }
 
-    pub fn render_elements<R: NiriRenderer>(
+    pub fn render<R: NiriRenderer>(
         &self,
         renderer: &mut R,
         view_rect: Rectangle<f64, Logical>,
         target: RenderTarget,
         focus_ring: bool,
-    ) -> Vec<FloatingSpaceRenderElement<R>> {
-        let mut rv = Vec::new();
-
+        push: &mut dyn FnMut(FloatingSpaceRenderElement<R>),
+    ) {
         let scale = Scale::from(self.scale);
 
         // Draw the closing windows on top of the other windows.
@@ -1069,7 +1068,7 @@ impl<W: LayoutElement> FloatingSpace<W> {
         // FIXME: I guess this should rather preserve the stacking order when the window is closed.
         for closing in self.closing_windows.iter().rev() {
             let elem = closing.render(renderer.as_gles_renderer(), view_rect, scale, target);
-            rv.push(elem.into());
+            push(elem.into());
         }
 
         let active = self.active_window_id.clone();
@@ -1077,13 +1076,10 @@ impl<W: LayoutElement> FloatingSpace<W> {
             // For the active tile, draw the focus ring.
             let focus_ring = focus_ring && Some(tile.window().id()) == active.as_ref();
 
-            rv.extend(
-                tile.render(renderer, tile_pos, focus_ring, target)
-                    .map(Into::into),
-            );
+            tile.render(renderer, tile_pos, focus_ring, target, &mut |elem| {
+                push(elem.into())
+            });
         }
-
-        rv
     }
 
     pub fn interactive_resize_begin(&mut self, window: W::Id, edges: ResizeEdge) -> bool {

src/layout/focus_ring.rs

@@ -1,6 +1,5 @@
 use std::iter::zip;
 
-use arrayvec::ArrayVec;
 use niri_config::{CornerRadius, Gradient, GradientRelativeTo};
 use smithay::backend::renderer::element::{Element as _, Kind};
 use smithay::utils::{Logical, Point, Rectangle, Size};
@@ -220,18 +219,17 @@ impl FocusRing {
         &self,
         renderer: &mut impl NiriRenderer,
         location: Point<f64, Logical>,
-    ) -> impl Iterator<Item = FocusRingRenderElement> {
-        let mut rv = ArrayVec::<_, 8>::new();
-
+        push: &mut dyn FnMut(FocusRingRenderElement),
+    ) {
         if self.config.off {
-            return rv.into_iter();
+            return;
         }
 
         let border_width = -self.locations[0].y;
 
         // If drawing as a border with width = 0, then there's nothing to draw.
         if self.is_border && border_width == 0. {
-            return rv.into_iter();
+            return;
         }
 
         let has_border_shader = BorderRenderElement::has_shader(renderer);
@@ -244,7 +242,7 @@ impl FocusRing {
                 SolidColorRenderElement::from_buffer(buffer, location, alpha, Kind::Unspecified)
                     .into()
             };
-            rv.push(elem);
+            push(elem);
         };
 
         if self.is_border {
@@ -258,8 +256,6 @@ impl FocusRing {
                 location + self.locations[0],
             );
         }
-
-        rv.into_iter()
     }
 
     pub fn width(&self) -> f64 {

src/layout/insert_hint_element.rs

@@ -59,7 +59,8 @@ impl InsertHintElement {
         &self,
         renderer: &mut impl NiriRenderer,
         location: Point<f64, Logical>,
-    ) -> impl Iterator<Item = FocusRingRenderElement> {
-        self.inner.render(renderer, location)
+        push: &mut dyn FnMut(FocusRingRenderElement),
+    ) {
+        self.inner.render(renderer, location, push)
     }
 }

src/layout/mod.rs

@@ -64,7 +64,7 @@ use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::snapshot::RenderSnapshot;
 use crate::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
 use crate::render_helpers::texture::TextureBuffer;
-use crate::render_helpers::{BakedBuffer, RenderTarget, SplitElements};
+use crate::render_helpers::{BakedBuffer, RenderTarget};
 use crate::rubber_band::RubberBand;
 use crate::utils::transaction::{Transaction, TransactionBlocker};
 use crate::utils::{
@@ -159,7 +159,11 @@ pub trait LayoutElement {
         scale: Scale<f64>,
         alpha: f32,
         target: RenderTarget,
-    ) -> SplitElements<LayoutElementRenderElement<R>>;
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
+        self.render_popups(renderer, location, scale, alpha, target, push);
+        self.render_normal(renderer, location, scale, alpha, target, push);
+    }
 
     /// Renders the non-popup parts of the element.
     fn render_normal<R: NiriRenderer>(
@@ -169,8 +173,9 @@ pub trait LayoutElement {
         scale: Scale<f64>,
         alpha: f32,
         target: RenderTarget,
-    ) -> Vec<LayoutElementRenderElement<R>> {
-        self.render(renderer, location, scale, alpha, target).normal
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
+        let _ = (renderer, location, scale, alpha, target, push);
     }
 
     /// Renders the popups of the element.
@@ -181,8 +186,9 @@ pub trait LayoutElement {
         scale: Scale<f64>,
         alpha: f32,
         target: RenderTarget,
-    ) -> Vec<LayoutElementRenderElement<R>> {
-        self.render(renderer, location, scale, alpha, target).popups
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
+        let _ = (renderer, location, scale, alpha, target, push);
     }
 
     /// Requests the element to change its size.
@@ -4713,38 +4719,37 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
-    pub fn render_interactive_move_for_output<'a, R: NiriRenderer + 'a>(
-        &'a self,
+    pub fn render_interactive_move_for_output<R: NiriRenderer>(
+        &self,
         renderer: &mut R,
         output: &Output,
         target: RenderTarget,
-    ) -> impl Iterator<Item = RescaleRenderElement<TileRenderElement<R>>> + 'a {
+        push: &mut dyn FnMut(RescaleRenderElement<TileRenderElement<R>>),
+    ) {
         if self.update_render_elements_time != self.clock.now() {
             error!("clock moved between updating render elements and rendering");
         }
 
-        let mut rv = None;
+        let Some(InteractiveMoveState::Moving(move_)) = &self.interactive_move else {
+            return;
+        };
 
-        if let Some(InteractiveMoveState::Moving(move_)) = &self.interactive_move {
-            if &move_.output == output {
-                let scale = Scale::from(move_.output.current_scale().fractional_scale());
-                let zoom = self.overview_zoom();
-                let location = move_.tile_render_location(zoom);
-                let iter = move_
-                    .tile
-                    .render(renderer, location, true, target)
-                    .map(move |elem| {
-                        RescaleRenderElement::from_element(
-                            elem,
-                            location.to_physical_precise_round(scale),
-                            zoom,
-                        )
-                    });
-                rv = Some(iter);
-            }
+        if &move_.output != output {
+            return;
         }
 
-        rv.into_iter().flatten()
+        let scale = Scale::from(move_.output.current_scale().fractional_scale());
+        let zoom = self.overview_zoom();
+        let location = move_.tile_render_location(zoom);
+        move_
+            .tile
+            .render(renderer, location, true, target, &mut |elem| {
+                push(RescaleRenderElement::from_element(
+                    elem,
+                    location.to_physical_precise_round(scale),
+                    zoom,
+                ));
+            });
     }
 
     pub fn refresh(&mut self, is_active: bool) {

src/layout/monitor.rs

@@ -1638,40 +1638,36 @@ impl<W: LayoutElement> Monitor<W> {
     pub fn render_insert_hint_between_workspaces<R: NiriRenderer>(
         &self,
         renderer: &mut R,
-    ) -> impl Iterator<Item = MonitorRenderElement<R>> {
-        let mut rv = None;
-
-        if !self.options.layout.insert_hint.off {
-            if let Some(render_loc) = self.insert_hint_render_loc {
-                if let InsertWorkspace::NewAt(_) = render_loc.workspace {
-                    let iter = self
-                        .insert_hint_element
-                        .render(renderer, render_loc.location)
-                        .map(MonitorInnerRenderElement::UncroppedInsertHint);
-                    rv = Some(iter);
-                }
-            }
+        push: &mut dyn FnMut(MonitorRenderElement<R>),
+    ) {
+        if self.options.layout.insert_hint.off {
+            return;
         }
+        let Some(render_loc) = self.insert_hint_render_loc else {
+            return;
+        };
+        let InsertWorkspace::NewAt(_) = render_loc.workspace else {
+            return;
+        };
 
-        rv.into_iter().flatten().map(|elem| {
-            let elem = RescaleRenderElement::from_element(elem, Point::default(), 1.);
-            RelocateRenderElement::from_element(elem, Point::default(), Relocate::Relative)
-        })
+        self.insert_hint_element
+            .render(renderer, render_loc.location, &mut |elem| {
+                let elem = MonitorInnerRenderElement::UncroppedInsertHint(elem);
+                let elem = RescaleRenderElement::from_element(elem, Point::default(), 1.);
+                let elem =
+                    RelocateRenderElement::from_element(elem, Point::default(), Relocate::Relative);
+                push(elem);
+            });
     }
 
-    pub fn render_elements<'a, R: NiriRenderer>(
-        &'a self,
-        renderer: &'a mut R,
+    pub fn render_workspaces<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
         target: RenderTarget,
         focus_ring: bool,
-    ) -> impl Iterator<
-        Item = (
-            Rectangle<f64, Logical>,
-            MonitorRenderElement<R>,
-            impl Iterator<Item = MonitorRenderElement<R>> + 'a,
-        ),
-    > {
-        let _span = tracy_client::span!("Monitor::render_elements");
+        push: &mut dyn FnMut(MonitorRenderElement<R>),
+    ) {
+        let _span = tracy_client::span!("Monitor::render_workspaces");
 
         let scale = self.scale.fractional_scale();
         // Ceil the height in physical pixels.
@@ -1701,95 +1697,77 @@ impl<W: LayoutElement> Monitor<W> {
 
         let zoom = self.overview_zoom();
 
-        // Draw the insert hint.
-        let mut insert_hint = None;
-        if !self.options.layout.insert_hint.off {
-            if let Some(render_loc) = self.insert_hint_render_loc {
-                if let InsertWorkspace::Existing(workspace_id) = render_loc.workspace {
-                    insert_hint = Some((
-                        workspace_id,
-                        self.insert_hint_element
-                            .render(renderer, render_loc.location),
-                    ));
+        let insert_hint_render_loc = self
+            .insert_hint_render_loc
+            .filter(|_| !self.options.layout.insert_hint.off);
+
+        let scale_relocate = move |geo: Rectangle<f64, Logical>, elem| {
+            let elem = RescaleRenderElement::from_element(elem, Point::from((0, 0)), zoom);
+            RelocateRenderElement::from_element(
+                elem,
+                // The offset we get from workspaces_with_render_geo() is already
+                // rounded to physical pixels, but it's in the logical coordinate
+                // space, so we need to convert it to physical.
+                geo.loc.to_physical_precise_round(scale),
+                Relocate::Relative,
+            )
+        };
+
+        for (ws, geo) in self.workspaces_with_render_geo() {
+            // Macro instead of closure because ws and insert hint have different elem types.
+            macro_rules! push {
+                () => {{
+                    &mut |elem| {
+                        let elem = CropRenderElement::from_element(elem, scale, crop_bounds);
+                        if let Some(elem) = elem {
+                            let elem = MonitorInnerRenderElement::from(elem);
+                            push(scale_relocate(geo, elem));
+                        }
+                    }
+                }};
+            }
+
+            ws.render_floating(renderer, target, focus_ring, push!());
+
+            if let Some(loc) = insert_hint_render_loc {
+                if loc.workspace == InsertWorkspace::Existing(ws.id()) {
+                    self.insert_hint_element
+                        .render(renderer, loc.location, push!());
                 }
             }
+
+            ws.render_scrolling(renderer, target, focus_ring, push!());
         }
-
-        self.workspaces_with_render_geo().map(move |(ws, geo)| {
-            let map_ws_contents = move |elem: WorkspaceRenderElement<R>| {
-                let elem = CropRenderElement::from_element(elem, scale, crop_bounds)?;
-                let elem = MonitorInnerRenderElement::Workspace(elem);
-                Some(elem)
-            };
-
-            let (floating, scrolling) = ws.render_elements(renderer, target, focus_ring);
-            let floating = floating.filter_map(map_ws_contents);
-            let scrolling = scrolling.filter_map(map_ws_contents);
-
-            let hint = if matches!(insert_hint, Some((hint_ws_id, _)) if hint_ws_id == ws.id()) {
-                let iter = insert_hint.take().unwrap().1;
-                let iter = iter.filter_map(move |elem| {
-                    let elem = CropRenderElement::from_element(elem, scale, crop_bounds)?;
-                    let elem = MonitorInnerRenderElement::InsertHint(elem);
-                    Some(elem)
-                });
-                Some(iter)
-            } else {
-                None
-            };
-            let hint = hint.into_iter().flatten();
-
-            let iter = floating.chain(hint).chain(scrolling);
-
-            let scale_relocate = move |elem| {
-                let elem = RescaleRenderElement::from_element(elem, Point::from((0, 0)), zoom);
-                RelocateRenderElement::from_element(
-                    elem,
-                    // The offset we get from workspaces_with_render_positions() is already
-                    // rounded to physical pixels, but it's in the logical coordinate
-                    // space, so we need to convert it to physical.
-                    geo.loc.to_physical_precise_round(scale),
-                    Relocate::Relative,
-                )
-            };
-
-            let iter = iter.map(scale_relocate);
-
-            let background = ws.render_background();
-            let background = scale_relocate(MonitorInnerRenderElement::SolidColor(background));
-
-            (geo, background, iter)
-        })
     }
 
-    pub fn render_workspace_shadows<'a, R: NiriRenderer>(
-        &'a self,
-        renderer: &'a mut R,
-    ) -> impl Iterator<Item = MonitorRenderElement<R>> + 'a {
+    pub fn render_workspace_shadows<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
+        push: &mut dyn FnMut(MonitorRenderElement<R>),
+    ) {
+        let Some(progress) = self.overview_progress.as_ref().map(|p| p.clamped_value()) else {
+            return;
+        };
+        let alpha = progress.clamp(0., 1.) as f32;
+
         let _span = tracy_client::span!("Monitor::render_workspace_shadows");
 
         let scale = self.scale.fractional_scale();
         let zoom = self.overview_zoom();
-        let overview_clamped_progress = self.overview_progress.as_ref().map(|p| p.clamped_value());
 
-        self.workspaces_with_render_geo()
-            .flat_map(move |(ws, geo)| {
-                let shadow = overview_clamped_progress.map(|value| {
-                    ws.render_shadow(renderer)
-                        .map(move |elem| elem.with_alpha(value.clamp(0., 1.) as f32))
-                        .map(MonitorInnerRenderElement::Shadow)
-                });
-                let iter = shadow.into_iter().flatten();
-
-                iter.map(move |elem| {
-                    let elem = RescaleRenderElement::from_element(elem, Point::from((0, 0)), zoom);
-                    RelocateRenderElement::from_element(
-                        elem,
-                        geo.loc.to_physical_precise_round(scale),
-                        Relocate::Relative,
-                    )
-                })
-            })
+        for (ws, geo) in self.workspaces_with_render_geo() {
+            ws.render_shadow(renderer, &mut |elem| {
+                let elem = elem.with_alpha(alpha);
+                let elem = MonitorInnerRenderElement::Shadow(elem);
+                let elem = RescaleRenderElement::from_element(elem, Point::from((0, 0)), zoom);
+                let elem = RelocateRenderElement::from_element(
+                    elem,
+                    geo.loc.to_physical_precise_round(scale),
+                    Relocate::Relative,
+                );
+                push(elem);
+            });
+        }
     }
 
     pub fn workspace_switch_gesture_begin(&mut self, is_touchpad: bool) {

src/layout/scrolling.rs

@@ -2897,25 +2897,24 @@ impl<W: LayoutElement> ScrollingSpace<W> {
             .is_fullscreen()
     }
 
-    pub fn render_elements<R: NiriRenderer>(
+    pub fn render<R: NiriRenderer>(
         &self,
         renderer: &mut R,
         target: RenderTarget,
         focus_ring: bool,
-    ) -> Vec<ScrollingSpaceRenderElement<R>> {
-        let mut rv = vec![];
-
+        push: &mut dyn FnMut(ScrollingSpaceRenderElement<R>),
+    ) {
         let scale = Scale::from(self.scale);
 
         // Draw the closing windows on top of the other windows.
         let view_rect = Rectangle::new(Point::from((self.view_pos(), 0.)), self.view_size);
         for closing in self.closing_windows.iter().rev() {
             let elem = closing.render(renderer.as_gles_renderer(), view_rect, scale, target);
-            rv.push(elem.into());
+            push(elem.into());
         }
 
         if self.columns.is_empty() {
-            return rv;
+            return;
         }
 
         let mut first = true;
@@ -2930,7 +2929,8 @@ impl<W: LayoutElement> ScrollingSpace<W> {
             {
                 let pos = view_off + col_off + col_render_off;
                 let pos = pos.to_physical_precise_round(scale).to_logical(scale);
-                rv.extend(col.tab_indicator.render(renderer, pos).map(Into::into));
+                col.tab_indicator
+                    .render(renderer, pos, &mut |elem| push(elem.into()));
             }
 
             for (tile, tile_off, visible) in col.tiles_in_render_order() {
@@ -2955,14 +2955,11 @@ impl<W: LayoutElement> ScrollingSpace<W> {
                     continue;
                 }
 
-                rv.extend(
-                    tile.render(renderer, tile_pos, focus_ring, target)
-                        .map(Into::into),
-                );
+                tile.render(renderer, tile_pos, focus_ring, target, &mut |elem| {
+                    push(elem.into())
+                });
             }
         }
-
-        rv
     }
 
     pub fn window_under(&self, pos: Point<f64, Logical>) -> Option<(&W, HitType)> {

src/layout/shadow.rs

@@ -166,19 +166,19 @@ impl Shadow {
         &self,
         renderer: &mut impl NiriRenderer,
         location: Point<f64, Logical>,
-    ) -> impl Iterator<Item = ShadowRenderElement> + '_ {
+        push: &mut dyn FnMut(ShadowRenderElement),
+    ) {
         if !self.config.on {
-            return None.into_iter().flatten();
+            return;
         }
 
         let has_shadow_shader = ShadowRenderElement::has_shader(renderer);
         if !has_shadow_shader {
-            return None.into_iter().flatten();
+            return;
         }
 
-        let rv = zip(&self.shaders, &self.shader_rects)
-            .map(move |(shader, rect)| shader.clone().with_location(location + rect.loc));
-
-        Some(rv).into_iter().flatten()
+        for (shader, rect) in zip(&self.shaders, &self.shader_rects) {
+            push(shader.clone().with_location(location + rect.loc));
+        }
     }
 }

src/layout/tab_indicator.rs

@@ -294,17 +294,17 @@ impl TabIndicator {
         &self,
         renderer: &mut impl NiriRenderer,
         pos: Point<f64, Logical>,
-    ) -> impl Iterator<Item = TabIndicatorRenderElement> + '_ {
+        push: &mut dyn FnMut(TabIndicatorRenderElement),
+    ) {
         let has_border_shader = BorderRenderElement::has_shader(renderer);
         if !has_border_shader {
-            return None.into_iter().flatten();
+            return;
         }
 
-        let rv = zip(&self.shaders, &self.shader_locs)
-            .map(move |(shader, loc)| shader.clone().with_location(pos + *loc))
-            .map(TabIndicatorRenderElement::from);
-
-        Some(rv).into_iter().flatten()
+        for (shader, loc) in zip(&self.shaders, &self.shader_locs) {
+            let elem = shader.clone().with_location(pos + *loc);
+            push(TabIndicatorRenderElement::from(elem));
+        }
     }
 
     /// Extra size occupied by the tab indicator.

src/layout/tests.rs

@@ -166,17 +166,6 @@ impl LayoutElement for TestWindow {
         false
     }
 
-    fn render<R: NiriRenderer>(
-        &self,
-        _renderer: &mut R,
-        _location: Point<f64, Logical>,
-        _scale: Scale<f64>,
-        _alpha: f32,
-        _target: RenderTarget,
-    ) -> SplitElements<LayoutElementRenderElement<R>> {
-        SplitElements::default()
-    }
-
     fn request_size(
         &mut self,
         size: Size<i32, Logical>,

src/layout/tile.rs

@@ -1007,13 +1007,14 @@ impl<W: LayoutElement> Tile<W> {
         Point::from((0., y))
     }
 
-    fn render_inner<'a, R: NiriRenderer + 'a>(
-        &'a self,
+    fn render_inner<R: NiriRenderer>(
+        &self,
         renderer: &mut R,
         location: Point<f64, Logical>,
         focus_ring: bool,
         target: RenderTarget,
-    ) -> impl Iterator<Item = TileRenderElement<R>> + 'a {
+        push: &mut dyn FnMut(TileRenderElement<R>),
+    ) {
         let _span = tracy_client::span!("Tile::render_inner");
 
         let scale = Scale::from(self.scale);
@@ -1056,29 +1057,31 @@ impl<W: LayoutElement> Tile<W> {
             .unwrap_or_default()
             .scaled_by(1. - expanded_progress as f32);
 
+        // Popups go on top, whether it's resize or not.
+        self.window.render_popups(
+            renderer,
+            window_render_loc,
+            scale,
+            win_alpha,
+            target,
+            &mut |elem| push(elem.into()),
+        );
+
         // If we're resizing, try to render a shader, or a fallback.
-        let mut resize_shader = None;
-        let mut resize_popups = None;
-        let mut resize_fallback = None;
-
+        let mut pushed_resize = false;
         if let Some(resize) = &self.resize_animation {
-            resize_popups = Some(
-                self.window
-                    .render_popups(renderer, window_render_loc, scale, win_alpha, target)
-                    .into_iter()
-                    .map(Into::into),
-            );
-
             if ResizeRenderElement::has_shader(renderer) {
                 let gles_renderer = renderer.as_gles_renderer();
 
                 if let Some(texture_from) = resize.snapshot.texture(gles_renderer, scale, target) {
-                    let window_elements = self.window.render_normal(
+                    let mut window_elements = Vec::new();
+                    self.window.render_normal(
                         gles_renderer,
                         Point::from((0., 0.)),
                         scale,
                         1.,
                         target,
+                        &mut |elem| window_elements.push(elem),
                     );
 
                     let current = resize
@@ -1125,46 +1128,33 @@ impl<W: LayoutElement> Tile<W> {
                         // This is not a problem for split popups as the code will look for them by
                         // original id when it doesn't find them on the offscreen.
                         self.window.set_offscreen_data(Some(data));
-                        resize_shader = Some(elem.into());
+                        push(elem.into());
+                        pushed_resize = true;
                     }
                 }
             }
 
-            if resize_shader.is_none() {
+            if !pushed_resize {
                 let fallback_buffer = SolidColorBuffer::new(area.size, [1., 0., 0., 1.]);
-                resize_fallback = Some(
-                    SolidColorRenderElement::from_buffer(
-                        &fallback_buffer,
-                        area.loc,
-                        win_alpha,
-                        Kind::Unspecified,
-                    )
-                    .into(),
+                let elem = SolidColorRenderElement::from_buffer(
+                    &fallback_buffer,
+                    area.loc,
+                    win_alpha,
+                    Kind::Unspecified,
                 );
+                push(elem.into());
+                pushed_resize = true;
             }
         }
 
         // If we're not resizing, render the window itself.
-        let mut window_surface = None;
-        let mut window_popups = None;
-        let mut rounded_corner_damage = None;
         let has_border_shader = BorderRenderElement::has_shader(renderer);
-        if resize_shader.is_none() && resize_fallback.is_none() {
-            let window = self
-                .window
-                .render(renderer, window_render_loc, scale, win_alpha, target);
-
+        if !pushed_resize {
             let geo = Rectangle::new(window_render_loc, window_size);
             let radius = radius.fit_to(window_size.w as f32, window_size.h as f32);
 
             let clip_shader = ClippedSurfaceRenderElement::shader(renderer).cloned();
-
-            if clip_to_geometry && clip_shader.is_some() {
-                let damage = self.rounded_corner_damage.element();
-                rounded_corner_damage = Some(damage.with_location(window_render_loc).into());
-            }
-
-            window_surface = Some(window.normal.into_iter().map(move |elem| match elem {
+            let clip = |elem| match elem {
                 LayoutElementRenderElement::Wayland(elem) => {
                     // If we should clip to geometry, render a clipped window.
                     if clip_to_geometry {
@@ -1213,21 +1203,24 @@ impl<W: LayoutElement> Tile<W> {
                     // Otherwise, render the solid color as is.
                     LayoutElementRenderElement::SolidColor(elem).into()
                 }
-            }));
+            };
 
-            window_popups = Some(window.popups.into_iter().map(Into::into));
+            if clip_to_geometry && clip_shader.is_some() {
+                let damage = self.rounded_corner_damage.element();
+                push(damage.with_location(window_render_loc).into());
+            }
+
+            self.window.render_normal(
+                renderer,
+                window_render_loc,
+                scale,
+                win_alpha,
+                target,
+                &mut |elem| push(clip(elem)),
+            );
         }
 
-        let rv = resize_popups
-            .into_iter()
-            .flatten()
-            .chain(resize_shader)
-            .chain(resize_fallback)
-            .chain(window_popups.into_iter().flatten())
-            .chain(rounded_corner_damage)
-            .chain(window_surface.into_iter().flatten());
-
-        let elem = (fullscreen_progress > 0.).then(|| {
+        if fullscreen_progress > 0. {
             let alpha = fullscreen_progress as f32;
 
             // During the un/fullscreen animation, render a border element in order to use the
@@ -1243,7 +1236,7 @@ impl<W: LayoutElement> Tile<W> {
 
                 let size = self.fullscreen_backdrop.size();
                 let color = self.fullscreen_backdrop.color();
-                BorderRenderElement::new(
+                let elem = BorderRenderElement::new(
                     size,
                     Rectangle::from_size(size),
                     GradientInterpolation::default(),
@@ -1256,47 +1249,50 @@ impl<W: LayoutElement> Tile<W> {
                     scale.x as f32,
                     alpha,
                 )
-                .with_location(location)
-                .into()
+                .with_location(location);
+                push(elem.into());
             } else {
-                SolidColorRenderElement::from_buffer(
+                let elem = SolidColorRenderElement::from_buffer(
                     &self.fullscreen_backdrop,
                     location,
                     alpha,
                     Kind::Unspecified,
-                )
-                .into()
+                );
+                push(elem.into());
             }
-        });
-        let rv = rv.chain(elem);
+        }
 
-        let elem = self.visual_border_width().map(|width| {
-            self.border
-                .render(renderer, location + Point::from((width, width)))
-                .map(Into::into)
-        });
-        let rv = rv.chain(elem.into_iter().flatten());
+        if let Some(width) = self.visual_border_width() {
+            self.border.render(
+                renderer,
+                location + Point::from((width, width)),
+                &mut |elem| push(elem.into()),
+            );
+        }
 
         // Hide the focus ring when maximized/fullscreened. It's not normally visible anyway due to
         // being outside the monitor or obscured by a solid colored bar, but it is visible under
         // semitransparent bars in maximized state (which is a bit weird) and in the overview (also
         // a bit weird).
-        let elem = (focus_ring && expanded_progress < 1.)
-            .then(|| self.focus_ring.render(renderer, location).map(Into::into));
-        let rv = rv.chain(elem.into_iter().flatten());
+        if focus_ring && expanded_progress < 1. {
+            self.focus_ring
+                .render(renderer, location, &mut |elem| push(elem.into()));
+        }
 
-        let elem = (expanded_progress < 1.)
-            .then(|| self.shadow.render(renderer, location).map(Into::into));
-        rv.chain(elem.into_iter().flatten())
+        if expanded_progress < 1. {
+            self.shadow
+                .render(renderer, location, &mut |elem| push(elem.into()));
+        }
     }
 
-    pub fn render<'a, R: NiriRenderer + 'a>(
-        &'a self,
+    pub fn render<R: NiriRenderer>(
+        &self,
         renderer: &mut R,
         location: Point<f64, Logical>,
         focus_ring: bool,
         target: RenderTarget,
-    ) -> impl Iterator<Item = TileRenderElement<R>> + 'a {
+        push: &mut dyn FnMut(TileRenderElement<R>),
+    ) {
         let _span = tracy_client::span!("Tile::render");
 
         let scale = Scale::from(self.scale);
@@ -1306,16 +1302,19 @@ impl<W: LayoutElement> Tile<W> {
             .as_ref()
             .map_or(1., |alpha| alpha.anim.clamped_value()) as f32;
 
-        let mut open_anim_elem = None;
-        let mut alpha_anim_elem = None;
-        let mut window_elems = None;
-
+        let mut pushed = false;
         self.window().set_offscreen_data(None);
 
         if let Some(open) = &self.open_animation {
             let renderer = renderer.as_gles_renderer();
-            let elements = self.render_inner(renderer, Point::from((0., 0.)), focus_ring, target);
-            let elements = elements.collect::<Vec<TileRenderElement<_>>>();
+            let mut elements = Vec::new();
+            self.render_inner(
+                renderer,
+                Point::from((0., 0.)),
+                focus_ring,
+                target,
+                &mut |elem| elements.push(elem),
+            );
             match open.render(
                 renderer,
                 &elements,
@@ -1326,7 +1325,8 @@ impl<W: LayoutElement> Tile<W> {
             ) {
                 Ok((elem, data)) => {
                     self.window().set_offscreen_data(Some(data));
-                    open_anim_elem = Some(elem.into());
+                    push(elem.into());
+                    pushed = true;
                 }
                 Err(err) => {
                     warn!("error rendering window opening animation: {err:?}");
@@ -1334,15 +1334,22 @@ impl<W: LayoutElement> Tile<W> {
             }
         } else if let Some(alpha) = &self.alpha_animation {
             let renderer = renderer.as_gles_renderer();
-            let elements = self.render_inner(renderer, Point::from((0., 0.)), focus_ring, target);
-            let elements = elements.collect::<Vec<TileRenderElement<_>>>();
+            let mut elements = Vec::new();
+            self.render_inner(
+                renderer,
+                Point::from((0., 0.)),
+                focus_ring,
+                target,
+                &mut |elem| elements.push(elem),
+            );
             match alpha.offscreen.render(renderer, scale, &elements) {
                 Ok((elem, _sync, data)) => {
                     let offset = elem.offset();
                     let elem = elem.with_alpha(tile_alpha).with_offset(location + offset);
 
                     self.window().set_offscreen_data(Some(data));
-                    alpha_anim_elem = Some(elem.into());
+                    push(elem.into());
+                    pushed = true;
                 }
                 Err(err) => {
                     warn!("error rendering tile to offscreen for alpha animation: {err:?}");
@@ -1350,14 +1357,11 @@ impl<W: LayoutElement> Tile<W> {
             }
         }
 
-        if open_anim_elem.is_none() && alpha_anim_elem.is_none() {
-            window_elems = Some(self.render_inner(renderer, location, focus_ring, target));
+        if !pushed {
+            self.render_inner(renderer, location, focus_ring, target, &mut |elem| {
+                push(elem)
+            });
         }
-
-        open_anim_elem
-            .into_iter()
-            .chain(alpha_anim_elem)
-            .chain(window_elems.into_iter().flatten())
     }
 
     pub fn store_unmap_snapshot_if_empty(&mut self, renderer: &mut GlesRenderer) {
@@ -1371,19 +1375,28 @@ impl<W: LayoutElement> Tile<W> {
     fn render_snapshot(&self, renderer: &mut GlesRenderer) -> TileRenderSnapshot {
         let _span = tracy_client::span!("Tile::render_snapshot");
 
-        let contents = self.render(renderer, Point::from((0., 0.)), false, RenderTarget::Output);
+        let mut contents = Vec::new();
+        self.render(
+            renderer,
+            Point::from((0., 0.)),
+            false,
+            RenderTarget::Output,
+            &mut |elem| contents.push(elem),
+        );
 
         // A bit of a hack to render blocked out as for screencast, but I think it's fine here.
-        let blocked_out_contents = self.render(
+        let mut blocked_out_contents = Vec::new();
+        self.render(
             renderer,
             Point::from((0., 0.)),
             false,
             RenderTarget::Screencast,
+            &mut |elem| blocked_out_contents.push(elem),
         );
 
         RenderSnapshot {
-            contents: contents.collect(),
-            blocked_out_contents: blocked_out_contents.collect(),
+            contents,
+            blocked_out_contents,
             block_out_from: self.window.rules().block_out_from,
             size: self.animated_tile_size(),
             texture: Default::default(),

src/layout/workspace.rs

@@ -1624,39 +1624,48 @@ impl<W: LayoutElement> Workspace<W> {
         }
     }
 
-    pub fn render_elements<R: NiriRenderer>(
+    pub fn render_scrolling<R: NiriRenderer>(
         &self,
         renderer: &mut R,
         target: RenderTarget,
         focus_ring: bool,
-    ) -> (
-        impl Iterator<Item = WorkspaceRenderElement<R>>,
-        impl Iterator<Item = WorkspaceRenderElement<R>>,
+        push: &mut dyn FnMut(WorkspaceRenderElement<R>),
     ) {
         let scrolling_focus_ring = focus_ring && !self.floating_is_active();
-        let scrolling = self
-            .scrolling
-            .render_elements(renderer, target, scrolling_focus_ring);
-        let scrolling = scrolling.into_iter().map(WorkspaceRenderElement::from);
+        self.scrolling
+            .render(renderer, target, scrolling_focus_ring, &mut |elem| {
+                push(elem.into())
+            });
+    }
 
+    pub fn render_floating<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
+        target: RenderTarget,
+        focus_ring: bool,
+        push: &mut dyn FnMut(WorkspaceRenderElement<R>),
+    ) {
+        if !self.is_floating_visible() {
+            return;
+        }
+
+        let view_rect = Rectangle::from_size(self.view_size);
         let floating_focus_ring = focus_ring && self.floating_is_active();
-        let floating = self.is_floating_visible().then(|| {
-            let view_rect = Rectangle::from_size(self.view_size);
-            let floating =
-                self.floating
-                    .render_elements(renderer, view_rect, target, floating_focus_ring);
-            floating.into_iter().map(WorkspaceRenderElement::from)
-        });
-        let floating = floating.into_iter().flatten();
-
-        (floating, scrolling)
+        self.floating.render(
+            renderer,
+            view_rect,
+            target,
+            floating_focus_ring,
+            &mut |elem| push(elem.into()),
+        );
     }
 
     pub fn render_shadow<R: NiriRenderer>(
         &self,
         renderer: &mut R,
-    ) -> impl Iterator<Item = ShadowRenderElement> + '_ {
-        self.shadow.render(renderer, Point::from((0., 0.)))
+        push: &mut dyn FnMut(ShadowRenderElement),
+    ) {
+        self.shadow.render(renderer, Point::from((0., 0.)), push);
     }
 
     pub fn render_background(&self) -> SolidColorRenderElement {

src/niri.rs

@@ -22,9 +22,7 @@ use smithay::backend::allocator::Fourcc;
 use smithay::backend::input::Keycode;
 use smithay::backend::renderer::damage::OutputDamageTracker;
 use smithay::backend::renderer::element::memory::MemoryRenderBufferRenderElement;
-use smithay::backend::renderer::element::surface::{
-    render_elements_from_surface_tree, WaylandSurfaceRenderElement,
-};
+use smithay::backend::renderer::element::surface::WaylandSurfaceRenderElement;
 use smithay::backend::renderer::element::utils::{
     select_dmabuf_feedback, CropRenderElement, Relocate, RelocateRenderElement,
     RescaleRenderElement,
@@ -157,10 +155,11 @@ use crate::render_helpers::debug::draw_opaque_regions;
 use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
+use crate::render_helpers::surface::push_elements_from_surface_tree;
 use crate::render_helpers::texture::TextureBuffer;
 use crate::render_helpers::{
     encompassing_geo, render_to_dmabuf, render_to_encompassing_texture, render_to_shm,
-    render_to_texture, render_to_vec, shaders, RenderTarget, SplitElements,
+    render_to_texture, render_to_vec, shaders, RenderTarget,
 };
 use crate::ui::config_error_notification::ConfigErrorNotification;
 use crate::ui::exit_confirm_dialog::{ExitConfirmDialog, ExitConfirmDialogRenderElement};
@@ -2075,9 +2074,8 @@ impl State {
 
                 self.backend.with_primary_renderer(|renderer| {
                     // FIXME: pointer.
-                    let elements = mapped
-                        .render_for_screen_cast(renderer, scale)
-                        .collect::<Vec<_>>();
+                    let mut elements = Vec::new();
+                    mapped.render_for_screen_cast(renderer, scale, &mut |elem| elements.push(elem));
 
                     if cast.dequeue_buffer_and_render(renderer, &elements, bbox.size, scale) {
                         cast.last_frame_time = get_monotonic_time();
@@ -3865,13 +3863,14 @@ impl Niri {
         }
     }
 
-    pub fn pointer_element<R: NiriRenderer>(
+    pub fn render_pointer<R: NiriRenderer>(
         &self,
         renderer: &mut R,
         output: &Output,
-    ) -> Vec<OutputRenderElements<R>> {
+        push: &mut dyn FnMut(PointerRenderElements<R>),
+    ) {
         if !self.pointer_visibility.is_visible() {
-            return vec![];
+            return;
         }
 
         let _span = tracy_client::span!("Niri::pointer_element");
@@ -3890,20 +3889,21 @@ impl Niri {
 
         let output_scale = Scale::from(output.current_scale().fractional_scale());
 
-        let mut pointer_elements = match render_cursor {
-            RenderCursor::Hidden => vec![],
+        match render_cursor {
+            RenderCursor::Hidden => (),
             RenderCursor::Surface { surface, hotspot } => {
                 let pointer_pos =
                     (pointer_pos - hotspot.to_f64()).to_physical_precise_round(output_scale);
 
-                render_elements_from_surface_tree(
+                push_elements_from_surface_tree(
                     renderer,
                     &surface,
                     pointer_pos,
                     output_scale,
                     1.,
                     Kind::Cursor,
-                )
+                    &mut |elem| push(elem.into()),
+                );
             }
             RenderCursor::Named {
                 icon,
@@ -3916,8 +3916,7 @@ impl Niri {
                     (pointer_pos - hotspot.to_f64()).to_physical_precise_round(output_scale);
 
                 let texture = self.cursor_texture_cache.get(icon, scale, &cursor, idx);
-                let mut pointer_elements = vec![];
-                let pointer_element = match MemoryRenderBufferRenderElement::from_buffer(
+                match MemoryRenderBufferRenderElement::from_buffer(
                     renderer,
                     pointer_pos,
                     &texture,
@@ -3926,34 +3925,27 @@ impl Niri {
                     None,
                     Kind::Cursor,
                 ) {
-                    Ok(element) => Some(element),
+                    Ok(element) => push(element.into()),
                     Err(err) => {
                         warn!("error importing a cursor texture: {err:?}");
-                        None
                     }
-                };
-                if let Some(element) = pointer_element {
-                    pointer_elements.push(OutputRenderElements::NamedPointer(element));
                 }
-
-                pointer_elements
             }
-        };
+        }
 
         if let Some(dnd_icon) = self.dnd_icon.as_ref() {
             let pointer_pos =
                 (pointer_pos + dnd_icon.offset.to_f64()).to_physical_precise_round(output_scale);
-            pointer_elements.extend(render_elements_from_surface_tree(
+            push_elements_from_surface_tree(
                 renderer,
                 &dnd_icon.surface,
                 pointer_pos,
                 output_scale,
                 1.,
                 Kind::ScanoutCandidate,
-            ));
+                &mut |elem| push(elem.into()),
+            );
         }
-
-        pointer_elements
     }
 
     pub fn refresh_pointer_outputs(&mut self) {
@@ -4317,7 +4309,7 @@ impl Niri {
         // The pointer goes on the top.
         let mut elements = vec![];
         if include_pointer {
-            elements = self.pointer_element(renderer, output);
+            self.render_pointer(renderer, output, &mut |elem| elements.push(elem.into()));
         }
 
         // Next, the screen transition texture.
@@ -4329,12 +4321,8 @@ impl Niri {
         }
 
         // Next, the exit confirm dialog.
-        elements.extend(
-            self.exit_confirm_dialog
-                .render(renderer, output)
-                .into_iter()
-                .map(OutputRenderElements::from),
-        );
+        self.exit_confirm_dialog
+            .render(renderer, output, &mut |elem| elements.push(elem.into()));
 
         // Next, the config error notification too.
         if let Some(element) = self.config_error_notification.render(renderer, output) {
@@ -4345,14 +4333,15 @@ impl Niri {
         if self.is_locked() {
             let state = self.output_state.get(output).unwrap();
             if let Some(surface) = state.lock_surface.as_ref() {
-                elements.extend(render_elements_from_surface_tree(
+                push_elements_from_surface_tree(
                     renderer,
                     surface.wl_surface(),
-                    (0, 0),
+                    Point::new(0, 0),
                     output_scale,
                     1.,
                     Kind::ScanoutCandidate,
-                ));
+                    &mut |elem| elements.push(elem.into()),
+                );
             }
 
             // Draw the solid color background.
@@ -4384,12 +4373,8 @@ impl Niri {
 
         // If the screenshot UI is open, draw it.
         if self.screenshot_ui.is_open() {
-            elements.extend(
-                self.screenshot_ui
-                    .render_output(output, target)
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
+            self.screenshot_ui
+                .render_output(output, target, &mut |elem| elements.push(elem.into()));
 
             // Add the backdrop for outputs that were connected while the screenshot UI was open.
             elements.push(backdrop);
@@ -4406,13 +4391,10 @@ impl Niri {
         }
 
         // Then, the Alt-Tab switcher.
-        let mru_elements = self
-            .window_mru_ui
-            .render_output(self, output, renderer, target)
-            .into_iter()
-            .flatten()
-            .map(OutputRenderElements::from);
-        elements.extend(mru_elements);
+        self.window_mru_ui
+            .render_output(self, output, renderer, target, &mut |elem| {
+                elements.push(elem.into())
+            });
 
         // Don't draw the focus ring on the workspaces while interactively moving above those
         // workspaces, since the interactively-moved window already has a focus ring.
@@ -4421,129 +4403,119 @@ impl Niri {
         // Get monitor elements.
         let mon = self.layout.monitor_for_output(output).unwrap();
         let zoom = mon.overview_zoom();
-        let monitor_elements = Vec::from_iter(
-            mon.render_elements(renderer, target, focus_ring)
-                .map(|(geo, bg, iter)| (geo, bg, Vec::from_iter(iter))),
-        );
-        let workspace_shadow_elements = Vec::from_iter(mon.render_workspace_shadows(renderer));
-        let insert_hint_elements = mon.render_insert_hint_between_workspaces(renderer);
-        let int_move_elements: Vec<_> = self
-            .layout
-            .render_interactive_move_for_output(renderer, output, target)
-            .collect();
 
         // Get layer-shell elements.
         let layer_map = layer_map_for_output(output);
-        let mut extend_from_layer =
-            |elements: &mut SplitElements<LayerSurfaceRenderElement<R>>, layer, for_backdrop| {
-                self.render_layer(renderer, target, &layer_map, layer, elements, for_backdrop);
-            };
+
+        // We use macros instead of closures to avoid borrowing issues (renderer and elements go
+        // into different functions).
+        macro_rules! push_popups_from_layer {
+            ($layer:expr, $backdrop:expr, $push:expr) => {{
+                self.render_layer_popups(renderer, target, &layer_map, $layer, $backdrop, $push);
+            }};
+            ($layer:expr, true) => {{
+                push_popups_from_layer!($layer, true, &mut |elem| elements.push(elem.into()));
+            }};
+            ($layer:expr, $push:expr) => {{
+                push_popups_from_layer!($layer, false, $push);
+            }};
+            ($layer:expr) => {{
+                push_popups_from_layer!($layer, false, &mut |elem| elements.push(elem.into()));
+            }};
+        }
+        macro_rules! push_normal_from_layer {
+            ($layer:expr, $backdrop:expr, $push:expr) => {{
+                self.render_layer_normal(renderer, target, &layer_map, $layer, $backdrop, $push);
+            }};
+            ($layer:expr, true) => {{
+                push_normal_from_layer!($layer, true, &mut |elem| elements.push(elem.into()));
+            }};
+            ($layer:expr, $push:expr) => {{
+                push_normal_from_layer!($layer, false, $push);
+            }};
+            ($layer:expr) => {{
+                push_normal_from_layer!($layer, false, &mut |elem| elements.push(elem.into()));
+            }};
+        }
 
         // The overlay layer elements go next.
-        let mut layer_elems = SplitElements::default();
-        extend_from_layer(&mut layer_elems, Layer::Overlay, false);
-        elements.extend(layer_elems.into_iter().map(OutputRenderElements::from));
-
-        // Collect the top layer elements.
-        let mut layer_elems = SplitElements::default();
-        extend_from_layer(&mut layer_elems, Layer::Top, false);
-        let top_layer = layer_elems;
+        push_popups_from_layer!(Layer::Overlay);
+        push_normal_from_layer!(Layer::Overlay);
 
         // When rendering above the top layer, we put the regular monitor elements first.
         // Otherwise, we will render all layer-shell pop-ups and the top layer on top.
         if mon.render_above_top_layer() {
-            // Collect all other layer-shell elements.
-            let mut layer_elems = SplitElements::default();
-            extend_from_layer(&mut layer_elems, Layer::Bottom, false);
-            extend_from_layer(&mut layer_elems, Layer::Background, false);
+            self.layout
+                .render_interactive_move_for_output(renderer, output, target, &mut |elem| {
+                    elements.push(elem.into())
+                });
 
-            elements.extend(
-                int_move_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
-            elements.extend(
-                insert_hint_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
+            mon.render_insert_hint_between_workspaces(renderer, &mut |elem| {
+                elements.push(elem.into())
+            });
 
-            let mut ws_background = None;
-            elements.extend(
-                monitor_elements
-                    .into_iter()
-                    .flat_map(|(_ws_geo, ws_bg, iter)| {
-                        ws_background = Some(ws_bg);
-                        iter
-                    })
-                    .map(OutputRenderElements::from),
-            );
+            mon.render_workspaces(renderer, target, focus_ring, &mut |elem| {
+                elements.push(elem.into())
+            });
 
-            elements.extend(top_layer.into_iter().map(OutputRenderElements::from));
-            elements.extend(layer_elems.into_iter().map(OutputRenderElements::from));
+            push_popups_from_layer!(Layer::Top);
+            push_normal_from_layer!(Layer::Top);
 
-            if let Some(ws_background) = ws_background {
-                elements.push(OutputRenderElements::from(ws_background));
+            push_popups_from_layer!(Layer::Bottom);
+            push_popups_from_layer!(Layer::Background);
+            push_normal_from_layer!(Layer::Bottom);
+            push_normal_from_layer!(Layer::Background);
+
+            // We don't expect more than one workspace when render_above_top_layer().
+            if let Some((ws, _geo)) = mon.workspaces_with_render_geo().next() {
+                elements.push(ws.render_background().into());
             }
-
-            elements.extend(
-                workspace_shadow_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
         } else {
-            elements.extend(top_layer.into_iter().map(OutputRenderElements::from));
+            push_popups_from_layer!(Layer::Top);
+            push_normal_from_layer!(Layer::Top);
 
-            elements.extend(
-                int_move_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
+            self.layout
+                .render_interactive_move_for_output(renderer, output, target, &mut |elem| {
+                    elements.push(elem.into())
+                });
 
-            elements.extend(
-                insert_hint_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
+            mon.render_insert_hint_between_workspaces(renderer, &mut |elem| {
+                elements.push(elem.into())
+            });
 
-            for (ws_geo, ws_background, ws_elements) in monitor_elements {
-                // Collect all other layer-shell elements.
-                let mut layer_elems = SplitElements::default();
-                extend_from_layer(&mut layer_elems, Layer::Bottom, false);
-                extend_from_layer(&mut layer_elems, Layer::Background, false);
-
-                elements.extend(
-                    layer_elems
-                        .popups
-                        .into_iter()
-                        .filter_map(|elem| scale_relocate_crop(elem, output_scale, zoom, ws_geo))
-                        .map(OutputRenderElements::from),
-                );
-
-                elements.extend(ws_elements.into_iter().map(OutputRenderElements::from));
-
-                elements.extend(
-                    layer_elems
-                        .normal
-                        .into_iter()
-                        .filter_map(|elem| scale_relocate_crop(elem, output_scale, zoom, ws_geo))
-                        .map(OutputRenderElements::from),
-                );
-
-                elements.push(OutputRenderElements::from(ws_background));
+            // Macro instead of closure to avoid borrowing elements.
+            macro_rules! process {
+                ($geo:expr) => {{
+                    &mut |elem| {
+                        if let Some(elem) = scale_relocate_crop(elem, output_scale, zoom, $geo) {
+                            elements.push(elem.into());
+                        }
+                    }
+                }};
             }
 
-            elements.extend(
-                workspace_shadow_elements
-                    .into_iter()
-                    .map(OutputRenderElements::from),
-            );
+            for (_ws, geo) in mon.workspaces_with_render_geo() {
+                push_popups_from_layer!(Layer::Bottom, process!(geo));
+                push_popups_from_layer!(Layer::Background, process!(geo));
+            }
+
+            mon.render_workspaces(renderer, target, focus_ring, &mut |elem| {
+                elements.push(elem.into())
+            });
+
+            for (ws, geo) in mon.workspaces_with_render_geo() {
+                push_normal_from_layer!(Layer::Bottom, process!(geo));
+                push_normal_from_layer!(Layer::Background, process!(geo));
+
+                process!(geo)(ws.render_background());
+            }
         }
 
+        mon.render_workspace_shadows(renderer, &mut |elem| elements.push(elem.into()));
+
         // Then the backdrop.
-        let mut layer_elems = SplitElements::default();
-        extend_from_layer(&mut layer_elems, Layer::Background, true);
-        elements.extend(layer_elems.into_iter().map(OutputRenderElements::from));
+        push_popups_from_layer!(Layer::Background, true);
+        push_normal_from_layer!(Layer::Background, true);
 
         elements.push(backdrop);
 
@@ -4554,18 +4526,14 @@ impl Niri {
         elements
     }
 
-    #[allow(clippy::too_many_arguments)]
-    fn render_layer<R: NiriRenderer>(
-        &self,
-        renderer: &mut R,
-        target: RenderTarget,
-        layer_map: &LayerMap,
+    fn layers_in_render_order<'a>(
+        &'a self,
+        layer_map: &'a LayerMap,
         layer: Layer,
-        elements: &mut SplitElements<LayerSurfaceRenderElement<R>>,
         for_backdrop: bool,
-    ) {
+    ) -> impl Iterator<Item = (&'a MappedLayer, Rectangle<i32, Logical>)> {
         // LayerMap returns layers in reverse stacking order.
-        let iter = layer_map.layers_on(layer).rev().filter_map(|surface| {
+        layer_map.layers_on(layer).rev().filter_map(move |surface| {
             let mapped = self.mapped_layer_surfaces.get(surface)?;
 
             if for_backdrop != mapped.place_within_backdrop() {
@@ -4574,9 +4542,34 @@ impl Niri {
 
             let geo = layer_map.layer_geometry(surface)?;
             Some((mapped, geo))
-        });
-        for (mapped, geo) in iter {
-            elements.extend(mapped.render(renderer, geo.loc.to_f64(), target));
+        })
+    }
+
+    fn render_layer_normal<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
+        target: RenderTarget,
+        layer_map: &LayerMap,
+        layer: Layer,
+        for_backdrop: bool,
+        push: &mut dyn FnMut(LayerSurfaceRenderElement<R>),
+    ) {
+        for (mapped, geo) in self.layers_in_render_order(layer_map, layer, for_backdrop) {
+            mapped.render_normal(renderer, geo.loc.to_f64(), target, push);
+        }
+    }
+
+    fn render_layer_popups<R: NiriRenderer>(
+        &self,
+        renderer: &mut R,
+        target: RenderTarget,
+        layer_map: &LayerMap,
+        layer: Layer,
+        for_backdrop: bool,
+        push: &mut dyn FnMut(LayerSurfaceRenderElement<R>),
+    ) {
+        for (mapped, geo) in self.layers_in_render_order(layer_map, layer, for_backdrop) {
+            mapped.render_popups(renderer, geo.loc.to_f64(), target, push);
         }
     }
 
@@ -5291,7 +5284,8 @@ impl Niri {
             }
 
             // FIXME: pointer.
-            let elements: Vec<_> = mapped.render_for_screen_cast(renderer, scale).collect();
+            let mut elements = Vec::new();
+            mapped.render_for_screen_cast(renderer, scale, &mut |elem| elements.push(elem));
 
             if cast.dequeue_buffer_and_render(renderer, &elements, bbox.size, scale) {
                 cast.last_frame_time = target_presentation_time;
@@ -5594,7 +5588,8 @@ impl Niri {
                 }
                 let res_output = res.ok();
 
-                let pointer = self.pointer_element(renderer, &output);
+                let mut pointer = Vec::new();
+                self.render_pointer(renderer, &output, &mut |elem| pointer.push(elem));
                 let res_pointer = if pointer.is_empty() {
                     None
                 } else {
@@ -5685,12 +5680,14 @@ impl Niri {
                 mapped.rules().opacity.unwrap_or(1.).clamp(0., 1.)
             };
         // FIXME: pointer.
-        let elements = mapped.render(
+        let mut elements = Vec::new();
+        mapped.render(
             renderer,
             mapped.window.geometry().loc.to_f64(),
             scale,
             alpha,
             RenderTarget::ScreenCapture,
+            &mut |elem| elements.push(elem),
         );
         let geo = encompassing_geo(scale, elements.iter());
         let elements = elements.iter().rev().map(|elem| {
@@ -6556,6 +6553,13 @@ fn scale_relocate_crop<E: Element>(
     CropRenderElement::from_element(elem, output_scale, ws_geo)
 }
 
+niri_render_elements! {
+    PointerRenderElements<R> => {
+        Wayland = WaylandSurfaceRenderElement<R>,
+        NamedPointer = MemoryRenderBufferRenderElement<R>,
+    }
+}
+
 niri_render_elements! {
     OutputRenderElements<R> => {
         Monitor = MonitorRenderElement<R>,
@@ -6564,8 +6568,11 @@ niri_render_elements! {
         RelocatedLayerSurface = CropRenderElement<RelocateRenderElement<RescaleRenderElement<
             LayerSurfaceRenderElement<R>
         >>>,
+        RelocatedColor = CropRenderElement<RelocateRenderElement<RescaleRenderElement<
+            SolidColorRenderElement
+        >>>,
+        Pointer = PointerRenderElements<R>,
         Wayland = WaylandSurfaceRenderElement<R>,
-        NamedPointer = MemoryRenderBufferRenderElement<R>,
         SolidColor = SolidColorRenderElement,
         ScreenshotUi = ScreenshotUiRenderElement,
         WindowMruUi = WindowMruUiRenderElement<R>,

src/render_helpers/mod.rs

@@ -58,13 +58,6 @@ pub struct BakedBuffer<B> {
     pub dst: Option<Size<i32, Logical>>,
 }
 
-/// Render elements split into normal and popup.
-#[derive(Debug)]
-pub struct SplitElements<E> {
-    pub normal: Vec<E>,
-    pub popups: Vec<E>,
-}
-
 pub trait ToRenderElement {
     type RenderElement;
 
@@ -87,41 +80,6 @@ impl RenderTarget {
     }
 }
 
-impl<E> Default for SplitElements<E> {
-    fn default() -> Self {
-        Self {
-            normal: Vec::new(),
-            popups: Vec::new(),
-        }
-    }
-}
-
-impl<E> IntoIterator for SplitElements<E> {
-    type Item = E;
-    type IntoIter = std::iter::Chain<std::vec::IntoIter<E>, std::vec::IntoIter<E>>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.popups.into_iter().chain(self.normal)
-    }
-}
-
-impl<E> SplitElements<E> {
-    pub fn iter(&self) -> std::iter::Chain<std::slice::Iter<'_, E>, std::slice::Iter<'_, E>> {
-        self.popups.iter().chain(&self.normal)
-    }
-
-    pub fn into_vec(self) -> Vec<E> {
-        let Self { normal, mut popups } = self;
-        popups.extend(normal);
-        popups
-    }
-
-    pub fn extend(&mut self, other: SplitElements<E>) {
-        self.popups.extend(other.popups);
-        self.normal.extend(other.normal);
-    }
-}
-
 impl ToRenderElement for BakedBuffer<TextureBuffer<GlesTexture>> {
     type RenderElement = PrimaryGpuTextureRenderElement;
 

src/render_helpers/surface.rs

@@ -1,8 +1,10 @@
+use smithay::backend::renderer::element::surface::WaylandSurfaceRenderElement;
+use smithay::backend::renderer::element::Kind;
 use smithay::backend::renderer::gles::{GlesRenderer, GlesTexture};
 use smithay::backend::renderer::utils::{import_surface, RendererSurfaceStateUserData};
-use smithay::backend::renderer::Renderer as _;
+use smithay::backend::renderer::{ImportAll, Renderer};
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
-use smithay::utils::{Logical, Point};
+use smithay::utils::{Logical, Physical, Point, Scale};
 use smithay::wayland::compositor::{with_surface_tree_downward, TraversalAction};
 
 use super::texture::TextureBuffer;
@@ -78,3 +80,67 @@ pub fn render_snapshot_from_surface_tree(
         |_, _, _| true,
     );
 }
+
+pub fn push_elements_from_surface_tree<R>(
+    renderer: &mut R,
+    surface: &WlSurface,
+    // Fractional scale expects surface buffers to be aligned to physical pixels.
+    location: Point<i32, Physical>,
+    scale: Scale<f64>,
+    alpha: f32,
+    kind: Kind,
+    push: &mut dyn FnMut(WaylandSurfaceRenderElement<R>),
+) where
+    R: Renderer + ImportAll,
+    R::TextureId: Clone + 'static,
+{
+    let _span = tracy_client::span!("push_elements_from_surface_tree");
+
+    let location = location.to_f64();
+
+    with_surface_tree_downward(
+        surface,
+        location,
+        |_, states, location| {
+            let mut location = *location;
+            let data = states.data_map.get::<RendererSurfaceStateUserData>();
+
+            if let Some(data) = data {
+                if let Some(view) = data.lock().unwrap().view() {
+                    location += view.offset.to_f64().to_physical(scale);
+                    TraversalAction::DoChildren(location)
+                } else {
+                    TraversalAction::SkipChildren
+                }
+            } else {
+                TraversalAction::SkipChildren
+            }
+        },
+        |surface, states, location| {
+            let mut location = *location;
+            let data = states.data_map.get::<RendererSurfaceStateUserData>();
+
+            if let Some(data) = data {
+                let has_view = if let Some(view) = data.lock().unwrap().view() {
+                    location += view.offset.to_f64().to_physical(scale);
+                    true
+                } else {
+                    false
+                };
+
+                if has_view {
+                    match WaylandSurfaceRenderElement::from_surface(
+                        renderer, surface, states, location, alpha, kind,
+                    ) {
+                        Ok(Some(surface)) => push(surface),
+                        Ok(None) => {} // surface is not mapped
+                        Err(err) => {
+                            warn!("failed to import surface: {}", err);
+                        }
+                    };
+                }
+            }
+        },
+        |_, _, _| true,
+    );
+}

src/ui/exit_confirm_dialog.rs

@@ -3,7 +3,6 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::Mutex;
 
-use arrayvec::ArrayVec;
 use niri_config::Config;
 use ordered_float::NotNan;
 use pangocairo::cairo::{self, ImageSurface};
@@ -151,14 +150,15 @@ impl ExitConfirmDialog {
         &self,
         renderer: &mut R,
         output: &Output,
-    ) -> ArrayVec<ExitConfirmDialogRenderElement, 2> {
-        let mut rv = ArrayVec::new();
-
+        push: &mut dyn FnMut(ExitConfirmDialogRenderElement),
+    ) {
         let (value, clamped_value) = match &self.state {
-            State::Hidden => return rv,
+            State::Hidden => return,
             State::Showing(anim) | State::Hiding(anim) => (anim.value(), anim.clamped_value()),
             State::Visible => (1., 1.),
         };
+        let _span = tracy_client::span!("ExitConfirmDialog::render");
+
         // Can be out of range when starting from past 0. or 1. from a spring bounce.
         let clamped_value = clamped_value.clamp(0., 1.);
 
@@ -168,7 +168,7 @@ impl ExitConfirmDialog {
         let mut buffers = self.buffers.borrow_mut();
         let Some(fallback) = buffers[&NotNan::new(1.).unwrap()].clone() else {
             error!("exit confirm dialog opened without fallback buffer");
-            return rv;
+            return;
         };
 
         let buffer = buffers
@@ -179,7 +179,7 @@ impl ExitConfirmDialog {
         let size = buffer.logical_size();
         let Ok(buffer) = TextureBuffer::from_memory_buffer(renderer.as_gles_renderer(), buffer)
         else {
-            return rv;
+            return;
         };
 
         let location = (output_size.to_point() - size.to_point()).downscale(2.);
@@ -201,7 +201,7 @@ impl ExitConfirmDialog {
             (location + size.downscale(2.)).to_physical_precise_round(scale),
             value.max(0.) * 0.2 + 0.8,
         );
-        rv.push(ExitConfirmDialogRenderElement::Texture(elem));
+        push(ExitConfirmDialogRenderElement::Texture(elem));
 
         // Backdrop.
         let data = output.user_data().get_or_insert(|| {
@@ -218,9 +218,7 @@ impl ExitConfirmDialog {
             clamped_value as f32,
             Kind::Unspecified,
         );
-        rv.push(ExitConfirmDialogRenderElement::SolidColor(elem));
-
-        rv
+        push(ExitConfirmDialogRenderElement::SolidColor(elem));
     }
 }
 

src/ui/mru.rs

@@ -346,7 +346,8 @@ impl Thumbnail {
         is_active: bool,
         bob_y: f64,
         target: RenderTarget,
-    ) -> impl Iterator<Item = WindowMruUiRenderElement<R>> {
+        push: &mut dyn FnMut(WindowMruUiRenderElement<R>),
+    ) {
         let _span = tracy_client::span!("Thumbnail::render");
 
         let round = move |logical: f64| round_logical_in_physical(scale, logical);
@@ -368,11 +369,6 @@ impl Thumbnail {
         };
         let bob_offset = Point::new(0., bob_y);
 
-        // FIXME: this could use mipmaps, for that it should be rendered through an offscreen.
-        let elems = mapped
-            .render_normal(renderer, Point::new(0., 0.), s, preview_alpha, target)
-            .into_iter();
-
         // Clip thumbnails to their geometry.
         let radius = if mapped.sizing_mode().is_normal() {
             mapped.rules().geometry_corner_radius
@@ -385,7 +381,7 @@ impl Thumbnail {
         let clip_shader = ClippedSurfaceRenderElement::shader(renderer).cloned();
         let geo = Rectangle::from_size(self.size.to_f64());
         // FIXME: deduplicate code with Tile::render_inner()
-        let elems = elems.map(move |elem| match elem {
+        let clip = move |elem| match elem {
             LayoutElementRenderElement::Wayland(elem) => {
                 if let Some(shader) = clip_shader.clone() {
                     if ClippedSurfaceRenderElement::will_clip(&elem, s, geo, radius) {
@@ -426,9 +422,9 @@ impl Thumbnail {
                 // Otherwise, render the solid color as is.
                 LayoutElementRenderElement::SolidColor(elem).into()
             }
-        });
+        };
 
-        let elems = elems.map(move |elem| {
+        let downscale = move |elem| {
             let thumb_scale = Scale {
                 x: preview_geo.size.w / geo.size.w,
                 y: preview_geo.size.h / geo.size.h,
@@ -445,7 +441,21 @@ impl Thumbnail {
                 Relocate::Relative,
             );
             WindowMruUiRenderElement::Thumbnail(elem)
-        });
+        };
+
+        // FIXME: this could use mipmaps, for that it should be rendered through an offscreen.
+        mapped.render_normal(
+            renderer,
+            Point::new(0., 0.),
+            s,
+            preview_alpha,
+            target,
+            &mut |elem| {
+                let elem = clip(elem);
+                let elem = downscale(elem);
+                push(elem)
+            },
+        );
 
         let mut title_size = None;
         let title_texture = self.title_texture(renderer.as_gles_renderer(), mapped, scale);
@@ -462,7 +472,7 @@ impl Thumbnail {
         let should_block_out = target.should_block_out(mapped.rules().block_out_from);
         let title_texture = title_texture.filter(|_| !should_block_out);
 
-        let title_elems = title_texture.map(|(texture, size)| {
+        if let Some((texture, size)) = title_texture {
             // Clip from the right if it doesn't fit.
             let src = Rectangle::from_size(size);
 
@@ -484,15 +494,15 @@ impl Thumbnail {
             let renderer = renderer.as_gles_renderer();
             if let Some(program) = GradientFadeTextureRenderElement::shader(renderer) {
                 let elem = GradientFadeTextureRenderElement::new(texture, program);
-                WindowMruUiRenderElement::GradientFadeElem(elem)
+                push(WindowMruUiRenderElement::GradientFadeElem(elem));
             } else {
                 let elem = PrimaryGpuTextureRenderElement(texture);
-                WindowMruUiRenderElement::TextureElement(elem)
+                push(WindowMruUiRenderElement::TextureElement(elem));
             }
-        });
+        }
 
         let is_urgent = mapped.is_urgent();
-        let background_elems = (is_active || is_urgent).then(|| {
+        if is_active || is_urgent {
             let padding = Point::new(padding, padding);
 
             let mut size = preview_geo.size;
@@ -532,9 +542,9 @@ impl Thumbnail {
                 scale,
                 0.5,
             );
-            let bg_elems = background
-                .render(renderer, loc)
-                .map(WindowMruUiRenderElement::FocusRing);
+            background.render(renderer, loc, &mut |elem| {
+                push(WindowMruUiRenderElement::FocusRing(elem))
+            });
 
             let mut border = self.border.borrow_mut();
             let mut config = *border.config();
@@ -554,15 +564,10 @@ impl Thumbnail {
                 1.,
             );
 
-            let border_elems = border
-                .render(renderer, loc)
-                .map(WindowMruUiRenderElement::FocusRing);
-
-            bg_elems.chain(border_elems)
-        });
-        let background_elems = background_elems.into_iter().flatten();
-
-        elems.chain(title_elems).chain(background_elems)
+            border.render(renderer, loc, &mut |elem| {
+                push(WindowMruUiRenderElement::FocusRing(elem))
+            });
+        }
     }
 }
 
@@ -1091,26 +1096,27 @@ impl WindowMruUi {
         }
     }
 
-    pub fn render_output<'a, R: NiriRenderer>(
-        &'a self,
-        niri: &'a Niri,
+    pub fn render_output<R: NiriRenderer>(
+        &self,
+        niri: &Niri,
         output: &Output,
-        renderer: &'a mut R,
+        renderer: &mut R,
         target: RenderTarget,
-    ) -> Option<impl Iterator<Item = WindowMruUiRenderElement<R>> + 'a> {
+        push: &mut dyn FnMut(WindowMruUiRenderElement<R>),
+    ) {
         let (inner, progress) = match &self.state {
-            UiState::Closed { .. } => return None,
+            UiState::Closed { .. } => return,
             UiState::Closing { inner, anim } => (inner, anim.clamped_value()),
             UiState::Open(inner) => {
                 if inner.is_fully_open() {
                     (inner, 1.)
                 } else {
-                    return None;
+                    return;
                 }
             }
         };
 
-        let span = tracy_client::span!("mru render");
+        let _span = tracy_client::span!("WindowMruUi::render_output");
 
         let alpha = progress.clamp(0., 1.) as f32;
 
@@ -1131,9 +1137,12 @@ impl WindowMruUi {
         };
 
         // During the closing fade, use an offscreen to avoid transparent compositing artifacts.
-        let offscreen_elem = if *output == inner.output && alpha < 1. {
+        let mut pushed_offscreen = false;
+        if *output == inner.output && alpha < 1. {
             let renderer = renderer.as_gles_renderer();
-            let mut elems = Vec::from_iter(inner.render(niri, renderer, target));
+
+            let mut elems = Vec::new();
+            inner.render(niri, renderer, target, &mut |elem| elems.push(elem));
             elems.push(WindowMruUiRenderElement::SolidColor(render_backdrop(1.)));
 
             let scale = output.current_scale().fractional_scale();
@@ -1149,39 +1158,27 @@ impl WindowMruUi {
                     // itself).
                     //
                     // Anyhow, this is not very noticeable since Alt-Tab closing happens quickly.
-                    Some(WindowMruUiRenderElement::Offscreen(elem.with_alpha(alpha)))
+                    push(WindowMruUiRenderElement::Offscreen(elem.with_alpha(alpha)));
+                    pushed_offscreen = true;
                 }
                 Err(err) => {
                     warn!("error rendering MRU to offscreen for fade-out: {err:?}");
-                    None
                 }
             }
-        } else {
-            None
-        };
+        }
 
         // When alpha is 1., render everything directly, without an offscreen.
         //
         // This is not used as fallback when offscreen fails to render because it looks better to
         // hide the previews immediately than to render them with alpha = 1. during a fade-out.
-        let normal_elems =
-            (*output == inner.output && alpha == 1.).then(|| inner.render(niri, renderer, target));
-        let normal_elems = normal_elems.into_iter().flatten();
+        if *output == inner.output && alpha == 1. {
+            inner.render(niri, renderer, target, &mut |elem| push(elem));
+        }
 
         // This is used for both normal elems and for other outputs.
-        let backdrop_elem = (offscreen_elem.is_none())
-            .then(|| WindowMruUiRenderElement::SolidColor(render_backdrop(alpha)));
-
-        // Make sure the span includes consuming the iterator.
-        let drop_span = std::iter::once(span).filter_map(|_| None);
-
-        Some(
-            offscreen_elem
-                .into_iter()
-                .chain(normal_elems)
-                .chain(backdrop_elem)
-                .chain(drop_span),
-        )
+        if !pushed_offscreen {
+            push(WindowMruUiRenderElement::SolidColor(render_backdrop(alpha)));
+        }
     }
 
     pub fn are_animations_ongoing(&self) -> bool {
@@ -1554,12 +1551,13 @@ impl Inner {
         })
     }
 
-    fn render<'a, R: NiriRenderer>(
-        &'a self,
-        niri: &'a Niri,
-        renderer: &'a mut R,
+    fn render<R: NiriRenderer>(
+        &self,
+        niri: &Niri,
+        renderer: &mut R,
         target: RenderTarget,
-    ) -> impl Iterator<Item = WindowMruUiRenderElement<R>> + 'a {
+        push: &mut dyn FnMut(WindowMruUiRenderElement<R>),
+    ) {
         let output_size = output_size(&self.output);
         let scale = self.output.current_scale().fractional_scale();
 
@@ -1567,7 +1565,7 @@ impl Inner {
             self.scope_panel
                 .borrow_mut()
                 .get(renderer.as_gles_renderer(), scale, self.wmru.scope);
-        let panel = panel_texture.map(move |texture| {
+        if let Some(texture) = panel_texture {
             let padding = round_logical_in_physical(scale, f64::from(PANEL_PADDING));
 
             let size = texture.logical_size();
@@ -1580,9 +1578,8 @@ impl Inner {
                 None,
                 Kind::Unspecified,
             ));
-            WindowMruUiRenderElement::TextureElement(elem)
-        });
-        let panel = panel.into_iter();
+            push(WindowMruUiRenderElement::TextureElement(elem));
+        }
 
         let current_id = self.wmru.current_id;
 
@@ -1591,26 +1588,20 @@ impl Inner {
 
         let config = self.config.borrow();
 
-        let thumbnails = self
-            .thumbnails_in_view_render()
-            .filter_map(move |(thumbnail, geo)| {
-                let id = thumbnail.id;
-                let Some((_, mapped)) = niri.layout.windows().find(|(_, m)| m.id() == id) else {
-                    error!("window in the MRU must be present in the layout");
-                    return None;
-                };
+        for (thumbnail, geo) in self.thumbnails_in_view_render() {
+            let id = thumbnail.id;
+            let Some((_, mapped)) = niri.layout.windows().find(|(_, m)| m.id() == id) else {
+                error!("window in the MRU must be present in the layout");
+                continue;
+            };
 
-                let config = &config.recent_windows;
+            let config = &config.recent_windows;
 
-                let is_active = Some(id) == current_id;
-                let elems = thumbnail.render(
-                    renderer, config, mapped, geo, scale, is_active, bob_y, target,
-                );
-                Some(elems)
-            });
-        let thumbnails = thumbnails.flatten();
-
-        panel.chain(thumbnails)
+            let is_active = Some(id) == current_id;
+            thumbnail.render(
+                renderer, config, mapped, geo, scale, is_active, bob_y, target, push,
+            );
+        }
     }
 
     fn thumbnail_under(&self, pos: Point<f64, Logical>) -> Option<MappedId> {

src/ui/screenshot_ui.rs

@@ -623,7 +623,8 @@ impl ScreenshotUi {
         &self,
         output: &Output,
         target: RenderTarget,
-    ) -> ArrayVec<ScreenshotUiRenderElement, 11> {
+        push: &mut dyn FnMut(ScreenshotUiRenderElement),
+    ) {
         let _span = tracy_client::span!("ScreenshotUi::render_output");
 
         let Self::Open {
@@ -637,10 +638,8 @@ impl ScreenshotUi {
             panic!("screenshot UI must be open to render it");
         };
 
-        let mut elements = ArrayVec::new();
-
         let Some(output_data) = output_data.get(output) else {
-            return elements;
+            return;
         };
 
         let scale = output_data.scale;
@@ -666,19 +665,18 @@ impl ScreenshotUi {
                 None,
                 Kind::Unspecified,
             ));
-            elements.push(elem.into());
+            push(elem.into());
         }
 
-        let buf_loc = zip(&output_data.buffers, &output_data.locations);
-        elements.extend(buf_loc.map(|(buffer, loc)| {
-            SolidColorRenderElement::from_buffer(
+        for (buffer, loc) in zip(&output_data.buffers, &output_data.locations) {
+            let elem = SolidColorRenderElement::from_buffer(
                 buffer,
                 loc.to_f64().to_logical(scale),
                 progress,
                 Kind::Unspecified,
-            )
-            .into()
-        }));
+            );
+            push(elem.into());
+        }
 
         // The screenshot itself goes last.
         let index = match target {
@@ -690,12 +688,10 @@ impl ScreenshotUi {
 
         if *show_pointer {
             if let Some(pointer) = screenshot.pointer.clone() {
-                elements.push(pointer.into());
+                push(pointer.into());
             }
         }
-        elements.push(screenshot.buffer.clone().into());
-
-        elements
+        push(screenshot.buffer.clone().into());
     }
 
     pub fn capture(

src/window/mapped.rs

@@ -2,7 +2,7 @@ use std::cell::{Cell, Ref, RefCell};
 use std::time::Duration;
 
 use niri_config::{Color, CornerRadius, GradientInterpolation, WindowRule};
-use smithay::backend::renderer::element::surface::render_elements_from_surface_tree;
+use smithay::backend::renderer::element::surface::WaylandSurfaceRenderElement;
 use smithay::backend::renderer::element::Kind;
 use smithay::backend::renderer::gles::GlesRenderer;
 use smithay::desktop::space::SpaceElement as _;
@@ -33,8 +33,10 @@ use crate::render_helpers::offscreen::OffscreenData;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::snapshot::RenderSnapshot;
 use crate::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
-use crate::render_helpers::surface::render_snapshot_from_surface_tree;
-use crate::render_helpers::{BakedBuffer, RenderTarget, SplitElements};
+use crate::render_helpers::surface::{
+    push_elements_from_surface_tree, render_snapshot_from_surface_tree,
+};
+use crate::render_helpers::{BakedBuffer, RenderTarget};
 use crate::utils::id::IdCounter;
 use crate::utils::transaction::Transaction;
 use crate::utils::{
@@ -472,7 +474,8 @@ impl Mapped {
         &self,
         renderer: &mut R,
         scale: Scale<f64>,
-    ) -> impl DoubleEndedIterator<Item = WindowCastRenderElements<R>> {
+        push: &mut dyn FnMut(WindowCastRenderElements<R>),
+    ) {
         let bbox = self.window.bbox_with_popups().to_physical_precise_up(scale);
 
         let has_border_shader = BorderRenderElement::has_shader(renderer);
@@ -484,11 +487,9 @@ impl Mapped {
             .to_physical_precise_round(scale)
             .to_logical(scale);
         let radius = radius.fit_to(window_size.w as f32, window_size.h as f32);
-
         let location = self.window.geometry().loc.to_f64() - bbox.loc.to_logical(scale);
-        let elements = self.render(renderer, location, scale, 1., RenderTarget::Screencast);
 
-        elements.into_iter().map(move |elem| {
+        let use_border = |elem| {
             if let LayoutElementRenderElement::SolidColor(elem) = &elem {
                 // In this branch we're rendering a blocked-out window with a solid color. We need
                 // to render it with a rounded corner shader even if clip_to_geometry is false,
@@ -516,7 +517,16 @@ impl Mapped {
             }
 
             WindowCastRenderElements::from(elem)
-        })
+        };
+
+        self.render(
+            renderer,
+            location,
+            scale,
+            1.,
+            RenderTarget::Screencast,
+            &mut |elem| push(use_border(elem)),
+        );
     }
 
     pub fn get_focus_timestamp(&self) -> Option<Duration> {
@@ -601,52 +611,6 @@ impl LayoutElement for Mapped {
         self.window.is_in_input_region(&surface_local)
     }
 
-    fn render<R: NiriRenderer>(
-        &self,
-        renderer: &mut R,
-        location: Point<f64, Logical>,
-        scale: Scale<f64>,
-        alpha: f32,
-        target: RenderTarget,
-    ) -> SplitElements<LayoutElementRenderElement<R>> {
-        let mut rv = SplitElements::default();
-
-        if target.should_block_out(self.rules.block_out_from) {
-            let mut buffer = self.block_out_buffer.borrow_mut();
-            buffer.resize(self.window.geometry().size.to_f64());
-            let elem =
-                SolidColorRenderElement::from_buffer(&buffer, location, alpha, Kind::Unspecified);
-            rv.normal.push(elem.into());
-        } else {
-            let buf_pos = location - self.window.geometry().loc.to_f64();
-
-            let surface = self.toplevel().wl_surface();
-            for (popup, popup_offset) in PopupManager::popups_for_surface(surface) {
-                let offset = self.window.geometry().loc + popup_offset - popup.geometry().loc;
-
-                rv.popups.extend(render_elements_from_surface_tree(
-                    renderer,
-                    popup.wl_surface(),
-                    (buf_pos + offset.to_f64()).to_physical_precise_round(scale),
-                    scale,
-                    alpha,
-                    Kind::ScanoutCandidate,
-                ));
-            }
-
-            rv.normal = render_elements_from_surface_tree(
-                renderer,
-                surface,
-                buf_pos.to_physical_precise_round(scale),
-                scale,
-                alpha,
-                Kind::ScanoutCandidate,
-            );
-        }
-
-        rv
-    }
-
     fn render_normal<R: NiriRenderer>(
         &self,
         renderer: &mut R,
@@ -654,23 +618,26 @@ impl LayoutElement for Mapped {
         scale: Scale<f64>,
         alpha: f32,
         target: RenderTarget,
-    ) -> Vec<LayoutElementRenderElement<R>> {
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
         if target.should_block_out(self.rules.block_out_from) {
             let mut buffer = self.block_out_buffer.borrow_mut();
             buffer.resize(self.window.geometry().size.to_f64());
             let elem =
                 SolidColorRenderElement::from_buffer(&buffer, location, alpha, Kind::Unspecified);
-            vec![elem.into()]
+            push(elem.into());
         } else {
             let buf_pos = location - self.window.geometry().loc.to_f64();
             let surface = self.toplevel().wl_surface();
-            render_elements_from_surface_tree(
+            let mut push = |elem: WaylandSurfaceRenderElement<R>| push(elem.into());
+            push_elements_from_surface_tree(
                 renderer,
                 surface,
                 buf_pos.to_physical_precise_round(scale),
                 scale,
                 alpha,
                 Kind::ScanoutCandidate,
+                &mut push,
             )
         }
     }
@@ -682,28 +649,27 @@ impl LayoutElement for Mapped {
         scale: Scale<f64>,
         alpha: f32,
         target: RenderTarget,
-    ) -> Vec<LayoutElementRenderElement<R>> {
+        push: &mut dyn FnMut(LayoutElementRenderElement<R>),
+    ) {
         if target.should_block_out(self.rules.block_out_from) {
-            vec![]
-        } else {
-            let mut rv = vec![];
+            return;
+        }
 
-            let buf_pos = location - self.window.geometry().loc.to_f64();
-            let surface = self.toplevel().wl_surface();
-            for (popup, popup_offset) in PopupManager::popups_for_surface(surface) {
-                let offset = self.window.geometry().loc + popup_offset - popup.geometry().loc;
+        let buf_pos = location - self.window.geometry().loc.to_f64();
+        let surface = self.toplevel().wl_surface();
+        let mut push = |elem: WaylandSurfaceRenderElement<R>| push(elem.into());
+        for (popup, popup_offset) in PopupManager::popups_for_surface(surface) {
+            let offset = self.window.geometry().loc + popup_offset - popup.geometry().loc;
 
-                rv.extend(render_elements_from_surface_tree(
-                    renderer,
-                    popup.wl_surface(),
-                    (buf_pos + offset.to_f64()).to_physical_precise_round(scale),
-                    scale,
-                    alpha,
-                    Kind::ScanoutCandidate,
-                ));
-            }
-
-            rv
+            push_elements_from_surface_tree(
+                renderer,
+                popup.wl_surface(),
+                (buf_pos + offset.to_f64()).to_physical_precise_round(scale),
+                scale,
+                alpha,
+                Kind::ScanoutCandidate,
+                &mut push,
+            );
         }
     }