henad-compute 0.3.0

//! The engine that runs a [`GpuAgentModel`] as a [`GpuSimState`], the counterpart of [`crate::cpu::agent_engine`].

use std::marker::PhantomData;
use std::sync::Arc;

use henad_core::authoring::model::binding::{BindingDecl, buffer_target};
use henad_core::authoring::model::field::Extent;
use henad_core::authoring::model::gpu_agent_model::{Geometry, GpuAgentModel, PassCtx, PassId};
use henad_core::model::SimState;
use henad_core::params::ParamValue;
use henad_core::view::{StatEntry, stat_entries};

use crate::display_scale::display_dims;
use crate::gpu::capacity::{Demand, layout_entry, storage_bindings};
use crate::gpu::contracts::{assert_buffer_labels, assert_workgroup_size};
use crate::gpu::primitives::dispatch::{WORKGROUP, linear_dispatch};
use crate::gpu::primitives::pipeline::{compute_pipeline, lane_buffer, storage_buffer, uniform_buffer};
use crate::gpu::primitives::readback::{CounterReadback, StatsPoll};
use crate::gpu::primitives::reduce::GpuLaneReduce;
use crate::gpu::primitives::spatial_hash::{GpuSpatialHash, HashGrid};
use crate::gpu::sim_thread::GpuSimState;
use crate::gpu::view::agents::GpuAgents;
use crate::gpu::view::display::{DisplayTarget, GpuDisplay, build_display_target};
use crate::gpu::{GpuContext, MAX_STEPS_PER_SUBMISSION};
use crate::snapshot::GpuSnapshot;

/// A resource that exists once, or once per ping-ponged side.
struct Sides<T> {
    a: T,
    b: Option<T>,
}

impl<T> Sides<T> {
    fn pick(&self, a_is_current: bool) -> &T {
        if a_is_current {
            &self.a
        } else {
            self.b.as_ref().unwrap_or(&self.a)
        }
    }
}

struct BufferSides {
    a: wgpu::Buffer,
    b: Option<wgpu::Buffer>,
}

impl BufferSides {
    /// Returns `(current, next)`, the same buffer twice when this one is written in place.
    fn sides(&self, a_is_current: bool) -> (&wgpu::Buffer, &wgpu::Buffer) {
        match &self.b {
            None => (&self.a, &self.a),
            Some(b) if a_is_current => (&self.a, b),
            Some(b) => (b, &self.a),
        }
    }
}

struct EncodedPass {
    label: String,
    pipeline: wgpu::ComputePipeline,
    binds: Sides<wgpu::BindGroup>,
    groups: (u32, u32),
    /// Uniform of the pass, kept so an action can reseed it before a press. Unused for a pass that is not an action.
    uniform: wgpu::Buffer,
}

impl EncodedPass {
    fn encode(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        a_is_current: bool,
        timestamps: Option<wgpu::ComputePassTimestampWrites<'_>>,
    ) {
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some(&self.label),
            timestamp_writes: timestamps,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, self.binds.pick(a_is_current), &[]);
        pass.dispatch_workgroups(self.groups.0, self.groups.1, 1);
    }
}

/// Returns the timestamp writes of a pass, `None` for a pass that neither opens nor closes a batch.
///
/// A `ComputePassTimestampWrites` needs at least one index set.
fn stamps(
    query_set: Option<&wgpu::QuerySet>,
    opening: bool,
    closing: bool,
) -> Option<wgpu::ComputePassTimestampWrites<'_>> {
    query_set
        .filter(|_| opening || closing)
        .map(|query_set| wgpu::ComputePassTimestampWrites {
            query_set,
            beginning_of_pass_write_index: opening.then_some(0),
            end_of_pass_write_index: closing.then_some(1),
        })
}

/// GPU-resident state for a [`GpuAgentModel`], with every buffer, pass and bind group its model declares.
pub struct GpuAgentState<M: GpuAgentModel> {
    geom: Geometry,
    tick: u64,

    device: wgpu::Device,
    queue: wgpu::Queue,

    buffers: Vec<BufferSides>,
    /// `true` when the `a` side of every double buffered buffer holds the current state.
    current_is_a: bool,
    /// Set when some buffer requests double buffering. Nothing flips otherwise.
    ping_pong: bool,

