Fixed incorrect behaviour of headless_renderer depending on image dim…

…ensions (#13388)

# Objective

- Fixes #13384 .

## Solution

- If the image became wider when copying from the texture to the buffer,
then the data is reduced to its original size when copying from the
buffer to the image.

## Testing

- Ran example with 1919x1080 resolution

![000](https://github.com/bevyengine/bevy/assets/17225606/47d95ed7-1c8c-4be4-a45a-1f485a3d6aa7)

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: François Mockers <francois.mockers@vleue.com>

examples/app/headless_renderer.rs

@@ -13,16 +13,15 @@ use bevy::{
     prelude::*,
     render::{
         camera::RenderTarget,
-        render_asset::RenderAssetUsages,
-        render_asset::RenderAssets,
+        render_asset::{RenderAssetUsages, RenderAssets},
         render_graph::{self, NodeRunError, RenderGraph, RenderGraphContext, RenderLabel},
         render_resource::{
             Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d,
             ImageCopyBuffer, ImageDataLayout, Maintain, MapMode, TextureDimension, TextureFormat,
             TextureUsages,
         },
         renderer::{RenderContext, RenderDevice, RenderQueue},
-        texture::BevyDefault,
+        texture::{BevyDefault, TextureFormatPixelInfo},
         Extract, Render, RenderApp, RenderSet,
     },
 };
@@ -361,6 +360,10 @@ impl render_graph::Node for ImageCopyDriver {
             let block_dimensions = src_image.texture_format.block_dimensions();
             let block_size = src_image.texture_format.block_copy_size(None).unwrap();
 
+            // Calculating correct size of image row because
+            // copy_texture_to_buffer can copy image only by rows aligned wgpu::COPY_BYTES_PER_ROW_ALIGNMENT
+            // That's why image in buffer can be little bit wider
+            // This should be taken into account at copy from buffer stage
             let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(
                 (src_image.size.x as usize / block_dimensions.0 as usize) * block_size as usize,
             );
@@ -492,7 +495,24 @@ fn update(
                 for image in images_to_save.iter() {
                     // Fill correct data from channel to image
                     let img_bytes = images.get_mut(image.id()).unwrap();
-                    img_bytes.data.clone_from(&image_data);
+
+                    // We need to ensure that this works regardless of the image dimensions
+                    // If the image became wider when copying from the texture to the buffer,
+                    // then the data is reduced to its original size when copying from the buffer to the image.
+                    let row_bytes = img_bytes.width() as usize
+                        * img_bytes.texture_descriptor.format.pixel_size();
+                    let aligned_row_bytes = RenderDevice::align_copy_bytes_per_row(row_bytes);
+                    if row_bytes == aligned_row_bytes {
+                        img_bytes.data.clone_from(&image_data);
+                    } else {
+                        // shrink data to original image size
+                        img_bytes.data = image_data
+                            .chunks(aligned_row_bytes)
+                            .take(img_bytes.height() as usize)
+                            .flat_map(|row| &row[..row_bytes.min(row.len())])
+                            .cloned()
+                            .collect();
+                    }
 
                     // Create RGBA Image Buffer
                     let img = match img_bytes.clone().try_into_dynamic() {