mesh_bindings.rs

//! Bind group layout related definitions for the mesh pipeline.

use bevy_asset::AssetId;
use bevy_ecs::prelude::Resource;
use bevy_math::Mat4;
use bevy_render::mesh::{morph::MAX_MORPH_WEIGHTS, Mesh};
use bevy_render::{render_resource::*, renderer::RenderDevice};
use bevy_utils::HashMap;
use smallvec::SmallVec;

use crate::render::skin::MAX_JOINTS;
use crate::MeshUniform;

const MORPH_WEIGHT_SIZE: usize = std::mem::size_of::<f32>();
pub const MORPH_BUFFER_SIZE: usize = MAX_MORPH_WEIGHTS * MORPH_WEIGHT_SIZE;

const JOINT_SIZE: usize = std::mem::size_of::<Mat4>();
pub(crate) const JOINT_BUFFER_SIZE: usize = MAX_JOINTS * JOINT_SIZE;

// --- Individual layout entries ---

fn buffer_layout(binding: u32, size: u64, visibility: ShaderStages) -> BindGroupLayoutEntry {
    BindGroupLayoutEntry {
        binding,
        visibility,
        count: None,
        ty: BindingType::Buffer {
            ty: BufferBindingType::Uniform,
            has_dynamic_offset: true,
            min_binding_size: BufferSize::new(size),
        },
    }
}
fn model_layout(render_device: &RenderDevice, binding: u32) -> BindGroupLayoutEntry {
    GpuArrayBuffer::<MeshUniform>::binding_layout(
        binding,
        ShaderStages::VERTEX_FRAGMENT,
        render_device,
    )
}
fn skinning_layout(binding: u32) -> BindGroupLayoutEntry {
    buffer_layout(binding, JOINT_BUFFER_SIZE as u64, ShaderStages::VERTEX)
}
fn weights_layout(binding: u32) -> BindGroupLayoutEntry {
    buffer_layout(binding, MORPH_BUFFER_SIZE as u64, ShaderStages::VERTEX)
}
fn targets_layout(binding: u32) -> BindGroupLayoutEntry {
    BindGroupLayoutEntry {
        binding,
        visibility: ShaderStages::VERTEX,
        ty: BindingType::Texture {
            view_dimension: TextureViewDimension::D3,
            sample_type: TextureSampleType::Float { filterable: false },
            multisampled: false,
        },
        count: None,
    }
}

// --- Individual bind group entries ---

fn entry(binding: u32, size: u64, buffer: &Buffer) -> BindGroupEntry {
    BindGroupEntry {
        binding,
        resource: BindingResource::Buffer(BufferBinding {
            buffer,
            offset: 0,
            size: Some(BufferSize::new(size).unwrap()),
        }),
    }
}
fn model_entry(binding: u32, resource: BindingResource) -> BindGroupEntry {
    BindGroupEntry { binding, resource }
}
fn skinning_entry(binding: u32, buffer: &Buffer) -> BindGroupEntry {
    entry(binding, JOINT_BUFFER_SIZE as u64, buffer)
}
fn weights_entry(binding: u32, buffer: &Buffer) -> BindGroupEntry {
    entry(binding, MORPH_BUFFER_SIZE as u64, buffer)
}
fn targets_entry(binding: u32, texture: &TextureView) -> BindGroupEntry {
    BindGroupEntry {
        binding,
        resource: BindingResource::TextureView(texture),
    }
}

// --- Layout entry lists & bind group entry lists ---

const MOTION_VECTORS_OFFSET: u32 = 3;

fn model_layout_entries(device: &RenderDevice) -> [BindGroupLayoutEntry; 1] {
    [model_layout(device, 0)]
}
fn model_entries(model: BindingResource) -> [BindGroupEntry; 1] {
    [model_entry(0, model)]
}

fn morph_layout_entries() -> [BindGroupLayoutEntry; 2] {
    [weights_layout(2), targets_layout(3)]
}
fn morph_entries<'a>(weights: &'a Buffer, targets: &'a TextureView) -> [BindGroupEntry<'a>; 2] {
    [weights_entry(2, weights), targets_entry(3, targets)]
}

