fidget-wgpu 0.5.1

//! GPU-accelerated 3D rendering
//!
//! # Theory
//!
//! This module implements an algorithm similar to the one described in
//! [Massively Parallel Rendering of Complex Closed-Form Implicit Surfaces (Keeter '20)](https://www.mattkeeter.com/research/mpr/).
//! The rest of this section is intended for people who have read that paper
//! ("MPR" for short).
//!
//! We use interval arithmetic on a high-fanout hierarchy of tiles (64³, 16³,
//! 4³), followed by voxel and normal evaluation.  At each stage of interval
//! arithmetic, we compute a **simplified tape** for each tile containing only
//! portions of the expression which are active.
//!
//! After the root tile evaluation, tiles are sparse.  Tiles and tapes use an
//! atomic bump allocator to claim portions of a fixed buffer.  The tile buffer
//! is always sized to fit all possible tiles; the tape buffer can run out of
//! space, in which case we fall back to the previous (unsimplified) tape.
//!
//! ## Changes versus MPR
//!
//! There are a few notable changes compared to the MPR paper and reference
//! implementation.
//!
//! First, modern GPU APIs support indirect dispatch based on buffers on the GPU
//! itself.  This saves a round-trip: the 64³ shader can compute a dispatch size
//! for the 16³ shader and store it in a buffer (and so on for subsequent
//! stages).
//!
//! In a more significant change, evaluation is broken into **strata**:
//!
//! - The initial pass of 64³ tiles renders all of those tiles, any which are
//!   active are accumulated into a set of `depth / 64` strata
//! - Strata are evaluated one at a time in z-sorted order; this is where 16³,
//!   4³, voxel, and normal evaluation happens.  You can think of this as doing
//!   raymarching on 64³ voxels at a time.
//!
//! Strata-sorted evaluation has a few advantages:
//!
//! - We can statically allocate enough space for all tiles: all 64³ in the
//!   image, then all 16³ and 4³ tiles in a single strata.  It would be
//!   prohibitive to allocate storage for all 4³ tiles in the entire volume, but
//!   doing per-strata evaluation reduces the memory scaling from N³ to N².
//! - We get some amount of Z culling, because each pass can bail out if the
//!   result in the heightmap fully covers the tile
//!
//! # Practice
//! There are four core objects, each with different lifetimes
//!
//! - [`Context`] contains all of the pipelines used for 3D rendering.  It
//!   is very expensive to build and should be constructed once per thread /
//!   worker.
//! - [`RenderShape`] contains serialized bytecode to render a particular shape.
//!   Best practice is to rebuild it only when a shape changes (i.e. not once
//!   per frame), although in practice it's pretty fast to construct.
//! - [`Workspace`] contains GPU buffers needed for rendering at a particular
//!   image size.  It is primarily expensive in GPU memory, as it contains
//!   several full-frame buffers.  Best practice is to construct one
//!   [`Workspace`] object per worker context (or per simultaneous render); it
//!   will be automatically resized when used.  Systems with high variability in
//!   image size may want to periodically compare [`size`](Workspace::size)
//!   versus [`capacity`](Workspace::capacity) and fully reallocate buffers (by
//!   constructing a new `Workspace` object) if they get too out of whack.
//! - [`RenderConfig`] sets the transform matrix for rendering.  This is cheap
//!   to construct and could be built once per frame
//!
//! With all that out of the way, usage is pretty simple:
//! - Build a [`Context`]
//! - Use [`RenderShape::new`] to convert from a
//!   [`VmShape`](fidget_core::vm::VmShape) to a [`RenderShape`]
//! - Use [`Context::workspace`] to get [`Workspace`]
//! - Use [`Gpu::read_buffer_for(buffers.output())`](Gpu::read_buffer_for) to
//!   get an output buffer
//! - Call [`Context::run`] or [`Context::run_async`] to get an image
//!
//! ## Sync and async operation
//!
//! GPU operations are asynchronous; operations are submitted to a queue, and
//! are completed at some point in the future.  [`Context::run`] blocks until
//! operations are complete, but is only valid on the desktop (and is disabled
//! by feature flag on the web); it uses [`wgpu::Device::poll`], which is a
//! no-op on the web. [`Context::run_async`] is the async equivalent, and can be
//! used either natively or on the web.
//!
//! ## Low-level building blocks
//!
//! [`Context::run`] and `run_async` do four things:
//!
//! - Run the GPU kernels to produce an output image, which is a
//!   [`GeometryPixel`] array in a GPU storage buffer
//! - Copy from that GPU storage buffer to a mappable buffer (for read-back)
//! - Map that buffer into a [`MappedImage`](crate::buf::MappedImage)
//! - Read image data back to the CPU
//!
//! Lower-level building blocks are also available: [`Context::submit`] submits
//! the render operations to the GPU, and [`Workspace::output`] returns the
//! output buffer.

use crate::{
    CopyVarsChanged, CopyVarsError, Gpu, RegPipeline, RenderShape,
    TAPE_DATA_CAPACITY, TapeWord,
    buf::{
        BufferSizeError, BufferType, FlexBuffer, ReadBuffer, buffer_ro,
        buffer_ro_dyn, buffer_rw,
    },
    shaders, tag,
};
use fidget_core::{
    render::{ImageSize, VoxelSize},
    shape::{MissingVar, ShapeVars},
};
use fidget_raster::voxel::{GeometryPixel, Image};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

pub mod effects;

pub use fidget_raster::voxel::{RenderConfig, RenderSize};

const COMMON_SHADER: &str = include_str!("shaders/common.wgsl");
const INTERVAL_INPUT: &str = include_str!("shaders/interval_input.wgsl");
const TRANSFORM_INPUT: &str = include_str!("shaders/transform_input.wgsl");
const VOXEL_TILES_SHADER: &str = include_str!("shaders/voxel_tiles.wgsl");
const INTERVAL_TILES_SHADER: &str = include_str!("shaders/interval_tiles.wgsl");
const REPACK_SHADER: &str = include_str!("shaders/repack.wgsl");
const SORT_SHADER: &str = include_str!("shaders/sort.wgsl");
const INTERVAL_ROOT_SHADER: &str = include_str!("shaders/interval_root.wgsl");
const CLEAR_SHADER: &str = include_str!("shaders/clear.wgsl");
const MERGE_SHADER: &str = include_str!("shaders/merge.wgsl");
const NORMALS_SHADER: &str = include_str!("shaders/normals.wgsl");

/// Error type when resizing intermediate tile buffers
#[derive(Debug, thiserror::Error)]
#[error("failed to resize `{buf}` tile buffer")]
pub struct TileBuffersError {
    /// Buffer which failed to resize
    pub buf: TileBufferName,
    /// Error returned by buffer resizing
    #[source]
    pub err: BufferSizeError,
}

/// Names of buffers used by the intermediate tile rendering pass
///
/// This is only used for error reporting
#[derive(Debug)]
#[expect(missing_docs)]
pub enum TileBufferName {
    Tiles,
    Sorted,
    Zmin,
}

impl std::fmt::Display for TileBufferName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            TileBufferName::Tiles => "tiles",
            TileBufferName::Sorted => "sorted",
            TileBufferName::Zmin => "zmin",
        };
        s.fmt(f)
    }
}

/// Error type when resizing root tile buffers
#[derive(Debug, thiserror::Error)]
#[error("failed to resize `{buf}` root tile buffer")]
pub struct RootTileBuffersError {
    /// Buffer which failed to resize
    pub buf: RootTileBufferName,
    /// Error returned by buffer resizing
    #[source]
    pub err: BufferSizeError,
}

/// Names of buffers used by the root tile rendering pass (for error reporting)
#[derive(Debug)]
#[expect(missing_docs)]
pub enum RootTileBufferName {
    Tiles,
    Strata,
    Zmin,
    Zmax,
}

impl std::fmt::Display for RootTileBufferName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            RootTileBufferName::Tiles => "tiles",
            RootTileBufferName::Strata => "strata",
            RootTileBufferName::Zmin => "zmin",
            RootTileBufferName::Zmax => "zmax",
        };
        s.fmt(f)
    }
}

/// Names of all buffers, used for error reporting
#[derive(Debug)]
#[expect(missing_docs)]
pub enum BufferName {
    /// Tiles from the 64³ root tile pass
    Tile64(RootTileBufferName),
    /// Tiles from the 16³ intermediate tile pass
    Tile16(TileBufferName),
    /// Tiles from the 4³ intermediate tile pass
    Tile4(TileBufferName),
    TileTapes,
    Voxels,
    Heightmap,
    Geom,
    Image,
    Vars,
}

impl std::fmt::Display for BufferName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BufferName::Tile64(buf) => write!(f, "`{buf}` tile64"),
            BufferName::Tile16(buf) => write!(f, "`{buf}` tile16"),
            BufferName::Tile4(buf) => write!(f, "`{buf}` tile4"),
            BufferName::TileTapes => write!(f, "`tile tapes`"),
            BufferName::Voxels => write!(f, "`voxels`"),
            BufferName::Heightmap => write!(f, "`heightmap`"),
            BufferName::Geom => write!(f, "`geom`"),
            BufferName::Image => write!(f, "`image`"),
            BufferName::Vars => write!(f, "`vars`"),
        }
    }
}

/// Error returned when resizing buffers in a [`Workspace`] object
#[derive(Debug, thiserror::Error)]
#[error("failed to build {buf} buffer")]
pub struct BufferError {
    /// Buffer which failed to resize
    pub buf: BufferName,
    /// Error returned by buffer resizing
    #[source]
    pub err: BufferSizeError,
}

