effectcraft/crates/gpu/src/particles.rs

//! GPU particles (`particles.wgsl`): the stepped particle systems of `effectcraft-effects`
//! ([`effectcraft_effects::psim`]) simulated with one invocation per particle. States stay on
//! the GPU as checkpoints (per system key, like the CPU's `SimCache`), so playback advances from
//! the previous frame's state instead of simulating from layer time 0.

use std::sync::{Mutex, OnceLock};

use effectcraft_effects::psim::{ParticleSystem, SimParticle, SimRequest};
use wgpu::util::DeviceExt;

use crate::context::GpuContext;

/// Bytes of one particle state (`PState` in the shader).
const STATE_BYTES: u64 = 48;
/// Checkpoints kept (all keys).
const CHECKPOINTS: usize = 24;

pub(crate) struct Pipes {
    bgl: wgpu::BindGroupLayout,
    pipe: wgpu::ComputePipeline,
}

struct Checkpoint {
    key: u64,
    steps: u64,
    n: u32,
    buf: wgpu::Buffer,
    used: u64,
}

#[derive(Default)]
pub(crate) struct Checkpoints {
    list: Vec<Checkpoint>,
    clock: u64,
}

fn pipes(g: &GpuContext) -> &Pipes {
    g.particles.get_or_init(|| {
        let d = &g.device;
        let src = [include_str!("shaders/common.wgsl"), include_str!("shaders/kernels.wgsl"), include_str!("shaders/particles.wgsl")].concat();
        let module =
            d.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("effectcraft particles"), source: wgpu::ShaderSource::Wgsl(src.into()) });
        let buffer = |binding: u32, ty: wgpu::BufferBindingType| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::COMPUTE,
            ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None },
            count: None,
        };
        let bgl = d.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("effectcraft particles"),
            entries: &[
                buffer(0, wgpu::BufferBindingType::Uniform),
                buffer(1, wgpu::BufferBindingType::Storage { read_only: true }),
                buffer(2, wgpu::BufferBindingType::Storage { read_only: true }),
                buffer(3, wgpu::BufferBindingType::Storage { read_only: false }),
            ],
        });
        let layout = d.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("effectcraft particles"),
            bind_group_layouts: &[None, Some(&bgl)],
            immediate_size: 0,
        });
        let pipe = d.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("particles"),
            layout: Some(&layout),
            module: &module,
            entry_point: Some("particles"),
            compilation_options: Default::default(),
            cache: None,
        });
        Pipes { bgl, pipe }
    })
}

/// The shader's parameter block.
fn uniforms(req: &SimRequest, n: u32, s0: u64, prev_n: u32) -> Vec<u8> {
    let dt = (1.0 / effectcraft_effects::psim::SPS) as f32;
    let mut u = [n, s0 as u32, req.steps as u32, 0];
    let u2;
    let mut f = [[0.0f32; 4]; 6];
    match &req.system {
        ParticleSystem::Engine(e) => {
            u2 = [e.anim, e.flat as u32, e.seed, prev_n];
            f[0] = [e.life, e.speed, e.extra, e.cone];
            f[1] = [e.producer[0], e.producer[1], e.producer[2], 0.0];
            f[2] = [e.radius[0], e.radius[1], e.radius[2], 0.0];
            f[3] = [e.gravity[0], e.gravity[1], e.gravity[2], 0.0];
            f[4] = [e.axis[0], e.axis[1], e.axis[2], 0.0];
            f[5] = [e.damp, dt, 0.0, 0.0];
        }
        ParticleSystem::Cannon(c) => {
            u[3] = 1;
            u2 = [0, 1, c.seed, prev_n];
            f[0] = [c.pos[0], c.pos[1], c.barrel_angle, c.barrel_radius];
            f[1] = [c.dir_spread, c.vel, c.vel_spread, c.grav_spread];
            f[2] = [c.gravity[0], c.gravity[1], dt, 0.0];
            f[3] = c.bounds;
        }
    }
    let mut out = Vec::with_capacity(128);
    for x in u.iter().chain(&u2) {
        out.extend_from_slice(&x.to_le_bytes());
    }
    for x in f.iter().flatten() {
        out.extend_from_slice(&x.to_le_bytes());
    }
    out
}