fn skin_layout_entries() -> [BindGroupLayoutEntry; 1] {
    [skinning_layout(1)]
}
fn skin_entries(skin: &Buffer) -> [BindGroupEntry; 1] {
    [skinning_entry(1, skin)]
}

fn motion_vectors_morph_layout_entries() -> [BindGroupLayoutEntry; 1] {
    [weights_layout(2 + MOTION_VECTORS_OFFSET)]
}
fn motion_vectors_morph_entries(weights: &Buffer) -> [BindGroupEntry; 1] {
    [weights_entry(2 + MOTION_VECTORS_OFFSET, weights)]
}

fn motion_vectors_skin_layout_entries() -> [BindGroupLayoutEntry; 1] {
    [skinning_layout(1 + MOTION_VECTORS_OFFSET)]
}
fn motion_vectors_skin_entries(skin: &Buffer) -> [BindGroupEntry; 1] {
    [skinning_entry(1 + MOTION_VECTORS_OFFSET, skin)]
}

/// Create a layout for a specific [`ActiveVariant`] by combining the layouts of
/// all the active shader features.
fn variant_layout(active_variant: &ActiveVariant, device: &RenderDevice) -> BindGroupLayout {
    let mut layout_entries = SmallVec::<[BindGroupLayoutEntry; 6]>::new();

    layout_entries.extend(model_layout_entries(device));

    if active_variant.morph {
        layout_entries.extend(morph_layout_entries());
    }
    if active_variant.skin {
        layout_entries.extend(skin_layout_entries());
    }
    if active_variant.morph && active_variant.motion_vectors {
        layout_entries.extend(motion_vectors_morph_layout_entries());
    }
    if active_variant.skin && active_variant.motion_vectors {
        layout_entries.extend(motion_vectors_skin_layout_entries());
    }

    device.create_bind_group_layout(&BindGroupLayoutDescriptor {
        entries: &layout_entries,
        label: Some(LAYOUT_LABELS.get_variant(*active_variant)),
    })
}

/// Create [`BindGroup`]s and add them to a [`MeshBindGroups`].
///
/// Use [`MeshBindGroups::bind_group_builder`] to get a [`MeshBindGroupBuilder`],
/// and use [`MeshBindGroupBuilder::add_variant`] to add new bind group variants.
pub struct MeshBindGroupBuilder<'a> {
    layouts: &'a MeshLayouts,
    device: &'a RenderDevice,
    model: BindingResource<'a>,
    pub bind_groups: &'a mut MeshBindGroups,
}
impl<'a> MeshBindGroupBuilder<'a> {
    /// Create a [`BindGroup`] with the provided active features, and add it to [`Self::bind_groups`].
    ///
    /// Each parameter to `add_variant` is a shader feature. When the parameter
    /// is `None`, it means that the shader feature is disabled.
    pub fn add_variant(
        &mut self,
        morph: Option<(AssetId<Mesh>, &Buffer, &TextureView)>,
        skin: Option<&Buffer>,
        previous_weights: Option<&Buffer>,
        previous_skin: Option<&Buffer>,
    ) {
        let mut bind_group_entries = SmallVec::<[BindGroupEntry; 6]>::new();

        bind_group_entries.extend(model_entries(self.model.clone()));

        if let Some((_, weights, targets)) = morph {
            bind_group_entries.extend(morph_entries(weights, targets));
        }
        if let Some(skin) = skin {
            bind_group_entries.extend(skin_entries(skin));
        }
        if let Some(previous_weights) = previous_weights {
            bind_group_entries.extend(motion_vectors_morph_entries(previous_weights));
        }
        if let Some(previous_skin) = previous_skin {
            bind_group_entries.extend(motion_vectors_skin_entries(previous_skin));
        }

        let skin = skin.is_some();
        let motion_vectors = previous_weights.is_some() || previous_skin.is_some();

        let active_variant = ActiveVariant::new(morph.is_some(), skin, motion_vectors);

        let bind_group = self.device.create_bind_group(&BindGroupDescriptor {
            entries: &bind_group_entries,
            label: Some(BIND_GROUP_LABELS.get_variant(active_variant)),
            layout: self.layouts.0.get_variant(active_variant),
        });
        self.bind_groups
            .insert(morph.map(|m| m.0), skin, motion_vectors, bind_group);
    }
}