/// Error returned when submitting a voxel rasterization job to the GPU
#[derive(Debug, thiserror::Error)]
pub enum SubmitError {
    /// Missing variable when evaluating
    #[error(transparent)]
    MissingVar(#[from] MissingVar),
    /// Error while resizing buffers
    #[error(transparent)]
    Buffer(#[from] BufferError),
}

impl From<CopyVarsError> for SubmitError {
    fn from(value: CopyVarsError) -> Self {
        match value {
            CopyVarsError::MissingVar(v) => Self::MissingVar(v),
            CopyVarsError::BufferSize(err) => Self::Buffer(BufferError {
                buf: BufferName::Vars,
                err,
            }),
        }
    }
}

////////////////////////////////////////////////////////////////////////////////

/// Doppelganger of the WGSL `struct Config`
///
/// Fields are carefully ordered to require no internal padding (enforced by
/// `zerocopy` derives)
#[derive(Debug, IntoBytes, Immutable, FromBytes, KnownLayout)]
#[cfg_attr(test, derive(facet::Facet))]
#[repr(C)]
struct Config {
    /// Screen-to-model transform matrix
    mat: [f32; 16],

    /// Input index of X, Y, Z axes
    ///
    /// `u32::MAX` is used as a marker if an axis is unused
    axes: [u32; 3],

    /// Initial offset in `tape_data`
    tape_data_offset: u32,

    /// Render size, rounded up to the nearest multiple of 64
    render_size: [u32; 3],

    /// Number of words in the trailing tape buffer
    tape_data_capacity: u32,

    /// Image size (not rounded)
    image_size: [u32; 3],

    /// Length of the root tape
    root_tape_len: u32,
    // This is followed by a flexible array member containing tape data
}

/// A render size is rounded up to the next multiple of 64 on every axis
///
/// The internal `VoxelSize` stores divided-by-64 values, so that the render
/// size cannot be constructed with an invalid state.
#[derive(Copy, Clone, Debug)]
struct TileRenderSize(VoxelSize);

impl From<VoxelSize> for TileRenderSize {
    fn from(image_size: VoxelSize) -> Self {
        let nx = image_size.width().div_ceil(64);
        let ny = image_size.height().div_ceil(64);
        let nz = image_size.depth().div_ceil(64);
        Self(VoxelSize::new(nx, ny, nz))
    }
}

impl TileRenderSize {
    /// Number of tiles in the X axis
    fn nx(&self) -> u32 {
        self.0.width()
    }

    /// Number of tiles in the Y axis
    fn ny(&self) -> u32 {
        self.0.height()
    }

    /// Number of tiles in the Z axis
    fn nz(&self) -> u32 {
        self.0.depth()
    }

    /// Number of voxels in the X axis (always a multiple of 64)
    fn width(&self) -> u32 {
        self.0.width() * 64
    }

    /// Number of voxels in the Y axis (always a multiple of 64)
    fn height(&self) -> u32 {
        self.0.height() * 64
    }

    /// Number of voxels in the Z axis (always a multiple of 64)
    fn depth(&self) -> u32 {
        self.0.depth() * 64
    }
}

/// Returns a shader for interval root tiles
fn interval_root_shader(reg_count: u8) -> String {
    let mut shader_code = shaders::opcode_constants();
    shader_code += &format!("const REG_COUNT: u32 = {reg_count};");
    shader_code += COMMON_SHADER;
    shader_code += INTERVAL_ROOT_SHADER;
    shader_code += INTERVAL_INPUT;
    shader_code += TRANSFORM_INPUT;
    shader_code += shaders::INTERVAL_OPS;
    shader_code += shaders::COMMON;
    shader_code += shaders::TAPE_INTERPRETER;
    shader_code += shaders::STACK;
    shader_code += shaders::TAPE_SIMPLIFY;
    shader_code
}

/// Returns a shader for interval root tile repacking
fn repack_shader() -> String {
    let mut shader_code = String::new();
    shader_code += REPACK_SHADER;
    shader_code += COMMON_SHADER;
    shader_code += shaders::COMMON;
    shader_code
}

/// Returns a shader for interval tile sorting
fn sort_shader() -> String {
    let mut shader_code = String::new();
    shader_code += SORT_SHADER;
    shader_code += COMMON_SHADER;
    shader_code += shaders::COMMON;
    shader_code
}

/// Returns a shader for interval tile evaluation
fn interval_tiles_shader(reg_count: u8) -> String {
    let mut shader_code = shaders::opcode_constants();
    shader_code += &format!("const REG_COUNT: u32 = {reg_count};");
    shader_code += INTERVAL_TILES_SHADER;
    shader_code += INTERVAL_INPUT;
    shader_code += TRANSFORM_INPUT;
    shader_code += COMMON_SHADER;
    shader_code += shaders::INTERVAL_OPS;
    shader_code += shaders::COMMON;
    shader_code += shaders::TAPE_INTERPRETER;
    shader_code += shaders::STACK;
    shader_code += shaders::TAPE_SIMPLIFY;
    shader_code
}

/// Returns a shader for voxel tile evaluation
fn voxel_tiles_shader(reg_count: u8) -> String {
    let mut shader_code = shaders::opcode_constants();
    shader_code += &format!("const REG_COUNT: u32 = {reg_count};");
    shader_code += VOXEL_TILES_SHADER;
    shader_code += TRANSFORM_INPUT;
    shader_code += COMMON_SHADER;
    shader_code += shaders::FLOAT_OPS;
    shader_code += shaders::COMMON;
    shader_code += shaders::TAPE_INTERPRETER;
    shader_code += shaders::DUMMY_STACK;
    shader_code
}

/// Returns a shader for normals evaluation
fn normals_shader(reg_count: u8) -> String {
    let mut shader_code = shaders::opcode_constants();
    shader_code += &format!("const REG_COUNT: u32 = {reg_count};");
    shader_code += NORMALS_SHADER;
    shader_code += TRANSFORM_INPUT;
    shader_code += COMMON_SHADER;
    shader_code += shaders::GRAD_OPS;
    shader_code += shaders::COMMON;
    shader_code += shaders::TAPE_INTERPRETER;
    shader_code += shaders::DUMMY_STACK;
    shader_code
}

/// Returns a shader for merging images
fn merge_shader() -> String {
    MERGE_SHADER.to_owned() + COMMON_SHADER + shaders::COMMON
}

/// Returns a shader for clearing counters in between strata passes
fn clear_shader() -> String {
    CLEAR_SHADER.to_owned() + COMMON_SHADER + shaders::COMMON
}

////////////////////////////////////////////////////////////////////////////////

/// Root context, which produces a list of 64³ tiles
struct RootContext {
    /// Pipelines for 64³ tile evaluation
    root_pipeline: RegPipeline,

    /// Bind group layout
    bind_group_layout: wgpu::BindGroupLayout,
}

/// Per-strata offset in the root tiles list
///
/// This must be equivalent to `strata_size_bytes` in the interval root shader
fn strata_size_bytes(render_size: TileRenderSize) -> usize {
    let nx = usize::try_from(render_size.nx()).unwrap();
    let ny = usize::try_from(render_size.ny()).unwrap();
    // Snap to `min_storage_buffer_offset_alignment`
    ((nx * ny + 4) * std::mem::size_of::<u32>()).next_multiple_of(256)
}

impl RootContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("root"),
                entries: &[
                    buffer_rw(0), // tiles_out
                    buffer_rw(1), // tile64_zmax
                ],
            });

        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("root pipeline"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });
        let device_ = device.clone();
        let root_pipeline = RegPipeline::build(Box::new(move |reg_count| {
            let shader_code = interval_root_shader(reg_count);
            let shader_module =
                device_.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: Some("interval root"),
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                });
            device_.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(&format!("interval root ({reg_count})")),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("interval_root_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        }));

        Self {
            bind_group_layout,
            root_pipeline,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        reg_count: u8,
        render_size: TileRenderSize,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let bind_group = workspace.bind_groups.root(ctx, workspace);
        compute_pass.set_pipeline(self.root_pipeline.get(reg_count));
        compute_pass.set_bind_group(1, bind_group, &[]);

        // Workgroup is 4x4x4, so we divide by 4 here on each axis
        let nx = render_size.nx().div_ceil(4);
        let ny = render_size.ny().div_ceil(4);
        let nz = render_size.nz().div_ceil(4);
        compute_pass.dispatch_workgroups(nx, ny, nz);
    }
}

/// Repack context, which strata-sorts a list of 64³ tiles
struct RepackContext {
    /// Pipeline for 64³ tile packing
    repack_pipeline: wgpu::ComputePipeline,

    /// Bind group layout
    bind_group_layout: wgpu::BindGroupLayout,
}

impl RepackContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("repack"),
                entries: &[
                    buffer_ro(0), // tiles_out
                    buffer_ro(1), // tile64_zmin
                    buffer_rw(2), // strata_tiles
                ],
            });

        let repack_pipeline = {
            let shader_code = repack_shader();
            let pipeline_layout = device.create_pipeline_layout(
                &wgpu::PipelineLayoutDescriptor {
                    label: Some("repack"),
                    bind_group_layouts: &[
                        Some(common_bind_group_layout),
                        Some(&bind_group_layout),
                    ],
                    immediate_size: 0u32,
                },
            );
            let shader_module =
                device.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: Some("repack"),
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                });
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("repack"),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("repack_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        };

        Self {
            bind_group_layout,
            repack_pipeline,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        render_size: TileRenderSize,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let bind_group = workspace.bind_groups.repack(ctx, workspace);

        compute_pass.set_pipeline(&self.repack_pipeline);
        compute_pass.set_bind_group(1, bind_group, &[]);

        // Workgroup is 64x1x1, so we divide on the X axis.  It doesn't matter
        // much; we just need one thread per possible output tile from the
        // previous stage, i.e. `(nx * ny * nz)` total threads.  This could be
        // optimized further with indirect dispatch, but ehhhhhhh
        let nx = render_size.nx().div_ceil(64);
        let ny = render_size.ny();
        let nz = render_size.nz();
        compute_pass.dispatch_workgroups(nx, ny, nz);
    }
}

////////////////////////////////////////////////////////////////////////////////

struct IntervalContext {
    /// Pipeline for 64³ -> 16³ tile evaluation
    interval64_pipeline: RegPipeline,

