Skip to main content

bevy_render/
uniform.rs

1use crate::{
2    render_resource::{encase::internal::WriteInto, DynamicUniformBuffer, ShaderType},
3    renderer::{RenderDevice, RenderQueue},
4    GpuResourceAppExt, Render, RenderApp, RenderSystems,
5};
6use bevy_app::{App, Plugin};
7use bevy_ecs::{component::Component, prelude::*};
8use core::{marker::PhantomData, ops::Deref};
9
10/// Stores the index of a uniform inside of [`ComponentUniforms`].
11#[derive(Component)]
12pub struct DynamicUniformIndex<C: Component> {
13    index: u32,
14    marker: PhantomData<C>,
15}
16
17impl<C: Component> DynamicUniformIndex<C> {
18    #[inline]
19    pub fn index(&self) -> u32 {
20        self.index
21    }
22}
23
24/// This plugin prepares the components of the corresponding type for the GPU
25/// by transforming them into uniforms.
26///
27/// They can then be accessed from the [`ComponentUniforms`] resource.
28/// For referencing the newly created uniforms a [`DynamicUniformIndex`] is inserted
29/// for every processed entity.
30///
31/// Therefore it sets up the [`RenderSystems::Prepare`] step
32/// for the specified [`ExtractComponent`](`crate::extract_component::ExtractComponent`).
33pub struct UniformComponentPlugin<C>(PhantomData<fn() -> C>);
34
35impl<C> Default for UniformComponentPlugin<C> {
36    fn default() -> Self {
37        Self(PhantomData)
38    }
39}
40
41impl<C: Component + ShaderType + WriteInto + Clone> Plugin for UniformComponentPlugin<C> {
42    fn build(&self, app: &mut App) {
43        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
44            render_app
45                .init_gpu_resource::<ComponentUniforms<C>>()
46                .add_systems(
47                    Render,
48                    prepare_uniform_components::<C>.in_set(RenderSystems::PrepareResources),
49                );
50        }
51    }
52}
53
54/// Stores all uniforms of the component type.
55#[derive(Resource)]
56pub struct ComponentUniforms<C: Component + ShaderType> {
57    uniforms: DynamicUniformBuffer<C>,
58}
59
60impl<C: Component + ShaderType> Deref for ComponentUniforms<C> {
61    type Target = DynamicUniformBuffer<C>;
62
63    #[inline]
64    fn deref(&self) -> &Self::Target {
65        &self.uniforms
66    }
67}
68
69impl<C: Component + ShaderType> ComponentUniforms<C> {
70    #[inline]
71    pub fn uniforms(&self) -> &DynamicUniformBuffer<C> {
72        &self.uniforms
73    }
74}
75
76impl<C: Component + ShaderType> Default for ComponentUniforms<C> {
77    fn default() -> Self {
78        Self {
79            uniforms: Default::default(),
80        }
81    }
82}
83
84/// This system prepares all components of the corresponding component type.
85/// They are transformed into uniforms and stored in the [`ComponentUniforms`] resource.
86fn prepare_uniform_components<C>(
87    mut commands: Commands,
88    render_device: Res<RenderDevice>,
89    render_queue: Res<RenderQueue>,
90    mut component_uniforms: ResMut<ComponentUniforms<C>>,
91    components: Query<(Entity, &C)>,
92) where
93    C: Component + ShaderType + WriteInto + Clone,
94{
95    let components_iter = components.iter();
96    let count = components_iter.len();
97    let Some(mut writer) =
98        component_uniforms
99            .uniforms
100            .get_writer(count, &render_device, &render_queue)
101    else {
102        return;
103    };
104    let entities = components_iter
105        .map(|(entity, component)| {
106            (
107                entity,
108                DynamicUniformIndex::<C> {
109                    index: writer.write(component),
110                    marker: PhantomData,
111                },
112            )
113        })
114        .collect::<Vec<_>>();
115    commands.try_insert_batch(entities);
116}