codecraft 0.1.1

//! The cluster grid: which lights reach which part of the frustum.
//!
//! A forward pass that loops every light for every fragment is quadratic in
//! the thing a scene most wants to add. Clustering cuts the frustum into
//! cells once a frame, asks which lights touch each cell, and leaves the
//! fragment with a handful to answer to instead of all of them. It is the
//! same trick as a broad phase in physics, and it is what "forward+" names.
//!
//! The culling runs on the GPU: one invocation per cell, every light tested
//! against the cell's bounds. See `clustered.wgsl` for the geometry.
//!
//! Directional lights are not here. The key and the rim are infinitely far
//! away and reach every cell, so there is nothing to cull and they stay in
//! the camera uniform where the shading pass can always find them.

use glam::Mat4;

use crate::sceneobjects::lights::{GpuLight, MAX_LIGHTS};

/// Cells across, down, and into the frame. Sixteen by nine matches the shape
/// of a window closely enough that a cell stays roughly square, and
/// twenty-four slices is the number the technique's own literature settled
/// on: enough that a slice is thin where lights overlap, few enough that the
/// grid stays a few thousand cells rather than a few hundred thousand.
pub const GRID: [u32; 3] = [16, 9, 24];

/// How many cells there are.
pub const CELLS: u32 = GRID[0] * GRID[1] * GRID[2];

/// How many lights one cell can hold. A cell that wants more drops the rest,
/// which is a dimmer cell and not a broken one -- and sixty-four lights
/// overlapping one froxel is a lighting problem before it is a renderer one.
pub const MAX_LIGHTS_PER_CELL: u32 = 64;

/// What the culling pass needs to know about the frame.
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ClusterUniform {
    inv_proj: [f32; 16],
    /// x, y: the frame in pixels. z, w: near and far.
    screen: [f32; 4],
    /// x, y, z: the grid. w: how many lights there are this frame.
    grid: [u32; 4],
    /// x, y: scale and bias taking `log(view depth)` to a slice.
    slice: [f32; 4],
}

/// Scale and bias mapping a view-space depth to a slice, so the shading pass
/// and the culling pass agree about which cell a fragment is in.
///
/// `slice = log(depth) * scale + bias`, which inverts
/// `depth = near * (far/near)^(slice/slices)`.
pub fn slice_scale_bias(near: f32, far: f32) -> (f32, f32) {
    let slices = GRID[2] as f32;
    let scale = slices / (far / near).ln();
    (scale, -near.ln() * scale)
}

/// The grid, the lights in it, and the pass that fills it.
pub struct Clusters {
    pipeline: wgpu::ComputePipeline,
    uniform: wgpu::Buffer,
    view: wgpu::Buffer,
    /// Every local light in the scene, which the shading pass reads too.
    pub lights: wgpu::Buffer,
    /// How many lights each cell ended up with.
    pub counts: wgpu::Buffer,
    /// `MAX_LIGHTS_PER_CELL` slots per cell, indices into `lights`.
    pub indices: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    /// How many lights the last [`prepare`](Self::prepare) was given.
    count: u32,
}

impl Clusters {
    pub fn new(device: &wgpu::Device) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("cluster culling shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("clustered.wgsl").into()),
        });

        let uniform = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster uniform"),
            size: std::mem::size_of::<ClusterUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let view = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster view matrix"),
            size: std::mem::size_of::<[f32; 16]>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let lights = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("local lights"),
            size: (MAX_LIGHTS * std::mem::size_of::<GpuLight>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let counts = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster light counts"),
            size: (CELLS as usize * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        let indices = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("cluster light indices"),
            size: ((CELLS * MAX_LIGHTS_PER_CELL) as usize * std::mem::size_of::<u32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });

        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("cluster bind group layout"),
            entries: &[
                uniform_entry(0),
                storage_entry(1, true),
                storage_entry(2, false),
                storage_entry(3, false),
                uniform_entry(4),
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("cluster bind group"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: lights.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: counts.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: indices.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: view.as_entire_binding(),
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("cluster pipeline layout"),
            bind_group_layouts: &[Some(&layout)],
            immediate_size: 0,
        });

        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("cluster culling pipeline"),
            layout: Some(&pipeline_layout),
            module: &shader,
            entry_point: Some("cull_cs"),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            cache: None,
        });

        Self {
            pipeline,
            uniform,
            view,
            lights,
            counts,
            indices,
            bind_group,
            count: 0,
        }
    }

    /// Send this frame's lights and the grid they will be sorted into.
    #[allow(clippy::too_many_arguments)]
    pub fn prepare(
        &mut self,
        queue: &wgpu::Queue,
        lights: &[GpuLight],
        view: Mat4,
        proj: Mat4,
        width: u32,
        height: u32,
        near: f32,
        far: f32,
    ) {
        self.count = lights.len() as u32;
        if !lights.is_empty() {
            queue.write_buffer(&self.lights, 0, bytemuck::cast_slice(lights));
        }
        queue.write_buffer(&self.view, 0, bytemuck::cast_slice(&view.to_cols_array()));

        let (scale, bias) = slice_scale_bias(near, far);
        let uniform = ClusterUniform {
            inv_proj: proj.inverse().to_cols_array(),
            screen: [width.max(1) as f32, height.max(1) as f32, near, far],
            grid: [GRID[0], GRID[1], GRID[2], self.count],
            slice: [scale, bias, 0.0, 0.0],
        };
        queue.write_buffer(&self.uniform, 0, bytemuck::bytes_of(&uniform));
    }

    /// Sort this frame's lights into the grid.
    pub fn cull(&self, encoder: &mut wgpu::CommandEncoder) {
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("cluster culling pass"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.bind_group, &[]);
        // One invocation per cell, in groups of sixty-four.
        pass.dispatch_workgroups(CELLS.div_ceil(64), 1, 1);
    }
}

fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_grid_is_a_few_thousand_cells() {
        assert_eq!(CELLS, 16 * 9 * 24);
        assert_eq!(CELLS, 3456);
    }

    #[test]
    fn a_depth_maps_to_the_slice_it_belongs_in() {
        let (near, far) = (0.1, 100.0);
        let (scale, bias) = slice_scale_bias(near, far);
        let slice = |depth: f32| (depth.ln() * scale + bias).floor() as i32;

        // The near plane is the front of the first slice, and the far plane
        // is the back of the last one -- which is one past the last index, so
        // the shader clamps it. Rounding can leave it a hair short of the
        // boundary, which is why this is a range and not an equality.
        assert_eq!(slice(near), 0);
        assert!(
            (GRID[2] as i32 - 1..=GRID[2] as i32).contains(&slice(far)),
            "the far plane should land at the back of the grid, got {}",
            slice(far),
        );
        assert_eq!(slice(far).min(GRID[2] as i32 - 1), GRID[2] as i32 - 1);

        // And slices go back exponentially: each one covers the same *ratio*
        // of depths, so they get thicker with distance.
        let boundary = |k: u32| near * (far / near).powf(k as f32 / GRID[2] as f32);
        assert_eq!(slice(boundary(5) + 1.0e-4), 5);
        assert!(
            boundary(20) - boundary(19) > boundary(2) - boundary(1),
            "a far slice covers more depth than a near one",
        );
    }
}