bevy_render/
extract_component.rs1use 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#[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;
56
57 fn extract_component(item: QueryItem<'_, Self::QueryData>) -> Option<Self::Out>;
62}
63
64pub 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#[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
124fn 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
158pub 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
208fn 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
224fn 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}