1use std::collections::HashMap;
2
3use crate::{
4 app::{App, SystemStage},
5 assets::{handle::Handle, storage::RawAssetHandle},
6 ecs::{plugin::Plugin, system::{Commands, Local, Query, Res, ResMut}},
7};
8
9use super::{
10 backend::WGPUBackend,
11 binding::{BindGroupLayout, BindGroupLayoutBuilder, BindingKind},
12 buffer::Buffer,
13 buffers::{BindGroup, BindGroupBuilder, BufferBuilder},
14 flags::ShaderStages,
15 layout::GlobalLayoutPool,
16 material::Material,
17 player::AnimationPlayer,
18 skinned_mesh::SkinnedMesh,
19};
20
21struct SkinningBatch {
24 matrices: Buffer, _info: Buffer, bind_group: BindGroup,
27 joint_count: u32,
28 capacity: u32,
29}
30
31impl SkinningBatch {
32 fn new(backend: &WGPUBackend, layout: &BindGroupLayout, joint_count: u32, capacity: u32) -> Self {
33 let matrices = BufferBuilder::empty((capacity as u64) * (joint_count as u64) * 64)
34 .with_label("pebble_skinning:joint_matrices")
35 .with_storage()
36 .build(backend);
37 let info = BufferBuilder::with_data(bytemuck::cast_slice(&[joint_count, 0u32, 0u32, 0u32]))
38 .with_label("pebble_skinning:info")
39 .with_uniform()
40 .build(backend);
41 let bind_group = BindGroupBuilder::new(layout)
42 .with_buffer(&matrices)
43 .with_buffer(&info)
44 .build(backend);
45 Self { matrices, _info: info, bind_group, joint_count, capacity }
46 }
47}
48
49pub struct SkinnedBatchRenderer {
63 layout: BindGroupLayout,
64 batches: HashMap<(RawAssetHandle, RawAssetHandle), SkinningBatch>,
65}
66
67impl SkinnedBatchRenderer {
68 fn new(backend: &WGPUBackend) -> Self {
69 let layout = BindGroupLayoutBuilder::new()
70 .with_label("pebble_skinning")
71 .with_entry("joint_matrices", 0, BindingKind::storage_buffer_read_only(ShaderStages::VERTEX))
72 .with_entry("skin_info", 1, BindingKind::uniform_buffer(ShaderStages::VERTEX))
73 .build(backend);
74 Self { layout, batches: HashMap::new() }
75 }
76
77 pub fn bind_group(&self, material: RawAssetHandle, mesh: RawAssetHandle) -> Option<&BindGroup> {
80 self.batches.get(&(material, mesh)).map(|b| &b.bind_group)
81 }
82
83 fn prepare(&mut self, key: (RawAssetHandle, RawAssetHandle), joint_count: u32, needed: u32, backend: &WGPUBackend) {
84 let layout = self.layout.clone();
85 let batch = self.batches.entry(key).or_insert_with(|| {
86 SkinningBatch::new(backend, &layout, joint_count, needed.max(256))
87 });
88 if needed > batch.capacity {
89 *batch = SkinningBatch::new(backend, &layout, batch.joint_count, (batch.capacity * 2).max(needed));
90 }
91 }
92}
93
94pub struct SkinnedBatchUnit {
98 pub material: RawAssetHandle,
99 pub mesh: RawAssetHandle,
100 pub instance_count: u32,
101}
102
103#[derive(Default)]
113pub struct SkinnedBatchStorage {
114 pub batches: Vec<SkinnedBatchUnit>,
115}
116
117fn init_skinned_batching(
120 mut commands: Commands,
121 backend: Res<WGPUBackend>,
122 mut pool: ResMut<GlobalLayoutPool>,
123) -> Option<()> {
124 let renderer = SkinnedBatchRenderer::new(&backend);
125 pool.register("pebble_skinning", renderer.layout.clone());
126 commands.insert_resource(renderer);
127 commands.insert_resource(SkinnedBatchStorage::default());
128 Some(())
129}
130
131struct GroupData {
132 joint_count: u32,
133 matrices: Vec<glam::Mat4>, count: u32,
135}
136
137fn batch_skinned_entities(
138 backend: Res<WGPUBackend>,
139 mut renderer: ResMut<SkinnedBatchRenderer>,
140 mut storage: ResMut<SkinnedBatchStorage>,
141 mut query: Query<(&Handle<Material>, &Handle<SkinnedMesh>, &AnimationPlayer)>,
142 mut groups: Local<HashMap<(RawAssetHandle, RawAssetHandle), GroupData>>,
143) {
144 storage.batches.clear();
145 groups.clear();
146
147 for (mat, mesh, player) in query.iter() {
148 let entry = groups.entry((mat.id, mesh.id)).or_insert_with(|| GroupData {
149 joint_count: player.joint_count() as u32,
150 matrices: Vec::new(),
151 count: 0,
152 });
153 entry.matrices.extend(player.compute_matrices());
154 entry.count += 1;
155 }
156
157 for ((mat_id, mesh_id), group) in groups.iter() {
158 renderer.prepare((*mat_id, *mesh_id), group.joint_count, group.count, &backend);
159 let batch = renderer.batches.get(&(*mat_id, *mesh_id)).unwrap();
160 batch.matrices.write(bytemuck::cast_slice(&group.matrices));
161 storage.batches.push(SkinnedBatchUnit {
162 material: *mat_id,
163 mesh: *mesh_id,
164 instance_count: group.count,
165 });
166 }
167}
168
169pub struct SkinnedBatchingPlugin;
176
177impl Plugin for SkinnedBatchingPlugin {
178 fn build(&self, app: &mut App) {
179 app.add_system(SystemStage::Startup, init_skinned_batching);
180 app.add_system(SystemStage::PreRender, batch_skinned_entities);
181 }
182}