viewport-lib 0.17.0

//! GPU skinning sidecar storage.
//!
//! Holds the per-mesh skin-weight storage buffers and the per-(mesh, instance)
//! joint palette storage buffers used by the skinned pipeline variants.
//!
//! Static meshes pay zero overhead: nothing is allocated until
//! [`crate::ViewportGpuResources::set_skin_weights`] is called for a mesh,
//! which is what marks that mesh as skinnable.
//!
//! See `docs/plans/skeletal-animation-plan.md`.

use std::collections::HashMap;

use wgpu::util::DeviceExt;

use crate::resources::SkinWeights;
use crate::resources::ViewportGpuResources;
use crate::resources::mesh_store::MeshId;

/// Packed per-vertex skin data uploaded to the GPU sidecar storage buffer.
///
/// Layout matches the `SkinVertex` WGSL struct in `mesh_skinned.wgsl`:
/// `weights` first (vec4 aligned to 16), then the two packed joint-index
/// u32s, then 8 bytes of trailing padding. WGSL rounds the struct up to its
/// alignment (16) for array stride, so the stride is 32 bytes even though
/// only 24 bytes carry data. The padding makes Rust agree.
///
/// Total: 32 bytes per vertex (24 bytes data + 8 bytes std430 padding).
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub(crate) struct PackedSkinVertex {
    pub weights: [f32; 4],
    /// `joints[0]` in the low 16 bits, `joints[1]` in the high 16 bits.
    pub joints_01: u32,
    /// `joints[2]` in the low 16 bits, `joints[3]` in the high 16 bits.
    pub joints_23: u32,
    /// Trailing padding so the struct stride matches WGSL std430.
    pub _pad: [u32; 2],
}

/// Per-instance joint palette: storage buffer plus its bind group.
pub(crate) struct InstancePalette {
    /// Storage buffer of `mat4x4<f32>` joint matrices.
    pub buffer: wgpu::Buffer,
    /// Bind group binding both the mesh's skin weights and this instance's
    /// palette buffer.
    pub bind_group: wgpu::BindGroup,
    /// Number of `Mat4` slots allocated in `buffer`. Used to decide when a
    /// realloc is required.
    pub joint_capacity: u32,
}

/// Per-mesh skinning data: the weights storage buffer plus a map of instance
/// palettes keyed by `instance_id`.
pub(crate) struct MeshSkinning {
    pub weights_buffer: wgpu::Buffer,
    pub instances: HashMap<u32, InstancePalette>,
}

/// Renderer-side skinning state.
///
/// Owns the bind group layout used by skinned pipeline variants and a per-mesh
/// map of skinning sidecars.
pub(crate) struct SkinningState {
    pub bind_group_layout: wgpu::BindGroupLayout,
    pub meshes: HashMap<MeshId, MeshSkinning>,
}

