nightshade-renderer 0.57.0

use crate::meshlet::gpu_mesh::MeshletMeshStreams;
use crate::wgpu::passes::geometry::meshlet::types::{
    MESHLET_MAX_TRIANGLES, MESHLET_SOFTWARE_RASTER_MAX_PIXELS, MeshletCullViewUniform,
    MeshletDispatchIndirectArgs, MeshletDrawIndirectArgs, MeshletInstanceUniform,
};

/// Clusters the raster list is sized for, as a multiple of the instance count.
/// Only the instances nearest the camera contribute more than a handful, so the
/// list is sized well under the worst case and the cull reports an overflow by
/// counting past it.
const CLUSTERS_RESERVED_PER_INSTANCE: u64 = 32;
const CLUSTER_CAPACITY_FLOOR: u64 = 1 << 16;
const CLUSTER_CAPACITY_CEILING: u64 = 1 << 22;

/// What the cull needs to test a cluster against the depth pyramid: the matrix
/// that projects a world point into the pyramid's screen, the pyramid's size and
/// mip count, and whether there is a pyramid at all. Empty on the first frame,
/// before any depth has been reduced.
#[derive(Default)]
pub struct MeshletOcclusionInputs {
    pub occluder_from_world: [[f32; 4]; 4],
    pub screen_size: (f32, f32),
    pub mip_count: u32,
    pub enabled: bool,
}

/// The frame-varying inputs the cut is chosen from: the view its frustum and
/// error project against, the pixels of error a cluster may carry, and the
/// occlusion pyramid. Bundled so the sync signature stays one argument as this
/// grows.
pub struct MeshletCullInputs<'a> {
    pub view: &'a crate::config::RenderView,
    pub lod_error_threshold: f32,
    pub occlusion: &'a MeshletOcclusionInputs,
}

/// Where the instances come from: the scene world holding the placements, the
/// baked assets they name, and the counter the scene sync bumps whenever it
/// wrote or retired one.
///
/// The counter is what keeps a still scene still. Placements are per-entity
/// state that mostly never changes, and reading a quarter of a million of them
/// back to rebuild an identical instance list, then uploading twenty megabytes
/// of it, is a frame's worth of work to arrive where the gpu already was.
pub struct MeshletInstanceInputs<'a> {
    pub scene_world: &'a nightshade_ecs::dynamic::DynWorld,
    pub assets: &'a crate::config::MeshletAssetCache,
    pub generation: u64,
}

/// What the scene needs to turn the frame's resolved materials into the table
/// the resolve shades from.
pub struct MeshletMaterialInputs<'a> {
    pub render_materials: &'a crate::config::RenderMaterials,
    /// Where each texture landed in the shared arrays. Owned by the pass,
    /// because it is filled as textures upload rather than per frame.
    pub layer_map: &'a std::collections::HashMap<
        crate::asset_id::TextureId,
        crate::wgpu::material_texture_arrays::MaterialTextureLayer,
    >,
    /// Set when a texture moved layers, so materials naming it convert again
    /// even though the material table itself did not change.
    pub layers_changed: bool,
}

/// A storage buffer that grows to fit whatever the frame needs, reporting when
/// it reallocated so dependent bind groups can be rebuilt.
struct GrowableStorageBuffer {
    label: &'static str,
    buffer: wgpu::Buffer,
    capacity_in_bytes: u64,
}

impl GrowableStorageBuffer {
    fn new(device: &wgpu::Device, label: &'static str) -> Self {
        let capacity_in_bytes = 256;
        Self {
            label,
            buffer: Self::allocate(device, label, capacity_in_bytes),
            capacity_in_bytes,
        }
    }

    fn allocate(device: &wgpu::Device, label: &'static str, size: u64) -> wgpu::Buffer {
        device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    }

    /// Grows to hold `size` bytes without writing anything, for a buffer the
    /// gpu fills.
    fn reserve(&mut self, device: &wgpu::Device, size: u64) -> bool {
        if size <= self.capacity_in_bytes {
            return false;
        }
        self.capacity_in_bytes = size.next_power_of_two();
        self.buffer = Self::allocate(device, self.label, self.capacity_in_bytes);
        true
    }

    fn write(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, bytes: &[u8]) -> bool {
        let reallocated = self.reserve(device, bytes.len() as u64);
        if !bytes.is_empty() {
            queue.write_buffer(&self.buffer, 0, bytes);
        }
        reallocated
    }
}