#[derive(Clone)]
struct VariantsData<T> {
    model_only: T,
    morphed: T,
    skinned: T,
    morphed_skinned: T,
    morphed_motion_vectors: T,
    skinned_motion_vectors: T,
    morphed_skinned_motion_vectors: T,
}
impl<T> VariantsData<T> {
    fn get_variant(&self, variant: ActiveVariant) -> &T {
        match (variant.morph, variant.skin, variant.motion_vectors) {
            (false, false, _) => &self.model_only,
            (true, false, false) => &self.morphed,
            (false, true, false) => &self.skinned,
            (true, true, false) => &self.morphed_skinned,
            (true, false, true) => &self.morphed_motion_vectors,
            (false, true, true) => &self.skinned_motion_vectors,
            (true, true, true) => &self.morphed_skinned_motion_vectors,
        }
    }
}
const ACTIVE_VARIANTS: VariantsData<ActiveVariant> = VariantsData {
    model_only: ActiveVariant::new(false, false, false),
    morphed: ActiveVariant::new(true, false, false),
    skinned: ActiveVariant::new(false, true, false),
    morphed_skinned: ActiveVariant::new(true, true, false),
    morphed_motion_vectors: ActiveVariant::new(true, false, true),
    skinned_motion_vectors: ActiveVariant::new(false, true, true),
    morphed_skinned_motion_vectors: ActiveVariant::new(true, true, true),
};
const BIND_GROUP_LABELS: VariantsData<&'static str> = VariantsData {
    model_only: "mesh_bind_group",
    morphed: "morphed_mesh_bind_group",
    skinned: "skinned_mesh_bind_group",
    morphed_skinned: "morphed_skinned_mesh_bind_group",
    morphed_motion_vectors: "morphed_motion_vectors_mesh_bind_group",
    skinned_motion_vectors: "skinned_motion_vectors_mesh_bind_group",
    morphed_skinned_motion_vectors: "morphed_skinned_motion_vectors_mesh_bind_group",
};
const LAYOUT_LABELS: VariantsData<&'static str> = VariantsData {
    model_only: "mesh_layout",
    morphed: "morphed_mesh_layout",
    skinned: "skinned_mesh_layout",
    morphed_skinned: "morphed_skinned_mesh_layout",
    morphed_motion_vectors: "morphed_motion_vectors_mesh_layout",
    skinned_motion_vectors: "skinned_motion_vectors_mesh_layout",
    morphed_skinned_motion_vectors: "morphed_skinned_motion_vectors_mesh_layout",
};

/// The [`BindGroup`]s for individual existing mesh shader variants.
///
/// Morph targets allow several different bind groups, because individual mesh
/// may have a different [`TextureView`] that represents the morph target's pose
/// vertex attribute values.
///
/// Non-morph target bind groups are optional. We don't know at compile time
/// whether motion vectors or skinned meshes will be used.
#[derive(Default, Resource)]
pub struct MeshBindGroups {
    pub model_only: Option<BindGroup>,
    pub morphed: HashMap<AssetId<Mesh>, BindGroup>,
    pub skinned: Option<BindGroup>,
    pub morphed_motion_vectors: HashMap<AssetId<Mesh>, BindGroup>,
    pub skinned_motion_vectors: Option<BindGroup>,
}