    /// Pipeline to sort 16³ tiles
    sort16_pipeline: wgpu::ComputePipeline,

    /// Pipeline for 16³ -> 4³ tile evaluation
    interval16_pipeline: RegPipeline,

    /// Pipeline to sort 4³ tiles
    sort4_pipeline: wgpu::ComputePipeline,

    /// Bind group layout for interval pipelines
    interval_bind_group_layout: wgpu::BindGroupLayout,

    /// Bind group layout for sort pipelines
    sort_bind_group_layout: wgpu::BindGroupLayout,
}

impl IntervalContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let interval_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("interval"),
                entries: &[
                    buffer_ro_dyn(0), // tiles_in
                    buffer_ro(1),     // tile_zmin
                    buffer_rw(2),     // subtiles_out
                    buffer_rw(3),     // subtile_zmin
                    buffer_rw(4),     // subtile_zhist
                ],
            });

        let interval_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("interval pipeline layout"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&interval_bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        let device_ = device.clone();
        let interval_pipeline_layout_ = interval_pipeline_layout.clone();
        let interval64_pipeline =
            RegPipeline::build(Box::new(move |reg_count| {
                let shader_code = interval_tiles_shader(reg_count);
                // SAFETY: the shader is carefully written
                let shader_module = unsafe {
                    device_.create_shader_module_trusted(
                        wgpu::ShaderModuleDescriptor {
                            label: Some(&format!(
                                "interval64 tiles shader ({reg_count})"
                            )),
                            source: wgpu::ShaderSource::Wgsl(
                                shader_code.into(),
                            ),
                        },
                        wgpu::ShaderRuntimeChecks::unchecked(),
                    )
                };
                device_.create_compute_pipeline(
                    &wgpu::ComputePipelineDescriptor {
                        label: Some(&format!("interval64 ({reg_count})")),
                        layout: Some(&interval_pipeline_layout_),
                        module: &shader_module,
                        entry_point: Some("interval_tile_main"),
                        compilation_options: wgpu::PipelineCompilationOptions {
                            constants: &[
                                ("TILE_SIZE", 64.0),
                                ("SUBTILE_SIZE", 16.0),
                            ],
                            ..Default::default()
                        },
                        cache: None,
                    },
                )
            }));

        let device_ = device.clone();
        let interval16_pipeline =
            RegPipeline::build(Box::new(move |reg_count| {
                let shader_code = interval_tiles_shader(reg_count);
                // SAFETY: the shader is carefully written
                let shader_module = unsafe {
                    device_.create_shader_module_trusted(
                        wgpu::ShaderModuleDescriptor {
                            label: Some(&format!(
                                "interval16 tiles shader ({reg_count})"
                            )),
                            source: wgpu::ShaderSource::Wgsl(
                                shader_code.into(),
                            ),
                        },
                        wgpu::ShaderRuntimeChecks::unchecked(),
                    )
                };
                device_.create_compute_pipeline(
                    &wgpu::ComputePipelineDescriptor {
                        label: Some(&format!("interval16 ({reg_count})")),
                        layout: Some(&interval_pipeline_layout),
                        module: &shader_module,
                        entry_point: Some("interval_tile_main"),
                        compilation_options: wgpu::PipelineCompilationOptions {
                            constants: &[
                                ("TILE_SIZE", 16.0),
                                ("SUBTILE_SIZE", 4.0),
                            ],
                            ..Default::default()
                        },
                        cache: None,
                    },
                )
            }));

        let sort_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("sort bind group layout"),
                entries: &[
                    buffer_ro(0), // subtiles_out
                    buffer_rw(1), // z_hist
                    buffer_rw(2), // sorted_subtiles
                ],
            });
        let sort_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("sort pipeline layout"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&sort_bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        let shader_code = sort_shader();
        // SAFETY: the shader is carefully written
        let shader_module = unsafe {
            device.create_shader_module_trusted(
                wgpu::ShaderModuleDescriptor {
                    label: Some("sort shader module"),
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                },
                wgpu::ShaderRuntimeChecks::unchecked(),
            )
        };
        let sort16_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("sort16"),
                layout: Some(&sort_pipeline_layout),
                module: &shader_module,
                entry_point: Some("sort_main"),
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &[("SUBTILE_SIZE", 16.0)],
                    ..Default::default()
                },
                cache: None,
            });
        let sort4_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("sort4"),
                layout: Some(&sort_pipeline_layout),
                module: &shader_module,
                entry_point: Some("sort_main"),
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &[("SUBTILE_SIZE", 4.0)],
                    ..Default::default()
                },
                cache: None,
            });

        Self {
            interval_bind_group_layout,
            sort_bind_group_layout,
            interval64_pipeline,
            sort16_pipeline,
            interval16_pipeline,
            sort4_pipeline,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        strata: u64,
        reg_count: u8,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let strata_bytes =
            u64::try_from(workspace.strata_size_bytes()).unwrap();
        let offset_bytes = strata * strata_bytes;
        let bind_group16 = workspace.bind_groups.interval16(ctx, workspace);
        compute_pass.set_pipeline(self.interval64_pipeline.get(reg_count));
        compute_pass.set_bind_group(
            1,
            bind_group16,
            &[u32::try_from(offset_bytes).unwrap()],
        );
        compute_pass.dispatch_workgroups_indirect(
            workspace.tile64.strata.data(),
            offset_bytes,
        );

        let bind_group_sort16 = workspace.bind_groups.sort16(ctx, workspace);
        compute_pass.set_pipeline(&self.sort16_pipeline);
        compute_pass.set_bind_group(1, bind_group_sort16, &[]);
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile16.tiles.data(), 0);

        let bind_group4 = workspace.bind_groups.interval4(ctx, workspace);
        compute_pass.set_pipeline(self.interval16_pipeline.get(reg_count));
        compute_pass.set_bind_group(1, bind_group4, &[0]);
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile16.sorted.data(), 0);

        let bind_group_sort4 = workspace.bind_groups.sort4(ctx, workspace);
        compute_pass.set_pipeline(&self.sort4_pipeline);
        compute_pass.set_bind_group(1, bind_group_sort4, &[]);
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile4.tiles.data(), 0);
    }
}

////////////////////////////////////////////////////////////////////////////////

struct VoxelContext {
    /// Bind group layout
    bind_group_layout: wgpu::BindGroupLayout,

    /// Pipeline for interpreted voxel evaluation
    voxel_pipeline: RegPipeline,
}