/// The gpu-resident meshlet scene: the shared mesh streams, the instance list,
/// and the cluster list plus indirect draw the cull writes into.
///
/// The cpu never walks the scene here. It uploads instances, reserves room for
/// the clusters, and resets the draw. Which clusters survive, and how many the
/// raster draws, are decided on the gpu.
///
/// `generation` changes whenever any buffer was reallocated, which is the
/// signal to rebuild bind groups.
pub struct MeshletScene {
    pub meshes: MeshletMeshStreams,
    instances: GrowableStorageBuffer,
    clusters: GrowableStorageBuffer,
    materials: GrowableStorageBuffer,
    draw_args: wgpu::Buffer,
    dispatch_args: wgpu::Buffer,
    cull_view: wgpu::Buffer,
    pub instance_count: u32,
    pub cluster_capacity: u32,
    /// Whether the cull may route clusters to the compute rasterizer at all,
    /// which it may not where the device has no 64 bit texture atomics.
    pub software_raster_enabled: bool,
    pub generation: u64,
    /// The material table's generation as of the last conversion, so the
    /// per-material work only repeats when the table actually changed.
    materials_generation: Option<u64>,
    /// The placement generation the instance list was built from, so it only
    /// rebuilds when the scene sync actually moved a placement.
    instances_generation: Option<u64>,
}

impl MeshletScene {
    pub fn new(device: &wgpu::Device, software_raster_enabled: bool) -> Self {
        Self {
            meshes: MeshletMeshStreams::new(device),
            instances: GrowableStorageBuffer::new(device, "meshlet instances"),
            clusters: GrowableStorageBuffer::new(device, "meshlet raster clusters"),
            materials: GrowableStorageBuffer::new(device, "meshlet materials"),
            draw_args: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("meshlet draw args"),
                size: std::mem::size_of::<MeshletDrawIndirectArgs>() as u64,
                usage: wgpu::BufferUsages::INDIRECT
                    | wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
            dispatch_args: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("meshlet software raster dispatch args"),
                size: std::mem::size_of::<MeshletDispatchIndirectArgs>() as u64,
                usage: wgpu::BufferUsages::INDIRECT
                    | wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
            cull_view: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("meshlet cull view"),
                size: std::mem::size_of::<MeshletCullViewUniform>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
            instance_count: 0,
            cluster_capacity: 0,
            software_raster_enabled,
            generation: 0,
            materials_generation: None,
            instances_generation: None,
        }
    }

    pub fn dispatch_args_buffer(&self) -> &wgpu::Buffer {
        &self.dispatch_args
    }

    pub fn materials_buffer(&self) -> &wgpu::Buffer {
        &self.materials.buffer
    }

    pub fn instances_buffer(&self) -> &wgpu::Buffer {
        &self.instances.buffer
    }

    pub fn clusters_buffer(&self) -> &wgpu::Buffer {
        &self.clusters.buffer
    }

    pub fn draw_args_buffer(&self) -> &wgpu::Buffer {
        &self.draw_args
    }

    pub fn cull_view_buffer(&self) -> &wgpu::Buffer {
        &self.cull_view
    }

    /// Converts the scene's resolved materials into the form the resolve reads,
    /// keyed the same way the mesh pass keys them so an entity shades the same
    /// whichever pass draws it. Slot zero is the built-in default, which is what
    /// an entity without a material lands on.
    fn sync_materials(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        materials: &MeshletMaterialInputs<'_>,
    ) {
        let render_materials = materials.render_materials;
        let layer_map = materials.layer_map;
        if self.materials_generation == Some(render_materials.generation)
            && !materials.layers_changed
        {
            return;
        }
        self.materials_generation = Some(render_materials.generation);

        let mut materials_data = vec![
            crate::wgpu::passes::geometry::material_gpu::default_material_data([
                0.7, 0.7, 0.7, 1.0,
            ]),
        ];
        for entry in &render_materials.entries {
            materials_data.push(
                crate::wgpu::passes::geometry::material_gpu::convert_material_to_gpu_data(
                    &entry.material,
                    &entry.texture_ids,
                    layer_map,
                ),
            );
        }

        if self
            .materials
            .write(device, queue, bytemuck::cast_slice(&materials_data))
        {
            self.generation = self.generation.wrapping_add(1);
        }
    }

