nannou_draw 0.20.0

//! A shader that renders a mesh multiple times in one draw call.

use crate::render::{ShaderBufferHandle, ShaderModelAsset};
use crate::{
    draw::{Draw, DrawCommand, drawing::Drawing, primitive::Primitive},
    render::{PreparedShaderModel, ShaderModel, queue_shader_model},
};
use bevy::pbr::{MATERIAL_BIND_GROUP_INDEX, SetMeshViewBindingArrayBindGroup};
use bevy::render::extract_instances::ExtractedInstances;
use bevy::{
    core_pipeline::core_3d::Transparent3d,
    ecs::system::{SystemParamItem, lifetimeless::*},
    pbr::{RenderMeshInstances, SetMeshBindGroup, SetMeshViewBindGroup},
    prelude::*,
    render::{
        Render, RenderApp, RenderSystems,
        extract_component::ExtractComponent,
        mesh::{RenderMesh, RenderMeshBufferInfo, allocator::MeshAllocator},
        render_asset::{RenderAssets, prepare_assets},
        render_phase::{
            AddRenderCommand, PhaseItem, RenderCommand, RenderCommandResult, SetItemPipeline,
            TrackedRenderPass,
        },
        storage::{GpuShaderBuffer, ShaderBuffer},
    },
};
use std::{hash::Hash, marker::PhantomData};

pub struct Indirect<'a> {
    draw: &'a Draw,
    primitive_index: Option<usize>,
    indirect_buffer: Option<Handle<ShaderBuffer>>,
}

impl<'a> Drop for Indirect<'a> {
    fn drop(&mut self) {
        if let Some((index, ssbo)) = self.primitive_index.take().zip(self.indirect_buffer.take()) {
            self.insert_indirect_draw_command(index, ssbo);
        }
    }
}

pub fn new(draw: &Draw) -> Indirect<'_> {
    Indirect {
        draw,
        primitive_index: None,
        indirect_buffer: None,
    }
}

impl<'a> Indirect<'a> {
    pub fn primitive<T>(mut self, drawing: Drawing<T>) -> Indirect<'a>
    where
        T: Into<Primitive>,
    {
        self.draw
            .state
            .write()
            .unwrap()
            .ignored_drawings
            .insert(drawing.index);
        self.primitive_index = Some(drawing.index);
        self
    }

    pub fn buffer(mut self, ssbo: Handle<ShaderBuffer>) -> Indirect<'a> {
        self.indirect_buffer = Some(ssbo);
        self
    }

    fn insert_indirect_draw_command(&self, index: usize, indirect_buffer: Handle<ShaderBuffer>) {
        let mut state = self.draw.state.write().unwrap();
        let primitive = state.drawing.remove(&index).unwrap();
        state
            .draw_commands
            .push(Some(DrawCommand::Indirect(primitive, indirect_buffer)));
    }
}

#[derive(Component, ExtractComponent, Clone)]
pub struct IndirectMesh;

pub struct IndirectShaderModelPlugin<SM>(PhantomData<SM>);

impl<SM> Default for IndirectShaderModelPlugin<SM>
where
    SM: Default,
{
    fn default() -> Self {
        IndirectShaderModelPlugin(PhantomData)
    }
}

impl<SM> Plugin for IndirectShaderModelPlugin<SM>
where
    SM: ShaderModel,
    SM::Data: PartialEq + Eq + Hash + Clone,
{
    fn build(&self, app: &mut App) {
        app.sub_app_mut(RenderApp)
            .add_render_command::<Transparent3d, DrawIndirectShaderModel<SM>>()
            .add_systems(
                Render,
                queue_shader_model::<SM, With<IndirectMesh>, DrawIndirectShaderModel<SM>>
                    .after(prepare_assets::<PreparedShaderModel<SM>>)
                    .in_set(RenderSystems::QueueMeshes),
            );
    }
}

type DrawIndirectShaderModel<SM> = (
    SetItemPipeline,
    SetMeshViewBindGroup<0>,
    SetMeshViewBindingArrayBindGroup<1>,
    SetMeshBindGroup<2>,
    SetShaderModelBindGroup<SM, MATERIAL_BIND_GROUP_INDEX>,
    DrawMeshIndirect,
);

