bevy_render/
extract_component.rs

1use crate::{
2    render_resource::{encase::internal::WriteInto, DynamicUniformBuffer, ShaderType},
3    renderer::{RenderDevice, RenderQueue},
4    view::ViewVisibility,
5    Extract, ExtractSchedule, Render, RenderApp, RenderSet,
6};
7use bevy_app::{App, Plugin};
8use bevy_asset::{Asset, Handle};
9use bevy_ecs::{
10    component::Component,
11    prelude::*,
12    query::{QueryFilter, QueryItem, ReadOnlyQueryData},
13    system::lifetimeless::Read,
14};
15use std::{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 can allow early optimization,
46    /// for example if there is an `enabled: bool` field on `Self`, or by only accepting
47    /// values within certain thresholds.
48    ///
49    /// The output may be different from the queried component.
50    /// This can be useful for example if only a subset of the fields are useful
51    /// in the render world.
52    ///
53    /// `Out` has a [`Bundle`] trait bound instead of a [`Component`] trait bound in order to allow use cases
54    /// such as tuples of components as output.
55    type Out: Bundle;
56
57    // TODO: https://github.com/rust-lang/rust/issues/29661
58    // type Out: Component = Self;
59
60    /// Defines how the component is transferred into the "render world".
61    fn extract_component(item: QueryItem<'_, Self::QueryData>) -> Option<Self::Out>;
62}
63
64/// This plugin prepares the components of the corresponding type for the GPU
65/// by transforming them into uniforms.
66///
67/// They can then be accessed from the [`ComponentUniforms`] resource.
68/// For referencing the newly created uniforms a [`DynamicUniformIndex`] is inserted
69/// for every processed entity.
70///
71/// Therefore it sets up the [`RenderSet::Prepare`] step
72/// for the specified [`ExtractComponent`].
73pub struct UniformComponentPlugin<C>(PhantomData<fn() -> C>);
74
75impl<C> Default for UniformComponentPlugin<C> {
76    fn default() -> Self {
77        Self(PhantomData)
78    }
79}
80
81impl<C: Component + ShaderType + WriteInto + Clone> Plugin for UniformComponentPlugin<C> {
82    fn build(&self, app: &mut App) {
83        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
84            render_app
85                .insert_resource(ComponentUniforms::<C>::default())
86                .add_systems(
87                    Render,
88                    prepare_uniform_components::<C>.in_set(RenderSet::PrepareResources),
89                );
90        }
91    }
92}
93
94/// Stores all uniforms of the component type.
95#[derive(Resource)]
96pub struct ComponentUniforms<C: Component + ShaderType> {
97    uniforms: DynamicUniformBuffer<C>,
98}
99
100impl<C: Component + ShaderType> Deref for ComponentUniforms<C> {
101    type Target = DynamicUniformBuffer<C>;
102
103    #[inline]
104    fn deref(&self) -> &Self::Target {
105        &self.uniforms
106    }
107}
108
109impl<C: Component + ShaderType> ComponentUniforms<C> {
110    #[inline]
111    pub fn uniforms(&self) -> &DynamicUniformBuffer<C> {
112        &self.uniforms
113    }
114}
115
116impl<C: Component + ShaderType> Default for ComponentUniforms<C> {
117    fn default() -> Self {
118        Self {
119            uniforms: Default::default(),
120        }
121    }
122}
123
124/// This system prepares all components of the corresponding component type.
125/// They are transformed into uniforms and stored in the [`ComponentUniforms`] resource.
126fn prepare_uniform_components<C>(
127    mut commands: Commands,
128    render_device: Res<RenderDevice>,
129    render_queue: Res<RenderQueue>,
130    mut component_uniforms: ResMut<ComponentUniforms<C>>,
131    components: Query<(Entity, &C)>,
132) where
133    C: Component + ShaderType + WriteInto + Clone,
134{
135    let components_iter = components.iter();
136    let count = components_iter.len();
137    let Some(mut writer) =
138        component_uniforms
139            .uniforms
140            .get_writer(count, &render_device, &render_queue)
141    else {
142        return;
143    };
144    let entities = components_iter
145        .map(|(entity, component)| {
146            (
147                entity,
148                DynamicUniformIndex::<C> {
149                    index: writer.write(component),
150                    marker: PhantomData,
151                },
152            )
153        })
154        .collect::<Vec<_>>();
155    commands.insert_or_spawn_batch(entities);
156}
157
158/// This plugin extracts the components into the "render world".
159///
160/// Therefore it sets up the [`ExtractSchedule`] step
161/// 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        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
188            if self.only_extract_visible {
189                render_app.add_systems(ExtractSchedule, extract_visible_components::<C>);
190            } else {
191                render_app.add_systems(ExtractSchedule, extract_components::<C>);
192            }
193        }
194    }
195}
196
197impl<T: Asset> ExtractComponent for Handle<T> {
198    type QueryData = Read<Handle<T>>;
199    type QueryFilter = ();
200    type Out = Handle<T>;
201
202    #[inline]
203    fn extract_component(handle: QueryItem<'_, Self::QueryData>) -> Option<Self::Out> {
204        Some(handle.clone_weak())
205    }
206}
207
208/// This system extracts all components of the corresponding [`ExtractComponent`] type.
209fn extract_components<C: ExtractComponent>(
210    mut commands: Commands,
211    mut previous_len: Local<usize>,
212    query: Extract<Query<(Entity, C::QueryData), C::QueryFilter>>,
213) {
214    let mut values = Vec::with_capacity(*previous_len);
215    for (entity, query_item) in &query {
216        if let Some(component) = C::extract_component(query_item) {
217            values.push((entity, component));
218        }
219    }
220    *previous_len = values.len();
221    commands.insert_or_spawn_batch(values);
222}
223
224/// This system extracts all visible components of the corresponding [`ExtractComponent`] type.
225fn extract_visible_components<C: ExtractComponent>(
226    mut commands: Commands,
227    mut previous_len: Local<usize>,
228    query: Extract<Query<(Entity, &ViewVisibility, C::QueryData), C::QueryFilter>>,
229) {
230    let mut values = Vec::with_capacity(*previous_len);
231    for (entity, view_visibility, query_item) in &query {
232        if view_visibility.get() {
233            if let Some(component) = C::extract_component(query_item) {
234                values.push((entity, component));
235            }
236        }
237    }
238    *previous_len = values.len();
239    commands.insert_or_spawn_batch(values);
240}