nightshade-renderer 0.57.0

use std::collections::HashMap;

use crate::meshlet::asset::{BvhNode, Meshlet, MeshletMesh};

const BUFFER_GROWTH_FACTOR: f32 = 2.0;
const INITIAL_BUFFER_BYTES: u64 = 256;

/// Where a single uploaded [`MeshletMesh`] lives inside the shared GPU
/// storage buffers. Offsets are in elements of each stream: `u32` words for
/// the position and normal streams, `[f32; 2]` for uvs, bytes for indices,
/// and struct counts for the bvh, meshlet, and cull-data streams.
pub struct MeshletMeshEntry {
    pub vertex_position_word_offset: u64,
    pub vertex_normal_word_offset: u64,
    pub vertex_uv_offset: u64,
    pub index_byte_offset: u64,
    pub bvh_node_offset: u64,
    pub meshlet_offset: u64,
    pub meshlet_cull_data_offset: u64,
    pub root_bvh_node_index: u32,
}

struct StreamBuffer {
    label: &'static str,
    element_size: u64,
    buffer: wgpu::Buffer,
    size_in_bytes: u64,
    element_count: u64,
}

impl StreamBuffer {
    fn new(device: &wgpu::Device, label: &'static str, element_size: u64) -> Self {
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: INITIAL_BUFFER_BYTES,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        Self {
            label,
            element_size,
            buffer,
            size_in_bytes: INITIAL_BUFFER_BYTES,
            element_count: 0,
        }
    }

    fn append(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, bytes: &[u8]) -> u64 {
        let element_offset = self.element_count;
        let offset_bytes = element_offset * self.element_size;
        let padded_length = (bytes.len() as u64).div_ceil(wgpu::COPY_BUFFER_ALIGNMENT)
            * wgpu::COPY_BUFFER_ALIGNMENT;
        let required_size = offset_bytes + padded_length;

        if required_size > self.size_in_bytes {
            let new_size = (required_size as f32 * BUFFER_GROWTH_FACTOR).ceil() as u64;
            let new_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(self.label),
                size: new_size,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });

            if offset_bytes > 0 {
                let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some(self.label),
                });
                encoder.copy_buffer_to_buffer(&self.buffer, 0, &new_buffer, 0, offset_bytes);
                queue.submit(std::iter::once(encoder.finish()));
            }

            self.buffer = new_buffer;
            self.size_in_bytes = new_size;
        }

        if padded_length == bytes.len() as u64 {
            queue.write_buffer(&self.buffer, offset_bytes, bytes);
        } else {
            let mut padded = bytes.to_vec();
            padded.resize(padded_length as usize, 0);
            queue.write_buffer(&self.buffer, offset_bytes, &padded);
        }

        self.element_count += padded_length / self.element_size;
        element_offset
    }
}

/// Owns the persistent, grow-on-demand GPU storage buffers that back
/// meshlet-based rendering. Each mesh stream is suballocated into one large
/// buffer, and per-mesh offsets are rebased to global buffer positions at
/// upload time so the bvh and meshlet records reference absolute locations.
pub struct MeshletMeshStreams {
    vertex_positions: StreamBuffer,
    vertex_normals: StreamBuffer,
    vertex_uvs: StreamBuffer,
    indices: StreamBuffer,
    bvh_nodes: StreamBuffer,
    meshlets: StreamBuffer,
    meshlet_cull_data: StreamBuffer,
    entries: HashMap<u64, MeshletMeshEntry>,
}

impl MeshletMeshStreams {
    pub fn new(device: &wgpu::Device) -> Self {
        Self {
            vertex_positions: StreamBuffer::new(
                device,
                "meshlet vertex positions",
                std::mem::size_of::<u32>() as u64,
            ),
            vertex_normals: StreamBuffer::new(
                device,
                "meshlet vertex normals",
                std::mem::size_of::<u32>() as u64,
            ),
            vertex_uvs: StreamBuffer::new(
                device,
                "meshlet vertex uvs",
                std::mem::size_of::<[f32; 2]>() as u64,
            ),
            indices: StreamBuffer::new(device, "meshlet indices", 1),
            bvh_nodes: StreamBuffer::new(
                device,
                "meshlet bvh nodes",
                std::mem::size_of::<BvhNode>() as u64,
            ),
            meshlets: StreamBuffer::new(device, "meshlets", std::mem::size_of::<Meshlet>() as u64),
            meshlet_cull_data: StreamBuffer::new(
                device,
                "meshlet cull data",
                std::mem::size_of::<crate::meshlet::asset::MeshletCullData>() as u64,
            ),
            entries: HashMap::new(),
        }
    }