struct SetShaderModelBindGroup<SM: ShaderModel, const I: usize>(PhantomData<SM>);
impl<P: PhaseItem, SM: ShaderModel, const I: usize> RenderCommand<P>
    for SetShaderModelBindGroup<SM, I>
{
    type Param = (
        SRes<RenderAssets<PreparedShaderModel<SM>>>,
        SRes<ExtractedInstances<ShaderModelAsset<SM>>>,
    );
    type ViewQuery = ();
    type ItemQuery = ();

    #[inline]
    fn render<'w>(
        item: &P,
        _view: (),
        _item_query: Option<()>,
        (models, instances): SystemParamItem<'w, '_, Self::Param>,
        pass: &mut TrackedRenderPass<'w>,
    ) -> RenderCommandResult {
        let models = models.into_inner();
        let instances = instances.into_inner();

        let Some(model_asset) = instances.get(&item.main_entity()) else {
            return RenderCommandResult::Skip;
        };
        let Some(model) = models.get(model_asset.0) else {
            return RenderCommandResult::Skip;
        };
        pass.set_bind_group(I, &model.bind_group, &[]);
        RenderCommandResult::Success
    }
}

struct DrawMeshIndirect;
impl<P: PhaseItem> RenderCommand<P> for DrawMeshIndirect {
    type Param = (
        SRes<RenderAssets<RenderMesh>>,
        SRes<RenderMeshInstances>,
        SRes<MeshAllocator>,
        SRes<RenderAssets<GpuShaderBuffer>>,
    );
    type ViewQuery = ();
    type ItemQuery = Read<ShaderBufferHandle>;

    #[inline]
    fn render<'w>(
        item: &P,
        _view: (),
        indirect_buffer: Option<&'w ShaderBufferHandle>,
        (meshes, render_mesh_instances, mesh_allocator, ssbos): SystemParamItem<
            'w,
            '_,
            Self::Param,
        >,
        pass: &mut TrackedRenderPass<'w>,
    ) -> RenderCommandResult {
        let mesh_allocator = mesh_allocator.into_inner();

        let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(item.main_entity())
        else {
            return RenderCommandResult::Skip;
        };
        let Some(gpu_mesh) = meshes.into_inner().get(mesh_instance.mesh_asset_id()) else {
            return RenderCommandResult::Skip;
        };
        let Some(indirect_buffer) = indirect_buffer else {
            return RenderCommandResult::Skip;
        };
        let Some(indirect_buffer) = ssbos.into_inner().get(&indirect_buffer.0) else {
            return RenderCommandResult::Skip;
        };
        let Some(vertex_buffer_slice) =
            mesh_allocator.mesh_vertex_slice(&mesh_instance.mesh_asset_id())
        else {
            return RenderCommandResult::Skip;
        };

        // The mesh shares a `MeshAllocator` slab with other meshes, so its vertex and
        // index data generally sit at a non-zero offset within the slab buffers. The
        // indirect args buffer (written on the GPU / CPU without knowledge of those
        // offsets) uses `first_index`/`base_vertex`/`first_vertex` of `0`, so we slice
        // the bound buffers to start at the mesh's data. That makes the zero-based args
        // address this mesh's geometry rather than whatever happens to sit at slab
        // offset 0. (`MeshBufferSlice::range` is measured in elements, not bytes.)
        let vertex_stride = gpu_mesh.layout.0.layout().array_stride;
        let vertex_offset = vertex_buffer_slice.range.start as u64 * vertex_stride;
        pass.set_vertex_buffer(0, vertex_buffer_slice.buffer.slice(vertex_offset..));

        match &gpu_mesh.buffer_info {
            RenderMeshBufferInfo::Indexed { index_format, .. } => {
                let Some(index_buffer_slice) =
                    mesh_allocator.mesh_index_slice(&mesh_instance.mesh_asset_id())
                else {
                    return RenderCommandResult::Skip;
                };

                let index_offset =
                    index_buffer_slice.range.start as u64 * index_format.byte_size() as u64;
                pass.set_index_buffer(
                    index_buffer_slice.buffer.slice(index_offset..),
                    *index_format,
                );
                pass.draw_indexed_indirect(&indirect_buffer.buffer, 0);
            }
            RenderMeshBufferInfo::NonIndexed => {
                pass.draw_indirect(&indirect_buffer.buffer, 0);
            }
        }
        RenderCommandResult::Success
    }
}