impl VoxelContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("voxel bind group layout"),
                entries: &[
                    buffer_ro(0), // tiles4_in
                    buffer_ro(1), // tile4_zmin
                    buffer_rw(2), // result
                ],
            });

        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("voxel pipeline layout"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });
        let device_ = device.clone();
        let voxel_pipeline = RegPipeline::build(Box::new(move |reg_count| {
            let shader_code = voxel_tiles_shader(reg_count);
            // SAFETY: The shader is careful, good luck
            let shader_module = unsafe {
                device_.create_shader_module_trusted(
                    wgpu::ShaderModuleDescriptor {
                        label: Some("voxel shader module"),
                        source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                    },
                    wgpu::ShaderRuntimeChecks::unchecked(),
                )
            };
            device_.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(&format!("voxels ({reg_count})")),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("voxel_ray_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        }));

        Self {
            bind_group_layout,
            voxel_pipeline,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        reg_count: u8,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let bind_group = workspace.bind_groups.voxel(ctx, workspace);
        compute_pass.set_pipeline(self.voxel_pipeline.get(reg_count));
        compute_pass.set_bind_group(1, bind_group, &[]);

        // Each workgroup is 4x4x4, i.e. covering a 4x4 splat of pixels with 4x
        // workers in the Z direction.
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile4.sorted.data(), 0);
    }
}

struct NormalsContext {
    /// Bind group layout
    bind_group_layout: wgpu::BindGroupLayout,

    /// Pipeline for normal evaluation
    normals_pipeline: RegPipeline,
}

impl NormalsContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("normals bind group layout"),
                entries: &[
                    buffer_ro(0), // image_heightmap
                    buffer_rw(1), // image_out
                ],
            });

        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("normals pipeline"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });
        let device_ = device.clone();
        let normals_pipeline = RegPipeline::build(Box::new(move |reg_count| {
            let shader_code = normals_shader(reg_count);
            let shader_module =
                device_.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: Some("normals shader module"),
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                });
            device_.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(&format!("normals ({reg_count})")),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("normals_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        }));

        Self {
            bind_group_layout,
            normals_pipeline,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        reg_count: u8,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let bind_group = workspace.bind_groups.normals(ctx, workspace);
        compute_pass.set_pipeline(self.normals_pipeline.get(reg_count));
        compute_pass.set_bind_group(1, bind_group, &[]);

        compute_pass.dispatch_workgroups(
            workspace.image_size.width().div_ceil(8),
            workspace.image_size.height().div_ceil(8),
            1,
        );
    }
}

////////////////////////////////////////////////////////////////////////////////

/// Context for 3D (combined heightmap and normal) rendering
pub struct Context {
    gpu: Gpu,

    /// Bind group layout for the common bind group (used by all stages)
    common_bind_group_layout: wgpu::BindGroupLayout,

    root_ctx: RootContext,
    repack_ctx: RepackContext,
    interval_ctx: IntervalContext,
    voxel_ctx: VoxelContext,
    normals_ctx: NormalsContext,
    merge_ctx: MergeContext,
    reset_ctx: ResetContext,
    clear_ctx: ClearContext,
}

tag!(TilesBufferTag, u32, usize, STORAGE | INDIRECT);
tag!(SortedBufferTag, u32, usize, STORAGE | INDIRECT);
tag!(ZminBufferTag, u32, ImageSize, STORAGE | COPY_DST);

struct TileBuffers<const N: u64> {
    /// Tiles written by the stage outputting N³ tiles
    tiles: FlexBuffer<TilesBufferTag>,
    /// Sorted version of [`tiles`](Self::tiles)
    sorted: FlexBuffer<SortedBufferTag>,
    /// Minimum Z height at each XY tile
    zmin: FlexBuffer<ZminBufferTag>,
}

impl<const N: u64> TileBuffers<N> {
    /// Returns a new `TileBuffers` object
    fn new(
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<Self, TileBuffersError> {
        let tile_buf_size = Self::tile_buf_size(render_size);
        let tiles =
            FlexBuffer::new(device, format!("active_tile{N}"), tile_buf_size)
                .map_err(|err| TileBuffersError {
                buf: TileBufferName::Tiles,
                err,
            })?;
        let sorted =
            FlexBuffer::new(device, format!("sorted_tile{N}"), tile_buf_size)
                .map_err(|err| TileBuffersError {
                buf: TileBufferName::Sorted,
                err,
            })?;
        let zmin = FlexBuffer::new(
            device,
            format!("tile{N}_zmin"),
            Self::zmin_buf_size(render_size),
        )
        .map_err(|err| TileBuffersError {
            buf: TileBufferName::Zmin,
            err,
        })?;

        Ok(Self {
            tiles,
            sorted,
            zmin,
        })
    }

    fn tile_buf_size(render_size: TileRenderSize) -> usize {
        let n = usize::try_from(N).unwrap();
        let nx = usize::try_from(render_size.width()).unwrap() / n;
        let ny = usize::try_from(render_size.height()).unwrap() / n;
        let nz = 64 / n;
        // wg_dispatch: [u32; 3]
        // count: u32,
        4 + nx * ny * nz
    }

    fn zmin_buf_size(render_size: TileRenderSize) -> ImageSize {
        ImageSize::new(
            render_size.width() / u32::try_from(N).unwrap(),
            render_size.height() / u32::try_from(N).unwrap(),
        )
    }

    fn grow_to_fit(
        &mut self,
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<(), TileBuffersError> {
        let TileBuffers {
            tiles,
            sorted,
            zmin,
        } = self;
        let tile_buf_size = Self::tile_buf_size(render_size);
        tiles.grow_to_fit(device, tile_buf_size).map_err(|err| {
            TileBuffersError {
                buf: TileBufferName::Tiles,
                err,
            }
        })?;
        sorted.grow_to_fit(device, tile_buf_size).map_err(|err| {
            TileBuffersError {
                buf: TileBufferName::Sorted,
                err,
            }
        })?;
        zmin.grow_to_fit(device, Self::zmin_buf_size(render_size))
            .map_err(|err| TileBuffersError {
                buf: TileBufferName::Zmin,
                err,
            })?;

        Ok(())
    }

    /// Returns the number of bytes in use by these buffers
    ///
    /// See [`self.capacity`](Self::capacity) for total bytes allocated
    pub fn size(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let TileBuffers {
            tiles,
            sorted,
            zmin,
        } = self;
        tiles.size_bytes() + sorted.size_bytes() + zmin.size_bytes()
    }

    /// Returns the number of bytes allocated by these buffers
    pub fn capacity(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let TileBuffers {
            tiles,
            sorted,
            zmin,
        } = self;
        tiles.capacity() + sorted.capacity() + zmin.capacity()
    }
}

tag!(RootTilesBufferTag, u32, usize, STORAGE | COPY_DST);
tag!(
    RootStrataBufferTag,
    u8,
    usize,
    STORAGE | INDIRECT | COPY_DST
);
tag!(RootZminBufferTag, u32, ImageSize, STORAGE | COPY_DST);
tag!(RootZmaxBufferTag, u32, ImageSize, STORAGE | COPY_DST);

/// Root tile buffers store strata-packed tile lists
struct RootTileBuffers {
    /// Initial output tiles
    tiles: FlexBuffer<RootTilesBufferTag>,
    /// Strata-sorted output tiles
    strata: FlexBuffer<RootStrataBufferTag>,
    zmin: FlexBuffer<RootZminBufferTag>,
    zmax: FlexBuffer<RootZmaxBufferTag>,
}

impl RootTileBuffers {
    /// Build a new root tiles buffer, which stores strata-packed tile lists
    fn new(
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<Self, RootTileBuffersError> {
        // Root tile buffers are always 64³ voxels
        const N: usize = 64;

        // Allocate enough words to write all of the output tiles
        let tiles = FlexBuffer::new(
            device,
            format!("tiles_out{N}"),
            Self::tiles_buf_size(render_size),
        )
        .map_err(|err| RootTileBuffersError {
            buf: RootTileBufferName::Tiles,
            err,
        })?;

        let strata = FlexBuffer::new(
            device,
            format!("strata_tile{N}"),
            Self::strata_buf_size(render_size),
        )
        .map_err(|err| RootTileBuffersError {
            buf: RootTileBufferName::Strata,
            err,
        })?;

        let z_buf_size = Self::z_buf_size(render_size);
        let zmin = FlexBuffer::new(device, format!("tile{N}_zmin"), z_buf_size)
            .map_err(|err| RootTileBuffersError {
                buf: RootTileBufferName::Zmin,
                err,
            })?;
        let zmax = FlexBuffer::new(device, format!("tile{N}_zmax"), z_buf_size)
            .map_err(|err| RootTileBuffersError {
                buf: RootTileBufferName::Zmax,
                err,
            })?;
        Ok(Self {
            tiles,
            strata,
            zmin,
            zmax,
        })
    }

    fn tiles_buf_size(render_size: TileRenderSize) -> usize {
        let nx = usize::try_from(render_size.nx()).unwrap();
        let ny = usize::try_from(render_size.ny()).unwrap();
        let nz = usize::try_from(render_size.nz()).unwrap();
        // wg_dispatch: [u32; 3] (unused)
        // count: u32,
        4 + nx * ny * nz
    }

    fn strata_buf_size(render_size: TileRenderSize) -> usize {
        let nz = usize::try_from(render_size.nz()).unwrap();
        let strata_size = strata_size_bytes(render_size);
        strata_size * nz
    }

    fn z_buf_size(render_size: TileRenderSize) -> ImageSize {
        ImageSize::new(render_size.nx(), render_size.ny())
    }

    /// Grows all of the buffers to fit a particular render size
    fn grow_to_fit(
        &mut self,
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<(), RootTileBuffersError> {
        // Destructure to make sure we take all members into account
        let RootTileBuffers {
            tiles,
            strata,
            zmin,
            zmax,
        } = self;
        tiles
            .grow_to_fit(device, Self::tiles_buf_size(render_size))
            .map_err(|err| RootTileBuffersError {
                buf: RootTileBufferName::Tiles,
                err,
            })?;
        strata
            .grow_to_fit(device, Self::strata_buf_size(render_size))
            .map_err(|err| RootTileBuffersError {
                buf: RootTileBufferName::Strata,
                err,
            })?;

        let z_buf_size = Self::z_buf_size(render_size);
        zmin.grow_to_fit(device, z_buf_size).map_err(|err| {
            RootTileBuffersError {
                buf: RootTileBufferName::Zmin,
                err,
            }
        })?;
        zmax.grow_to_fit(device, z_buf_size).map_err(|err| {
            RootTileBuffersError {
                buf: RootTileBufferName::Zmax,
                err,
            }
        })?;

        Ok(())
    }

    /// Returns the number of bytes in use by buffers
    pub fn size(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let RootTileBuffers {
            tiles,
            strata,
            zmin,
            zmax,
        } = self;
        tiles.size_bytes()
            + strata.size_bytes()
            + zmin.size_bytes()
            + zmax.size_bytes()
    }

    /// Returns the number of bytes allocated to buffers
    pub fn capacity(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let RootTileBuffers {
            tiles,
            strata,
            zmin,
            zmax,
        } = self;
        tiles.capacity() + strata.capacity() + zmin.capacity() + zmax.capacity()
    }
}

tag!(TileTapesBufferTag, u32, usize, STORAGE | COPY_DST);
tag!(VoxelsBufferTag, u32, VoxelSize, STORAGE | COPY_DST);
tag!(pub GeomBufferTag, GeometryPixel, VoxelSize, STORAGE | COPY_SRC | COPY_DST,
    "Tag for a on-GPU buffer storing [`GeometryPixel`] values");
tag!(pub(crate) VarsBufferTag, f32, usize, STORAGE | COPY_DST);

/// Workspace for rendering
///
/// This object is constructed by [`Context::workspace`] and may only be used with
/// that particular [`Context`].
pub struct Workspace {
    /// Image render size
    ///
    /// Note that the tile buffers below round up to the nearest root tile
    /// (64³ voxels).
    image_size: VoxelSize,

    /// Config and tape data buffer (constant size)
    config_buf: wgpu::Buffer,

    /// Scratch space to upload variable values
    vars_buf: FlexBuffer<VarsBufferTag>,

    /// Buffer for z histogram counters
    ///
    /// This is laid out as follows:
    ///
    /// - 4 `u32` words (for the 16³ pass)
    /// - 240 bytes of padding
    /// - 16 `u32` words (for the 4³ pass)
    z_hist_buf: wgpu::Buffer,

    /// Map from tile to the relevant tape (as a start index)
    tile_tapes: FlexBuffer<TileTapesBufferTag>,

    /// Root tile Z heights (64³)
    tile64: RootTileBuffers,

    /// Z heights for filled first-stage tiles (16³)
    tile16: TileBuffers<16>,

    /// Z heights for filled second-stage tiles (4³)
    tile4: TileBuffers<4>,

    /// Z heights for voxel tile evaluation (rounded up)
    voxels: FlexBuffer<VoxelsBufferTag>,

    /// Buffer of [`GeometryPixel`] data, generated by the normal pass
    ///
    /// This is at the original image size
    geom: FlexBuffer<GeomBufferTag>,

    /// Cached bind groups
    bind_groups: BindGroups,
}

/// Cached bind groups (constructed on-demand)
#[derive(Default)]
struct BindGroups {
    common: std::cell::OnceCell<wgpu::BindGroup>,
    merge: std::cell::OnceCell<wgpu::BindGroup>,
    root: std::cell::OnceCell<wgpu::BindGroup>,
    repack: std::cell::OnceCell<wgpu::BindGroup>,
    interval16: std::cell::OnceCell<wgpu::BindGroup>,
    sort16: std::cell::OnceCell<wgpu::BindGroup>,
    interval4: std::cell::OnceCell<wgpu::BindGroup>,
    sort4: std::cell::OnceCell<wgpu::BindGroup>,
    voxel: std::cell::OnceCell<wgpu::BindGroup>,
    normals: std::cell::OnceCell<wgpu::BindGroup>,
    clear: std::cell::OnceCell<wgpu::BindGroup>,
}

impl BindGroups {
    fn common(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.common.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("common bind group"),
                    layout: &ctx.common_bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.config_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile_tapes.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.vars_buf.bind_active(),
                        },
                    ],
                })
        })
    }

    fn clear(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.clear.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("clear bind group"),
                    layout: &ctx.clear_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace
                                .tile16
                                .tiles
                                .data()
                                .slice(0..16)
                                .try_into()
                                .unwrap(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace
                                .tile16
                                .sorted
                                .data()
                                .slice(0..16)
                                .try_into()
                                .unwrap(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace
                                .tile4
                                .tiles
                                .data()
                                .slice(0..16)
                                .try_into()
                                .unwrap(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 3,
                            resource: workspace
                                .tile4
                                .sorted
                                .data()
                                .slice(0..16)
                                .try_into()
                                .unwrap(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 4,
                            resource: workspace.z_hist_buf.as_entire_binding(),
                        },
                    ],
                })
        })
    }

    fn merge(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.merge.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("merge bind group"),
                    layout: &ctx.merge_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.tile64.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile16.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.tile4.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 3,
                            resource: workspace.voxels.bind_active(),
                        },
                    ],
                })
        })
    }

    fn root(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.root.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("interval root bind group"),
                    layout: &ctx.root_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.tile64.tiles.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile64.zmax.bind_active(),
                        },
                    ],
                })
        })
    }

    fn repack(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.repack.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("repack bind group"),
                    layout: &ctx.repack_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.tile64.tiles.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile64.zmax.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.tile64.strata.bind_active(),
                        },
                    ],
                })
        })
    }

    fn interval16(
        &self,
        ctx: &Context,
        workspace: &Workspace,
    ) -> &wgpu::BindGroup {
        let strata_bytes =
            u64::try_from(workspace.strata_size_bytes()).unwrap();
        self.interval16.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("interval16 bind group"),
                    layout: &ctx.interval_ctx.interval_bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace
                                .tile64
                                .strata
                                .data()
                                .slice(0..strata_bytes) // dynamic offset!
                                .try_into()
                                .unwrap(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile64.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.tile16.tiles.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 3,
                            resource: workspace.tile16.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 4,
                            resource: workspace
                                .z_hist_buf
                                .slice(0..16)
                                .try_into()
                                .unwrap(),
                        },
                    ],
                })
        })
    }

    fn sort16(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.sort16.get_or_init(|| {
            Self::sort_bind_group(
                ctx,
                &workspace.tile16,
                workspace.z_hist_buf.slice(0..16).try_into().unwrap(),
            )
        })
    }

    fn sort4(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.sort4.get_or_init(|| {
            Self::sort_bind_group(
                ctx,
                &workspace.tile4,
                workspace.z_hist_buf.slice(256..320).try_into().unwrap(),
            )
        })
    }

    fn sort_bind_group<const N: u64>(
        ctx: &Context,
        tile_bufs: &TileBuffers<N>,
        z_hist: wgpu::BindingResource,
    ) -> wgpu::BindGroup {
        ctx.gpu
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some(&format!("sort{N} bind group")),
                layout: &ctx.interval_ctx.sort_bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: tile_bufs.tiles.bind_active(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: z_hist,
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: tile_bufs.sorted.bind_active(),
                    },
                ],
            })
    }

    fn interval4(
        &self,
        ctx: &Context,
        workspace: &Workspace,
    ) -> &wgpu::BindGroup {
        self.interval4.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("interval4 bind group"),
                    layout: &ctx.interval_ctx.interval_bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.tile16.sorted.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile16.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.tile4.tiles.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 3,
                            resource: workspace.tile4.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 4,
                            resource: workspace
                                .z_hist_buf
                                .slice(256..320)
                                .try_into()
                                .unwrap(),
                        },
                    ],
                })
        })
    }

    fn voxel(&self, ctx: &Context, workspace: &Workspace) -> &wgpu::BindGroup {
        self.voxel.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("voxel bind group"),
                    layout: &ctx.voxel_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.tile4.sorted.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile4.zmin.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.voxels.bind_active(),
                        },
                    ],
                })
        })
    }

    fn normals(
        &self,
        ctx: &Context,
        workspace: &Workspace,
    ) -> &wgpu::BindGroup {
        self.normals.get_or_init(|| {
            ctx.gpu
                .device
                .create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("normals bind group"),
                    layout: &ctx.normals_ctx.bind_group_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: workspace.voxels.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.geom.bind_active(),
                        },
                    ],
                })
        })
    }
}