    /// Spatial hash, if the model declares one.
    index: Option<GpuSpatialHash>,
    index_binds: Option<Sides<wgpu::BindGroup>>,

    steps: Vec<EncodedPass>,
    display: Option<(EncodedPass, Arc<GpuDisplay>)>,

    reduce: GpuLaneReduce,
    reduce_pass: EncodedPass,
    counters: Option<CounterReadback>,

    actions: Vec<EncodedPass>,
    /// Seed of the next action press, advanced per press so that pressing twice draws twice.
    action_seed: u32,
    /// Parameter values, kept to rewrite the action uniforms on each press.
    params: Vec<ParamValue>,

    agents: Sides<Arc<GpuAgents>>,

    _marker: PhantomData<M>,
}

impl<M: GpuAgentModel> std::fmt::Debug for GpuAgentState<M> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuAgentState")
            .field("model", &M::ID)
            .field("tick", &self.tick)
            .finish_non_exhaustive()
    }
}

impl<M: GpuAgentModel> GpuAgentState<M> {
    /// Builds the state on `ctx`, seeded with the model's default seed.
    ///
    /// # Panics
    ///
    /// Panics as [`Self::new_seeded`] does.
    pub fn new(ctx: &GpuContext, params: &[ParamValue]) -> Self {
        Self::new_seeded(ctx, params, None)
    }

    /// Returns the geometry for `params`, without touching a device. `limits` sets only the display cap.
    pub fn geometry_for(params: &[ParamValue], limits: &wgpu::Limits) -> Geometry {
        let (num_agents, extent) = M::dims(params);
        let num_agents = num_agents.max(1);
        let extent = Extent {
            w: extent.w.max(1.0),
            h: extent.h.max(1.0),
        };
        let (width, height) = extent.cells();
        Geometry {
            num_agents,
            extent,
            width,
            height,
            n_cells: width * height,
            display: display_dims(width, height, limits.max_texture_dimension_2d),
            index: M::INDEX.then(|| HashGrid::new(extent, M::index_cell_size(params))),
        }
    }

    /// Resources that would be allocated for this model based on `params`.
    pub fn demand(params: &[ParamValue], limits: &wgpu::Limits) -> Demand {
        let geom = Self::geometry_for(params, limits);
        let mut demand = Demand::default();
        for (spec, len) in M::BUFFERS.iter().zip(M::buffer_lens(&geom)) {
            demand.push_sides(&format!("{}_{}", M::ID, spec.label), len, spec.double_buffered);
        }
        if let Some(grid) = geom.index {
            demand.push_index(M::ID, grid.num_cells(), geom.num_agents);
        }
        if M::DISPLAY.is_some() {
            demand.set_display(geom.width, geom.height, limits);
        }

        for (label, storage) in Self::declared_passes() {
            demand.push_pass(label, storage);
        }
        demand
    }

    /// Storage buffers each declared pass binds. Read by both [`Self::demand`] and
    /// [`Self::max_storage_bindings`], so the device that a host requests and the shortfall that the UI
    /// reports cannot disagree.
    fn declared_passes() -> Vec<(String, u32)> {
        let mut passes: Vec<(String, u32)> = M::STEP_PASSES
            .iter()
            .map(|spec| (format!("{}_{}", M::ID, spec.label), storage_bindings(spec.bindings)))
            .collect();
        if let Some(spec) = &M::DISPLAY {
            passes.push((format!("{}_display", M::ID), storage_bindings(spec.bindings)));
        }
        passes.push((format!("{}_reduce_leaf", M::ID), storage_bindings(M::REDUCE.bindings)));
        for action in M::ACTIONS {
            passes.push((
                format!("{}_action_{}", M::ID, action.desc.id),
                storage_bindings(action.pass.bindings),
            ));
        }
        passes
    }

    /// Returns the most storage buffers any declared pass binds. It does not depend on params, so a host can query it
    /// before it has a device.
    pub fn max_storage_bindings() -> u32 {
        Self::declared_passes()
            .into_iter()
            .map(|(_, storage)| storage)
            .max()
            .unwrap_or(0)
    }

