pebble/wgpu/
skinning.rs

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
21// ── internal ──────────────────────────────────────────────────────────────────
22
23struct SkinningBatch {
24    matrices:    Buffer,   // storage: capacity * joint_count * mat4
25    _info:       Buffer,   // uniform: SkinningInfo { joint_count: u32, _pad: [u32; 3] }
26    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
49// ── public ────────────────────────────────────────────────────────────────────
50
51/// Holds one GPU buffer pair per unique `(material, mesh)` batch — one per
52/// skeleton type in practice. Retrieve the bind group for a batch from your
53/// render system via [`bind_group`](Self::bind_group).
54///
55/// WGSL layout for `GroupEntry::Global("pebble_skinning")`:
56/// ```wgsl
57/// struct SkinningInfo { joint_count: u32 }
58/// @group(N) @binding(0) var<storage, read> joint_matrices: array<mat4x4<f32>>;
59/// @group(N) @binding(1) var<uniform>        skin_info:     SkinningInfo;
60/// // in vs_main: let base = instance_index * skin_info.joint_count;
61/// ```
62pub 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    /// The bind group for the `(material, mesh)` batch — pass to
78    /// `set_bind_group` immediately before `draw_indexed(0..index_count, 0, 0..instance_count)`.
79    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
94/// One entry per unique `(material, mesh)` pair — produced by
95/// [`batch_skinned_entities`] each `PreRender` tick. Iterate in your render
96/// system alongside [`SkinnedBatchRenderer`] to drive draw calls.
97pub struct SkinnedBatchUnit {
98    pub material:       RawAssetHandle,
99    pub mesh:           RawAssetHandle,
100    pub instance_count: u32,
101}
102
103/// The frame output of the skinning batch system.
104///
105/// ```ignore
106/// for batch in storage.batches.iter() {
107///     let Some(bind_group) = renderer.bind_group(batch.material, batch.mesh) else { continue };
108///     pass.set_bind_group(0, bind_group, &[]);
109///     pass.draw_indexed(0..mesh.index_count, 0, 0..batch.instance_count);
110/// }
111/// ```
112#[derive(Default)]
113pub struct SkinnedBatchStorage {
114    pub batches: Vec<SkinnedBatchUnit>,
115}
116
117// ── systems ───────────────────────────────────────────────────────────────────
118
119fn init_skinned_batching(
120    mut commands: Commands,
121    backend:      Option<Res<WGPUBackend>>,
122    mut pool:     ResMut<GlobalLayoutPool>,
123) -> Option<()> {
124    let backend = backend?;
125    let renderer = SkinnedBatchRenderer::new(&backend);
126    pool.register("pebble_skinning", renderer.layout.clone());
127    commands.insert_resource(renderer);
128    commands.insert_resource(SkinnedBatchStorage::default());
129    Some(())
130}
131
132struct GroupData {
133    joint_count: u32,
134    matrices:    Vec<glam::Mat4>, // flat: entity0 joints, entity1 joints, ...
135    count:       u32,
136}
137
138fn batch_skinned_entities(
139    backend:      Option<Res<WGPUBackend>>,
140    renderer:     Option<ResMut<SkinnedBatchRenderer>>,
141    storage:      Option<ResMut<SkinnedBatchStorage>>,
142    mut query:    Query<(&Handle<Material>, &Handle<SkinnedMesh>, &AnimationPlayer)>,
143    mut groups:   Local<HashMap<(RawAssetHandle, RawAssetHandle), GroupData>>,
144) {
145    let (Some(backend), Some(mut renderer), Some(mut storage)) = (backend, renderer, storage) else {
146        return;
147    };
148
149    storage.batches.clear();
150    groups.clear();
151
152    for (mat, mesh, player) in query.iter() {
153        let entry = groups.entry((mat.id, mesh.id)).or_insert_with(|| GroupData {
154            joint_count: player.joint_count() as u32,
155            matrices:    Vec::new(),
156            count:       0,
157        });
158        entry.matrices.extend(player.compute_matrices());
159        entry.count += 1;
160    }
161
162    for ((mat_id, mesh_id), group) in groups.iter() {
163        renderer.prepare((*mat_id, *mesh_id), group.joint_count, group.count, &backend);
164        let batch = renderer.batches.get(&(*mat_id, *mesh_id)).unwrap();
165        batch.matrices.write(bytemuck::cast_slice(&group.matrices));
166        storage.batches.push(SkinnedBatchUnit {
167            material:       *mat_id,
168            mesh:           *mesh_id,
169            instance_count: group.count,
170        });
171    }
172}
173
174/// Registers the `"pebble_skinning"` bind group layout, creates
175/// [`SkinnedBatchRenderer`] and [`SkinnedBatchStorage`] resources, and runs
176/// [`batch_skinned_entities`] in [`SystemStage::PreRender`].
177///
178/// The plugin does **not** advance animation time — add your own system that
179/// calls [`AnimationPlayer::advance`](super::player::AnimationPlayer::advance).
180pub struct SkinnedBatchingPlugin;
181
182impl Plugin for SkinnedBatchingPlugin {
183    fn build(&self, app: &mut App) {
184        app.add_system(SystemStage::Startup, init_skinned_batching);
185        app.add_system(SystemStage::PreRender, batch_skinned_entities);
186    }
187}
pebble/wgpu/skinning.rs

pebble/wgpu/
skinning.rs
