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)]
@@ -379,9 +380,11 @@ const QUAD_INDICES: [usize; 6] = [0, 2, 3, 0, 1, 2];
 pub struct UiBatch {
     pub range: Range<u32>,
     pub image: Handle<Image>,
-    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 +401,27 @@ pub fn prepare_uinodes(
 
     let mut start = 0;
     let mut end = 0;
-    let mut current_batch_handle = Default::default();
-    let mut last_z = 0.0;
+    let mut stored_batch_handle = DEFAULT_IMAGE_HANDLE.typed();
+
+    #[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 is_textured(&stored_batch_handle) {
+                if stored_batch_handle.id() != extracted_uinode.image.id() {
+                    commands.spawn(UiBatch {
+                        range: start..end,
+                        image: stored_batch_handle,
+                    });
+                    start = end;
+                }
             }
-            current_batch_handle = extracted_uinode.image.clone_weak();
-        }
+            stored_batch_handle = extracted_uinode.image.clone_weak();
+            TEXTURED_QUAD
+        } else {
+            UNTEXTURED_QUAD
+        };
 
         let mut uinode_rect = extracted_uinode.rect;
 
@@ -469,7 +479,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,19 +524,18 @@ pub fn prepare_uinodes(
                 position: positions_clipped[i].into(),
                 uv: uvs[i].into(),
                 color,
+                mode,
             });
         }
 
-        last_z = extracted_uinode.transform.w_axis[2];
         end += QUAD_INDICES.len() as u32;
     }
 
     // if start != end, there is one last batch to process
     if start != end {
         commands.spawn(UiBatch {
             range: start..end,
-            image: current_batch_handle,
-            z: last_z,
+            image: stored_batch_handle,
         });
     }
 
@@ -604,7 +613,7 @@ pub fn queue_uinodes(
                     draw_function: draw_ui_function,
                     pipeline,
                     entity,
-                    sort_key: FloatOrd(batch.z),
+                    sort_key: FloatOrd(0.),
                 });
             }
         }

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;
 }