impl SkinningState {
    pub fn new(device: &wgpu::Device) -> Self {
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("skin_bgl"),
            entries: &[
                // binding 0: skin weights storage buffer (per-vertex joint indices + weights)
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 1: joint palette storage buffer (mat4x4<f32> per joint)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });
        Self {
            bind_group_layout,
            meshes: HashMap::new(),
        }
    }

    pub(crate) fn pack(weights: &SkinWeights) -> Vec<PackedSkinVertex> {
        weights
            .joint_indices
            .iter()
            .zip(weights.joint_weights.iter())
            .map(|(j, w)| {
                let j0 = j[0] as u32;
                let j1 = j[1] as u32;
                let j2 = j[2] as u32;
                let j3 = j[3] as u32;
                PackedSkinVertex {
                    weights: *w,
                    joints_01: j0 | (j1 << 16),
                    joints_23: j2 | (j3 << 16),
                    _pad: [0, 0],
                }
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

impl ViewportGpuResources {
    /// Always `true`. GPU skinning ships unconditionally as part of the
    /// renderer's capability set.
    ///
    /// Plugins read this to decide whether to emit
    /// [`crate::SkinnedPoseUpdate`] (GPU path) or
    /// [`crate::SkinnedMeshUpdate`] (CPU path) each frame. The flag stays in
    /// the API so a future build that intentionally drops the skinned
    /// pipelines (memory-constrained target, headless test mode) can return
    /// `false` here without changing plugin code.
    pub fn supports_gpu_skinning(&self) -> bool {
        true
    }

    /// Attach per-vertex skin weights to an uploaded mesh.
    ///
    /// Creates a sidecar storage buffer holding packed joint indices + weights
    /// (24 bytes per vertex) and the bind group used by the skinned pipeline
    /// variants. The mesh's vertex buffer is not modified.
    ///
    /// Calling this on a `mesh_id` marks the mesh as skinnable: subsequent
    /// draws of any scene node referencing this mesh are routed through the
    /// skinned pipeline variant once a joint palette has been uploaded via
    /// [`Self::set_skin_palette`].
    ///
    /// Calling this again on the same `mesh_id` replaces the weights buffer
    /// and invalidates all previously created instance bind groups; consumers
    /// should re-upload palettes for any active instances afterwards.
    pub fn set_skin_weights(
        &mut self,
        device: &wgpu::Device,
        mesh_id: MeshId,
        weights: &SkinWeights,
    ) {
        let packed = SkinningState::pack(weights);
        self.install_skin_weights(device, mesh_id, &packed);
    }

    /// Start an asynchronous skin-weights upload.
    ///
    /// Returns a `JobId` immediately. The packed weight stream (joint
    /// indices and blend weights, four per vertex) is computed on a worker
    /// thread; buffer creation and the `mesh.skinning` insert run on the
    /// main thread during the next `process_uploads` call after the worker
    /// finishes.
    ///
    /// `weights` transfers into the worker; clone at the call site to
    /// retain ownership.
    ///
    /// As with the synchronous `set_skin_weights`, this replaces any prior
    /// weights for the same `mesh_id` and invalidates active per-instance
    /// bind groups. Re-upload palettes via `set_skin_palette` afterwards.
    pub fn begin_upload_skin_weights(
        &mut self,
        device: &wgpu::Device,
        mesh_id: MeshId,
        weights: SkinWeights,
    ) -> crate::resources::JobId {
        let device_for_apply = device.clone();
        let mut runner = self.jobs.lock().expect("upload job runner poisoned");
        runner.submit_cpu(move |progress| {
            progress.set(0.2);
            let packed = SkinningState::pack(&weights);
            progress.set(0.9);
            Ok(crate::resources::upload_jobs::JobProduct::with_apply(
                Box::new(move |resources: &mut ViewportGpuResources| {
                    resources.install_skin_weights(&device_for_apply, mesh_id, &packed);
                }),
            ))
        })
    }

    /// Shared apply step for both `set_skin_weights` and
    /// `begin_upload_skin_weights`: build the storage buffer and register
    /// the mesh as skinnable.
    fn install_skin_weights(
        &mut self,
        device: &wgpu::Device,
        mesh_id: MeshId,
        packed: &[PackedSkinVertex],
    ) {
        let weights_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("skin_weights_buffer"),
            contents: bytemuck::cast_slice(packed),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });
        self.skinning.meshes.insert(
            mesh_id,
            MeshSkinning {
                weights_buffer,
                instances: HashMap::new(),
            },
        );
    }

    /// Upload the joint palette for one instance of a skinned mesh.
    ///
    /// `instance_id` lets multiple skinned instances of one bind-pose mesh
    /// coexist (the crowd case). For single-instance meshes pass `0`.
    ///
    /// Allocates or grows the per-instance palette storage buffer as needed.
    /// Re-binds the per-instance bind group when the buffer is reallocated.
    /// `set_skin_weights` must have been called for `mesh_id` first.
    ///
    /// `palette[i]` is `world_transform[i] * inverse_bind[i]` for joint `i`,
    /// the LBS-ready matrix produced by
    /// [`crate::JointMatrices::compute`].
    pub fn set_skin_palette(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        mesh_id: MeshId,
        instance_id: u32,
        palette: &[glam::Mat4],
    ) -> bool {
        let bgl = &self.skinning.bind_group_layout;
        let mesh = match self.skinning.meshes.get_mut(&mesh_id) {
            Some(m) => m,
            None => return false,
        };

        let joints_needed = palette.len() as u32;
        let needs_realloc = match mesh.instances.get(&instance_id) {
            Some(inst) => inst.joint_capacity < joints_needed,
            None => true,
        };

        if needs_realloc {
            let capacity = joints_needed.max(1);
            let buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("skin_palette_buffer"),
                size: (capacity as u64) * std::mem::size_of::<[[f32; 4]; 4]>() as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("skin_bind_group"),
                layout: bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: mesh.weights_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: buffer.as_entire_binding(),
                    },
                ],
            });
            mesh.instances.insert(
                instance_id,
                InstancePalette {
                    buffer,
                    bind_group,
                    joint_capacity: capacity,
                },
            );
        }

        let inst = mesh.instances.get(&instance_id).unwrap();
        let bytes: Vec<[[f32; 4]; 4]> = palette.iter().map(|m| m.to_cols_array_2d()).collect();
        queue.write_buffer(&inst.buffer, 0, bytemuck::cast_slice(&bytes));
        true
    }

    /// Whether `mesh_id` has been marked as skinnable via
    /// [`Self::set_skin_weights`].
    pub fn is_skinned_mesh(&self, mesh_id: MeshId) -> bool {
        self.skinning.meshes.contains_key(&mesh_id)
    }

    /// Whether the given `(mesh_id, instance_id)` has a palette uploaded and
    /// is ready to be drawn through the skinned pipeline.
    pub(crate) fn skin_instance_bind_group(
        &self,
        mesh_id: MeshId,
        instance_id: u32,
    ) -> Option<&wgpu::BindGroup> {
        self.skinning
            .meshes
            .get(&mesh_id)?
            .instances
            .get(&instance_id)
            .map(|p| &p.bind_group)
    }
}

