Improved UI render batching

crates/bevy_ui/src/render/mod.rs

@@ -349,6 +349,7 @@ struct UiVertex {
     pub position: [f32; 3],
     pub uv: [f32; 2],
     pub color: [f32; 4],
+    pub mode: u32,
 }
 
 #[derive(Resource)]
@@ -382,6 +383,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>,
@@ -398,20 +402,31 @@ pub fn prepare_uinodes(
 
     let mut start = 0;
     let mut end = 0;
-    let mut current_batch_handle = Default::default();
+    let mut stored_batch_handle = DEFAULT_IMAGE_HANDLE.typed();
     let mut last_z = 0.0;
+    let default_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 extracted_uinode.image.id() != default_id {
+            if stored_batch_handle.id() != default_id {
+                if stored_batch_handle.id() != extracted_uinode.image.id() {
+                    if start != end {
+                        commands.spawn(UiBatch {
+                            range: start..end,
+                            image: stored_batch_handle,
+                            z: last_z,
+                        });
+                        start = end;
+                    }
+                    stored_batch_handle = extracted_uinode.image.clone_weak();
+                }
+            } else {
+                stored_batch_handle = extracted_uinode.image.clone_weak();
             }
-            current_batch_handle = extracted_uinode.image.clone_weak();
-        }
+            TEXTURED_QUAD
+        } else {
+            UNTEXTURED_QUAD
+        };
 
         let mut uinode_rect = extracted_uinode.rect;
 
@@ -469,7 +484,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);
@@ -514,6 +529,7 @@ pub fn prepare_uinodes(
                 position: positions_clipped[i].into(),
                 uv: uvs[i].into(),
                 color,
+                mode,
             });
         }
 
@@ -525,7 +541,7 @@ pub fn prepare_uinodes(
     if start != end {
         commands.spawn(UiBatch {
             range: start..end,
-            image: current_batch_handle,
+            image: stored_batch_handle,
             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

@@ -6,6 +6,7 @@ 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 +15,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;
 }
 
@@ -30,6 +33,10 @@ var sprite_sampler: sampler;
 @fragment
 fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
     var color = textureSample(sprite_texture, sprite_sampler, in.uv);
-    color = in.color * color;
+    if in.mode == 0u {
+        color = in.color * color;
+    } else {
+        color = in.color;   
+    }
     return color;
 }