Improved UI render batching

crates/bevy_ui/src/render/mod.rs

@@ -554,6 +554,7 @@ struct UiVertex {
     pub position: [f32; 3],
     pub uv: [f32; 2],
     pub color: [f32; 4],
+    pub mode: u32,
 }
 
 #[derive(Resource)]
@@ -587,6 +588,9 @@ pub struct UiBatch {
     pub z: f32,
 }
 
+const TEXTURED_QUAD: u32 = 0;
+const UNTEXTURED_QUAD: u32 = 1;
+
 pub fn prepare_uinodes(
     mut commands: Commands,
     render_device: Res<RenderDevice>,
@@ -603,20 +607,33 @@ pub fn prepare_uinodes(
 
     let mut start = 0;
     let mut end = 0;
-    let mut current_batch_handle = Default::default();
+    let mut current_batch_image = DEFAULT_IMAGE_HANDLE.typed();
     let mut last_z = 0.0;
+
+    #[inline]
+    fn is_textured(image: &Handle<Image>) -> bool {
+        image.id() != DEFAULT_IMAGE_HANDLE.id()
+    }
+
     for extracted_uinode in &extracted_uinodes.uinodes {
-        if current_batch_handle != extracted_uinode.image {
-            if start != end {
-                commands.spawn(UiBatch {
-                    range: start..end,
-                    image: current_batch_handle,
-                    z: last_z,
-                });
-                start = end;
+        let mode = if is_textured(&extracted_uinode.image) {
+            if current_batch_image.id() != extracted_uinode.image.id() {
+                if is_textured(&current_batch_image) && start != end {
+                    commands.spawn(UiBatch {
+                        range: start..end,
+                        image: current_batch_image,
+                        z: last_z,
+                    });
+                    start = end;
+                }
+                current_batch_image = extracted_uinode.image.clone_weak();
             }
-            current_batch_handle = extracted_uinode.image.clone_weak();
-        }
+            TEXTURED_QUAD
+        } else {
+            // Untextured `UiBatch`es are never spawned within the loop.
+            // If all the `extracted_uinodes` are untextured a single untextured UiBatch will be spawned after the loop terminates.
+            UNTEXTURED_QUAD
+        };
 
         let mut uinode_rect = extracted_uinode.rect;
 
@@ -674,7 +691,7 @@ pub fn prepare_uinodes(
                 continue;
             }
         }
-        let uvs = if current_batch_handle.id() == DEFAULT_IMAGE_HANDLE.id() {
+        let uvs = if mode == UNTEXTURED_QUAD {
             [Vec2::ZERO, Vec2::X, Vec2::ONE, Vec2::Y]
         } else {
             let atlas_extent = extracted_uinode.atlas_size.unwrap_or(uinode_rect.max);
@@ -719,6 +736,7 @@ pub fn prepare_uinodes(
                 position: positions_clipped[i].into(),
                 uv: uvs[i].into(),
                 color,
+                mode,
             });
         }
 
@@ -730,7 +748,7 @@ pub fn prepare_uinodes(
     if start != end {
         commands.spawn(UiBatch {
             range: start..end,
-            image: current_batch_handle,
+            image: current_batch_image,
             z: last_z,
         });
     }

crates/bevy_ui/src/render/pipeline.rs

@@ -77,6 +77,8 @@ impl SpecializedRenderPipeline for UiPipeline {
                 VertexFormat::Float32x2,
                 // color
                 VertexFormat::Float32x4,
+                // mode
+                VertexFormat::Uint32,
             ],
         );
         let shader_defs = Vec::new();

crates/bevy_ui/src/render/ui.wgsl

@@ -1,11 +1,14 @@
 #import bevy_render::view
 
+const TEXTURED_QUAD: u32 = 0u;
+
 @group(0) @binding(0)
 var<uniform> view: View;
 
 struct VertexOutput {
     @location(0) uv: vec2<f32>,
     @location(1) color: vec4<f32>,
+    @location(3) mode: u32,
     @builtin(position) position: vec4<f32>,
 };
 
@@ -14,11 +17,13 @@ fn vertex(
     @location(0) vertex_position: vec3<f32>,
     @location(1) vertex_uv: vec2<f32>,
     @location(2) vertex_color: vec4<f32>,
+    @location(3) mode: u32,
 ) -> VertexOutput {
     var out: VertexOutput;
     out.uv = vertex_uv;
     out.position = view.view_proj * vec4<f32>(vertex_position, 1.0);
     out.color = vertex_color;
+    out.mode = mode;
     return out;
 }
 
@@ -29,7 +34,12 @@ var sprite_sampler: sampler;
 
 @fragment
 fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
+    // textureSample can only be called in unform control flow, not inside an if branch.
     var color = textureSample(sprite_texture, sprite_sampler, in.uv);
-    color = in.color * color;
+    if in.mode == TEXTURED_QUAD {
+        color = in.color * color;
+    } else {
+        color = in.color;   
+    }
     return color;
 }