    /// Builds the state on `ctx`, seeded with `seed`, or with the model's own default seed when it is `None`.
    ///
    /// # Panics
    ///
    /// Panics if the device cannot hold the model. A host calls [`Self::demand`] first, and this assert is the
    /// backstop. Also panics if a shader declares a workgroup size other than the one its pass dispatches, or a
    /// buffer label is reserved or ends in `_in` or `_out`.
    #[expect(clippy::too_many_lines, reason = "one linear construction of every wgpu object")]
    #[cfg_attr(
        all(target_arch = "wasm32", target_feature = "atomics"),
        expect(
            clippy::arc_with_non_send_sync,
            reason = "the agent layers hold wgpu buffers, which atomics leave unsendable"
        )
    )]
    pub fn new_seeded(ctx: &GpuContext, params: &[ParamValue], seed: Option<u64>) -> Self {
        let device = &ctx.device;
        let queue = &ctx.queue;

        let limits = device.limits();
        let shortfalls = Self::demand(params, &limits).shortfalls(&limits);
        assert!(
            shortfalls.is_empty(),
            "{} does not fit this device at these params: {}",
            M::ID,
            shortfalls.join("; ")
        );
        assert_buffer_labels(M::ID, M::BUFFERS.iter().map(|spec| spec.label));
        // Every pass but the display folds a linear domain onto workgroups of `WORKGROUP`.
        let linear = [WORKGROUP, 1, 1];
        for (pass, shader) in M::STEP_PASSES
            .iter()
            .map(|spec| (spec.label, spec.shader))
            .chain([("reduce leaf", M::REDUCE.shader)])
            .chain(M::ACTIONS.iter().map(|action| (action.desc.id, action.pass.shader)))
        {
            assert_workgroup_size(M::ID, pass, shader, linear);
        }
        if let Some(spec) = &M::DISPLAY {
            assert_workgroup_size(M::ID, "display", spec.shader, [spec.workgroup, spec.workgroup, 1]);
        }

        let mut geom = Self::geometry_for(params, &limits);
        let (num_agents, extent) = (geom.num_agents, geom.extent);
        let (width, height) = (geom.width, geom.height);

        // The hash fits its own grid to the index cell size. That size can differ from the field's cell size.
        let index =
            M::INDEX.then(|| GpuSpatialHash::new(device, queue, M::ID, extent, M::index_cell_size(params), num_agents));
        geom.index = index.as_ref().map(GpuSpatialHash::grid);

        // Buffers.
        let lens = M::buffer_lens(&geom);
        assert_eq!(
            lens.len(),
            M::BUFFERS.len(),
            "{}: buffer_lens must return one length per BUFFERS entry",
            M::ID
        );
        let seeds = M::seed_buffers(&geom, params, seed);
        assert_eq!(
            seeds.len(),
            M::BUFFERS.len(),
            "{}: seed_buffers must return one vector per BUFFERS entry",
            M::ID
        );

        let buffers: Vec<BufferSides> = M::BUFFERS
            .iter()
            .zip(&lens)
            .map(|(spec, &len)| {
                let make_buffer = |side: &str| {
                    let label = format!("{}_{}_{side}", M::ID, spec.label);
                    if spec.drawable {
                        lane_buffer(device, &label, len)
                    } else {
                        storage_buffer(device, &label, len)
                    }
                };
                BufferSides {
                    a: make_buffer("a"),
                    // Only the current side is seeded, so the other side is written by the first step.
                    b: spec.double_buffered.then(|| make_buffer("b")),
                }
            })
            .collect();

        let mut clear = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some(&format!("{}_seed", M::ID)),
        });
        for (k, (buffer, bytes)) in buffers.iter().zip(&seeds).enumerate() {
            if bytes.is_empty() {
                clear.clear_buffer(&buffer.a, 0, None);
                continue;
            }
            assert_eq!(
                bytes.len(),
                lens[k] * std::mem::size_of::<u32>(),
                "{}: seed for buffer '{}' must be buffer_lens[{k}] words of bytes",
                M::ID,
                M::BUFFERS[k].label
            );
            queue.write_buffer(&buffer.a, 0, bytes);
        }
        queue.submit(Some(clear.finish()));

        let ping_pong = M::BUFFERS.iter().any(|spec| spec.double_buffered);
        let index_binds = index.as_ref().map(|hash| {
            let pos = &buffers[M::POS_BUFFER];
            Sides {
                a: hash.bind_positions(device, &format!("{}_hash_bind_a", M::ID), &pos.a),
                b: pos
                    .b
                    .as_ref()
                    .map(|b| hash.bind_positions(device, &format!("{}_hash_bind_b", M::ID), b)),
            }
        });

        let counters =
            (M::COUNTERS > 0).then(|| CounterReadback::new(device, &format!("{}_counters", M::ID), M::COUNTERS));

        let display_spec = M::DISPLAY;
        // The display pass binds the view, and the snapshot carries the handle.
        let (display_view, display_handle) = match display_spec
            .as_ref()
            .map(|_| build_display_target(device, ctx.target_format, width, height))
        {
            Some(DisplayTarget { view, display, .. }) => (Some(view), Some(display)),
            None => (None, None),
        };

        let reduce_domain = M::REDUCE.domain.invocations(&geom);
        let reduce = GpuLaneReduce::new(device, queue, M::ID, M::REDUCE.lanes, reduce_domain);

        let build = PassBuilder {
            device,
            queue,
            geom: &geom,
            params,
            buffers: &buffers,
            ping_pong,
            index: index.as_ref(),
            counters: counters.as_ref(),
            display_view: display_view.as_ref(),
            reduce: &reduce,
        };

        let steps: Vec<EncodedPass> = M::STEP_PASSES
            .iter()
            .enumerate()
            .map(|(i, spec)| {
                let invocations = spec.domain.invocations(&geom);
                build.linear_pass::<M>(PassId::Step(i), spec.label, spec.shader, spec.bindings, invocations)
            })
            .collect();

        // One invocation per texel, as the grid engine's display pass does.
        let display = display_spec.as_ref().zip(display_handle).map(|(spec, handle)| {
            let (tex_w, tex_h) = geom.display;
            let groups = (tex_w.div_ceil(spec.workgroup), tex_h.div_ceil(spec.workgroup));
            let pass = build.pass::<M>(
                PassId::Display,
                "display",
                spec.shader,
                spec.bindings,
                groups,
                tex_w * tex_h,
            );
            (pass, handle)
        });

        let reduce_groups = reduce.agent_groups();
        let reduce_pass = build.pass::<M>(
            PassId::Reduce,
            "reduce_leaf",
            M::REDUCE.shader,
            M::REDUCE.bindings,
            reduce_groups,
            reduce_domain,
        );

        // Truncated from the same stream the CPU engines use, since the WGSL generator is 32 bit.
        let action_seed = henad_core::action::action_seed(seed) as u32;
        let actions: Vec<EncodedPass> = M::ACTIONS
            .iter()
            .enumerate()
            .map(|(i, action)| {
                let spec = &action.pass;
                let invocations = spec.domain.invocations(&geom);
                build.pass_in_place::<M>(
                    PassId::Action(i),
                    &format!("action_{}", action.desc.id),
                    spec.shader,
                    spec.bindings,
                    linear_dispatch(invocations),
                    invocations,
                    true,
                    action_seed,
                )
            })
            .collect();

        let make_agents = |a_is_current: bool| {
            let (pos, _) = buffers[M::POS_BUFFER].sides(a_is_current);
            let (color, _) = buffers[M::COLOR_BUFFER].sides(a_is_current);
            Arc::new(GpuAgents {
                pos: pos.clone(),
                color: color.clone(),
                count: num_agents,
                world_w: extent.w,
                world_h: extent.h,
            })
        };
        let agents = Sides {
            a: make_agents(true),
            b: ping_pong.then(|| make_agents(false)),
        };

        Self {
            geom,
            tick: 0,
            device: device.clone(),
            queue: queue.clone(),
            buffers,
            current_is_a: true,
            ping_pong,
            index,
            index_binds,
            steps,
            display,
            reduce,
            reduce_pass,
            counters,
            actions,
            action_seed,
            params: params.to_vec(),
            agents,
            _marker: PhantomData,
        }
    }

    /// Submits `steps` steps, at most [`MAX_STEPS_PER_SUBMISSION`] per command buffer, as the runner does. Nothing
    /// waits on the GPU.
    pub fn run_batched(&mut self, steps: u32) {
        let mut remaining = steps;
        while remaining > 0 {
            let batch = remaining.min(MAX_STEPS_PER_SUBMISSION);
            let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("henad_gpu_agent_batch"),
            });
            self.encode_steps(&mut encoder, batch, None);
            self.queue.submit(Some(encoder.finish()));
            remaining -= batch;
        }
    }

    /// Runs the snapshot passes and waits for the readback, so `stats()` reports the current tick.
    pub fn refresh_stats(&mut self) {
        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("henad_gpu_agent_snapshot"),
        });
        self.encode_snapshot_passes(&mut encoder);
        let device = self.device.clone();
        self.queue.submit(Some(encoder.finish()));
        self.begin_stats_readback();
        self.poll_stats_readback(&device, true);
    }

    /// Returns the current side of buffer `index` as raw words, blocking on the GPU.
    ///
    /// # Panics
    ///
    /// Panics if the readback fails.
    pub fn read_buffer(&self, index: usize) -> Vec<u32> {
        let (buffer, _) = self.buffers[index].sides(self.current_is_a);
        let size = buffer.size();
        let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("henad_gpu_agent_readback"),
            size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("henad_gpu_agent_readback"),
        });
        encoder.copy_buffer_to_buffer(buffer, 0, &staging, 0, size);
        self.queue.submit(Some(encoder.finish()));

        let (tx, rx) = flume::bounded(1);
        staging
            .slice(..)
            .map_async(wgpu::MapMode::Read, move |r| drop(tx.send(r)));
        self.device
            .poll(wgpu::PollType::wait_indefinitely())
            .expect("readback poll");
        rx.recv().expect("readback channel").expect("readback map");
        let data = staging.slice(..).get_mapped_range().expect("readback range");
        let out = bytemuck::cast_slice::<u8, u32>(&data).to_vec();
        drop(data);
        staging.unmap();
        out
    }

    /// Geometry the state was built at.
    pub fn geometry(&self) -> &Geometry {
        &self.geom
    }
}

