Make TrackedRenderPass::set_vertex_buffer aware of slice size (#14916)

# Objective

- Fixes #14841

## Solution

- Compute BufferSlice size manually and use it for comparison in
`TrackedRenderPass`

## Testing

- Gizmo example does not crash with #14721 (without system ordering),
and `slice` computes correct size there

---

## Migration Guide

- `TrackedRenderPass::set_vertex_buffer` function has been modified to
update vertex buffers when the same buffer with the same offset is
provided, but its size has changed. Some existing code may rely on the
previous behavior, which did not update the vertex buffer in this
scenario.

---------

Co-authored-by: Zachary Harrold <zac@harrold.com.au>

crates/bevy_gizmos/src/pipeline_2d.rs

@@ -54,9 +54,7 @@ impl Plugin for LineGizmo2dPlugin {
             )
             .add_systems(
                 Render,
-                // FIXME: added `chain()` to workaround vertex buffer being not updated when sliced size changed
                 (queue_line_gizmos_2d, queue_line_joint_gizmos_2d)
-                    .chain()
                     .in_set(GizmoRenderSystem::QueueLineGizmos2d)
                     .after(prepare_assets::<GpuLineGizmo>),
             );

crates/bevy_gizmos/src/pipeline_3d.rs

@@ -53,9 +53,7 @@ impl Plugin for LineGizmo3dPlugin {
             )
             .add_systems(
                 Render,
-                // FIXME: added `chain()` to workaround vertex buffer being not updated when sliced size changed
                 (queue_line_gizmos_3d, queue_line_joint_gizmos_3d)
-                    .chain()
                     .in_set(GizmoRenderSystem::QueueLineGizmos3d)
                     .after(prepare_assets::<GpuLineGizmo>),
             );

crates/bevy_render/src/render_phase/draw_state.rs

@@ -21,13 +21,14 @@ use wgpu::{IndexFormat, QuerySet, RenderPass};
 struct DrawState {
     pipeline: Option<RenderPipelineId>,
     bind_groups: Vec<(Option<BindGroupId>, Vec<u32>)>,
-    vertex_buffers: Vec<Option<(BufferId, u64)>>,
+    /// List of vertex buffers by [`BufferId`], offset, and size. See [`DrawState::buffer_slice_key`]
+    vertex_buffers: Vec<Option<(BufferId, u64, u64)>>,
     index_buffer: Option<(BufferId, u64, IndexFormat)>,
 }
 
 impl DrawState {
     /// Marks the `pipeline` as bound.
-    pub fn set_pipeline(&mut self, pipeline: RenderPipelineId) {
+    fn set_pipeline(&mut self, pipeline: RenderPipelineId) {
         // TODO: do these need to be cleared?
         // self.bind_groups.clear();
         // self.vertex_buffers.clear();
@@ -36,17 +37,12 @@ impl DrawState {
     }
 
     /// Checks, whether the `pipeline` is already bound.
-    pub fn is_pipeline_set(&self, pipeline: RenderPipelineId) -> bool {
+    fn is_pipeline_set(&self, pipeline: RenderPipelineId) -> bool {
         self.pipeline == Some(pipeline)
     }
 
     /// Marks the `bind_group` as bound to the `index`.
-    pub fn set_bind_group(
-        &mut self,
-        index: usize,
-        bind_group: BindGroupId,
-        dynamic_indices: &[u32],
-    ) {
+    fn set_bind_group(&mut self, index: usize, bind_group: BindGroupId, dynamic_indices: &[u32]) {
         let group = &mut self.bind_groups[index];
         group.0 = Some(bind_group);
         group.1.clear();
@@ -54,7 +50,7 @@ impl DrawState {
     }
 
     /// Checks, whether the `bind_group` is already bound to the `index`.
-    pub fn is_bind_group_set(
+    fn is_bind_group_set(
         &self,
         index: usize,
         bind_group: BindGroupId,
@@ -68,26 +64,35 @@ impl DrawState {
     }
 
     /// Marks the vertex `buffer` as bound to the `index`.
-    pub fn set_vertex_buffer(&mut self, index: usize, buffer: BufferId, offset: u64) {
-        self.vertex_buffers[index] = Some((buffer, offset));
+    fn set_vertex_buffer(&mut self, index: usize, buffer_slice: BufferSlice) {
+        self.vertex_buffers[index] = Some(self.buffer_slice_key(&buffer_slice));
     }
 
     /// Checks, whether the vertex `buffer` is already bound to the `index`.
-    pub fn is_vertex_buffer_set(&self, index: usize, buffer: BufferId, offset: u64) -> bool {
+    fn is_vertex_buffer_set(&self, index: usize, buffer_slice: &BufferSlice) -> bool {
         if let Some(current) = self.vertex_buffers.get(index) {
-            *current == Some((buffer, offset))
+            *current == Some(self.buffer_slice_key(buffer_slice))
         } else {
             false
         }
     }
 
+    /// Returns the value used for checking whether `BufferSlice`s are equivalent.
+    fn buffer_slice_key(&self, buffer_slice: &BufferSlice) -> (BufferId, u64, u64) {
+        (
+            buffer_slice.id(),
+            buffer_slice.offset(),
+            buffer_slice.size(),
+        )
+    }
+
     /// Marks the index `buffer` as bound.
-    pub fn set_index_buffer(&mut self, buffer: BufferId, offset: u64, index_format: IndexFormat) {
+    fn set_index_buffer(&mut self, buffer: BufferId, offset: u64, index_format: IndexFormat) {
         self.index_buffer = Some((buffer, offset, index_format));
     }
 
     /// Checks, whether the index `buffer` is already bound.
-    pub fn is_index_buffer_set(
+    fn is_index_buffer_set(
         &self,
         buffer: BufferId,
         offset: u64,
@@ -188,30 +193,27 @@ impl<'a> TrackedRenderPass<'a> {
     /// [`draw`]: TrackedRenderPass::draw
     /// [`draw_indexed`]: TrackedRenderPass::draw_indexed
     pub fn set_vertex_buffer(&mut self, slot_index: usize, buffer_slice: BufferSlice<'a>) {
-        let offset = buffer_slice.offset();
-        if self
-            .state
-            .is_vertex_buffer_set(slot_index, buffer_slice.id(), offset)
-        {
+        if self.state.is_vertex_buffer_set(slot_index, &buffer_slice) {
             detailed_trace!(
-                "set vertex buffer {} (already set): {:?} ({})",
+                "set vertex buffer {} (already set): {:?} (offset = {}, size = {})",
                 slot_index,
                 buffer_slice.id(),
-                offset
+                buffer_slice.offset(),
+                buffer_slice.size(),
             );
             return;
         }
         detailed_trace!(
-            "set vertex buffer {}: {:?} ({})",
+            "set vertex buffer {}: {:?} (offset = {}, size = {})",
             slot_index,
             buffer_slice.id(),
-            offset
+            buffer_slice.offset(),
+            buffer_slice.size(),
         );
 
         self.pass
             .set_vertex_buffer(slot_index as u32, *buffer_slice);
-        self.state
-            .set_vertex_buffer(slot_index, buffer_slice.id(), offset);
+        self.state.set_vertex_buffer(slot_index, buffer_slice);
     }
 
     /// Sets the active index buffer.

crates/bevy_render/src/render_resource/buffer.rs

@@ -8,6 +8,7 @@ render_resource_wrapper!(ErasedBuffer, wgpu::Buffer);
 pub struct Buffer {
     id: BufferId,
     value: ErasedBuffer,
+    size: wgpu::BufferAddress,
 }
 
 impl Buffer {
@@ -17,14 +18,21 @@ impl Buffer {
     }
 
     pub fn slice(&self, bounds: impl RangeBounds<wgpu::BufferAddress>) -> BufferSlice {
-        BufferSlice {
-            id: self.id,
-            // need to compute and store this manually because wgpu doesn't export offset on wgpu::BufferSlice
-            offset: match bounds.start_bound() {
+        // need to compute and store this manually because wgpu doesn't export offset and size on wgpu::BufferSlice
+        let offset = match bounds.start_bound() {
             Bound::Included(&bound) => bound,
             Bound::Excluded(&bound) => bound + 1,
             Bound::Unbounded => 0,
-            },
+        };
+        let size = match bounds.end_bound() {
+            Bound::Included(&bound) => bound + 1,
+            Bound::Excluded(&bound) => bound,
+            Bound::Unbounded => self.size,
+        } - offset;
+        BufferSlice {
+            id: self.id,
+            offset,
+            size,
             value: self.value.slice(bounds),
         }
     }
@@ -39,6 +47,7 @@ impl From<wgpu::Buffer> for Buffer {
     fn from(value: wgpu::Buffer) -> Self {
         Buffer {
             id: BufferId::new(),
+            size: value.size(),
             value: ErasedBuffer::new(value),
         }
     }
@@ -58,6 +67,7 @@ pub struct BufferSlice<'a> {
     id: BufferId,
     offset: wgpu::BufferAddress,
     value: wgpu::BufferSlice<'a>,
+    size: wgpu::BufferAddress,
 }
 
 impl<'a> BufferSlice<'a> {
@@ -70,6 +80,11 @@ impl<'a> BufferSlice<'a> {
     pub fn offset(&self) -> wgpu::BufferAddress {
         self.offset
     }
+
+    #[inline]
+    pub fn size(&self) -> wgpu::BufferAddress {
+        self.size
+    }
 }
 
 impl<'a> Deref for BufferSlice<'a> {