path: root/src/renderer/instances.rs

use fxhash::FxHashMap;

use crate::{
    camera::Camera,
    config::{CELLS_IN_CHUNK, CHUNK_SIZE},
    sim::{cell_manager::manager::CellManager, entity::EntityId, sim_manager::SimManager},
};

// TODO derive from shared chunk config
const CHUNK_SLOTS: usize = 11 * 200;

fn create_instance_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("Instance buffer"),
        size: (capacity * size_of::<RendererInstance>()) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

fn create_cell_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("Cell buffer"),
        size: (capacity * CELLS_IN_CHUNK) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct RendererInstance {
    centre: [f32; 2],
    cos_sin: [f32; 2],
    half_size: [f32; 2],
    dims: [u32; 2],
    cell_offset: u32,
    _padding: u32,
}

fn create_instance_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    camera_uniform_buffer: &wgpu::Buffer,
    palette_buffer: &wgpu::Buffer,
    instance_buffer: &wgpu::Buffer,
    cell_buffer: &wgpu::Buffer,
    lighting_texture: &wgpu::TextureView,
    sampler: &wgpu::Sampler,
    lighting_uniform_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("Instance bind group"),
        layout,
        entries: &[
            wgpu::BindGroupEntry {
                binding: 0,
                resource: camera_uniform_buffer.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 1,
                resource: palette_buffer.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 2,
                resource: instance_buffer.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 3,
                resource: cell_buffer.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 4,
                resource: wgpu::BindingResource::TextureView(lighting_texture),
            },
            wgpu::BindGroupEntry {
                binding: 5,
                resource: wgpu::BindingResource::Sampler(sampler),
            },
            wgpu::BindGroupEntry {
                binding: 6,
                resource: lighting_uniform_buffer.as_entire_binding(),
            },
        ],
    })
}

pub struct InstancesRendererState {
    instance_pipeline: wgpu::RenderPipeline,
    instance_bind_group_layout: wgpu::BindGroupLayout,
    camera_uniform_buffer: wgpu::Buffer,
    palette_buffer: wgpu::Buffer,
    instance_buffer: wgpu::Buffer,
    cell_buffer: wgpu::Buffer,
    instance_bind_group: wgpu::BindGroup,
    entity_capacity: usize,
    sampler: wgpu::Sampler,

    renderer_entities: FxHashMap<EntityId, usize>,
    instances_count: usize,
}

impl InstancesRendererState {
    /// Rebind the lighting resources, needed whenever the lighting texture is recreated.
    pub fn rebuild_bind_group(
        &mut self,
        device: &wgpu::Device,
        lighting_texture: &wgpu::TextureView,
        lighting_uniform_buffer: &wgpu::Buffer,
    ) {
        self.instance_bind_group = create_instance_bind_group(
            device,
            &self.instance_bind_group_layout,
            &self.camera_uniform_buffer,
            &self.palette_buffer,
            &self.instance_buffer,
            &self.cell_buffer,
            lighting_texture,
            &self.sampler,
            lighting_uniform_buffer,
        );
    }