/// Carries what every pass needs to resolve its bindings into wgpu objects.
struct PassBuilder<'a> {
    device: &'a wgpu::Device,
    queue: &'a wgpu::Queue,
    geom: &'a Geometry,
    params: &'a [ParamValue],
    buffers: &'a [BufferSides],
    ping_pong: bool,
    index: Option<&'a GpuSpatialHash>,
    counters: Option<&'a CounterReadback>,
    display_view: Option<&'a wgpu::TextureView>,
    reduce: &'a GpuLaneReduce,
}

impl PassBuilder<'_> {
    /// A pass over a linear invocation domain, folded onto the 2D workgroup grid.
    fn linear_pass<M: GpuAgentModel>(
        &self,
        id: PassId,
        label: &str,
        shader: &str,
        bindings: &[BindingDecl],
        invocations: u32,
    ) -> EncodedPass {
        self.pass::<M>(id, label, shader, bindings, linear_dispatch(invocations), invocations)
    }

    /// Builds a pass over `groups` workgroups.
    ///
    /// `groups.0` is the fold width that `henad::dispatch::linear_index` expects, and goes in the uniform as
    /// `groups_x`.
    fn pass<M: GpuAgentModel>(
        &self,
        id: PassId,
        label: &str,
        shader: &str,
        bindings: &[BindingDecl],
        groups: (u32, u32),
        invocations: u32,
    ) -> EncodedPass {
        self.pass_in_place::<M>(id, label, shader, bindings, groups, invocations, false, 0)
    }

    /// Builds a pass as [`Self::pass`] does, with `in_place` binding writes to the side that already holds the state.
    ///
    /// Nothing swaps after an action, so writing the far side would throw the work away.
    #[expect(
        clippy::too_many_arguments,
        reason = "one call site, and every argument is a pass fact"
    )]
    fn pass_in_place<M: GpuAgentModel>(
        &self,
        id: PassId,
        label: &str,
        shader: &str,
        bindings: &[BindingDecl],
        groups: (u32, u32),
        invocations: u32,
        in_place: bool,
        seed: u32,
    ) -> EncodedPass {
        let label = format!("{}_{label}", M::ID);

        let uniform = uniform_buffer(
            self.device,
            self.queue,
            &format!("{label}_params"),
            &M::pass_params_bytes(
                id,
                PassCtx {
                    geom: self.geom,
                    invocations,
                    groups_x: groups.0,
                    seed,
                },
                self.params,
            ),
        );

        let entries: Vec<wgpu::BindGroupLayoutEntry> = bindings
            .iter()
            .enumerate()
            .map(|(i, decl)| layout_entry(i as u32, decl))
            .collect();
        let layout = self.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some(&format!("{label}_layout")),
            entries: &entries,
        });

        let make_bind = |a_is_current: bool, side: &str| {
            let entries: Vec<wgpu::BindGroupEntry<'_>> = bindings
                .iter()
                .enumerate()
                .map(|(i, decl)| wgpu::BindGroupEntry {
                    binding: i as u32,
                    resource: self.resource::<M>(decl, a_is_current, in_place, &uniform),
                })
                .collect();
            self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some(&format!("{label}_bind_{side}")),
                layout: &layout,
                entries: &entries,
            })
        };

        let binds = Sides {
            a: make_bind(true, "a"),
            b: self.ping_pong.then(|| make_bind(false, "b")),
        };

        let pipeline = compute_pipeline(self.device, &label, shader, &layout);
        EncodedPass {
            label,
            pipeline,
            binds,
            groups,
            uniform,
        }
    }

    /// Returns the resource that a shader binding's name refers to. The name alone links the binding to its resource.
    ///
    /// Panics rather than returning an error, since a name that resolves to nothing is a shader
    /// and model that disagree, and no run can be correct after that.
    fn resource<'r, M: GpuAgentModel>(
        &'r self,
        decl: &BindingDecl,
        a_is_current: bool,
        in_place: bool,
        uniform: &'r wgpu::Buffer,
    ) -> wgpu::BindingResource<'r> {
        if let Some((label, writes)) = buffer_target(decl) {
            let k = M::BUFFERS
                .iter()
                .position(|spec| spec.label == label)
                .unwrap_or_else(|| panic!("{}: no buffer labelled `{label}`, wanted by `{}`", M::ID, decl.name));
            let (read, write) = self.buffers[k].sides(a_is_current);
            return if writes && !in_place { write } else { read }.as_entire_binding();
        }
        match decl.name {
            "params" => uniform.as_entire_binding(),
            "cell_start" => self.index.expect("INDEX is declared").cell_start_binding(),
            "sorted" => self.index.expect("INDEX is declared").sorted_binding(),
            "counters" => self.counters.expect("COUNTERS is non-zero").binding(),
            "partials" => self.reduce.partials_binding(),
            "output" => wgpu::BindingResource::TextureView(self.display_view.expect("DISPLAY is declared")),
            other => panic!("{}: `{other}` is reserved but the engine has no resource for it", M::ID),
        }
    }
}