impl Workspace {
    /// Returns the current image size
    pub fn image_size(&self) -> VoxelSize {
        self.image_size
    }

    /// Returns a handle to the image output buffer
    ///
    /// This is intended for subsequent shaders which want to use the
    /// [`GeometryPixel`] image data without copying to the CPU.  It requires a
    /// exclusive borrow of the `Workspace` object (and then extends that
    /// lifetime) so that other callers can't simultaneously touch the buffer.
    pub fn output(&mut self) -> &FlexBuffer<GeomBufferTag> {
        &self.geom
    }

    fn new(device: &wgpu::Device) -> Self {
        // The config buffer is statically sized, so we can check it here
        static_assertions::const_assert!(
            (std::mem::size_of::<Config>()
                + TAPE_DATA_CAPACITY * std::mem::size_of::<TapeWord>())
                as u64
                <= BufferType::Storage.max_size()
        );

        let image_size = 64.into();

        let config_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("config"),
            size: (std::mem::size_of::<Config>()
                + TAPE_DATA_CAPACITY * std::mem::size_of::<TapeWord>())
                as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let render_size = TileRenderSize::from(image_size);
        let voxels = FlexBuffer::new(
            device,
            "voxels",
            Self::voxels_buf_size(render_size),
        )
        .unwrap();
        let tile_tapes = FlexBuffer::new(
            device,
            "tile tape",
            Self::tile_tapes_buf_size(render_size),
        )
        .unwrap();

        let geom = FlexBuffer::new(device, "geom", image_size).unwrap();

        let tile64 = RootTileBuffers::new(device, render_size).unwrap();
        let tile16 = TileBuffers::new(device, render_size).unwrap();
        let tile4 = TileBuffers::new(device, render_size).unwrap();

        // z_hist_buf never changes size
        let z_hist_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tiles_zhist"),
            size: u64::try_from(
                (4 * std::mem::size_of::<u32>()).next_multiple_of(256)
                    + (16 * std::mem::size_of::<u32>()),
            )
            .unwrap(),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let vars_buf = FlexBuffer::new(device, "vars", 4).unwrap();

        Self {
            config_buf,
            image_size,
            tile_tapes,
            tile64,
            tile16,
            tile4,
            voxels,
            geom,
            vars_buf,
            z_hist_buf,
            bind_groups: Default::default(),
        }
    }

    fn render_size(&self) -> TileRenderSize {
        self.image_size.into()
    }

    /// Returns the size of one strata (in bytes)
    fn strata_size_bytes(&self) -> usize {
        strata_size_bytes(self.render_size())
    }

    /// Returns the number of bytes in the `tile_tapes` buffer
    ///
    /// The tile tape array is... complicated
    ///
    /// The first `nx * ny * nz` words are tape indices for the root tiles
    /// (64³), densely allocated in x / y / z order.  This is
    /// straight-forward.
    ///
    /// After that point, it gets weirder.  At any given point in time, we're
    /// evaluating a single strata (i.e. a 64-voxel deep slice of the image).
    /// We allocated enough tape words for that strata, also in x / y / z
    /// order, but z is limited to either 0..4 for the 16³ subtiles, or
    /// 0..16 for 4³ subtiles.
    ///
    /// In other words, it looks something like this:
    ///
    /// ```text
    /// | index | index | index | ... |     densely packed 64³ tape indices
    /// | index | index | index | ... |     16² XY tiles × 4  Z positions
    /// | index | index | index | ... |     4²  XY tiles × 16 Z positions
    /// ```
    fn tile_tapes_buf_size(render_size: TileRenderSize) -> usize {
        let nx = usize::try_from(render_size.nx()).unwrap();
        let ny = usize::try_from(render_size.ny()).unwrap();
        let nz = usize::try_from(render_size.nz()).unwrap();

        // Each tile contains 16³ and 4³ subtiles
        let xy_size = (64usize / 4).pow(3) + (64usize / 16).pow(3);

        // Total size computation:
        //    nx * ny * nz + (nx * ny * xy_size)
        // => nx * ny * (nz + xy_size)
        nx.checked_mul(ny)
            .unwrap()
            .checked_mul(nz.checked_add(xy_size).unwrap())
            .unwrap()
    }

    fn voxels_buf_size(render_size: TileRenderSize) -> VoxelSize {
        VoxelSize::new(
            render_size.width(),
            render_size.height(),
            render_size.depth(),
        )
    }