/// Simulate `req` (see [`effectcraft_effects::psim::ParticleSim`]).
pub(crate) fn simulate(g: &GpuContext, req: &SimRequest) -> Option<Vec<SimParticle>> {
    if !g.can_readback() || req.steps > u32::MAX as u64 {
        return None;
    }
    let n = req.births.len() as u32;
    if n == 0 {
        return Some(vec![]);
    }
    let size = n as u64 * STATE_BYTES;
    if size > g.device.limits().max_storage_buffer_binding_size {
        return None;
    }
    let p = pipes(g);
    // The latest checkpoint at or before the target step.
    let (s0, prev, prev_n) = {
        let mut c = g.particle_states.lock().ok()?;
        c.clock += 1;
        let now = c.clock;
        match c.list.iter_mut().filter(|k| k.key == req.key && k.steps <= req.steps).max_by_key(|k| k.steps) {
            Some(k) => {
                k.used = now;
                (k.steps, Some(k.buf.clone()), k.n)
            }
            None => (0, None, 0),
        }
    };
    let d = &g.device;
    let births: Vec<u8> = req.births.iter().flat_map(|b| b.to_le_bytes()).collect();
    let births = d.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("births"), contents: &births, usage: wgpu::BufferUsages::STORAGE });
    let prev = prev.unwrap_or_else(|| {
        d.create_buffer(&wgpu::BufferDescriptor { label: Some("particles"), size: STATE_BYTES, usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false })
    });
    let next = d.create_buffer(&wgpu::BufferDescriptor {
        label: Some("particles"),
        size,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let ub = d.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("particles"),
        contents: &uniforms(req, n, s0, prev_n),
        usage: wgpu::BufferUsages::UNIFORM,
    });
    let bg = d.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &p.bgl,
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: ub.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: births.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: prev.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: next.as_entire_binding() },
        ],
    });
    let read = d.create_buffer(&wgpu::BufferDescriptor {
        label: Some("particles readback"),
        size,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let mut enc = d.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("particles") });
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("particles"), timestamp_writes: None });
        pass.set_pipeline(&p.pipe);
        pass.set_bind_group(1, &bg, &[]);
        pass.dispatch_workgroups(n.div_ceil(64), 1, 1);
    }
    enc.copy_buffer_to_buffer(&next, 0, &read, 0, size);
    g.queue.submit([enc.finish()]);
    let (tx, rx) = std::sync::mpsc::channel();
    read.map_async(wgpu::MapMode::Read, .., move |r| {
        let _ = tx.send(r);
    });
    g.device.poll(wgpu::PollType::wait_indefinitely()).ok()?;
    rx.recv().ok()?.ok()?;
    let view = read.get_mapped_range(..).ok()?;
    let f: Vec<f32> = view.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect();
    drop(view);
    read.unmap();
    let parts = f
        .chunks_exact(12)
        .enumerate()
        .filter(|(_, s)| s[9] != 0.0)
        .map(|(id, s)| SimParticle { p: [s[0], s[1], s[2]], v: [s[4], s[5], s[6]], age: s[3], life: s[7], rnd: s[8], id: id as u32 })
        .collect();
    // Keep the new state as a checkpoint.
    if let Ok(mut c) = g.particle_states.lock() {
        let now = c.clock;
        c.list.retain(|k| !(k.key == req.key && k.steps == req.steps));
        c.list.push(Checkpoint { key: req.key, steps: req.steps, n, buf: next, used: now });
        while c.list.len() > CHECKPOINTS {
            let Some(i) = c.list.iter().enumerate().min_by_key(|(_, k)| k.used).map(|(i, _)| i) else { break };
            c.list.remove(i);
        }
    }
    Some(parts)
}

pub(crate) type PipesCell = OnceLock<Pipes>;
pub(crate) type StatesCell = Mutex<Checkpoints>;