impl<M: GpuAgentModel> SimState for GpuAgentState<M> {
    /// Steps once in its own submission, for a caller that holds only a `SimState`. The GPU runner batches
    /// instead.
    fn step(&mut self) {
        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("gpu_agent_single_step"),
        });
        self.encode_steps(&mut encoder, 1, None);
        self.queue.submit(Some(encoder.finish()));
    }

    fn tick(&self) -> u64 {
        self.tick
    }

    fn stats(&self) -> Vec<StatEntry> {
        let counters = self.counters.as_ref().map_or(&[][..], CounterReadback::values);
        stat_entries(M::STATS, M::stats(&self.reduce.sums(), counters, &self.geom))
    }

    /// Encodes the action into its own submission.
    ///
    /// The GPU runner calls [`GpuSimState::encode_action`] instead, and submits the action before
    /// its snapshot.
    fn act(&mut self, index: usize) -> bool {
        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("henad_gpu_agent_action"),
        });
        if !GpuSimState::encode_action(self, &mut encoder, index) {
            return false;
        }
        self.queue.submit(Some(encoder.finish()));
        true
    }

    /// Rejects every live edit. Its entry declares every parameter reload-only.
    fn set_param(&mut self, _index: usize, _value: &ParamValue) -> bool {
        false
    }

    fn population(&self) -> u64 {
        u64::from(self.geom.num_agents)
    }

    fn heap_bytes(&self) -> usize {
        let buffers: usize = self
            .buffers
            .iter()
            .map(|sides| (sides.a.size() + sides.b.as_ref().map_or(0, wgpu::Buffer::size)) as usize)
            .sum();
        let display = self.display.as_ref().map_or(0, |_| {
            (self.geom.display.0 as usize) * (self.geom.display.1 as usize) * 4
        });
        buffers
            + display
            + self.index.as_ref().map_or(0, GpuSpatialHash::heap_bytes)
            + self.reduce.heap_bytes()
            + M::COUNTERS * std::mem::size_of::<u32>()
    }
}