    /// Resizes to render the target image size
    ///
    /// Internal buffers are resized to fit (only getting larger)
    fn set_image_size(
        &mut self,
        device: &wgpu::Device,
        image_size: VoxelSize,
    ) -> Result<(), BufferError> {
        let render_size = TileRenderSize::from(image_size);
        let Workspace {
            image_size: image_size_ref,
            config_buf: _,
            z_hist_buf: _,
            vars_buf: _,
            tile_tapes,
            tile64,
            tile16,
            tile4,
            voxels,
            geom,
            bind_groups,
        } = self;
        // Clear our cached bind groups if the image sizes is changing
        if *image_size_ref != image_size {
            // TODO this also clears the common bind group, which isn't
            // necessary (since it's not size-dependent)
            *bind_groups = Default::default();
        }
        *image_size_ref = image_size;
        tile_tapes
            .grow_to_fit(device, Self::tile_tapes_buf_size(render_size))
            .map_err(|err| BufferError {
                buf: BufferName::TileTapes,
                err,
            })?;
        tile64
            .grow_to_fit(device, render_size)
            .map_err(|e| BufferError {
                buf: BufferName::Tile64(e.buf),
                err: e.err,
            })?;
        tile16
            .grow_to_fit(device, render_size)
            .map_err(|e| BufferError {
                buf: BufferName::Tile16(e.buf),
                err: e.err,
            })?;
        tile4
            .grow_to_fit(device, render_size)
            .map_err(|e| BufferError {
                buf: BufferName::Tile4(e.buf),
                err: e.err,
            })?;

        voxels
            .grow_to_fit(device, Self::voxels_buf_size(render_size))
            .map_err(|err| BufferError {
                buf: BufferName::Voxels,
                err,
            })?;
        geom.grow_to_fit(device, image_size)
            .map_err(|err| BufferError {
                buf: BufferName::Geom,
                err,
            })?;
        Ok(())
    }

    /// Returns total allocated size (in bytes)
    pub fn capacity(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let Workspace {
            image_size: _,
            config_buf,
            z_hist_buf,
            tile_tapes,
            vars_buf,
            tile64,
            tile16,
            tile4,
            voxels,
            geom,
            bind_groups: _,
        } = self;
        config_buf.size()
            + vars_buf.capacity()
            + z_hist_buf.size()
            + tile_tapes.capacity()
            + tile64.capacity()
            + tile16.capacity()
            + tile4.capacity()
            + voxels.capacity()
            + geom.capacity()
    }

    /// Returns total active size (in bytes)
    pub fn size(&self) -> u64 {
        // Destructure to make sure we take all members into account
        let Workspace {
            image_size: _,
            config_buf,
            vars_buf,
            z_hist_buf,
            tile_tapes,
            tile64,
            tile16,
            tile4,
            voxels,
            geom,
            bind_groups: _,
        } = self;
        config_buf.size()
            + vars_buf.size_bytes()
            + z_hist_buf.size()
            + tile_tapes.size_bytes()
            + tile64.size()
            + tile16.size()
            + tile4.size()
            + voxels.size_bytes()
            + geom.size_bytes()
    }
}

impl Context {
    /// Build a new 3D rendering context, given a device and queue
    pub fn new(gpu: &Gpu) -> Self {
        // Create bind group layout and bind group
        let common_bind_group_layout = gpu.device.create_bind_group_layout(
            &wgpu::BindGroupLayoutDescriptor {
                label: Some("common bind group layout"),
                entries: &[
                    buffer_rw(0), // config (including tape buffer)
                    buffer_rw(1), // tile_tape (hierarchical)
                    buffer_ro(2), // vars
                ],
            },
        );

        let root_ctx = RootContext::new(&gpu.device, &common_bind_group_layout);
        let repack_ctx =
            RepackContext::new(&gpu.device, &common_bind_group_layout);
        let interval_ctx =
            IntervalContext::new(&gpu.device, &common_bind_group_layout);
        let voxel_ctx =
            VoxelContext::new(&gpu.device, &common_bind_group_layout);
        let normals_ctx =
            NormalsContext::new(&gpu.device, &common_bind_group_layout);
        let merge_ctx =
            MergeContext::new(&gpu.device, &common_bind_group_layout);
        let clear_ctx =
            ClearContext::new(&gpu.device, &common_bind_group_layout);
        let reset_ctx = ResetContext;

        Self {
            gpu: gpu.clone(),
            common_bind_group_layout,
            root_ctx,
            repack_ctx,
            interval_ctx,
            voxel_ctx,
            normals_ctx,
            merge_ctx,
            reset_ctx,
            clear_ctx,
        }
    }

    /// Builds a new [`Workspace`] object for use in rendering
    ///
    /// The buffers are initialized with a dummy size and resized automatically
    /// when passed into any of the runner functions (e.g. [`run`](Self::run) or
    /// [`submit`](Self::submit)).
    pub fn workspace(&self) -> Workspace {
        Workspace::new(&self.gpu.device)
    }

    /// Renders the image, with a blocking wait to read pixel data from the GPU
    ///
    /// This function is not present when built for the `wasm32` target
    #[cfg(not(target_arch = "wasm32"))]
    pub fn run(
        &self,
        shape: &RenderShape,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<GeomBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.run_with_vars(shape, &Default::default(), workspace, out, settings)
    }

    /// Renders the image, with a blocking wait to read pixel data from the GPU
    ///
    /// This function is not present when built for the `wasm32` target
    #[cfg(not(target_arch = "wasm32"))]
    pub fn run_with_vars(
        &self,
        shape: &RenderShape,
        vars: &ShapeVars<f32>,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<GeomBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.submit_with_vars(shape, vars, workspace, &settings)?;
        self.gpu.copy(&workspace.geom, out);
        let image = self.gpu.map_image(out);
        Ok(image.image())
    }

    /// Renders the image, with an async wait to read pixel data from the GPU
    ///
    /// This can be called either natively or on the web
    pub async fn run_async(
        &self,
        shape: &RenderShape,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<GeomBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.run_with_vars_async(
            shape,
            &Default::default(),
            workspace,
            out,
            settings,
        )
        .await
    }

    /// Renders the image, with an async wait to read pixel data from the GPU
    ///
    /// This can be called either natively or on the web
    pub async fn run_with_vars_async(
        &self,
        shape: &RenderShape,
        vars: &ShapeVars<f32>,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<GeomBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.submit_with_vars(shape, vars, workspace, &settings)?;
        self.gpu.copy(&workspace.geom, out);
        let image = self.gpu.map_image_async(out).await;
        Ok(image.image())
    }

    /// Submits a single image to be rendered on the GPU
    ///
    /// The resulting image (as a buffer of [`GeometryPixel`] data) is available
    /// on the GPU in [`workspace.output()`](Workspace::output).
    pub fn submit(
        &self,
        shape: &RenderShape,
        workspace: &mut Workspace,
        settings: &RenderConfig,
    ) -> Result<(), SubmitError> {
        self.submit_with_vars(shape, &Default::default(), workspace, settings)
    }

    /// Submits a single image to be rendered on the GPU, with extra variables
    ///
    /// See [`submit`](Self::submit) for additional details.
    pub fn submit_with_vars(
        &self,
        shape: &RenderShape,
        vars: &ShapeVars<f32>,
        workspace: &mut Workspace,
        settings: &RenderConfig,
    ) -> Result<(), SubmitError> {
        // Create a command encoder and dispatch the compute work
        let mut encoder = self.gpu.device.create_command_encoder(
            &wgpu::CommandEncoderDescriptor { label: None },
        );
        self.encode_with_vars(shape, vars, workspace, settings, &mut encoder)?;
        self.gpu.queue.submit(Some(encoder.finish()));
        Ok(())
    }

    /// Encodes a single image to be rendered on the GPU, with extra variables
    ///
    /// See [`submit`](Self::submit) for additional details.
    pub fn encode_with_vars(
        &self,
        shape: &RenderShape,
        vars: &ShapeVars<f32>,
        workspace: &mut Workspace,
        settings: &RenderConfig,
        encoder: &mut wgpu::CommandEncoder,
    ) -> Result<(), SubmitError> {
        workspace.set_image_size(&self.gpu.device, settings.image_size)?;
        let render_size = TileRenderSize::from(workspace.image_size);

        let mat =
            settings.world_to_model * workspace.image_size.screen_to_world();

        // Divide by 2 to go from `u32` -> `TapeWord`
        let start_offset = u32::try_from(shape.bytecode.len()).unwrap() / 2;
        let config = Config {
            mat: mat.data.as_slice().try_into().unwrap(),
            axes: shape.axes(),
            render_size: [
                render_size.width(),
                render_size.height(),
                render_size.depth(),
            ],
            tape_data_capacity: TAPE_DATA_CAPACITY.try_into().unwrap(),
            image_size: [
                workspace.image_size.width(),
                workspace.image_size.height(),
                workspace.image_size.depth(),
            ],
            tape_data_offset: start_offset,
            root_tape_len: start_offset,
        };

        {
            // We load the `Config` and shape tape data.
            let config_len = std::mem::size_of_val(&config);
            let mut writer = self
                .gpu
                .queue
                .write_buffer_with(
                    &workspace.config_buf,
                    0,
                    ((config_len + shape.bytecode.as_bytes().len()) as u64)
                        .try_into()
                        .unwrap(),
                )
                .unwrap();
            writer
                .slice(..config_len)
                .copy_from_slice(config.as_bytes());
            writer
                .slice(config_len..)
                .copy_from_slice(shape.bytecode.as_bytes());
        }

        // Copy vars (if present), then reset relevant bind groups if the buffer
        // size has changed.
        if matches!(
            shape.copy_vars(&self.gpu, vars, &mut workspace.vars_buf)?,
            CopyVarsChanged::BufferChanged
        ) {
            workspace.bind_groups.common = Default::default();
        }

        // Initial buffer reset pass
        self.reset_ctx.run(encoder, workspace);

        let mut compute_pass =
            encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: None,
                timestamp_writes: None,
            });

        // Build the common config buffer
        let common_bind_group = workspace.bind_groups.common(self, workspace);
        compute_pass.set_bind_group(0, common_bind_group, &[]);

        // Populate root tiles (64x64x64, densely packed)
        self.root_ctx.run(
            self,
            workspace,
            shape.bytecode.reg_count(),
            render_size,
            &mut compute_pass,
        );
        // Repack root tiles into strata
        self.repack_ctx
            .run(self, workspace, render_size, &mut compute_pass);

        // Evaluate tiles in reverse-Z order by strata (64 voxels deep)
        let strata_count = u64::from(render_size.depth()).div_ceil(64);
        for strata in 0..strata_count {
            self.interval_ctx.run(
                self,
                workspace,
                strata,
                shape.bytecode.reg_count(),
                &mut compute_pass,
            );
            self.voxel_ctx.run(
                self,
                workspace,
                shape.bytecode.reg_count(),
                &mut compute_pass,
            );

            // Merge filled tiles from large -> small, populating the heightmap
            self.merge_ctx.run(self, workspace, &mut compute_pass);
            self.normals_ctx.run(
                self,
                workspace,
                shape.bytecode.reg_count(),
                &mut compute_pass,
            );

            self.clear_ctx.run(self, workspace, &mut compute_pass);
        }

        // Submit the commands and wait for the GPU to complete
        Ok(())
    }
}

