forestia/bevy
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
38c3423693
bevy/crates/bevy_render/src/view/window/screenshot.rs
Joona Aalto 38c3423693
Event Split: Event, EntityEvent, and BufferedEvent (#19647) 
# Objective

Closes #19564.

The current `Event` trait looks like this:

```rust
pub trait Event: Send + Sync + 'static {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
    
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}
```

The `Event` trait is used by both buffered events
(`EventReader`/`EventWriter`) and observer events. If they are observer
events, they can optionally be targeted at specific `Entity`s or
`ComponentId`s, and can even be propagated to other entities.

However, there has long been a desire to split the trait semantically
for a variety of reasons, see #14843, #14272, and #16031 for discussion.
Some reasons include:

- It's very uncommon to use a single event type as both a buffered event
and targeted observer event. They are used differently and tend to have
distinct semantics.
- A common footgun is using buffered events with observers or event
readers with observer events, as there is no type-level error that
prevents this kind of misuse.
- #19440 made `Trigger::target` return an `Option<Entity>`. This
*seriously* hurts ergonomics for the general case of entity observers,
as you need to `.unwrap()` each time. If we could statically determine
whether the event is expected to have an entity target, this would be
unnecessary.

There's really two main ways that we can categorize events: push vs.
pull (i.e. "observer event" vs. "buffered event") and global vs.
targeted:

|              | Push            | Pull                        |
| ------------ | --------------- | --------------------------- |
| **Global**   | Global observer | `EventReader`/`EventWriter` |
| **Targeted** | Entity observer | -                           |

There are many ways to approach this, each with their tradeoffs.
Ultimately, we kind of want to split events both ways:

- A type-level distinction between observer events and buffered events,
to prevent people from using the wrong kind of event in APIs
- A statically designated entity target for observer events to avoid
accidentally using untargeted events for targeted APIs

This PR achieves these goals by splitting event traits into `Event`,
`EntityEvent`, and `BufferedEvent`, with `Event` being the shared trait
implemented by all events.

## `Event`, `EntityEvent`, and `BufferedEvent`

`Event` is now a very simple trait shared by all events.

```rust
pub trait Event: Send + Sync + 'static {
    // Required for observer APIs
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}
```

You can call `trigger` for *any* event, and use a global observer for
listening to the event.

```rust
#[derive(Event)]
struct Speak {
    message: String,
}

// ...

app.add_observer(|trigger: On<Speak>| {
    println!("{}", trigger.message);
});

// ...

commands.trigger(Speak {
    message: "Y'all like these reworked events?".to_string(),
});
```

To allow an event to be targeted at entities and even propagated
further, you can additionally implement the `EntityEvent` trait:

```rust
pub trait EntityEvent: Event {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
}
```

This lets you call `trigger_targets`, and to use targeted observer APIs
like `EntityCommands::observe`:

```rust
#[derive(Event, EntityEvent)]
#[entity_event(traversal = &'static ChildOf, auto_propagate)]
struct Damage {
    amount: f32,
}

// ...

let enemy = commands.spawn((Enemy, Health(100.0))).id();

// Spawn some armor as a child of the enemy entity.
// When the armor takes damage, it will bubble the event up to the enemy.
let armor_piece = commands
    .spawn((ArmorPiece, Health(25.0), ChildOf(enemy)))
    .observe(|trigger: On<Damage>, mut query: Query<&mut Health>| {
        // Note: `On::target` only exists because this is an `EntityEvent`.
        let mut health = query.get(trigger.target()).unwrap();
        health.0 -= trigger.amount();
    });

commands.trigger_targets(Damage { amount: 10.0 }, armor_piece);
```

> [!NOTE]
> You *can* still also trigger an `EntityEvent` without targets using
`trigger`. We probably *could* make this an either-or thing, but I'm not
sure that's actually desirable.

To allow an event to be used with the buffered API, you can implement
`BufferedEvent`:

```rust
pub trait BufferedEvent: Event {}
```

The event can then be used with `EventReader`/`EventWriter`:

```rust
#[derive(Event, BufferedEvent)]
struct Message(String);

fn write_hello(mut writer: EventWriter<Message>) {
    writer.write(Message("I hope these examples are alright".to_string()));
}

fn read_messages(mut reader: EventReader<Message>) {
    // Process all buffered events of type `Message`.
    for Message(message) in reader.read() {
        println!("{message}");
    }
}
```

In summary:

- Need a basic event you can trigger and observe? Derive `Event`!
- Need the event to be targeted at an entity? Derive `EntityEvent`!
- Need the event to be buffered and support the
`EventReader`/`EventWriter` API? Derive `BufferedEvent`!

## Alternatives

I'll now cover some of the alternative approaches I have considered and
briefly explored. I made this section collapsible since it ended up
being quite long :P

<details>

<summary>Expand this to see alternatives</summary>

### 1. Unified `Event` Trait

One option is not to have *three* separate traits (`Event`,
`EntityEvent`, `BufferedEvent`), and to instead just use associated
constants on `Event` to determine whether an event supports targeting
and buffering or not:

```rust
pub trait Event: Send + Sync + 'static {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
    const TARGETED: bool = false;
    const BUFFERED: bool = false;
    
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}
```

Methods can then use bounds like `where E: Event<TARGETED = true>` or
`where E: Event<BUFFERED = true>` to limit APIs to specific kinds of
events.

This would keep everything under one `Event` trait, but I don't think
it's necessarily a good idea. It makes APIs harder to read, and docs
can't easily refer to specific types of events. You can also create
weird invariants: what if you specify `TARGETED = false`, but have
`Traversal` and/or `AUTO_PROPAGATE` enabled?

### 2. `Event` and `Trigger`

Another option is to only split the traits between buffered events and
observer events, since that is the main thing people have been asking
for, and they have the largest API difference.

If we did this, I think we would need to make the terms *clearly*
separate. We can't really use `Event` and `BufferedEvent` as the names,
since it would be strange that `BufferedEvent` doesn't implement
`Event`. Something like `ObserverEvent` and `BufferedEvent` could work,
but it'd be more verbose.

For this approach, I would instead keep `Event` for the current
`EventReader`/`EventWriter` API, and call the observer event a
`Trigger`, since the "trigger" terminology is already used in the
observer context within Bevy (both as a noun and a verb). This is also
what a long [bikeshed on
Discord](https://discord.com/channels/691052431525675048/749335865876021248/1298057661878898791)
seemed to land on at the end of last year.

```rust
// For `EventReader`/`EventWriter`
pub trait Event: Send + Sync + 'static {}

// For observers
pub trait Trigger: Send + Sync + 'static {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
    const TARGETED: bool = false;
    
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}
```

The problem is that "event" is just a really good term for something
that "happens". Observers are rapidly becoming the more prominent API,
so it'd be weird to give them the `Trigger` name and leave the good
`Event` name for the less common API.

So, even though a split like this seems neat on the surface, I think it
ultimately wouldn't really work. We want to keep the `Event` name for
observer events, and there is no good alternative for the buffered
variant. (`Message` was suggested, but saying stuff like "sends a
collision message" is weird.)

### 3. `GlobalEvent` + `TargetedEvent`

What if instead of focusing on the buffered vs. observed split, we
*only* make a distinction between global and targeted events?

```rust
// A shared event trait to allow global observers to work
pub trait Event: Send + Sync + 'static {
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}

// For buffered events and non-targeted observer events
pub trait GlobalEvent: Event {}

// For targeted observer events
pub trait TargetedEvent: Event {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
}
```

This is actually the first approach I implemented, and it has the neat
characteristic that you can only use non-targeted APIs like `trigger`
with a `GlobalEvent` and targeted APIs like `trigger_targets` with a
`TargetedEvent`. You have full control over whether the entity should or
should not have a target, as they are fully distinct at the type-level.

However, there's a few problems:

- There is no type-level indication of whether a `GlobalEvent` supports
buffered events or just non-targeted observer events
- An `Event` on its own does literally nothing, it's just a shared trait
required to make global observers accept both non-targeted and targeted
events
- If an event is both a `GlobalEvent` and `TargetedEvent`, global
observers again have ambiguity on whether an event has a target or not,
undermining some of the benefits
- The names are not ideal

### 4. `Event` and `EntityEvent`

We can fix some of the problems of Alternative 3 by accepting that
targeted events can also be used in non-targeted contexts, and simply
having the `Event` and `EntityEvent` traits:

```rust
// For buffered events and non-targeted observer events
pub trait Event: Send + Sync + 'static {
    fn register_component_id(world: &mut World) -> ComponentId { ... }
    fn component_id(world: &World) -> Option<ComponentId> { ... }
}

// For targeted observer events
pub trait EntityEvent: Event {
    type Traversal: Traversal<Self>;
    const AUTO_PROPAGATE: bool = false;
}
```

This is essentially identical to this PR, just without a dedicated
`BufferedEvent`. The remaining major "problem" is that there is still
zero type-level indication of whether an `Event` event *actually*
supports the buffered API. This leads us to the solution proposed in
this PR, using `Event`, `EntityEvent`, and `BufferedEvent`.

</details>

## Conclusion

The `Event` + `EntityEvent` + `BufferedEvent` split proposed in this PR
aims to solve all the common problems with Bevy's current event model
while keeping the "weirdness" factor minimal. It splits in terms of both
the push vs. pull *and* global vs. targeted aspects, while maintaining a
shared concept for an "event".

### Why I Like This

- The term "event" remains as a single concept for all the different
kinds of events in Bevy.
- Despite all event types being "events", they use fundamentally
different APIs. Instead of assuming that you can use an event type with
any pattern (when only one is typically supported), you explicitly opt
in to each one with dedicated traits.
- Using separate traits for each type of event helps with documentation
and clearer function signatures.
- I can safely make assumptions on expected usage.
- If I see that an event is an `EntityEvent`, I can assume that I can
use `observe` on it and get targeted events.
- If I see that an event is a `BufferedEvent`, I can assume that I can
use `EventReader` to read events.
- If I see both `EntityEvent` and `BufferedEvent`, I can assume that
both APIs are supported.

In summary: This allows for a unified concept for events, while limiting
the different ways to use them with opt-in traits. No more guess-work
involved when using APIs.

### Problems?

- Because `BufferedEvent` implements `Event` (for more consistent
semantics etc.), you can still use all buffered events for non-targeted
observers. I think this is fine/good. The important part is that if you
see that an event implements `BufferedEvent`, you know that the
`EventReader`/`EventWriter` API should be supported. Whether it *also*
supports other APIs is secondary.
- I currently only support `trigger_targets` for an `EntityEvent`.
However, you can technically target components too, without targeting
any entities. I consider that such a niche and advanced use case that
it's not a huge problem to only support it for `EntityEvent`s, but we
could also split `trigger_targets` into `trigger_entities` and
`trigger_components` if we wanted to (or implement components as
entities :P).
- You can still trigger an `EntityEvent` *without* targets. I consider
this correct, since `Event` implements the non-targeted behavior, and
it'd be weird if implementing another trait *removed* behavior. However,
it does mean that global observers for entity events can technically
return `Entity::PLACEHOLDER` again (since I got rid of the
`Option<Entity>` added in #19440 for ergonomics). I think that's enough
of an edge case that it's not a huge problem, but it is worth keeping in
mind.
- ~~Deriving both `EntityEvent` and `BufferedEvent` for the same type
currently duplicates the `Event` implementation, so you instead need to
manually implement one of them.~~ Changed to always requiring `Event` to
be derived.

## Related Work

There are plans to implement multi-event support for observers,
especially for UI contexts. [Cart's
example](https://github.com/bevyengine/bevy/issues/14649#issuecomment-2960402508)
API looked like this:

```rust
// Truncated for brevity
trigger: Trigger<(
    OnAdd<Pressed>,
    OnRemove<Pressed>,
    OnAdd<InteractionDisabled>,
    OnRemove<InteractionDisabled>,
    OnInsert<Hovered>,
)>,
```

I believe this shouldn't be in conflict with this PR. If anything, this
PR might *help* achieve the multi-event pattern for entity observers
with fewer footguns: by statically enforcing that all of these events
are `EntityEvent`s in the context of `EntityCommands::observe`, we can
avoid misuse or weird cases where *some* events inside the trigger are
targeted while others are not.
2025-06-15 16:46:34 +00:00

695 lines
26 KiB
Rust
Raw Blame History

use super::ExtractedWindows;
use crate::{
camera::{ManualTextureViewHandle, ManualTextureViews, NormalizedRenderTarget, RenderTarget},
gpu_readback,
prelude::Shader,
render_asset::{RenderAssetUsages, RenderAssets},
render_resource::{
binding_types::texture_2d, BindGroup, BindGroupEntries, BindGroupLayout,
BindGroupLayoutEntries, Buffer, BufferUsages, CachedRenderPipelineId, FragmentState,
PipelineCache, RenderPipelineDescriptor, SpecializedRenderPipeline,
SpecializedRenderPipelines, Texture, TextureUsages, TextureView, VertexState,
},
renderer::RenderDevice,
texture::{GpuImage, OutputColorAttachment},
view::{prepare_view_attachments, prepare_view_targets, ViewTargetAttachments, WindowSurfaces},
ExtractSchedule, MainWorld, Render, RenderApp, RenderSystems,
};
use alloc::{borrow::Cow, sync::Arc};
use bevy_app::{First, Plugin, Update};
use bevy_asset::{embedded_asset, load_embedded_asset, Handle};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
entity::EntityHashMap, event::event_update_system, prelude::*, system::SystemState,
};
use bevy_image::{Image, TextureFormatPixelInfo};
use bevy_platform::collections::HashSet;
use bevy_reflect::Reflect;
use bevy_tasks::AsyncComputeTaskPool;
use bevy_utils::default;
use bevy_window::{PrimaryWindow, WindowRef};
use core::ops::Deref;
use std::{
path::Path,
sync::{
mpsc::{Receiver, Sender},
Mutex,
},
};
use tracing::{error, info, warn};
use wgpu::{CommandEncoder, Extent3d, TextureFormat};
#[derive(Event, EntityEvent, Deref, DerefMut, Reflect, Debug)]
#[reflect(Debug)]
pub struct ScreenshotCaptured(pub Image);
/// 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 marker component that indicates that a screenshot is currently being captured.
#[derive(Component, Default)]
pub struct Capturing;
/// A marker component that indicates that a screenshot has been captured, the image is ready, and
/// the screenshot entity can be despawned.
#[derive(Component, Default)]
pub struct Captured;
impl Screenshot {
/// Capture a screenshot of the provided window entity.
pub fn window(window: Entity) -> Self {
Self(RenderTarget::Window(WindowRef::Entity(window)))
}
/// 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<Image>) -> Self {
Self(RenderTarget::Image(image.into()))
}
/// 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<Mutex<Receiver<(Entity, Image)>>>);
#[derive(Resource, Deref, DerefMut, Default)]
struct RenderScreenshotTargets(EntityHashMap<NormalizedRenderTarget>);
#[derive(Resource, Deref, DerefMut, Default)]
struct RenderScreenshotsPrepared(EntityHashMap<ScreenshotPreparedState>);
#[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<Path>) -> impl FnMut(On<ScreenshotCaptured>) {
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
// the screenshot looks right
let img = dyn_img.to_rgb8();
#[cfg(not(target_arch = "wasm32"))]
match img.save_with_format(&path, format) {
Ok(_) => info!("Screenshot saved to {}", path.display()),
Err(e) => error!("Cannot save screenshot, IO error: {e}"),
}
#[cfg(target_arch = "wasm32")]
{
let save_screenshot = || {
use image::EncodableLayout;
use wasm_bindgen::{JsCast, JsValue};
let mut image_buffer = std::io::Cursor::new(Vec::new());
img.write_to(&mut image_buffer, format)
.map_err(|e| JsValue::from_str(&format!("{e}")))?;
// SAFETY: `image_buffer` only exist in this closure, and is not used after this line
let parts = js_sys::Array::of1(&unsafe {
js_sys::Uint8Array::view(image_buffer.into_inner().as_bytes())
.into()
});
let blob = web_sys::Blob::new_with_u8_array_sequence(&parts)?;
let url = web_sys::Url::create_object_url_with_blob(&blob)?;
let window = web_sys::window().unwrap();
let document = window.document().unwrap();
let link = document.create_element("a")?;
link.set_attribute("href", &url)?;
link.set_attribute(
"download",
path.file_name()
.and_then(|filename| filename.to_str())
.ok_or_else(|| JsValue::from_str("Invalid filename"))?,
)?;
let html_element = link.dyn_into::<web_sys::HtmlElement>()?;
html_element.click();
web_sys::Url::revoke_object_url(&url)?;
Ok::<(), JsValue>(())
};
match (save_screenshot)() {
Ok(_) => info!("Screenshot saved to {}", path.display()),
Err(e) => error!("Cannot save screenshot, error: {e:?}"),
};
}
}
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<Entity, With<Captured>>) {
for entity in screenshots.iter() {
commands.entity(entity).despawn();
}
}
pub fn trigger_screenshots(
mut commands: Commands,
captured_screenshots: ResMut<CapturedScreenshots>,
) {
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<RenderScreenshotTargets>,
mut main_world: ResMut<MainWorld>,
mut system_state: Local<
Option<
SystemState<(
Commands,
Query<Entity, With<PrimaryWindow>>,
Query<(Entity, &Screenshot), Without<Capturing>>,
)>,
>,
>,
mut seen_targets: Local<HashSet<NormalizedRenderTarget>>,
) {
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();
continue;
}
seen_targets.insert(render_target.clone());
targets.insert(entity, render_target);
commands.entity(entity).insert(Capturing);
}
system_state.apply(&mut main_world);
}
fn prepare_screenshots(
targets: Res<RenderScreenshotTargets>,
mut prepared: ResMut<RenderScreenshotsPrepared>,
window_surfaces: Res<WindowSurfaces>,
render_device: Res<RenderDevice>,
screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
pipeline_cache: Res<PipelineCache>,
mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
images: Res<RenderAssets<GpuImage>>,
manual_texture_views: Res<ManualTextureViews>,
mut view_target_attachments: ResMut<ViewTargetAttachments>,
) {
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.handle) else {
warn!("Unknown image for screenshot, skipping: {:?}", image);
continue;
};
let format = gpu_image.texture_format;
let (texture_view, state) = prepare_screenshot_state(
gpu_image.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<ScreenshotToScreenPipeline>,
) -> (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: gpu_readback::get_aligned_size(size, 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;
impl Plugin for ScreenshotPlugin {
fn build(&self, app: &mut bevy_app::App) {
app.add_systems(
First,
clear_screenshots
.after(event_update_system)
.before(ApplyDeferred),
)
.register_type::<Screenshot>()
.register_type::<ScreenshotCaptured>();
embedded_asset!(app, "screenshot.wgsl");
}
fn finish(&self, app: &mut bevy_app::App) {
let (tx, rx) = std::sync::mpsc::channel();
app.add_systems(Update, trigger_screenshots)
.insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx))));
if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
render_app
.insert_resource(RenderScreenshotsSender(tx))
.init_resource::<RenderScreenshotTargets>()
.init_resource::<RenderScreenshotsPrepared>()
.init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>()
.init_resource::<ScreenshotToScreenPipeline>()
.add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all())
.add_systems(
Render,
prepare_screenshots
.after(prepare_view_attachments)
.before(prepare_view_targets)
.in_set(RenderSystems::ManageViews),
);
}
}
}
#[derive(Resource)]
pub struct ScreenshotToScreenPipeline {
pub bind_group_layout: BindGroupLayout,
pub shader: Handle<Shader>,
}
impl FromWorld for ScreenshotToScreenPipeline {
fn from_world(render_world: &mut World) -> Self {
let device = render_world.resource::<RenderDevice>();
let bind_group_layout = device.create_bind_group_layout(
"screenshot-to-screen-bgl",
&BindGroupLayoutEntries::single(
wgpu::ShaderStages::FRAGMENT,
texture_2d(wgpu::TextureSampleType::Float { filterable: false }),
),
);
let shader = load_embedded_asset!(render_world, "screenshot.wgsl");
Self {
bind_group_layout,
shader,
}
}
}
impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
type Key = TextureFormat;
fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
RenderPipelineDescriptor {
label: Some(Cow::Borrowed("screenshot-to-screen")),
layout: vec![self.bind_group_layout.clone()],
vertex: VertexState {
buffers: vec![],
shader_defs: vec![],
entry_point: Cow::Borrowed("vs_main"),
shader: self.shader.clone(),
},
primitive: wgpu::PrimitiveState {
cull_mode: Some(wgpu::Face::Back),
..Default::default()
},
depth_stencil: None,
multisample: Default::default(),
fragment: Some(FragmentState {
shader: self.shader.clone(),
entry_point: Cow::Borrowed("fs_main"),
shader_defs: vec![],
targets: vec![Some(wgpu::ColorTargetState {
format: key,
blend: None,
write_mask: wgpu::ColorWrites::ALL,
})],
}),
push_constant_ranges: Vec::new(),
zero_initialize_workgroup_memory: false,
}
}
}
pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) {
let targets = world.resource::<RenderScreenshotTargets>();
let prepared = world.resource::<RenderScreenshotsPrepared>();
let pipelines = world.resource::<PipelineCache>();
let gpu_images = world.resource::<RenderAssets<GpuImage>>();
let windows = world.resource::<ExtractedWindows>();
let manual_texture_views = world.resource::<ManualTextureViews>();
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,
texture_format,
&texture_view,
);
}
NormalizedRenderTarget::Image(image) => {
let Some(gpu_image) = gpu_images.get(&image.handle) else {
warn!("Unknown image for screenshot, skipping: {:?}", image);
continue;
};
let width = gpu_image.size.width;
let height = gpu_image.size.height;
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,
);
}
};
}
}
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) {
let extent = Extent3d {
width,
height,
depth_or_array_layers: 1,
};
encoder.copy_texture_to_buffer(
prepared_state.texture.as_image_copy(),
wgpu::TexelCopyBufferInfo {
buffer: &prepared_state.buffer,
layout: gpu_readback::layout_data(extent, texture_format),
},
extent,
);
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) {
#[cfg(feature = "trace")]
let _span = tracing::info_span!("collect_screenshots").entered();
let sender = world.resource::<RenderScreenshotsSender>().deref().clone();
let prepared = world.resource::<RenderScreenshotsPrepared>();
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();
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);
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 =
gpu_readback::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,
depth_or_array_layers: 1,
},
wgpu::TextureDimension::D2,
result,
texture_format,
RenderAssetUsages::RENDER_WORLD,
),
)) {
error!("Failed to send screenshot: {}", e);
}
};
AsyncComputeTaskPool::get().spawn(finish).detach();
}
}
Reference in New Issue View Git Blame Copy Permalink