bevy_render/
extract_component.rs1use 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#[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
33pub trait ExtractComponent: Component {
38 type QueryData: ReadOnlyQueryData;
40 type QueryFilter: QueryFilter;
42
43 type Out: Bundle;
57
58 fn extract_component(item: QueryItem<'_, Self::QueryData>) -> Option<Self::Out>;
63}
64
65pub 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#[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
125fn 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
159pub 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
199fn 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
217fn 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}