impl<M: GpuAgentModel> GpuSimState for GpuAgentState<M> {
    /// Records one compute pass per declared pass per step, all into one encoder.
    ///
    /// A pass is the synchronization boundary wgpu inserts barriers at, so a step's passes
    /// cannot be collapsed into dispatches inside one pass. The later ones would read stale
    /// data.
    fn encode_steps(&mut self, encoder: &mut wgpu::CommandEncoder, count: u32, timestamps: Option<&wgpu::QuerySet>) {
        if count == 0 {
            return;
        }
        let last_pass = self.steps.len().saturating_sub(1);

        for i in 0..count {
            let is_first = i == 0;
            let is_last = i == count - 1;

            // A batch begins with the index rebuild, so the opening stamp goes on its counting
            // pass. A stamp on an empty pass is silently never written.
            if let (Some(hash), Some(binds)) = (&self.index, &self.index_binds) {
                hash.encode_build(
                    encoder,
                    binds.pick(self.current_is_a),
                    timestamps.filter(|_| is_first).map(|query_set| (query_set, 0)),
                );
            }

            for (j, pass) in self.steps.iter().enumerate() {
                let opening = is_first && j == 0 && self.index.is_none();
                let closing = is_last && j == last_pass;
                pass.encode(encoder, self.current_is_a, stamps(timestamps, opening, closing));
            }

            if self.ping_pong {
                self.current_is_a = !self.current_is_a;
            }
        }

        self.tick += u64::from(count);
    }