    pub fn update_buffers(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        sim: &mut SimManager,
        camera: &Camera,
    ) {
        puffin::profile_function!();
        let mut cell_buffer: [u8; CELLS_IN_CHUNK] = [0; CELLS_IN_CHUNK];
        let mut instances: Vec<RendererInstance> = Vec::new();

        for &idx in sim.cell_manager.chunk_position_to_chunk_idx.values() {
            let chunk = &mut sim.cell_manager.chunks[idx];
            if !chunk.needs_texture_update {
                continue;
            }
            chunk.needs_texture_update = false;

            for (byte, cell) in cell_buffer.iter_mut().zip(chunk.cells.iter()) {
                *byte = cell.material as u8;
            }

            queue.write_buffer(
                &self.cell_buffer,
                (idx * CELLS_IN_CHUNK) as u64,
                &cell_buffer,
            );
        }

        // TODO optimize, this is very inefficient, just build it per frame with positions and skip the lookup?
        self.renderer_entities.clear();
        for entity in sim.entities.values() {
            // TODO add "needs texture update"?
            if let Some(cells) = &entity.data.cells {
                for i in 0..cells.cells.len() {
                    cell_buffer[i] = cells.cells[i].material as u8;
                }

                let next_slot = CHUNK_SLOTS + self.renderer_entities.len();
                let slot = *self
                    .renderer_entities
                    .entry(entity.data.id)
                    .or_insert(next_slot);

                queue.write_buffer(
                    &self.cell_buffer,
                    (slot * CHUNK_SIZE as usize * CHUNK_SIZE as usize) as u64,
                    &cell_buffer,
                );
            }
        }

        let (lower, upper) = camera.viewport_bounds_world();
        let ((cxl, cyl), _) =
            CellManager::split_game_position(lower.x.floor() as i32, lower.y.floor() as i32);
        let ((cxu, cyu), _) =
            CellManager::split_game_position(upper.x.floor() as i32, upper.y.floor() as i32);

        let half_size = [CHUNK_SIZE as f32 / 2.0, CHUNK_SIZE as f32 / 2.0];
        let dims = [CHUNK_SIZE as u32, CHUNK_SIZE as u32];

        for cx in cxl..=cxu {
            for cy in cyl..=cyu {
                if let Some(idx) = sim.cell_manager.chunk_position_to_chunk_idx.get(&(cx, cy)) {
                    instances.push(RendererInstance {
                        centre: [
                            (cx * CHUNK_SIZE) as f32 + half_size[0],
                            (cy * CHUNK_SIZE) as f32 + half_size[1],
                        ],
                        cos_sin: [1.0, 0.0],
                        half_size,
                        dims,
                        cell_offset: (idx * CELLS_IN_CHUNK) as u32,
                        _padding: 0,
                    });
                }
            }
        }

        let mut entities = 0;
        for (id, slot) in &self.renderer_entities {
            if let Some(entity) = sim.entities.get(id)
                && let Some(cells) = &entity.data.cells
                && let Some((pos, (cos, sin))) = entity.data._transform(&sim.rb_manager)
            {
                // TODO access this better
                // let sleeping = sim
                //     .rb_manager
                //     .physics_manager
                //     .world
                //     .bodies
                //     .get(entity.data.rb_h.unwrap())
                //     .unwrap()
                //     .is_sleeping();
                let sleeping = false;

                entities += 1;
                instances.push(RendererInstance {
                    centre: if sleeping {
                        pos.round().to_array()
                    } else {
                        pos.to_array()
                    },
                    cos_sin: [cos, sin],
                    half_size: (cells.size / 2).as_vec2().to_array(),
                    dims: cells.size.as_uvec2().to_array(),
                    cell_offset: (slot * CHUNK_SIZE as usize * CHUNK_SIZE as usize) as u32,
                    _padding: 0,
                });
            }
        }

        if entities > self.entity_capacity {
            self.entity_capacity = entities.next_power_of_two();
            self.instance_buffer =
                create_instance_buffer(device, self.entity_capacity + CHUNK_SLOTS);
            self.cell_buffer = create_cell_buffer(device, CHUNK_SLOTS + self.entity_capacity);
        }

        queue.write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(&instances));

        self.instances_count = instances.len();
    }

    pub fn render(&self, render_pass: &mut wgpu::RenderPass) {
        render_pass.set_pipeline(&self.instance_pipeline);
        render_pass.set_bind_group(0, &self.instance_bind_group, &[]);
        render_pass.draw(0..4, 0..self.instances_count as u32);
    }

    pub fn new(
        device: &wgpu::Device,
        surface_config: &wgpu::SurfaceConfiguration,
        camera_uniform_buffer: &wgpu::Buffer,
        palette_buffer: &wgpu::Buffer,
        lighting_texture: &wgpu::TextureView,
        lighting_uniform_buffer: &wgpu::Buffer,
    ) -> Self {
        let instance_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Instance shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shader/instance.wgsl").into()),
        });

        let instance_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &[
                    // camera uniform
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 0,
                        count: None,
                        visibility: wgpu::ShaderStages::VERTEX,
                    },
                    // palette
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 1,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                    // instance buffer
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 2,
                        count: None,
                        visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                    },
                    // cells
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 3,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                    // lighting
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Texture {
                            multisampled: false,
                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
                            view_dimension: wgpu::TextureViewDimension::D2Array,
                        },
                        binding: 4,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                    // sampler
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                        binding: 5,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                    // lighting uniform
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 6,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                ],
            });

        let instance_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                immediate_size: 0,
                bind_group_layouts: &[Some(&instance_bind_group_layout)],
            });

        let instance_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&instance_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &instance_shader,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &instance_shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_config.format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        let entity_capacity = 2048;
        let instance_buffer = create_instance_buffer(&device, CHUNK_SLOTS + entity_capacity);

        let cell_buffer = create_cell_buffer(&device, entity_capacity + CHUNK_SLOTS);

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            ..Default::default()
        });
        let instance_bind_group = create_instance_bind_group(
            device,
            &instance_bind_group_layout,
            camera_uniform_buffer,
            palette_buffer,
            &instance_buffer,
            &cell_buffer,
            lighting_texture,
            &sampler,
            lighting_uniform_buffer,
        );

        InstancesRendererState {
            instance_pipeline,
            instance_bind_group_layout,
            camera_uniform_buffer: camera_uniform_buffer.clone(),
            palette_buffer: palette_buffer.clone(),
            instance_buffer,
            cell_buffer,
            instance_bind_group,
            entity_capacity,
            sampler,

            renderer_entities: FxHashMap::default(),
            instances_count: 0,
        }
    }
}