1use std::cell::{Cell, Ref, RefCell};
2use std::ops::Range;
3
4use wgpu::util::DeviceExt;
5
6use crate::frame::Matrix;
7use crate::ranges::FreeRanges;
8
9pub const NO_SKIN: u32 = u32::MAX;
10pub const SKIN_BINDING: u32 = 4;
11
12pub const SKIN_WGSL: &str = r#"
13@group(0) @binding(4) var<storage, read> skin_data: array<vec4u>;
14fn skin_joint(palette: u32, joint: u32) -> mat4x4f {
15 let at = palette + joint * 4u;
16 return mat4x4f(
17 bitcast<vec4f>(skin_data[at]),
18 bitcast<vec4f>(skin_data[at + 1u]),
19 bitcast<vec4f>(skin_data[at + 2u]),
20 bitcast<vec4f>(skin_data[at + 3u]),
21 );
22}
23fn skin_blend(skin: vec4u, vertex: u32, palette: u32) -> mat4x4f {
24 let record = skin.x + vertex * 2u;
25 let packed = skin_data[record];
26 let weights = bitcast<vec4f>(skin_data[record + 1u]);
27 var blended = skin_joint(palette, packed.x & 0xffffu) * weights.x;
28 blended = blended + skin_joint(palette, packed.x >> 16u) * weights.y;
29 blended = blended + skin_joint(palette, packed.y & 0xffffu) * weights.z;
30 blended = blended + skin_joint(palette, packed.y >> 16u) * weights.w;
31 return blended;
32}
33fn skin_unit(v: vec3f) -> vec3f {
34 let length = sqrt(dot(v, v));
35 return select(v, v / length, length > 0.0);
36}
37fn skin_position(position: vec3f, vertex: u32, skin: vec4u) -> vec3f {
38 return (skin_blend(skin, vertex, skin.y) * vec4f(position, 1.0)).xyz;
39}
40fn skin_shape(position: vec3f, normal: vec3f, tangent: vec4f, vertex: u32, skin: vec4u) -> DeformedVertex {
41 let current = skin_blend(skin, vertex, skin.y);
42 let previous = skin_blend(skin, vertex, skin.z);
43 var out: DeformedVertex;
44 out.position = (current * vec4f(position, 1.0)).xyz;
45 out.normal = skin_unit((current * vec4f(normal, 0.0)).xyz);
46 out.tangent = vec4f(skin_unit((current * vec4f(tangent.xyz, 0.0)).xyz), tangent.w);
47 out.previous_position = (previous * vec4f(position, 1.0)).xyz;
48 return out;
49}
50"#;
51
52#[derive(Clone, Copy, Debug)]
53pub struct SkinData<'a> {
54 pub joints: &'a [[u16; 4]],
55 pub weights: &'a [[f32; 4]],
56}
57
58impl SkinData<'_> {
59 pub fn palette(&self) -> usize {
60 self.joints
61 .iter()
62 .zip(self.weights)
63 .flat_map(|(joints, weights)| {
64 joints
65 .iter()
66 .zip(weights)
67 .filter(|(_, weight)| **weight != 0.0)
68 .map(|(joint, _)| *joint as usize + 1)
69 })
70 .max()
71 .unwrap_or(1)
72 }
73}
74
75#[derive(Clone, Copy, Debug)]
76pub struct InstancePose<'a> {
77 pub instance: usize,
78 pub joints: &'a [Matrix],
79 pub previous: &'a [Matrix],
80}
81
82#[derive(Clone, Debug, PartialEq)]
83struct StoredPose {
84 instance: usize,
85 joints: Vec<Matrix>,
86 previous: Vec<Matrix>,
87}
88
89#[derive(Clone, Copy, Debug, PartialEq, Eq)]
90pub(crate) struct Placed {
91 pub current: u32,
92 pub previous: u32,
93 pub joints: u32,
94}
95
96pub(crate) struct PoseLayout {
97 pub placed: Vec<(usize, Placed)>,
98 pub words: Vec<[u32; 4]>,
99 pub base: u32,
100}
101
102impl PoseLayout {
103 pub fn get(&self, instance: usize) -> Option<Placed> {
104 self.placed
105 .binary_search_by_key(&instance, |(index, _)| *index)
106 .ok()
107 .map(|at| self.placed[at].1)
108 }
109}
110
111pub(crate) struct SkinStore {
112 records: Vec<[u32; 4]>,
113 free: FreeRanges,
114 palette_capacity: usize,
115 buffer: wgpu::Buffer,
116 poses: RefCell<Vec<StoredPose>>,
117 layout: RefCell<PoseLayout>,
118 pending: Cell<bool>,
119}
120
121fn words(matrices: &[Matrix]) -> impl Iterator<Item = [u32; 4]> + '_ {
122 matrices
123 .iter()
124 .flat_map(|matrix| matrix.iter().map(|column| column.map(f32::to_bits)))
125}
126
127fn layout(base: u32, poses: &[StoredPose]) -> PoseLayout {
128 let mut words_out = Vec::new();
129 let mut placed = Vec::with_capacity(poses.len());
130 for pose in poses {
131 let current = base + words_out.len() as u32;
132 words_out.extend(words(&pose.joints));
133 let previous = base + words_out.len() as u32;
134 words_out.extend(words(&pose.previous));
135 placed.push((
136 pose.instance,
137 Placed {
138 current,
139 previous,
140 joints: pose.joints.len() as u32,
141 },
142 ));
143 }
144 PoseLayout {
145 placed,
146 words: words_out,
147 base,
148 }
149}
150
151impl SkinStore {
152 pub fn new(device: &wgpu::Device) -> Self {
153 let palette_capacity = 4;
154 Self {
155 records: Vec::new(),
156 free: FreeRanges::default(),
157 palette_capacity,
158 buffer: skin_buffer(device, &[], palette_capacity),
159 poses: RefCell::new(Vec::new()),
160 layout: RefCell::new(layout(0, &[])),
161 pending: Cell::new(true),
162 }
163 }
164
165 pub fn buffer(&self) -> &wgpu::Buffer {
166 &self.buffer
167 }
168
169 fn rebuild(&mut self, device: &wgpu::Device) {
170 self.buffer = skin_buffer(device, &self.records, self.palette_capacity);
171 self.pending.set(true);
172 }
173
174 pub fn allocate(
175 &mut self,
176 device: &wgpu::Device,
177 skin: SkinData<'_>,
178 ) -> Result<Range<u32>, String> {
179 let words: Vec<[u32; 4]> = skin
180 .joints
181 .iter()
182 .zip(skin.weights)
183 .flat_map(|(joints, weights)| {
184 [
185 [
186 u32::from(joints[0]) | u32::from(joints[1]) << 16,
187 u32::from(joints[2]) | u32::from(joints[3]) << 16,
188 0,
189 0,
190 ],
191 weights.map(f32::to_bits),
192 ]
193 })
194 .collect();
195 let base = self
196 .free
197 .take(&mut self.records, words.len(), [0; 4], "skinned vertices")?;
198 self.records[base as usize..base as usize + words.len()].copy_from_slice(&words);
199 *self.layout.get_mut() = layout(self.records.len() as u32, &self.poses.borrow());
200 self.rebuild(device);
201 Ok(base..base + words.len() as u32)
202 }
203
204 pub fn release(&mut self, range: Range<u32>) {
205 self.free.give(range);
206 }
207
208 pub fn set_poses(&self, poses: &[InstancePose<'_>]) -> Result<(), String> {
209 let mut stored: Vec<StoredPose> = Vec::with_capacity(poses.len());
210 for pose in poses {
211 if pose.joints.is_empty() || pose.joints.len() != pose.previous.len() {
212 return Err(format!(
213 "the pose of instance {} needs joints and as many previous joints",
214 pose.instance
215 ));
216 }
217 if pose.joints.len() > pfx_core::anim::MAX_JOINTS {
218 return Err(format!(
219 "the pose of instance {} has {} joints, over {}",
220 pose.instance,
221 pose.joints.len(),
222 pfx_core::anim::MAX_JOINTS
223 ));
224 }
225 if !pose
226 .joints
227 .iter()
228 .chain(pose.previous)
229 .flatten()
230 .flatten()
231 .all(|value| value.is_finite())
232 {
233 return Err(format!(
234 "the pose of instance {} is not finite",
235 pose.instance
236 ));
237 }
238 stored.push(StoredPose {
239 instance: pose.instance,
240 joints: pose.joints.to_vec(),
241 previous: pose.previous.to_vec(),
242 });
243 }
244 stored.sort_by_key(|pose| pose.instance);
245 if stored
246 .windows(2)
247 .any(|pair| pair[0].instance == pair[1].instance)
248 {
249 return Err("two poses name the same instance".into());
250 }
251 *self.layout.borrow_mut() = layout(self.records.len() as u32, &stored);
252 *self.poses.borrow_mut() = stored;
253 self.pending.set(true);
254 Ok(())
255 }
256
257 pub fn same(&self, poses: &[InstancePose<'_>]) -> bool {
258 let stored = self.poses.borrow();
259 stored.len() == poses.len()
260 && poses.iter().all(|pose| {
261 stored
262 .binary_search_by_key(&pose.instance, |stored| stored.instance)
263 .is_ok_and(|at| {
264 stored[at].joints == pose.joints && stored[at].previous == pose.previous
265 })
266 })
267 }
268
269 pub fn layout(&self) -> Ref<'_, PoseLayout> {
270 self.layout.borrow()
271 }
272
273 pub fn upload(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) -> bool {
274 if !self.pending.replace(false) {
275 return false;
276 }
277 let needed = self.layout.get_mut().words.len();
278 let mut rebuilt = false;
279 if needed > self.palette_capacity {
280 self.palette_capacity = needed.next_power_of_two();
281 self.rebuild(device);
282 self.pending.set(false);
283 rebuilt = true;
284 }
285 let layout = self.layout.get_mut();
286 if !layout.words.is_empty() {
287 queue.write_buffer(
288 &self.buffer,
289 u64::from(layout.base) * 16,
290 bytemuck::cast_slice(&layout.words),
291 );
292 }
293 rebuilt
294 }
295}
296
297fn skin_buffer(device: &wgpu::Device, records: &[[u32; 4]], palette: usize) -> wgpu::Buffer {
298 let mut contents = records.to_vec();
299 contents.resize(records.len() + palette.max(1), [0; 4]);
300 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
301 label: Some("live skins"),
302 contents: bytemuck::cast_slice(&contents),
303 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
304 })
305}