    fn encode_action(&mut self, encoder: &mut wgpu::CommandEncoder, index: usize) -> bool {
        let Some(action) = self.actions.get(index) else {
            return false;
        };
        // A fresh seed per press, so pressing twice draws twice. The write is queued before the
        // encoder is submitted, so the pass reads the new value.
        self.action_seed = self.action_seed.wrapping_mul(747_796_405).wrapping_add(2_891_336_453);
        self.queue.write_buffer(
            &action.uniform,
            0,
            &M::pass_params_bytes(
                PassId::Action(index),
                PassCtx {
                    geom: &self.geom,
                    invocations: M::ACTIONS[index].pass.domain.invocations(&self.geom),
                    groups_x: action.groups.0,
                    seed: self.action_seed,
                },
                &self.params,
            ),
        );
        action.encode(encoder, self.current_is_a, None);
        true
    }

    fn encode_snapshot_passes(&mut self, encoder: &mut wgpu::CommandEncoder) {
        if let Some((pass, _)) = &self.display {
            pass.encode(encoder, self.current_is_a, None);
        }
        self.encode_stats_passes(encoder);
    }

    fn encode_stats_passes(&mut self, encoder: &mut wgpu::CommandEncoder) {
        self.reduce_pass.encode(encoder, self.current_is_a, None);
        self.reduce.encode(encoder);
        if let Some(counters) = &mut self.counters {
            counters.encode_copy(encoder);
        }
    }

    fn begin_stats_readback(&mut self) {
        self.reduce.begin_readback();
        if let Some(counters) = &mut self.counters {
            counters.begin_map();
        }
    }

    fn poll_stats_readback(&mut self, device: &wgpu::Device, block: bool) -> StatsPoll {
        let sums_poll = self.reduce.poll_readback(device, block);
        let counters_poll = match &mut self.counters {
            Some(readback) if block => readback.poll_blocking(device),
            Some(readback) => readback.poll(device),
            None => StatsPoll::Landed,
        };
        match (sums_poll, counters_poll) {
            (StatsPoll::Pending, _) | (_, StatsPoll::Pending) => StatsPoll::Pending,
            (StatsPoll::Failed, _) | (_, StatsPoll::Failed) => StatsPoll::Failed,
            (StatsPoll::Landed, StatsPoll::Landed) => StatsPoll::Landed,
        }
    }

    fn stats_readback_pending(&self) -> bool {
        self.reduce.readback_pending() || self.counters.as_ref().is_some_and(CounterReadback::is_pending)
    }

    fn view(&self) -> GpuSnapshot {
        GpuSnapshot {
            display: self.display.as_ref().map(|(_, display)| Arc::clone(display)),
            agents: Some(Arc::clone(self.agents.pick(self.current_is_a))),
        }
    }
}