struct ClearContext {
    bind_group_layout: wgpu::BindGroupLayout,
    pipeline: wgpu::ComputePipeline,
}

impl ClearContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("clear bind group layout"),
                entries: &[
                    buffer_rw(0), // tile16_count
                    buffer_rw(1), // tile16_sort
                    buffer_rw(2), // tile4_count
                    buffer_rw(3), // tile4_sort
                    buffer_rw(4), // zhist_buf
                ],
            });

        // Create the compute pipeline
        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("clear pipeline layout"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        // Compile the shader
        let shader_code = clear_shader();
        let shader_module =
            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("clear shader module"),
                source: wgpu::ShaderSource::Wgsl(shader_code.into()),
            });

        let pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("clear"),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("clear_main"),
                compilation_options: Default::default(),
                cache: None,
            });

        Self {
            pipeline,
            bind_group_layout,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let bind_group = workspace.bind_groups.clear(ctx, workspace);
        compute_pass.set_pipeline(&self.pipeline);
        compute_pass.set_bind_group(1, bind_group, &[]);
        compute_pass.dispatch_workgroups(1, 1, 1);
    }
}

struct MergeContext {
    bind_group_layout: wgpu::BindGroupLayout,
    pipeline: wgpu::ComputePipeline,
}

impl MergeContext {
    fn new(
        device: &wgpu::Device,
        common_bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        let shader_code = merge_shader();

        // Create bind group layout and bind group
        let bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("merge bind group layout"),
                entries: &[
                    buffer_rw(0), // tile64_zmin
                    buffer_rw(1), // tile16_zmin
                    buffer_rw(2), // tile4_zmin
                    buffer_rw(3), // voxels
                ],
            });

        // Create the compute pipeline
        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("merge pipeline layout"),
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        // Compile the shader
        let shader_module =
            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("merge shader module"),
                source: wgpu::ShaderSource::Wgsl(shader_code.into()),
            });

        let pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("merge"),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("merge_main"),
                compilation_options: Default::default(),
                cache: None,
            });

        Self {
            pipeline,
            bind_group_layout,
        }
    }

    fn run(
        &self,
        ctx: &Context,
        workspace: &Workspace,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        let render_size = workspace.render_size();
        let bind_group = workspace.bind_groups.merge(ctx, workspace);
        compute_pass.set_pipeline(&self.pipeline);
        compute_pass.set_bind_group(1, bind_group, &[]);
        compute_pass.dispatch_workgroups(
            render_size.width().div_ceil(8),
            render_size.height().div_ceil(8),
            1,
        );
    }
}

struct ResetContext;

impl ResetContext {
    fn run(&self, encoder: &mut wgpu::CommandEncoder, workspace: &Workspace) {
        // Clear only the `count` member of the tile64 `tiles_out` buffer
        encoder.clear_buffer(workspace.tile64.tiles.data(), 12, Some(4));

        // Per-strata counters may now be at a different location in memory if
        // we're using the buffers for multiple renders of different sizes!  To
        // be safe, we'll clear them here, rather than in a render pass.
        let strata_size_bytes = workspace.strata_size_bytes();
        for s in 0..workspace.render_size().nz() {
            encoder.clear_buffer(
                workspace.tile64.strata.data(),
                u64::from(s) * u64::try_from(strata_size_bytes).unwrap(),
                Some(16),
            );
        }

        // Clear all of the heightmaps and output maps
        workspace.tile64.zmin.clear(encoder);
        workspace.tile64.zmax.clear(encoder);
        workspace.tile16.zmin.clear(encoder);
        workspace.tile4.zmin.clear(encoder);
        workspace.voxels.clear(encoder);
        workspace.geom.clear(encoder);

        // Clear the whole tile tape map (TODO is this needed?)
        workspace.tile_tapes.clear(encoder);

        // tiles / sorted counters and z_hist are reset in clear shader
    }
}

#[cfg(test)]
mod test {
    use super::{effects::PackedVoxel, *};
    use crate::color::{ShapeColor, ShapeColorBuffers};
    use fidget_core::{context::Tree, vm::VmShape};
    use std::collections::HashSet;

    #[test]
    fn compile_interval_root_shader() {
        crate::compile_shader(&interval_root_shader(16), "interval root");
    }

    #[test]
    fn compile_interval_tiles_shader() {
        crate::compile_shader(&interval_tiles_shader(16), "interval tiles");
    }

    #[test]
    fn compile_voxel_tiles_shader() {
        crate::compile_shader(&voxel_tiles_shader(16), "voxel tiles");
    }

    #[test]
    fn compile_normals_shader() {
        crate::compile_shader(&normals_shader(16), "normals tiles");
    }

    #[test]
    fn compile_repack_shader() {
        crate::compile_shader(&repack_shader(), "repack");
    }

    #[test]
    fn compile_sort_shader() {
        crate::compile_shader(&sort_shader(), "sort");
    }

    #[test]
    fn compile_merge_shader() {
        crate::compile_shader(&merge_shader(), "merge");
    }

    #[test]
    fn compile_clear_shader() {
        crate::compile_shader(&clear_shader(), "clear");
    }

