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