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#[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
24pub 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#[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
84fn 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}