    #[derive(
        Copy,
        Clone,
        Debug,
        IntoBytes,
        Immutable,
        FromBytes,
        KnownLayout,
        PartialEq,
        Eq,
        Hash,
    )]
    #[repr(C)]
    struct Rgba {
        r: u8,
        g: u8,
        b: u8,
        a: u8,
    }

    struct RenderOutput {
        merged: Vec<PackedVoxel>,
        colors: Vec<Rgba>,
        shaded: fidget_raster::Image<u32, VoxelSize>,
        heightmap: fidget_raster::Image<u32, VoxelSize>,
    }

    fn render(
        shapes: &[(Tree, ShapeColor<Tree>)],
        render_config: RenderConfig,
    ) -> RenderOutput {
        let gpu = pollster::block_on(Gpu::init_basic()).unwrap();
        let voxel_ctx = Context::new(&gpu);
        let effects_ctx = effects::Context::new(&gpu);

        let mut buf = voxel_ctx.workspace();
        let mut merge_buf = effects_ctx.merge_workspace();
        let mut shade_buf = effects_ctx.shade_workspace();

        // Render and accumulate each shape
        for (shape, _) in shapes {
            let shape =
                RenderShape::new(&VmShape::from(shape.clone())).unwrap();
            voxel_ctx.submit(&shape, &mut buf, &render_config).unwrap();
            effects_ctx
                .submit_merge(buf.output(), Default::default(), &mut merge_buf)
                .unwrap();
        }
        let merged = gpu.read_vec(merge_buf.output());
        let shape_colors = shapes
            .iter()
            .map(|(_, c)| match c {
                ShapeColor::Rgb { r, g, b } => ShapeColor::Rgb {
                    r: VmShape::from(r.clone()),
                    g: VmShape::from(g.clone()),
                    b: VmShape::from(b.clone()),
                },
                ShapeColor::Hsl { h, s, l } => ShapeColor::Hsl {
                    h: VmShape::from(h.clone()),
                    s: VmShape::from(s.clone()),
                    l: VmShape::from(l.clone()),
                },
            })
            .collect::<Vec<_>>();
        let shape_colors = ShapeColorBuffers::new(&shape_colors).unwrap();
        let mut color_workspace = effects_ctx.color_workspace();

        // Compute SSAO buffer
        let mut ssao_buf = effects_ctx.ssao_workspace();
        effects_ctx.submit_ssao(&merge_buf, &mut ssao_buf).unwrap();

        // Compute per-pixel colors
        effects_ctx
            .submit_color(
                &merge_buf,
                &render_config.world_to_model,
                &shape_colors,
                &mut color_workspace,
                &mut shade_buf,
            )
            .unwrap();

        // At this point, we should have either transparent (white) or opaque
        // colored pixels in the color output buffer
        let raw_colors = gpu.read_vec(shade_buf.output());
        let colors = <[Rgba]>::ref_from_bytes(raw_colors.as_bytes())
            .unwrap()
            .to_vec();

        // Compute shaded image (with color, overwriting shade_buf)
        let mut shade_out = gpu.read_buffer_for(shade_buf.output());
        effects_ctx
            .submit_shade(&merge_buf, Some(&ssao_buf), &mut shade_buf)
            .unwrap();
        gpu.copy(shade_buf.output(), &mut shade_out);
        let shaded = gpu.map_image(&mut shade_out).image();

        // Recompute per-pixel colors
        effects_ctx
            .submit_color(
                &merge_buf,
                &render_config.world_to_model,
                &shape_colors,
                &mut color_workspace,
                &mut shade_buf,
            )
            .unwrap();
        // Compute heightmap
        effects_ctx
            .submit_heightmap(&merge_buf, &mut shade_buf)
            .unwrap();
        gpu.copy(shade_buf.output(), &mut shade_out);
        let heightmap = gpu.map_image(&mut shade_out).image();

        RenderOutput {
            merged,
            colors,
            shaded,
            heightmap,
        }
    }

    #[test]
    fn voxel_pipeline() {
        // We only run in CI if we're on MacOS (because other runners don't have
        // GPUs and will fail to build the context).
        #[cfg(not(target_os = "macos"))]
        if std::env::var("CI").is_ok() {
            return;
        }

        let size = 128;

        let (x, y, z) = Tree::axes();
        let x_ = x.clone() - 0.2;
        let sphere1 = (x_.square() + y.square() + z.square()).sqrt()
            - Tree::constant(0.5);
        let x_ = x + 0.2;
        let sphere2 = (x_.square() + y.square() + z.square()).sqrt()
            - Tree::constant(0.5);
        let spheres = sphere1.min(sphere2);
        let out = render(
            &[(
                spheres,
                ShapeColor::Rgb {
                    r: Tree::constant(0.0),
                    g: Tree::constant(1.0),
                    b: Tree::constant(0.25),
                },
            )],
            RenderConfig::from_size(size.into()),
        );

        let color_set = out.colors.iter().cloned().collect::<HashSet<_>>();
        assert!(
            color_set.contains(&Rgba {
                a: 0,
                r: 0xFF,
                g: 0xFF,
                b: 0xFF
            }) && color_set.contains(&Rgba {
                a: 0xFF,
                r: 0x00,
                g: 0xFF,
                b: 0x3F
            }),
            "invalid color set {color_set:X?} in {} pixels",
            out.colors.len()
        );

        let (shaded, img_size) = out.shaded.take();
        assert_eq!(img_size.width(), size);
        assert_eq!(img_size.height(), size);

        for (m, c) in out.merged.iter().zip(out.colors.iter()) {
            let p = m.z;
            if p == 0 {
                assert_eq!(c.a, 0);
            } else {
                assert_eq!(c.a, 0xFF);
            }
        }
        for (m, c) in out.merged.iter().zip(shaded.iter()) {
            let p = m.z;
            if p == 0 {
                assert_eq!(*c, 0);
            } else {
                // Check that the alpha channel is fully opaque
                assert_eq!(c & (0xFF << 24), (0xFF << 24));
                // Check that the other colors are approximately correct
                let g = (c >> 8) & 0xFF;
                let b = (c >> 16) & 0xFF;
                assert!(
                    (g / 4).abs_diff(b) < 2,
                    "invalid G/B ratio {g}/{b} (expected 4)"
                );
            }
        }
        for (m, c) in out.merged.iter().zip(out.heightmap.iter()) {
            let p = m.z;
            if p == 0 {
                assert_eq!(*c, 0);
            } else {
                // Check that the alpha channel is fully opaque
                assert_eq!(c & (0xFF << 24), (0xFF << 24));
                // Check that the other colors are approximately correct
                let g = (c >> 8) & 0xFF;
                let b = (c >> 16) & 0xFF;
                assert!(
                    (g / 4).abs_diff(b) < 2,
                    "invalid G/B ratio {g}/{b} (expected 4)"
                );
            }
        }
        // Make sure that the heightmap and shaded images are different!
        assert!(out.heightmap.iter().zip(shaded.iter()).any(|(a, b)| a != b));
    }

    #[test]
    fn voxel_transform() {
        // We only run in CI if we're on MacOS (because other runners don't have
        // GPUs and will fail to build the context).
        #[cfg(not(target_os = "macos"))]
        if std::env::var("CI").is_ok() {
            return;
        }

        let (x, y, z) = Tree::axes();
        let sphere =
            (x.square() + y.square() + z.square()).sqrt() - Tree::constant(0.5);

        let size = 64;
        let cfg = RenderConfig {
            image_size: size.into(),
            world_to_model: nalgebra::Translation3::new(-0.5, 0.0, 0.0)
                .to_homogeneous(),
        };
        let out = render(
            &[(
                sphere.clone(),
                ShapeColor::Rgb {
                    r: Tree::x() + 0.5, // offset to avoid saturation at 0
                    g: Tree::y() + 0.5, // offset to avoid saturation at 0
                    b: Tree::z(),
                },
            )],
            cfg,
        );

        let mut iter = out.colors.iter();
        let mut pixels = String::with_capacity(out.colors.len());
        for _y in 0..size {
            for _x in 0..size {
                pixels += if iter.next().unwrap().a == 0 {
                    "."
                } else {
                    "#"
                };
            }
            pixels += "\n";
        }
        assert_eq!(
            pixels,
            "\
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ...........................................###########..........
                .........................................###############........
                .......................................###################......
                ......................................#####################.....
                .....................................#######################....
                ....................................#########################...
                ...................................###########################..
                ...................................###########################..
                ..................................#############################.
                ..................................#############################.
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                .................................###############################
                ..................................#############################.
                ..................................#############################.
                ...................................###########################..
                ...................................###########################..
                ....................................#########################...
                .....................................#######################....
                ......................................#####################.....
                .......................................###################......
                .........................................###############........
                ...........................................###########..........
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
            "
            .replace(" ", "")
        );

        let (_out, img_size) = out.shaded.take();
        assert_eq!(img_size.width(), size);
        assert_eq!(img_size.height(), size);

        for (m, c) in out.merged.iter().zip(out.colors.iter()) {
            let p = m.z;
            if p == 0 {
                assert_eq!(c.a, 0);
            } else {
                assert_eq!(c.a, 0xFF);
            }
        }

        // Every pixel should have a unique color, plus the background color
        let color_set = out.colors.iter().cloned().collect::<HashSet<_>>();
        assert_eq!(
            color_set.len(),
            pixels.chars().filter(|c| *c == '#').count() + 1,
            "invalid color set {color_set:X?} in {} pixels",
            out.colors.len()
        );

        ////////////////////////////////////////////////////////////////////////
        // Now, run the same test with a different transform.  This lets us
        // confirm that colors are pinned to the correct space corodinates
        let cfg = RenderConfig {
            image_size: size.into(),
            world_to_model: nalgebra::Translation3::new(0.5, 0.0, 0.0)
                .to_homogeneous(),
        };
        let out = render(
            &[(
                sphere,
                ShapeColor::Rgb {
                    r: Tree::x() + 0.5, // offset to avoid saturation at 0
                    g: Tree::y() + 0.5, // offset to avoid saturation at 0
                    b: Tree::z(),
                },
            )],
            cfg,
        );

        let mut iter = out.colors.iter();
        let mut pixels = String::with_capacity(out.colors.len());
        for _y in 0..size {
            for _x in 0..size {
                pixels += if iter.next().unwrap().a == 0 {
                    "."
                } else {
                    "#"
                };
            }
            pixels += "\n";
        }
        assert_eq!(
            pixels,
            "\
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ...........###########..........................................
                .........###############........................................
                .......###################......................................
                ......#####################.....................................
                .....#######################....................................
                ....#########################...................................
                ...###########################..................................
                ...###########################..................................
                ..#############################.................................
                ..#############################.................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                .###############################................................
                ..#############################.................................
                ..#############################.................................
                ...###########################..................................
                ...###########################..................................
                ....#########################...................................
                .....#######################....................................
                ......#####################.....................................
                .......###################......................................
                .........###############........................................
                ...........###########..........................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
                ................................................................
            "
            .replace(" ", "")
        );

        let (_out, img_size) = out.shaded.take();
        assert_eq!(img_size.width(), size);
        assert_eq!(img_size.height(), size);

        for (m, c) in out.merged.iter().zip(out.colors.iter()) {
            let p = m.z;
            if p == 0 {
                assert_eq!(c.a, 0);
            } else {
                assert_eq!(c.a, 0xFF);
            }
        }

        // Every pixel should have a unique color, plus the background color
        let color_set = out.colors.iter().cloned().collect::<HashSet<_>>();
        assert_eq!(
            color_set.len(),
            pixels.chars().filter(|c| *c == '#').count() + 1,
            "invalid color set {color_set:X?} in {} pixels",
            out.colors.len()
        );
    }

    #[test]
    fn voxel_hsl() {
        // We only run in CI if we're on MacOS (because other runners don't have
        // GPUs and will fail to build the context).
        #[cfg(not(target_os = "macos"))]
        if std::env::var("CI").is_ok() {
            return;
        }

        let size = 128;
        let (x, y, z) = Tree::axes();
        let x_ = x.clone() - 0.2;
        let sphere1 = (x_.square() + y.square() + z.square()).sqrt()
            - Tree::constant(0.5);
        let x_ = x + 0.2;
        let sphere2 = (x_.square() + y.square() + z.square()).sqrt()
            - Tree::constant(0.5);
        let out = render(
            &[
                (
                    sphere1,
                    ShapeColor::Hsl {
                        h: Tree::constant(0.0),
                        s: Tree::constant(1.0),
                        l: Tree::constant(0.5),
                    },
                ),
                (
                    sphere2,
                    ShapeColor::Hsl {
                        h: Tree::constant(0.5),
                        s: Tree::constant(0.4),
                        l: Tree::constant(0.7),
                    },
                ),
            ],
            RenderConfig::from_size(size.into()),
        );

        let color_set = out.colors.iter().cloned().collect::<HashSet<_>>();
        assert_eq!(color_set.len(), 3);
        assert!(
            color_set.contains(&Rgba {
                r: 0xFF,
                g: 0xFF,
                b: 0xFF,
                a: 0x00,
            }) && color_set.contains(&Rgba {
                r: 0xFF,
                g: 0x00,
                b: 0x00,
                a: 0xFF,
            }) && color_set.contains(&Rgba {
                r: 0x93,
                g: 0xD1,
                b: 0xD1,
                a: 0xFF
            }),
            "invalid color set {color_set:X?} in {} pixels",
            out.colors.len()
        );
        println!("{color_set:#x?}");
    }

    #[test]
    fn voxel_config_layout() {
        // Pick any shader, since `struct Config` is in the common text
        crate::test::compare_struct_layout::<Config>(&merge_shader(), "Config");
    }
}