    /// Appends every stream of `asset` to its buffer (growing as needed) and
    /// returns the global index of the mesh's root bvh node. If the mesh was
    /// already uploaded, returns its existing root index without reuploading.
    pub fn queue_upload(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        asset_id: u64,
        asset: &MeshletMesh,
    ) -> u32 {
        if let Some(entry) = self.entries.get(&asset_id) {
            return entry.root_bvh_node_index;
        }

        let vertex_position_word_offset = self.vertex_positions.append(
            device,
            queue,
            bytemuck::cast_slice(&asset.vertex_positions),
        );
        let vertex_normal_word_offset =
            self.vertex_normals
                .append(device, queue, bytemuck::cast_slice(&asset.vertex_normals));
        let vertex_uv_offset =
            self.vertex_uvs
                .append(device, queue, bytemuck::cast_slice(&asset.vertex_uvs));
        let index_byte_offset = self.indices.append(device, queue, &asset.indices);

        let meshlet_offset = self.meshlets.element_count;
        let vertex_position_bit_base = (vertex_position_word_offset * 32) as u32;
        let vertex_attribute_base = vertex_normal_word_offset as u32;
        let index_base = index_byte_offset as u32;
        let mut rebased_meshlets: Vec<Meshlet> = asset.meshlets.to_vec();
        for meshlet in &mut rebased_meshlets {
            meshlet.start_vertex_position_bit += vertex_position_bit_base;
            meshlet.start_vertex_attribute_id += vertex_attribute_base;
            meshlet.start_index_id += index_base;
        }
        self.meshlets
            .append(device, queue, bytemuck::cast_slice(&rebased_meshlets));

        let meshlet_cull_data_offset = self.meshlet_cull_data.append(
            device,
            queue,
            bytemuck::cast_slice(&asset.meshlet_cull_data),
        );

        let bvh_node_offset = self.bvh_nodes.element_count;
        let base_bvh_node_index = bvh_node_offset as u32;
        let base_meshlet_index = meshlet_offset as u32;
        let mut rebased_bvh: Vec<BvhNode> = asset.bvh.to_vec();
        for node in &mut rebased_bvh {
            for child_index in 0..node.aabbs.len() {
                let base = if node.child_counts[child_index] == u8::MAX {
                    base_bvh_node_index
                } else {
                    base_meshlet_index
                };
                node.aabbs[child_index].child_offset += base;
            }
        }
        self.bvh_nodes
            .append(device, queue, bytemuck::cast_slice(&rebased_bvh));

        let root_bvh_node_index = base_bvh_node_index;
        self.entries.insert(
            asset_id,
            MeshletMeshEntry {
                vertex_position_word_offset,
                vertex_normal_word_offset,
                vertex_uv_offset,
                index_byte_offset,
                bvh_node_offset,
                meshlet_offset,
                meshlet_cull_data_offset,
                root_bvh_node_index,
            },
        );

        root_bvh_node_index
    }

    pub fn entry(&self, asset_id: u64) -> Option<&MeshletMeshEntry> {
        self.entries.get(&asset_id)
    }

    pub fn vertex_positions_buffer(&self) -> &wgpu::Buffer {
        &self.vertex_positions.buffer
    }

    pub fn vertex_normals_buffer(&self) -> &wgpu::Buffer {
        &self.vertex_normals.buffer
    }

    pub fn vertex_uvs_buffer(&self) -> &wgpu::Buffer {
        &self.vertex_uvs.buffer
    }

    pub fn indices_buffer(&self) -> &wgpu::Buffer {
        &self.indices.buffer
    }

    pub fn bvh_nodes_buffer(&self) -> &wgpu::Buffer {
        &self.bvh_nodes.buffer
    }

    pub fn meshlets_buffer(&self) -> &wgpu::Buffer {
        &self.meshlets.buffer
    }

    pub fn meshlet_cull_data_buffer(&self) -> &wgpu::Buffer {
        &self.meshlet_cull_data.buffer
    }
}