bevy_render/
extract_component.rs

1use crate::{
2    render_resource::{encase::internal::WriteInto, DynamicUniformBuffer, ShaderType},
3    renderer::{RenderDevice, RenderQueue},
4    sync_component::SyncComponentPlugin,
5    sync_world::RenderEntity,
6    view::ViewVisibility,
7    Extract, ExtractSchedule, Render, RenderApp, RenderSet,
8};
9use bevy_app::{App, Plugin};
10use bevy_ecs::{
11    component::Component,
12    prelude::*,
13    query::{QueryFilter, QueryItem, ReadOnlyQueryData},
14};
15use core::{marker::PhantomData, ops::Deref};
16
17pub use bevy_render_macros::ExtractComponent;
18
19/// Stores the index of a uniform inside of [`ComponentUniforms`].
20#[derive(Component)]
21pub struct DynamicUniformIndex<C: Component> {
22    index: u32,
23    marker: PhantomData<C>,
24}
25
26impl<C: Component> DynamicUniformIndex<C> {
27    #[inline]
28    pub fn index(&self) -> u32 {
29        self.index
30    }
31}
32
33/// Describes how a component gets extracted for rendering.
34///
35/// Therefore the component is transferred from the "app world" into the "render world"
36/// in the [`ExtractSchedule`] step.
37pub trait ExtractComponent: Component {
38    /// ECS [`ReadOnlyQueryData`] to fetch the components to extract.
39    type QueryData: ReadOnlyQueryData;
40    /// Filters the entities with additional constraints.
41    type QueryFilter: QueryFilter;
42
43    /// The output from extraction.
44    ///
45    /// Returning `None` based on the queried item will remove the component from the entity in
46    /// the render world. This can be used, for example, to conditionally extract camera settings
47    /// in order to disable a rendering feature on the basis of those settings, without removing
48    /// the component from the entity in the main world.
49    ///
50    /// The output may be different from the queried component.
51    /// This can be useful for example if only a subset of the fields are useful
52    /// in the render world.
53    ///
54    /// `Out` has a [`Bundle`] trait bound instead of a [`Component`] trait bound in order to allow use cases
55    /// such as tuples of components as output.
56    type Out: Bundle;
57
58    // TODO: https://github.com/rust-lang/rust/issues/29661
59    // type Out: Component = Self;
60
61    /// Defines how the component is transferred into the "render world".
62    fn extract_component(item: QueryItem<'_, Self::QueryData>) -> Option<Self::Out>;
63}
64
65/// This plugin prepares the components of the corresponding type for the GPU
66/// by transforming them into uniforms.
67///
68/// They can then be accessed from the [`ComponentUniforms`] resource.
69/// For referencing the newly created uniforms a [`DynamicUniformIndex`] is inserted
70/// for every processed entity.
71///
72/// Therefore it sets up the [`RenderSet::Prepare`] step
73/// for the specified [`ExtractComponent`].
74pub struct UniformComponentPlugin<C>(PhantomData<fn() -> C>);
75
76impl<C> Default for UniformComponentPlugin<C> {
77    fn default() -> Self {
78        Self(PhantomData)
79    }
80}
81
82impl<C: Component + ShaderType + WriteInto + Clone> Plugin for UniformComponentPlugin<C> {
83    fn build(&self, app: &mut App) {
84        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
85            render_app
86                .insert_resource(ComponentUniforms::<C>::default())
87                .add_systems(
88                    Render,
89                    prepare_uniform_components::<C>.in_set(RenderSet::PrepareResources),
90                );
91        }
92    }
93}
94
95/// Stores all uniforms of the component type.
96#[derive(Resource)]
97pub struct ComponentUniforms<C: Component + ShaderType> {
98    uniforms: DynamicUniformBuffer<C>,
99}
100
101impl<C: Component + ShaderType> Deref for ComponentUniforms<C> {
102    type Target = DynamicUniformBuffer<C>;
103
104    #[inline]
105    fn deref(&self) -> &Self::Target {
106        &self.uniforms
107    }
108}
109
110impl<C: Component + ShaderType> ComponentUniforms<C> {
111    #[inline]
112    pub fn uniforms(&self) -> &DynamicUniformBuffer<C> {
113        &self.uniforms
114    }
115}
116
117impl<C: Component + ShaderType> Default for ComponentUniforms<C> {
118    fn default() -> Self {
119        Self {
120            uniforms: Default::default(),
121        }
122    }
123}
124
125/// This system prepares all components of the corresponding component type.
126/// They are transformed into uniforms and stored in the [`ComponentUniforms`] resource.
127fn prepare_uniform_components<C>(
128    mut commands: Commands,
129    render_device: Res<RenderDevice>,
130    render_queue: Res<RenderQueue>,
131    mut component_uniforms: ResMut<ComponentUniforms<C>>,
132    components: Query<(Entity, &C)>,
133) where
134    C: Component + ShaderType + WriteInto + Clone,
135{
136    let components_iter = components.iter();
137    let count = components_iter.len();
138    let Some(mut writer) =
139        component_uniforms
140            .uniforms
141            .get_writer(count, &render_device, &render_queue)
142    else {
143        return;
144    };
145    let entities = components_iter
146        .map(|(entity, component)| {
147            (
148                entity,
149                DynamicUniformIndex::<C> {
150                    index: writer.write(component),
151                    marker: PhantomData,
152                },
153            )
154        })
155        .collect::<Vec<_>>();
156    commands.insert_or_spawn_batch(entities);
157}
158
159/// This plugin extracts the components into the render world for synced entities.
160///
161/// To do so, it sets up the [`ExtractSchedule`] step for the specified [`ExtractComponent`].
162pub struct ExtractComponentPlugin<C, F = ()> {
163    only_extract_visible: bool,
164    marker: PhantomData<fn() -> (C, F)>,
165}
166
167impl<C, F> Default for ExtractComponentPlugin<C, F> {
168    fn default() -> Self {
169        Self {
170            only_extract_visible: false,
171            marker: PhantomData,
172        }
173    }
174}
175
176impl<C, F> ExtractComponentPlugin<C, F> {
177    pub fn extract_visible() -> Self {
178        Self {
179            only_extract_visible: true,
180            marker: PhantomData,
181        }
182    }
183}
184
185impl<C: ExtractComponent> Plugin for ExtractComponentPlugin<C> {
186    fn build(&self, app: &mut App) {
187        app.add_plugins(SyncComponentPlugin::<C>::default());
188
189        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
190            if self.only_extract_visible {
191                render_app.add_systems(ExtractSchedule, extract_visible_components::<C>);
192            } else {
193                render_app.add_systems(ExtractSchedule, extract_components::<C>);
194            }
195        }
196    }
197}
198
199/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are synced via [`crate::sync_world::SyncToRenderWorld`].
200fn extract_components<C: ExtractComponent>(
201    mut commands: Commands,
202    mut previous_len: Local<usize>,
203    query: Extract<Query<(RenderEntity, C::QueryData), C::QueryFilter>>,
204) {
205    let mut values = Vec::with_capacity(*previous_len);
206    for (entity, query_item) in &query {
207        if let Some(component) = C::extract_component(query_item) {
208            values.push((entity, component));
209        } else {
210            commands.entity(entity).remove::<C::Out>();
211        }
212    }
213    *previous_len = values.len();
214    commands.insert_or_spawn_batch(values);
215}
216
217/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are visible and synced via [`crate::sync_world::SyncToRenderWorld`].
218fn extract_visible_components<C: ExtractComponent>(
219    mut commands: Commands,
220    mut previous_len: Local<usize>,
221    query: Extract<Query<(RenderEntity, &ViewVisibility, C::QueryData), C::QueryFilter>>,
222) {
223    let mut values = Vec::with_capacity(*previous_len);
224    for (entity, view_visibility, query_item) in &query {
225        if view_visibility.get() {
226            if let Some(component) = C::extract_component(query_item) {
227                values.push((entity, component));
228            } else {
229                commands.entity(entity).remove::<C::Out>();
230            }
231        }
232    }
233    *previous_len = values.len();
234    commands.insert_or_spawn_batch(values);
235}