diff --git a/crates/bevy_dev_tools/src/ci_testing/mod.rs b/crates/bevy_dev_tools/src/ci_testing/mod.rs index f6510f68a1..59949fa0bf 100644 --- a/crates/bevy_dev_tools/src/ci_testing/mod.rs +++ b/crates/bevy_dev_tools/src/ci_testing/mod.rs @@ -7,6 +7,7 @@ pub use self::config::*; use bevy_app::prelude::*; use bevy_ecs::schedule::IntoSystemConfigs; +use bevy_render::view::screenshot::trigger_screenshots; use bevy_time::TimeUpdateStrategy; use std::time::Duration; @@ -51,7 +52,9 @@ impl Plugin for CiTestingPlugin { .insert_resource(config) .add_systems( Update, - systems::send_events.before(bevy_window::close_when_requested), + systems::send_events + .before(trigger_screenshots) + .before(bevy_window::close_when_requested), ); } } diff --git a/crates/bevy_dev_tools/src/ci_testing/systems.rs b/crates/bevy_dev_tools/src/ci_testing/systems.rs index abb2b28bcc..2b652abf96 100644 --- a/crates/bevy_dev_tools/src/ci_testing/systems.rs +++ b/crates/bevy_dev_tools/src/ci_testing/systems.rs @@ -1,9 +1,8 @@ use super::config::*; use bevy_app::AppExit; use bevy_ecs::prelude::*; -use bevy_render::view::screenshot::ScreenshotManager; -use bevy_utils::tracing::{debug, info, warn}; -use bevy_window::PrimaryWindow; +use bevy_render::view::screenshot::{save_to_disk, Screenshot}; +use bevy_utils::tracing::{debug, info}; pub(crate) fn send_events(world: &mut World, mut current_frame: Local) { let mut config = world.resource_mut::(); @@ -23,21 +22,10 @@ pub(crate) fn send_events(world: &mut World, mut current_frame: Local) { info!("Exiting after {} frames. Test successful!", *current_frame); } CiTestingEvent::Screenshot => { - let mut primary_window_query = - world.query_filtered::>(); - let Ok(main_window) = primary_window_query.get_single(world) else { - warn!("Requesting screenshot, but PrimaryWindow is not available"); - continue; - }; - let Some(mut screenshot_manager) = world.get_resource_mut::() - else { - warn!("Requesting screenshot, but ScreenshotManager is not available"); - continue; - }; let path = format!("./screenshot-{}.png", *current_frame); - screenshot_manager - .save_screenshot_to_disk(main_window, path) - .unwrap(); + world + .spawn(Screenshot::primary_window()) + .observe(save_to_disk(path)); info!("Took a screenshot at frame {}.", *current_frame); } // Custom events are forwarded to the world. diff --git a/crates/bevy_render/src/view/mod.rs b/crates/bevy_render/src/view/mod.rs index f9c4a5c59b..3729f4048f 100644 --- a/crates/bevy_render/src/view/mod.rs +++ b/crates/bevy_render/src/view/mod.rs @@ -5,6 +5,7 @@ use bevy_asset::{load_internal_asset, Handle}; pub use visibility::*; pub use window::*; +use crate::camera::NormalizedRenderTarget; use crate::extract_component::ExtractComponentPlugin; use crate::{ camera::{ @@ -25,12 +26,13 @@ use crate::{ }; use bevy_app::{App, Plugin}; use bevy_color::LinearRgba; +use bevy_derive::{Deref, DerefMut}; use bevy_ecs::prelude::*; use bevy_math::{mat3, vec2, vec3, Mat3, Mat4, UVec4, Vec2, Vec3, Vec4, Vec4Swizzles}; use bevy_reflect::{std_traits::ReflectDefault, Reflect}; use bevy_render_macros::ExtractComponent; use bevy_transform::components::GlobalTransform; -use bevy_utils::HashMap; +use bevy_utils::{hashbrown::hash_map::Entry, HashMap}; use std::{ ops::Range, sync::{ @@ -119,6 +121,9 @@ impl Plugin for ViewPlugin { render_app.add_systems( Render, ( + prepare_view_attachments + .in_set(RenderSet::ManageViews) + .before(prepare_view_targets), prepare_view_targets .in_set(RenderSet::ManageViews) .after(prepare_windows) @@ -132,7 +137,9 @@ impl Plugin for ViewPlugin { fn finish(&self, app: &mut App) { if let Some(render_app) = app.get_sub_app_mut(RenderApp) { - render_app.init_resource::(); + render_app + .init_resource::() + .init_resource::(); } } } @@ -458,6 +465,13 @@ pub struct ViewTarget { out_texture: OutputColorAttachment, } +/// Contains [`OutputColorAttachment`] used for each target present on any view in the current +/// frame, after being prepared by [`prepare_view_attachments`]. Users that want to override +/// the default output color attachment for a specific target can do so by adding a +/// [`OutputColorAttachment`] to this resource before [`prepare_view_targets`] is called. +#[derive(Resource, Default, Deref, DerefMut)] +pub struct ViewTargetAttachments(HashMap); + pub struct PostProcessWrite<'a> { pub source: &'a TextureView, pub destination: &'a TextureView, @@ -794,11 +808,41 @@ struct MainTargetTextures { main_texture: Arc, } -#[allow(clippy::too_many_arguments)] -pub fn prepare_view_targets( - mut commands: Commands, +/// Prepares the view target [`OutputColorAttachment`] for each view in the current frame. +pub fn prepare_view_attachments( windows: Res, images: Res>, + manual_texture_views: Res, + cameras: Query<&ExtractedCamera>, + mut view_target_attachments: ResMut, +) { + view_target_attachments.clear(); + for camera in cameras.iter() { + let Some(target) = &camera.target else { + continue; + }; + + match view_target_attachments.entry(target.clone()) { + Entry::Occupied(_) => {} + Entry::Vacant(entry) => { + let Some(attachment) = target + .get_texture_view(&windows, &images, &manual_texture_views) + .cloned() + .zip(target.get_texture_format(&windows, &images, &manual_texture_views)) + .map(|(view, format)| { + OutputColorAttachment::new(view.clone(), format.add_srgb_suffix()) + }) + else { + continue; + }; + entry.insert(attachment); + } + }; + } +} + +pub fn prepare_view_targets( + mut commands: Commands, clear_color_global: Res, render_device: Res, mut texture_cache: ResMut, @@ -809,24 +853,16 @@ pub fn prepare_view_targets( &CameraMainTextureUsages, &Msaa, )>, - manual_texture_views: Res, + view_target_attachments: Res, ) { let mut textures = HashMap::default(); - let mut output_textures = HashMap::default(); for (entity, camera, view, texture_usage, msaa) in cameras.iter() { let (Some(target_size), Some(target)) = (camera.physical_target_size, &camera.target) else { continue; }; - let Some(out_texture) = output_textures.entry(target.clone()).or_insert_with(|| { - target - .get_texture_view(&windows, &images, &manual_texture_views) - .zip(target.get_texture_format(&windows, &images, &manual_texture_views)) - .map(|(view, format)| { - OutputColorAttachment::new(view.clone(), format.add_srgb_suffix()) - }) - }) else { + let Some(out_attachment) = view_target_attachments.get(target) else { continue; }; @@ -913,7 +949,7 @@ pub fn prepare_view_targets( main_texture: main_textures.main_texture.clone(), main_textures, main_texture_format, - out_texture: out_texture.clone(), + out_texture: out_attachment.clone(), }); } } diff --git a/crates/bevy_render/src/view/window/mod.rs b/crates/bevy_render/src/view/window/mod.rs index 027a15e467..816d8c4e8d 100644 --- a/crates/bevy_render/src/view/window/mod.rs +++ b/crates/bevy_render/src/view/window/mod.rs @@ -1,9 +1,6 @@ use crate::{ - render_resource::{ - BindGroupEntries, PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView, - }, + render_resource::{SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, - texture::TextureFormatPixelInfo, Extract, ExtractSchedule, Render, RenderApp, RenderSet, WgpuWrapper, }; use bevy_app::{App, Last, Plugin}; @@ -18,21 +15,16 @@ use bevy_winit::CustomCursorCache; use std::{ num::NonZeroU32, ops::{Deref, DerefMut}, - sync::PoisonError, }; use wgpu::{ - BufferUsages, SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, - TextureViewDescriptor, + SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, TextureViewDescriptor, }; pub mod cursor; pub mod screenshot; -use screenshot::{ - ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline, -}; - use self::cursor::update_cursors; +use screenshot::{ScreenshotPlugin, ScreenshotToScreenPipeline}; pub struct WindowRenderPlugin; @@ -78,11 +70,9 @@ pub struct ExtractedWindow { pub swap_chain_texture_view: Option, pub swap_chain_texture: Option, pub swap_chain_texture_format: Option, - pub screenshot_memory: Option, pub size_changed: bool, pub present_mode_changed: bool, pub alpha_mode: CompositeAlphaMode, - pub screenshot_func: Option, } impl ExtractedWindow { @@ -120,7 +110,6 @@ impl DerefMut for ExtractedWindows { fn extract_windows( mut extracted_windows: ResMut, - screenshot_manager: Extract>, mut closing: Extract>, windows: Extract)>>, mut removed: Extract>, @@ -149,8 +138,6 @@ fn extract_windows( swap_chain_texture_format: None, present_mode_changed: false, alpha_mode: window.composite_alpha_mode, - screenshot_func: None, - screenshot_memory: None, }); // NOTE: Drop the swap chain frame here @@ -189,20 +176,6 @@ fn extract_windows( extracted_windows.remove(&removed_window); window_surfaces.remove(&removed_window); } - // This lock will never block because `callbacks` is `pub(crate)` and this is the singular callsite where it's locked. - // Even if a user had multiple copies of this system, since the system has a mutable resource access the two systems would never run - // at the same time - // TODO: since this is guaranteed, should the lock be replaced with an UnsafeCell to remove the overhead, or is it minor enough to be ignored? - for (window, screenshot_func) in screenshot_manager - .callbacks - .lock() - .unwrap_or_else(PoisonError::into_inner) - .drain() - { - if let Some(window) = extracted_windows.get_mut(&window) { - window.screenshot_func = Some(screenshot_func); - } - } } struct SurfaceData { @@ -254,9 +227,6 @@ pub fn prepare_windows( mut windows: ResMut, mut window_surfaces: ResMut, render_device: Res, - screenshot_pipeline: Res, - pipeline_cache: Res, - mut pipelines: ResMut>, #[cfg(target_os = "linux")] render_instance: Res, ) { for window in windows.windows.values_mut() { @@ -340,53 +310,6 @@ pub fn prepare_windows( } }; window.swap_chain_texture_format = Some(surface_data.configuration.format); - - if window.screenshot_func.is_some() { - let texture = render_device.create_texture(&wgpu::TextureDescriptor { - label: Some("screenshot-capture-rendertarget"), - size: wgpu::Extent3d { - width: surface_data.configuration.width, - height: surface_data.configuration.height, - depth_or_array_layers: 1, - }, - mip_level_count: 1, - sample_count: 1, - dimension: wgpu::TextureDimension::D2, - format: surface_data.configuration.format.add_srgb_suffix(), - usage: TextureUsages::RENDER_ATTACHMENT - | TextureUsages::COPY_SRC - | TextureUsages::TEXTURE_BINDING, - view_formats: &[], - }); - let texture_view = texture.create_view(&Default::default()); - let buffer = render_device.create_buffer(&wgpu::BufferDescriptor { - label: Some("screenshot-transfer-buffer"), - size: screenshot::get_aligned_size( - window.physical_width, - window.physical_height, - surface_data.configuration.format.pixel_size() as u32, - ) as u64, - usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - let bind_group = render_device.create_bind_group( - "screenshot-to-screen-bind-group", - &screenshot_pipeline.bind_group_layout, - &BindGroupEntries::single(&texture_view), - ); - let pipeline_id = pipelines.specialize( - &pipeline_cache, - &screenshot_pipeline, - surface_data.configuration.format, - ); - window.swap_chain_texture_view = Some(texture_view); - window.screenshot_memory = Some(ScreenshotPreparedState { - texture, - buffer, - bind_group, - pipeline_id, - }); - } } } diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index 21f6f6a6bf..bc53fc9880 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -1,16 +1,13 @@ -use std::{borrow::Cow, path::Path, sync::PoisonError}; - -use bevy_app::Plugin; -use bevy_asset::{load_internal_asset, Handle}; -use bevy_ecs::{entity::EntityHashMap, prelude::*}; -use bevy_tasks::AsyncComputeTaskPool; -use bevy_utils::tracing::{error, info, info_span}; -use std::sync::Mutex; -use thiserror::Error; -use wgpu::{ - CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, +use super::ExtractedWindows; +use crate::camera::{ + ManualTextureViewHandle, ManualTextureViews, NormalizedRenderTarget, RenderTarget, +}; +use crate::render_asset::RenderAssets; +use crate::render_resource::{BindGroupEntries, BufferUsages, TextureUsages, TextureView}; +use crate::texture::{GpuImage, OutputColorAttachment}; +use crate::view::{ + prepare_view_attachments, prepare_view_targets, ViewTargetAttachments, WindowSurfaces, }; - use crate::{ prelude::{Image, Shader}, render_asset::RenderAssetUsages, @@ -21,51 +18,116 @@ use crate::{ }, renderer::RenderDevice, texture::TextureFormatPixelInfo, - RenderApp, + ExtractSchedule, MainWorld, Render, RenderApp, RenderSet, +}; +use bevy_app::{First, Plugin, Update}; +use bevy_asset::{load_internal_asset, Handle}; +use bevy_derive::{Deref, DerefMut}; +use bevy_ecs::event::event_update_system; +use bevy_ecs::system::SystemState; +use bevy_ecs::{entity::EntityHashMap, prelude::*}; +use bevy_hierarchy::DespawnRecursiveExt; +use bevy_reflect::Reflect; +use bevy_tasks::AsyncComputeTaskPool; +use bevy_utils::tracing::{error, info, warn}; +use bevy_utils::{default, HashSet}; +use bevy_window::{PrimaryWindow, WindowRef}; +use std::ops::Deref; +use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{Arc, Mutex}; +use std::{borrow::Cow, path::Path}; +use wgpu::{ + CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, }; -use super::ExtractedWindows; +#[derive(Event, Deref, DerefMut, Reflect, Debug)] +#[reflect(Debug)] +pub struct ScreenshotCaptured(pub Image); -pub type ScreenshotFn = Box; +/// A component that signals to the renderer to capture a screenshot this frame. +/// +/// This component should be spawned on a new entity with an observer that will trigger +/// with [`ScreenshotCaptured`] when the screenshot is ready. +/// +/// Screenshots are captured asynchronously and may not be available immediately after the frame +/// that the component is spawned on. The observer should be used to handle the screenshot when it +/// is ready. +/// +/// Note that the screenshot entity will be despawned after the screenshot is captured and the +/// observer is triggered. +/// +/// # Usage +/// +/// ``` +/// # use bevy_ecs::prelude::*; +/// # use bevy_render::view::screenshot::{save_to_disk, Screenshot}; +/// +/// fn take_screenshot(mut commands: Commands) { +/// commands.spawn(Screenshot::primary_window()) +/// .observe(save_to_disk("screenshot.png")); +/// } +/// ``` +#[derive(Component, Deref, DerefMut, Reflect, Debug)] +#[reflect(Component, Debug)] +pub struct Screenshot(pub RenderTarget); -/// A resource which allows for taking screenshots of the window. -#[derive(Resource, Default)] -pub struct ScreenshotManager { - // this is in a mutex to enable extraction with only an immutable reference - pub(crate) callbacks: Mutex>, -} +/// A marker component that indicates that a screenshot is currently being captured. +#[derive(Component)] +pub struct Capturing; -#[derive(Error, Debug)] -#[error("A screenshot for this window has already been requested.")] -pub struct ScreenshotAlreadyRequestedError; +/// A marker component that indicates that a screenshot has been captured, the image is ready, and +/// the screenshot entity can be despawned. +#[derive(Component)] +pub struct Captured; -impl ScreenshotManager { - /// Signals the renderer to take a screenshot of this frame. - /// - /// The given callback will eventually be called on one of the [`AsyncComputeTaskPool`]s threads. - pub fn take_screenshot( - &mut self, - window: Entity, - callback: impl FnOnce(Image) + Send + Sync + 'static, - ) -> Result<(), ScreenshotAlreadyRequestedError> { - self.callbacks - .get_mut() - .unwrap_or_else(PoisonError::into_inner) - .try_insert(window, Box::new(callback)) - .map(|_| ()) - .map_err(|_| ScreenshotAlreadyRequestedError) +impl Screenshot { + /// Capture a screenshot of the provided window entity. + pub fn window(window: Entity) -> Self { + Self(RenderTarget::Window(WindowRef::Entity(window))) } - /// Signals the renderer to take a screenshot of this frame. - /// - /// The screenshot will eventually be saved to the given path, and the format will be derived from the extension. - pub fn save_screenshot_to_disk( - &mut self, - window: Entity, - path: impl AsRef, - ) -> Result<(), ScreenshotAlreadyRequestedError> { - let path = path.as_ref().to_owned(); - self.take_screenshot(window, move |img| match img.try_into_dynamic() { + /// Capture a screenshot of the primary window, if one exists. + pub fn primary_window() -> Self { + Self(RenderTarget::Window(WindowRef::Primary)) + } + + /// Capture a screenshot of the provided render target image. + pub fn image(image: Handle) -> Self { + Self(RenderTarget::Image(image)) + } + + /// Capture a screenshot of the provided manual texture view. + pub fn texture_view(texture_view: ManualTextureViewHandle) -> Self { + Self(RenderTarget::TextureView(texture_view)) + } +} + +struct ScreenshotPreparedState { + pub texture: Texture, + pub buffer: Buffer, + pub bind_group: BindGroup, + pub pipeline_id: CachedRenderPipelineId, + pub size: Extent3d, +} + +#[derive(Resource, Deref, DerefMut)] +pub struct CapturedScreenshots(pub Arc>>); + +#[derive(Resource, Deref, DerefMut, Default)] +struct RenderScreenshotTargets(EntityHashMap); + +#[derive(Resource, Deref, DerefMut, Default)] +struct RenderScreenshotsPrepared(EntityHashMap); + +#[derive(Resource, Deref, DerefMut)] +struct RenderScreenshotsSender(Sender<(Entity, Image)>); + +/// Saves the captured screenshot to disk at the provided path. +pub fn save_to_disk(path: impl AsRef) -> impl FnMut(Trigger) { + let path = path.as_ref().to_owned(); + move |trigger| { + let img = trigger.event().deref().clone(); + match img.try_into_dynamic() { Ok(dyn_img) => match image::ImageFormat::from_path(&path) { Ok(format) => { // discard the alpha channel which stores brightness values when HDR is enabled to make sure @@ -118,17 +180,238 @@ impl ScreenshotManager { Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"), }, Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"), - }) + } } } +fn clear_screenshots(mut commands: Commands, screenshots: Query>) { + for entity in screenshots.iter() { + commands.entity(entity).despawn_recursive(); + } +} + +pub fn trigger_screenshots( + mut commands: Commands, + captured_screenshots: ResMut, +) { + let captured_screenshots = captured_screenshots.lock().unwrap(); + while let Ok((entity, image)) = captured_screenshots.try_recv() { + commands.entity(entity).insert(Captured); + commands.trigger_targets(ScreenshotCaptured(image), entity); + } +} + +fn extract_screenshots( + mut targets: ResMut, + mut main_world: ResMut, + mut system_state: Local< + Option< + SystemState<( + Commands, + Query>, + Query<(Entity, &Screenshot), Without>, + )>, + >, + >, + mut seen_targets: Local>, +) { + if system_state.is_none() { + *system_state = Some(SystemState::new(&mut main_world)); + } + let system_state = system_state.as_mut().unwrap(); + let (mut commands, primary_window, screenshots) = system_state.get_mut(&mut main_world); + + targets.clear(); + seen_targets.clear(); + + let primary_window = primary_window.iter().next(); + + for (entity, screenshot) in screenshots.iter() { + let render_target = screenshot.0.clone(); + let Some(render_target) = render_target.normalize(primary_window) else { + warn!( + "Unknown render target for screenshot, skipping: {:?}", + render_target + ); + continue; + }; + if seen_targets.contains(&render_target) { + warn!( + "Duplicate render target for screenshot, skipping entity {:?}: {:?}", + entity, render_target + ); + // If we don't despawn the entity here, it will be captured again in the next frame + commands.entity(entity).despawn_recursive(); + continue; + } + seen_targets.insert(render_target.clone()); + targets.insert(entity, render_target); + commands.entity(entity).insert(Capturing); + } + + system_state.apply(&mut main_world); +} + +#[allow(clippy::too_many_arguments)] +fn prepare_screenshots( + targets: Res, + mut prepared: ResMut, + window_surfaces: Res, + render_device: Res, + screenshot_pipeline: Res, + pipeline_cache: Res, + mut pipelines: ResMut>, + images: Res>, + manual_texture_views: Res, + mut view_target_attachments: ResMut, +) { + prepared.clear(); + for (entity, target) in targets.iter() { + match target { + NormalizedRenderTarget::Window(window) => { + let window = window.entity(); + let Some(surface_data) = window_surfaces.surfaces.get(&window) else { + warn!("Unknown window for screenshot, skipping: {:?}", window); + continue; + }; + let format = surface_data.configuration.format.add_srgb_suffix(); + let size = Extent3d { + width: surface_data.configuration.width, + height: surface_data.configuration.height, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + NormalizedRenderTarget::Image(image) => { + let Some(gpu_image) = images.get(image) else { + warn!("Unknown image for screenshot, skipping: {:?}", image); + continue; + }; + let format = gpu_image.texture_format; + let size = Extent3d { + width: gpu_image.size.x, + height: gpu_image.size.y, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + NormalizedRenderTarget::TextureView(texture_view) => { + let Some(manual_texture_view) = manual_texture_views.get(texture_view) else { + warn!( + "Unknown manual texture view for screenshot, skipping: {:?}", + texture_view + ); + continue; + }; + let format = manual_texture_view.format; + let size = Extent3d { + width: manual_texture_view.size.x, + height: manual_texture_view.size.y, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + } + } +} + +fn prepare_screenshot_state( + size: Extent3d, + format: TextureFormat, + render_device: &RenderDevice, + pipeline: &ScreenshotToScreenPipeline, + pipeline_cache: &PipelineCache, + pipelines: &mut SpecializedRenderPipelines, +) -> (TextureView, ScreenshotPreparedState) { + let texture = render_device.create_texture(&wgpu::TextureDescriptor { + label: Some("screenshot-capture-rendertarget"), + size, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format, + usage: TextureUsages::RENDER_ATTACHMENT + | TextureUsages::COPY_SRC + | TextureUsages::TEXTURE_BINDING, + view_formats: &[], + }); + let texture_view = texture.create_view(&Default::default()); + let buffer = render_device.create_buffer(&wgpu::BufferDescriptor { + label: Some("screenshot-transfer-buffer"), + size: get_aligned_size(size.width, size.height, format.pixel_size() as u32) as u64, + usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let bind_group = render_device.create_bind_group( + "screenshot-to-screen-bind-group", + &pipeline.bind_group_layout, + &BindGroupEntries::single(&texture_view), + ); + let pipeline_id = pipelines.specialize(pipeline_cache, pipeline, format); + + ( + texture_view, + ScreenshotPreparedState { + texture, + buffer, + bind_group, + pipeline_id, + size, + }, + ) +} + pub struct ScreenshotPlugin; const SCREENSHOT_SHADER_HANDLE: Handle = Handle::weak_from_u128(11918575842344596158); impl Plugin for ScreenshotPlugin { fn build(&self, app: &mut bevy_app::App) { - app.init_resource::(); + app.add_systems( + First, + clear_screenshots + .after(event_update_system) + .before(apply_deferred), + ) + .add_systems(Update, trigger_screenshots) + .register_type::() + .register_type::(); load_internal_asset!( app, @@ -139,8 +422,23 @@ impl Plugin for ScreenshotPlugin { } fn finish(&self, app: &mut bevy_app::App) { + let (tx, rx) = std::sync::mpsc::channel(); + app.insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx)))); + if let Some(render_app) = app.get_sub_app_mut(RenderApp) { - render_app.init_resource::>(); + render_app + .insert_resource(RenderScreenshotsSender(tx)) + .init_resource::() + .init_resource::() + .init_resource::>() + .add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all()) + .add_systems( + Render, + prepare_screenshots + .after(prepare_view_attachments) + .before(prepare_view_targets) + .in_set(RenderSet::ManageViews), + ); } } } @@ -221,114 +519,187 @@ impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { } } -pub struct ScreenshotPreparedState { - pub texture: Texture, - pub buffer: Buffer, - pub bind_group: BindGroup, - pub pipeline_id: CachedRenderPipelineId, -} - pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) { - let windows = world.resource::(); + let targets = world.resource::(); + let prepared = world.resource::(); let pipelines = world.resource::(); + let gpu_images = world.resource::>(); + let windows = world.resource::(); + let manual_texture_views = world.resource::(); - for window in windows.values() { - if let Some(memory) = &window.screenshot_memory { - let width = window.physical_width; - let height = window.physical_height; - let texture_format = window.swap_chain_texture_format.unwrap(); - - encoder.copy_texture_to_buffer( - memory.texture.as_image_copy(), - wgpu::ImageCopyBuffer { - buffer: &memory.buffer, - layout: layout_data(width, height, texture_format), - }, - Extent3d { + for (entity, render_target) in targets.iter() { + match render_target { + NormalizedRenderTarget::Window(window) => { + let window = window.entity(); + let Some(window) = windows.get(&window) else { + continue; + }; + let width = window.physical_width; + let height = window.physical_height; + let Some(texture_format) = window.swap_chain_texture_format else { + continue; + }; + let Some(swap_chain_texture) = window.swap_chain_texture.as_ref() else { + continue; + }; + let texture_view = swap_chain_texture.texture.create_view(&Default::default()); + render_screenshot( + encoder, + prepared, + pipelines, + entity, width, height, - ..Default::default() - }, - ); - if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) { - let true_swapchain_texture_view = window - .swap_chain_texture - .as_ref() - .unwrap() - .texture - .create_view(&Default::default()); - let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { - label: Some("screenshot_to_screen_pass"), - color_attachments: &[Some(wgpu::RenderPassColorAttachment { - view: &true_swapchain_texture_view, - resolve_target: None, - ops: wgpu::Operations { - load: wgpu::LoadOp::Load, - store: wgpu::StoreOp::Store, - }, - })], - depth_stencil_attachment: None, - timestamp_writes: None, - occlusion_query_set: None, - }); - pass.set_pipeline(pipeline); - pass.set_bind_group(0, &memory.bind_group, &[]); - pass.draw(0..3, 0..1); + texture_format, + &texture_view, + ); } + NormalizedRenderTarget::Image(image) => { + let Some(gpu_image) = gpu_images.get(image) else { + warn!("Unknown image for screenshot, skipping: {:?}", image); + continue; + }; + let width = gpu_image.size.x; + let height = gpu_image.size.y; + let texture_format = gpu_image.texture_format; + let texture_view = gpu_image.texture_view.deref(); + render_screenshot( + encoder, + prepared, + pipelines, + entity, + width, + height, + texture_format, + texture_view, + ); + } + NormalizedRenderTarget::TextureView(texture_view) => { + let Some(texture_view) = manual_texture_views.get(texture_view) else { + warn!( + "Unknown manual texture view for screenshot, skipping: {:?}", + texture_view + ); + continue; + }; + let width = texture_view.size.x; + let height = texture_view.size.y; + let texture_format = texture_view.format; + let texture_view = texture_view.texture_view.deref(); + render_screenshot( + encoder, + prepared, + pipelines, + entity, + width, + height, + texture_format, + texture_view, + ); + } + }; + } +} + +#[allow(clippy::too_many_arguments)] +fn render_screenshot( + encoder: &mut CommandEncoder, + prepared: &RenderScreenshotsPrepared, + pipelines: &PipelineCache, + entity: &Entity, + width: u32, + height: u32, + texture_format: TextureFormat, + texture_view: &wgpu::TextureView, +) { + if let Some(prepared_state) = &prepared.get(entity) { + encoder.copy_texture_to_buffer( + prepared_state.texture.as_image_copy(), + wgpu::ImageCopyBuffer { + buffer: &prepared_state.buffer, + layout: layout_data(width, height, texture_format), + }, + Extent3d { + width, + height, + ..Default::default() + }, + ); + + if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) { + let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("screenshot_to_screen_pass"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: texture_view, + resolve_target: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Load, + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + pass.set_pipeline(pipeline); + pass.set_bind_group(0, &prepared_state.bind_group, &[]); + pass.draw(0..3, 0..1); } } } pub(crate) fn collect_screenshots(world: &mut World) { - let _span = info_span!("collect_screenshots"); + #[cfg(feature = "trace")] + let _span = bevy_utils::tracing::info_span!("collect_screenshots").entered(); - let mut windows = world.resource_mut::(); - for window in windows.values_mut() { - if let Some(screenshot_func) = window.screenshot_func.take() { - let width = window.physical_width; - let height = window.physical_height; - let texture_format = window.swap_chain_texture_format.unwrap(); - let pixel_size = texture_format.pixel_size(); - let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap(); + let sender = world.resource::().deref().clone(); + let prepared = world.resource::(); - let finish = async move { - let (tx, rx) = async_channel::bounded(1); - let buffer_slice = buffer.slice(..); - // The polling for this map call is done every frame when the command queue is submitted. - buffer_slice.map_async(wgpu::MapMode::Read, move |result| { - let err = result.err(); - if err.is_some() { - panic!("{}", err.unwrap().to_string()); - } - tx.try_send(()).unwrap(); - }); - rx.recv().await.unwrap(); - let data = buffer_slice.get_mapped_range(); - // we immediately move the data to CPU memory to avoid holding the mapped view for long - let mut result = Vec::from(&*data); - drop(data); - drop(buffer); + for (entity, prepared) in prepared.iter() { + let entity = *entity; + let sender = sender.clone(); + let width = prepared.size.width; + let height = prepared.size.height; + let texture_format = prepared.texture.format(); + let pixel_size = texture_format.pixel_size(); + let buffer = prepared.buffer.clone(); - if result.len() != ((width * height) as usize * pixel_size) { - // Our buffer has been padded because we needed to align to a multiple of 256. - // We remove this padding here - let initial_row_bytes = width as usize * pixel_size; - let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize; - - let mut take_offset = buffered_row_bytes; - let mut place_offset = initial_row_bytes; - for _ in 1..height { - result.copy_within( - take_offset..take_offset + buffered_row_bytes, - place_offset, - ); - take_offset += buffered_row_bytes; - place_offset += initial_row_bytes; - } - result.truncate(initial_row_bytes * height as usize); + let finish = async move { + let (tx, rx) = async_channel::bounded(1); + let buffer_slice = buffer.slice(..); + // The polling for this map call is done every frame when the command queue is submitted. + buffer_slice.map_async(wgpu::MapMode::Read, move |result| { + let err = result.err(); + if err.is_some() { + panic!("{}", err.unwrap().to_string()); } + tx.try_send(()).unwrap(); + }); + rx.recv().await.unwrap(); + let data = buffer_slice.get_mapped_range(); + // we immediately move the data to CPU memory to avoid holding the mapped view for long + let mut result = Vec::from(&*data); + drop(data); - screenshot_func(Image::new( + if result.len() != ((width * height) as usize * pixel_size) { + // Our buffer has been padded because we needed to align to a multiple of 256. + // We remove this padding here + let initial_row_bytes = width as usize * pixel_size; + let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize; + + let mut take_offset = buffered_row_bytes; + let mut place_offset = initial_row_bytes; + for _ in 1..height { + result.copy_within(take_offset..take_offset + buffered_row_bytes, place_offset); + take_offset += buffered_row_bytes; + place_offset += initial_row_bytes; + } + result.truncate(initial_row_bytes * height as usize); + } + + if let Err(e) = sender.send(( + entity, + Image::new( Extent3d { width, height, @@ -338,10 +709,12 @@ pub(crate) fn collect_screenshots(world: &mut World) { result, texture_format, RenderAssetUsages::RENDER_WORLD, - )); - }; + ), + )) { + error!("Failed to send screenshot: {:?}", e); + } + }; - AsyncComputeTaskPool::get().spawn(finish).detach(); - } + AsyncComputeTaskPool::get().spawn(finish).detach(); } } diff --git a/examples/window/screenshot.rs b/examples/window/screenshot.rs index 66591ac053..1bf800d18c 100644 --- a/examples/window/screenshot.rs +++ b/examples/window/screenshot.rs @@ -1,29 +1,48 @@ //! An example showing how to save screenshots to disk use bevy::prelude::*; -use bevy::render::view::screenshot::ScreenshotManager; -use bevy::window::PrimaryWindow; +use bevy::window::SystemCursorIcon; +use bevy_render::view::cursor::CursorIcon; +use bevy_render::view::screenshot::{save_to_disk, Capturing, Screenshot}; fn main() { App::new() .add_plugins(DefaultPlugins) .add_systems(Startup, setup) - .add_systems(Update, screenshot_on_spacebar) + .add_systems(Update, (screenshot_on_spacebar, screenshot_saving)) .run(); } fn screenshot_on_spacebar( + mut commands: Commands, input: Res>, - main_window: Query>, - mut screenshot_manager: ResMut, mut counter: Local, ) { if input.just_pressed(KeyCode::Space) { let path = format!("./screenshot-{}.png", *counter); *counter += 1; - screenshot_manager - .save_screenshot_to_disk(main_window.single(), path) - .unwrap(); + commands + .spawn(Screenshot::primary_window()) + .observe(save_to_disk(path)); + } +} + +fn screenshot_saving( + mut commands: Commands, + screenshot_saving: Query>, + windows: Query>, +) { + let window = windows.single(); + match screenshot_saving.iter().count() { + 0 => { + commands.entity(window).remove::(); + } + x if x > 0 => { + commands + .entity(window) + .insert(CursorIcon::from(SystemCursorIcon::Progress)); + } + _ => {} } }