#[cfg(test)]
mod async_skin_tests {
    use crate::ViewportGpuResources;
    use crate::geometry::primitives;
    use crate::resources::{SkinWeights, UploadStatus};

    fn try_make_device() -> Option<(wgpu::Device, wgpu::Queue)> {
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::LowPower,
            compatible_surface: None,
            force_fallback_adapter: false,
        }))
        .ok()?;
        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()
    }

    fn unit_weights(vertex_count: usize) -> SkinWeights {
        SkinWeights {
            joint_indices: vec![[0u8; 4]; vertex_count],
            joint_weights: vec![[1.0, 0.0, 0.0, 0.0]; vertex_count],
        }
    }

    #[test]
    fn sync_set_skin_weights_marks_mesh_skinnable() {
        let Some((device, _queue)) = try_make_device() else {
            eprintln!("skipping: no wgpu adapter available");
            return;
        };
        let mut resources =
            ViewportGpuResources::new(&device, wgpu::TextureFormat::Rgba8UnormSrgb, 1);
        let plane = primitives::grid_plane(1.0, 1.0, 4, 4);
        let mesh_id = resources.upload_mesh_data(&device, &plane).unwrap();
        let weights = unit_weights(plane.positions.len());

        resources.set_skin_weights(&device, mesh_id, &weights);
        assert!(resources.is_skinned_mesh(mesh_id));
    }

    #[test]
    fn begin_upload_skin_weights_completes() {
        let Some((device, queue)) = try_make_device() else {
            eprintln!("skipping: no wgpu adapter available");
            return;
        };
        let mut resources =
            ViewportGpuResources::new(&device, wgpu::TextureFormat::Rgba8UnormSrgb, 1);
        let plane = primitives::grid_plane(1.0, 1.0, 4, 4);
        let mesh_id = resources.upload_mesh_data(&device, &plane).unwrap();
        let weights = unit_weights(plane.positions.len());

        assert!(!resources.is_skinned_mesh(mesh_id));
        let id = resources.begin_upload_skin_weights(&device, mesh_id, weights);

        for _ in 0..200 {
            resources.process_uploads(&device, &queue);
            match resources.upload_status(id) {
                UploadStatus::Ready => break,
                UploadStatus::Failed(e) => panic!("upload failed: {e:?}"),
                UploadStatus::Pending { .. } => {
                    std::thread::sleep(std::time::Duration::from_millis(5));
                }
                UploadStatus::Unknown => panic!("job disappeared"),
            }
        }

        assert!(
            resources.is_skinned_mesh(mesh_id),
            "skin weights did not install"
        );
    }
}