1use core::{iter, mem};
2
3use bevy_ecs::prelude::*;
4use bevy_render::sync_world::MainEntityHashMap;
5use bevy_render::{
6 batching::NoAutomaticBatching,
7 mesh::morph::{MeshMorphWeights, MAX_MORPH_WEIGHTS},
8 render_resource::{BufferUsages, RawBufferVec},
9 renderer::{RenderDevice, RenderQueue},
10 view::ViewVisibility,
11 Extract,
12};
13use bytemuck::NoUninit;
14
15#[derive(Component)]
16pub struct MorphIndex {
17 pub index: u32,
18}
19
20#[derive(Default, Resource)]
26pub struct MorphIndices {
27 pub current: MainEntityHashMap<MorphIndex>,
30
31 pub prev: MainEntityHashMap<MorphIndex>,
34}
35
36#[derive(Resource)]
43pub struct MorphUniforms {
44 pub current_buffer: RawBufferVec<f32>,
46 pub prev_buffer: RawBufferVec<f32>,
48}
49
50impl Default for MorphUniforms {
51 fn default() -> Self {
52 Self {
53 current_buffer: RawBufferVec::new(BufferUsages::UNIFORM),
54 prev_buffer: RawBufferVec::new(BufferUsages::UNIFORM),
55 }
56 }
57}
58
59pub fn prepare_morphs(
60 render_device: Res<RenderDevice>,
61 render_queue: Res<RenderQueue>,
62 mut uniform: ResMut<MorphUniforms>,
63) {
64 if uniform.current_buffer.is_empty() {
65 return;
66 }
67 let len = uniform.current_buffer.len();
68 uniform.current_buffer.reserve(len, &render_device);
69 uniform
70 .current_buffer
71 .write_buffer(&render_device, &render_queue);
72
73 }
76
77const fn can_align(step: usize, target: usize) -> bool {
78 step % target == 0 || target % step == 0
79}
80
81const WGPU_MIN_ALIGN: usize = 256;
82
83fn add_to_alignment<T: NoUninit + Default>(buffer: &mut RawBufferVec<T>) {
85 let n = WGPU_MIN_ALIGN;
86 let t_size = size_of::<T>();
87 if !can_align(n, t_size) {
88 panic!(
90 "RawBufferVec should contain only types with a size multiple or divisible by {n}, \
91 {} has a size of {t_size}, which is neither multiple or divisible by {n}",
92 core::any::type_name::<T>()
93 );
94 }
95
96 let buffer_size = buffer.len();
97 let byte_size = t_size * buffer_size;
98 let bytes_over_n = byte_size % n;
99 if bytes_over_n == 0 {
100 return;
101 }
102 let bytes_to_add = n - bytes_over_n;
103 let ts_to_add = bytes_to_add / t_size;
104 buffer.extend(iter::repeat_with(T::default).take(ts_to_add));
105}
106
107pub fn extract_morphs(
110 morph_indices: ResMut<MorphIndices>,
111 uniform: ResMut<MorphUniforms>,
112 query: Extract<Query<(Entity, &ViewVisibility, &MeshMorphWeights)>>,
113) {
114 let (morph_indices, uniform) = (morph_indices.into_inner(), uniform.into_inner());
116
117 mem::swap(&mut morph_indices.current, &mut morph_indices.prev);
120 mem::swap(&mut uniform.current_buffer, &mut uniform.prev_buffer);
121 morph_indices.current.clear();
122 uniform.current_buffer.clear();
123
124 for (entity, view_visibility, morph_weights) in &query {
125 if !view_visibility.get() {
126 continue;
127 }
128 let start = uniform.current_buffer.len();
129 let weights = morph_weights.weights();
130 let legal_weights = weights.iter().take(MAX_MORPH_WEIGHTS).copied();
131 uniform.current_buffer.extend(legal_weights);
132 add_to_alignment::<f32>(&mut uniform.current_buffer);
133
134 let index = (start * size_of::<f32>()) as u32;
135 morph_indices
136 .current
137 .insert(entity.into(), MorphIndex { index });
138 }
139}
140
141pub fn no_automatic_morph_batching(
144 mut commands: Commands,
145 query: Query<Entity, (With<MeshMorphWeights>, Without<NoAutomaticBatching>)>,
146) {
147 for entity in &query {
148 commands.entity(entity).try_insert(NoAutomaticBatching);
149 }
150}