impl MeshBindGroups {
    pub fn new() -> Self {
        MeshBindGroups::default()
    }
    /// Get a specific [`BindGroup`] that was previously added.
    pub fn get(
        &self,
        morph_id: Option<AssetId<Mesh>>,
        skin: bool,
        motion_vectors: bool,
    ) -> Option<&BindGroup> {
        match (morph_id, skin, motion_vectors) {
            (Some(id), _, true) => self.morphed_motion_vectors.get(&id),
            (Some(id), _, false) => self.morphed.get(&id),
            (None, false, _) => self.model_only.as_ref(),
            (None, true, true) => self.skinned_motion_vectors.as_ref(),
            (None, true, false) => self.skinned.as_ref(),
        }
    }
    fn insert(
        &mut self,
        morph_id: Option<AssetId<Mesh>>,
        skin: bool,
        motion_vectors: bool,
        bind_group: BindGroup,
    ) {
        let insert = |map: &mut HashMap<_, _>, id, bind_group| {
            map.insert(id, bind_group);
        };
        match (morph_id, skin, motion_vectors) {
            (Some(id), _, true) => insert(&mut self.morphed_motion_vectors, id, bind_group),
            (Some(id), _, false) => insert(&mut self.morphed, id, bind_group),
            (None, false, _) => self.model_only = Some(bind_group),
            (None, true, true) => self.skinned_motion_vectors = Some(bind_group),
            (None, true, false) => self.skinned = Some(bind_group),
        }
    }
    pub fn clear(&mut self) {
        self.model_only = None;
        self.morphed.clear();
        self.morphed_motion_vectors.clear();
        self.skinned = None;
        self.skinned_motion_vectors = None;
    }
    /// Clears `self` and returns a [`MeshBindGroupBuilder`].
    pub fn bind_group_builder<'a>(
        &'a mut self,
        device: &'a RenderDevice,
        model: BindingResource<'a>,
        layouts: &'a MeshLayouts,
    ) -> MeshBindGroupBuilder<'a> {
        self.clear();

        MeshBindGroupBuilder {
            layouts,
            device,
            model,
            bind_groups: self,
        }
    }
}

/// All possible [`BindGroupLayout`]s in bevy's default mesh shader (`mesh.wgsl`).
#[derive(Clone)]
pub struct MeshLayouts(VariantsData<BindGroupLayout>);

impl MeshLayouts {
    /// Prepare the layouts used by the default bevy [`Mesh`].
    ///
    /// [`Mesh`]: bevy_render::prelude::Mesh
    pub fn new(device: &RenderDevice) -> Self {
        let layout = |variant| variant_layout(&variant, device);
        let variants = ACTIVE_VARIANTS;

        MeshLayouts(VariantsData {
            model_only: layout(variants.model_only),
            morphed: layout(variants.morphed),
            skinned: layout(variants.skinned),
            morphed_skinned: layout(variants.morphed_skinned),
            morphed_motion_vectors: layout(variants.morphed_motion_vectors),
            skinned_motion_vectors: layout(variants.skinned_motion_vectors),
            morphed_skinned_motion_vectors: layout(variants.morphed_skinned_motion_vectors),
        })
    }
    pub fn get(&self, morph: bool, skin: bool, motion_vectors: bool) -> &BindGroupLayout {
        let active_variant = ActiveVariant::new(morph, skin, motion_vectors);
        self.0.get_variant(active_variant)
    }
}

/// The set of active features for a given mesh shader instance.
///
/// Individual meshes may have different features. For example,
/// one mesh may have morph targets, another skinning.
///
/// Each of those features enable different shader code and bind groups through
/// the naga C-like pre-processor (`#ifdef FOO`).
///
/// As a result, different meshes may use different variants of the shader and need
/// different bind group layouts.
#[derive(Default, Debug, Clone, Copy)]
pub struct ActiveVariant {
    /// Whether this mesh uses morph targets.
    ///
    /// This is determined by the `MeshPipelineKey::MORPH_TARGETS` pipeline key
    /// flag.
    ///
    /// This pipeline key flag is set whenever the mesh being rendered has the
    /// `morph_targets` field set to `true`.
    pub morph: bool,

    /// Whether this mesh uses skeletal skinning.
    ///
    /// It is true whenever the mesh being rendered has both `Mesh::ATTRIBUTE_JOINT_INDEX`
    /// and `Mesh::ATTRIBUTE_JOINT_WEIGHT` vertex attributes.
    pub skin: bool,

    /// Whether this mesh is being rendered with motion vectors.
    ///
    /// This is determined by the `MeshPipelineKey::MOTION_VECTOR_PREPASS` pipeline key
    /// flag.
    ///
    /// This pipeline key flag is set whenever the view rendering the mesh has the
    /// `MotionVectorPrepass` component, **and** the mesh, if it has morph targets
    /// or skinning, was rendered last frame.
    ///
    /// Note that the same mesh can be rendered by different views, some of them
    /// with or without motion vectors.
    pub motion_vectors: bool,
}
impl ActiveVariant {
    pub const fn new(morph: bool, skin: bool, motion_vectors: bool) -> Self {
        Self {
            morph,
            skin,
            motion_vectors,
        }
    }
}