Make the get function on InstanceInputUniformBuffer less error pr…

…one (#17131)

# Objective

the `get` function on [`InstanceInputUniformBuffer`] seems very
error-prone. This PR hopes to fix this.

## Solution

Do a few checks to ensure the index is in bounds and that the `BDI` is
not removed.
Return `Option<BDI>` instead of `BDI`. 

## Testing

- Did you test these changes? If so, how?
added a test to verify that the instance buffer works correctly

## Future Work
Performance decreases when using .binary_search(). However this is
likely due to the fact that [`InstanceInputUniformBuffer::get`] for now
is never used, and only get_unchecked.

## Migration Guide
`InstanceInputUniformBuffer::get` now returns `Option<BDI>` instead of
`BDI` to reduce panics. If you require the old functionality of
`InstanceInputUniformBuffer::get` consider using
`InstanceInputUniformBuffer::get_unchecked`.

---------

Co-authored-by: Tim Overbeek <oorbeck@gmail.com>

crates/bevy_pbr/src/render/mesh.rs

@@ -957,7 +957,8 @@ impl RenderMeshInstanceGpuBuilder {
 
                 // Save the old mesh input uniform. The mesh preprocessing
                 // shader will need it to compute motion vectors.
-                let previous_mesh_input_uniform = current_input_buffer.get(current_uniform_index);
+                let previous_mesh_input_uniform =
+                    current_input_buffer.get_unchecked(current_uniform_index);
                 let previous_input_index = previous_input_buffer.add(previous_mesh_input_uniform);
                 mesh_input_uniform.previous_input_index = previous_input_index;
 

crates/bevy_render/src/batching/gpu_preprocessing.rs

@@ -217,11 +217,31 @@ where
     }
 
     /// Returns the piece of buffered data at the given index.
-    pub fn get(&self, uniform_index: u32) -> BDI {
+    ///
+    /// Returns [`None`] if the index is out of bounds or the data is removed.
+    pub fn get(&self, uniform_index: u32) -> Option<BDI> {
+        if (uniform_index as usize) >= self.buffer.len()
+            || self.free_uniform_indices.contains(&uniform_index)
+        {
+            None
+        } else {
+            Some(self.get_unchecked(uniform_index))
+        }
+    }
+
+    /// Returns the piece of buffered data at the given index.
+    /// Can return data that has previously been removed.
+    ///
+    /// # Panics
+    /// if `uniform_index` is not in bounds of [`Self::buffer`].
+    pub fn get_unchecked(&self, uniform_index: u32) -> BDI {
         self.buffer.values()[uniform_index as usize]
     }
 
     /// Stores a piece of buffered data at the given index.
+    ///
+    /// # Panics
+    /// if `uniform_index` is not in bounds of [`Self::buffer`].
     pub fn set(&mut self, uniform_index: u32, element: BDI) {
         self.buffer.values_mut()[uniform_index as usize] = element;
     }
@@ -990,3 +1010,21 @@ pub fn write_indirect_parameters_buffer(
         .write_buffer(&render_device, &render_queue);
     indirect_parameters_buffer.buffer.clear();
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn instance_buffer_correct_behavior() {
+        let mut instance_buffer = InstanceInputUniformBuffer::new();
+
+        let index = instance_buffer.add(2);
+        instance_buffer.remove(index);
+        assert_eq!(instance_buffer.get_unchecked(index), 2);
+        assert_eq!(instance_buffer.get(index), None);
+
+        instance_buffer.add(5);
+        assert_eq!(instance_buffer.buffer().len(), 1);
+    }
+}