    /// Uploads any asset the scene references for the first time and rebuilds
    /// the instance list from the placements.
    ///
    /// The whole list, not a range of it: an instance's slot is its position in
    /// query order, so one placement arriving in the middle moves every slot
    /// after it. What makes that affordable is running it only when a placement
    /// actually moved.
    fn rebuild_instances(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        instances: &MeshletInstanceInputs<'_>,
    ) {
        for (_, placement) in instances
            .scene_world
            .query_ref::<&crate::config::MeshletPlacement>()
            .iter()
        {
            if self.meshes.entry(placement.asset_id).is_some() {
                continue;
            }
            let Some(asset) = instances.assets.get(&placement.asset_id) else {
                continue;
            };
            self.meshes
                .queue_upload(device, queue, placement.asset_id, asset);
            self.generation = self.generation.wrapping_add(1);
        }

        let mut instance_uniforms: Vec<MeshletInstanceUniform> = Vec::new();
        for (_, placement) in instances
            .scene_world
            .query_ref::<&crate::config::MeshletPlacement>()
            .iter()
        {
            let Some(entry) = self.meshes.entry(placement.asset_id) else {
                continue;
            };
            instance_uniforms.push(MeshletInstanceUniform {
                world_from_local: placement.transform.into(),
                root_bvh_node_index: entry.root_bvh_node_index,
                material_id: placement.material_id,
                padding: [0; 2],
            });
        }
        self.instance_count = instance_uniforms.len() as u32;

        let capacity = (self.instance_count as u64 * CLUSTERS_RESERVED_PER_INSTANCE)
            .clamp(CLUSTER_CAPACITY_FLOOR, CLUSTER_CAPACITY_CEILING);
        self.cluster_capacity = capacity as u32;

        let mut reallocated =
            self.instances
                .write(device, queue, bytemuck::cast_slice(&instance_uniforms));
        reallocated |=
            self.clusters.reserve(
                device,
                capacity
                    * std::mem::size_of::<
                        crate::wgpu::passes::geometry::meshlet::types::InstancedOffset,
                    >() as u64,
            );
        if reallocated {
            self.generation = self.generation.wrapping_add(1);
        }
    }

    /// Readies the buffers the cull writes into and the view it chooses the cut
    /// against, then rebuilds the instance list if a placement moved.
    pub fn sync(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        instances: &MeshletInstanceInputs<'_>,
        materials: &MeshletMaterialInputs<'_>,
        cull: &MeshletCullInputs<'_>,
    ) {
        let view = cull.view;
        let lod_error_threshold = cull.lod_error_threshold;
        let occlusion = cull.occlusion;
        self.sync_materials(device, queue, materials);
        if self.instances_generation != Some(instances.generation) {
            self.instances_generation = Some(instances.generation);
            self.rebuild_instances(device, queue, instances);
        }

        queue.write_buffer(
            &self.draw_args,
            0,
            bytemuck::bytes_of(&MeshletDrawIndirectArgs {
                vertex_count: MESHLET_MAX_TRIANGLES * 3,
                instance_count: 0,
                first_vertex: 0,
                first_instance: 0,
            }),
        );

        queue.write_buffer(
            &self.dispatch_args,
            0,
            bytemuck::bytes_of(&MeshletDispatchIndirectArgs {
                workgroup_count_x: 0,
                workgroup_count_y: 0,
                workgroup_count_z: 0,
                total_clusters: 0,
                software_clusters: 0,
            }),
        );

        let viewport_height = view.screen_size.1.max(1) as f32;
        let mut frustum_planes = [[0.0_f32; 4]; 6];
        for (slot, plane) in view.frustum_planes.iter().enumerate() {
            frustum_planes[slot] = [plane.x, plane.y, plane.z, plane.w];
        }
        queue.write_buffer(
            &self.cull_view,
            0,
            bytemuck::bytes_of(&MeshletCullViewUniform {
                frustum_planes,
                occluder_from_world: occlusion.occluder_from_world,
                camera_position: [
                    view.camera_position.x,
                    view.camera_position.y,
                    view.camera_position.z,
                    1.0,
                ],
                params: [
                    viewport_height * 0.5 / (view.y_fov_rad * 0.5).tan().max(1.0e-6),
                    view.z_near.max(1.0e-4),
                    view.orthographic
                        .map(|(_, height)| viewport_height / height.abs().max(1.0e-6))
                        .unwrap_or(0.0),
                    lod_error_threshold.max(1.0e-3),
                ],
                counts: [
                    self.instance_count,
                    self.cluster_capacity,
                    MESHLET_SOFTWARE_RASTER_MAX_PIXELS,
                    u32::from(self.software_raster_enabled),
                ],
                limits: [
                    device.limits().max_compute_workgroups_per_dimension,
                    u32::from(occlusion.enabled && occlusion.mip_count > 0),
                    occlusion.mip_count,
                    0,
                ],
                occluder_screen_size: [
                    occlusion.screen_size.0.max(1.0),
                    occlusion.screen_size.1.max(1.0),
                    0.0,
                    0.0,
                ],
            }),
        );
    }
}