fidget-wgpu 0.5.1

//! GPU-accelerated 2D rendering
//!
//! See the [`voxel`](crate::voxel) module for details docs; this module is
//! analogous (down to the naming of types).
use crate::{
    CopyVarsChanged, CopyVarsError, Gpu, RegPipeline, RenderShape,
    TAPE_DATA_CAPACITY, TapeWord,
    buf::{
        BufferSizeError, BufferType, FlexBuffer, ReadBuffer, buffer_ro,
        buffer_rw,
    },
    shaders, tag,
    voxel::VarsBufferTag,
};
use fidget_core::{
    render::ImageSize,
    shape::{MissingVar, ShapeVars},
};
use fidget_raster::pixel::{Image, RawDistancePixel};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

pub use fidget_raster::pixel::{RenderConfig, RenderSize};
pub mod effects;

const COMMON_SHADER: &str = include_str!("shaders/common.wgsl");
const DISTANCE_PIXEL_SHADER: &str = include_str!("shaders/distance_pixel.wgsl");
const INTERVAL_INPUT: &str = include_str!("shaders/interval_input.wgsl");
const TRANSFORM_INPUT: &str = include_str!("shaders/transform_input.wgsl");
const INTERVAL_ROOT_SHADER: &str = include_str!("shaders/interval_root.wgsl");
const INTERVAL_TILES_SHADER: &str = include_str!("shaders/interval_tiles.wgsl");
const PIXEL_TILES_SHADER: &str = include_str!("shaders/pixel_tiles.wgsl");
const MERGE_SHADER: &str = include_str!("shaders/merge.wgsl");

/// 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 += DISTANCE_PIXEL_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 tile -> subtile reduction
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 += COMMON_SHADER;
    shader_code += DISTANCE_PIXEL_SHADER;
    shader_code += INTERVAL_TILES_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 pixel tile evaluation
fn pixel_tiles_shader(reg_count: u8) -> String {
    let mut shader_code = shaders::opcode_constants();
    shader_code += &format!("const REG_COUNT: u32 = {reg_count};");
    shader_code += PIXEL_TILES_SHADER;
    shader_code += TRANSFORM_INPUT;
    shader_code += COMMON_SHADER;
    shader_code += DISTANCE_PIXEL_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 merging images
fn merge_shader() -> String {
    MERGE_SHADER.to_owned()
        + COMMON_SHADER
        + DISTANCE_PIXEL_SHADER
        + shaders::COMMON
}

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

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

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

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 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
    }
}

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

/// 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,
}

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: None,
                entries: &[
                    buffer_rw(0), // tiles_out
                    buffer_rw(1), // tile_values
                ],
            });
        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                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: None,
                    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_tiles(ctx, workspace);
        compute_pass.set_pipeline(self.root_pipeline.get(reg_count));
        compute_pass.set_bind_group(1, bind_group, &[]);

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

/// Intermediate tiles context, which produces a list of 8² tiles
struct IntervalTilesContext {
    /// Pipelines for 8² tile evaluation
    tiles_pipeline: RegPipeline,

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

impl IntervalTilesContext {
    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: None,
                entries: &[
                    buffer_ro(0), // tiles_in
                    buffer_rw(1), // subtiles_out
                    buffer_rw(2), // subtile_values
                ],
            });
        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        let device_ = device.clone();
        let tiles_pipeline = RegPipeline::build(Box::new(move |reg_count| {
            let shader_code = interval_tiles_shader(reg_count);
            let shader_module =
                device_.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: None,
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                });
            device_.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(&format!("interval tiles ({reg_count})")),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("interval_tiles_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        }));

        Self {
            bind_group_layout,
            tiles_pipeline,
        }
    }

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

        // Indirect dispatch based on previous tile output
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile64.tiles.data(), 0);
    }
}

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

/// Pixel tiles context, which outputs individual pixels
struct PixelTilesContext {
    /// Pipelines for 8² pixel tile evaluation
    tiles_pipeline: RegPipeline,

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

impl PixelTilesContext {
    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: None,
                entries: &[
                    buffer_ro(0), // tiles_in
                    buffer_rw(1), // result
                ],
            });
        let pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[
                    Some(common_bind_group_layout),
                    Some(&bind_group_layout),
                ],
                immediate_size: 0u32,
            });

        let device_ = device.clone();
        let tiles_pipeline = RegPipeline::build(Box::new(move |reg_count| {
            let shader_code = pixel_tiles_shader(reg_count);
            let shader_module =
                device_.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: None,
                    source: wgpu::ShaderSource::Wgsl(shader_code.into()),
                });
            device_.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(&format!("pixel tiles ({reg_count})")),
                layout: Some(&pipeline_layout),
                module: &shader_module,
                entry_point: Some("pixel_tiles_main"),
                compilation_options: Default::default(),
                cache: None,
            })
        }));

        Self {
            bind_group_layout,
            tiles_pipeline,
        }
    }

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

        // Indirect dispatch based on previous tile output
        compute_pass
            .dispatch_workgroups_indirect(workspace.tile8.tiles.data(), 0);
    }
}

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

/// 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
    ///
    /// This is a 3x3 matrix in WGSL, but each row is padded to 16 bytes, so
    /// it's a total of 12 floats.
    mat: [f32; 12],

    /// 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; 2],

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

    /// Z position at which to render
    z: f32,

    /// Flag indicating whether to recurse down to individual pixels
    pixel_perfect: u32,

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

    /// Padding
    _pad: u32,
    // This is followed by a flexible array member containing tape data
}

tag!(TileTapesBufferTag, u32, usize, STORAGE | COPY_DST);
tag!(pub PixelBufferTag, RawDistancePixel, ImageSize, STORAGE | COPY_SRC | COPY_DST,
    "Tag for a on-GPU buffer storing [`RawDistancePixel`] values");

/// 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: ImageSize,

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

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

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

    /// Root tile buffers (64²)
    tile64: TileBuffers<64>,

    /// Second-stage tile buffers (8²)
    tile8: TileBuffers<8>,

    /// Pixel data
    pixels: FlexBuffer<PixelBufferTag>,

    /// Cached bind groups
    bind_groups: BindGroups,
}

impl Workspace {
    /// Builds a new set of buffers with a default size
    ///
    /// It is expected that these will be resized before being used
    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()
        );

        // Dummy size for infallible construction
        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 tile_tapes = FlexBuffer::new(
            device,
            "tile tape",
            Self::tile_tapes_buf_size(render_size),
        )
        .unwrap();

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

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

        Self {
            config_buf,
            vars_buf,
            image_size,
            tile_tapes,
            tile64,
            tile8,
            pixels,
            bind_groups: Default::default(),
        }
    }

    /// Returns the number of bytes in the `tile_tapes` buffer
    ///
    /// This is two levels of densely-allocated tiles; the first set are
    /// 64² and the second are 8².
    ///
    /// In other words, it looks something like this:
    ///
    /// ```text
    /// | index | index | index | ... |     densely packed 64² tape indices
    /// | index | index | index | ... |     densely packed 8² tape indices
    /// ```
    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();

        // Total size computation:
        //    (nx * ny) + (nx * ny * 64)
        // => nx * ny * 65
        nx.checked_mul(ny).unwrap().checked_mul(65).unwrap()
    }

    /// 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: ImageSize,
    ) -> Result<(), BufferError> {
        let render_size = TileRenderSize::from(image_size);
        let Workspace {
            image_size: image_size_ref,
            tile_tapes,
            tile64,
            tile8,
            config_buf: _,
            vars_buf: _,
            pixels,
            bind_groups,
        } = self;
        // Clear our cached bind groups if the image sizes is changing
        if *image_size_ref != image_size {
            *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,
            })?;
        tile8
            .grow_to_fit(device, render_size)
            .map_err(|e| BufferError {
                buf: BufferName::Tile8(e.buf),
                err: e.err,
            })?;
        pixels
            .grow_to_fit(device, image_size)
            .map_err(|err| BufferError {
                buf: BufferName::Pixel,
                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,
            vars_buf,
            tile_tapes,
            tile64,
            tile8,
            pixels,
            bind_groups: _,
        } = self;
        config_buf.size()
            + vars_buf.capacity()
            + tile_tapes.capacity()
            + tile64.capacity()
            + tile8.capacity()
            + pixels.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: _,
            vars_buf,
            config_buf,
            tile_tapes,
            tile64,
            tile8,
            pixels,
            bind_groups: _,
        } = self;
        config_buf.size()
            + vars_buf.size_bytes()
            + tile_tapes.size_bytes()
            + tile64.size()
            + tile8.size()
            + pixels.size_bytes()
    }

    /// Returns a handle to the image storage buffer
    ///
    /// This is intended for subsequent shaders which want to use the
    /// [`RawDistancePixel`] image data without copying to the CPU.  It requires
    /// an 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<PixelBufferTag> {
        &self.pixels
    }
}

/// 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,
}

/// Names of all buffers, used for error reporting
#[derive(Debug)]
pub enum BufferName {
    /// Tiles from the 64² root tile pass
    Tile64(TileBufferName),
    /// Tiles from the 8² intermediate tile pass
    Tile8(TileBufferName),
    /// Buffer for tile tapes
    TileTapes,
    /// GPU-written image pixels (as [`RawDistancePixel`] values)
    Pixel,
    /// CPU-mappable image pixels (as [`RawDistancePixel`] values)
    Image,
    /// Variables
    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::Tile8(buf) => write!(f, "`{buf}` tile8"),
            BufferName::TileTapes => write!(f, "`tile tapes`"),
            BufferName::Pixel => write!(f, "`pixel`"),
            BufferName::Image => write!(f, "`image`"),
            BufferName::Vars => write!(f, "`vars`"),
        }
    }
}

/// 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)]
    Buffers(#[from] BufferError),
}

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

/// Cached bind groups (constructed on-demand)
#[derive(Default)]
struct BindGroups {
    common: std::cell::OnceCell<wgpu::BindGroup>,
    root_tiles: std::cell::OnceCell<wgpu::BindGroup>,
    interval_tiles: std::cell::OnceCell<wgpu::BindGroup>,
    pixel_tiles: std::cell::OnceCell<wgpu::BindGroup>,
    merge: 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 root_tiles(
        &self,
        ctx: &Context,
        workspace: &Workspace,
    ) -> &wgpu::BindGroup {
        self.root_tiles.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.values.bind_active(),
                        },
                    ],
                })
        })
    }

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

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

    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.values.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: workspace.tile8.values.bind_active(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: workspace.pixels.bind_active(),
                        },
                    ],
                })
        })
    }
}

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

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

/// Root tile buffers store strata-packed tile lists
struct TileBuffers<const N: usize> {
    /// Output tiles
    tiles: FlexBuffer<TilesBufferTag>,

    /// Tile values (empty / full)
    values: FlexBuffer<ValuesBufferTag>,
}

/// Error type when resizing root tile buffers
#[derive(Debug, thiserror::Error)]
#[error("failed to resize `{buf}` root tile buffer")]
pub struct TileBufferError {
    /// Buffer which failed to resize
    pub buf: TileBufferName,
    /// 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 TileBufferName {
    Tiles,
    Values,
}

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::Values => "values",
        };
        s.fmt(f)
    }
}

impl<const N: usize> TileBuffers<N> {
    /// Build a new root tiles buffer, which stores strata-packed tile lists
    fn new(
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<Self, TileBufferError> {
        // 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| TileBufferError {
            buf: TileBufferName::Tiles,
            err,
        })?;

        let values_buf_size = Self::values_buf_size(render_size);
        let values =
            FlexBuffer::new(device, format!("tile{N}_values"), values_buf_size)
                .map_err(|err| TileBufferError {
                    buf: TileBufferName::Values,
                    err,
                })?;
        Ok(Self { tiles, values })
    }

    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();
        // wg_dispatch: [u32; 3] (unused)
        // count: u32,
        4 + nx
            .checked_mul(ny)
            .unwrap()
            .checked_mul((64 / N) * (64 / N))
            .unwrap()
    }

    fn values_buf_size(render_size: TileRenderSize) -> ImageSize {
        ImageSize::new(
            render_size.nx().checked_mul(64 / N as u32).unwrap(),
            render_size.ny().checked_mul(64 / N as u32).unwrap(),
        )
    }

    /// Grows all of the buffers to fit a particular render size
    fn grow_to_fit(
        &mut self,
        device: &wgpu::Device,
        render_size: TileRenderSize,
    ) -> Result<(), TileBufferError> {
        // Destructure to make sure we take all members into account
        let TileBuffers { tiles, values } = self;
        tiles
            .grow_to_fit(device, Self::tiles_buf_size(render_size))
            .map_err(|err| TileBufferError {
                buf: TileBufferName::Tiles,
                err,
            })?;
        values
            .grow_to_fit(device, Self::values_buf_size(render_size))
            .map_err(|err| TileBufferError {
                buf: TileBufferName::Values,
                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 TileBuffers { tiles, values } = self;
        tiles.size_bytes() + values.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 TileBuffers { tiles, values } = self;
        tiles.capacity() + values.capacity()
    }
}

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

/// Context for 2D (distance field) rendering
pub struct Context {
    gpu: Gpu,

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

    /// Context for root tile evaluation (64²)
    root_ctx: RootContext,

    /// Context for second-stage tile evaluation (generating 8² tiles)
    tiles_ctx: IntervalTilesContext,

    /// Context for per-pixel evaluation (taking 8² tiles)
    pixels_ctx: PixelTilesContext,

    /// Context which resets buffers before evaluation
    reset_ctx: ResetContext,

    /// Context to merge tile and pixel images
    merge_ctx: MergeContext,
}

impl Context {
    /// Build a new 2D rendering context, given a device and queue
    ///
    /// If render timestamps are desirable, then the device should be
    /// initialized with [`wgpu::Features::TIMESTAMP_QUERY`].
    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 tiles_ctx =
            IntervalTilesContext::new(&gpu.device, &common_bind_group_layout);
        let pixels_ctx =
            PixelTilesContext::new(&gpu.device, &common_bind_group_layout);
        let merge_ctx =
            MergeContext::new(&gpu.device, &common_bind_group_layout);

        Self {
            gpu: gpu.clone(),
            common_bind_group_layout,
            root_ctx,
            tiles_ctx,
            pixels_ctx,
            merge_ctx,
            reset_ctx: ResetContext,
        }
    }

    /// 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<PixelBufferTag>,
        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<PixelBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.submit_with_vars(shape, vars, workspace, &settings)?;
        self.gpu.copy(workspace.output(), 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
    pub async fn run_async(
        &self,
        shape: &RenderShape,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<PixelBufferTag>,
        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
    pub async fn run_with_vars_async(
        &self,
        shape: &RenderShape,
        vars: &ShapeVars<f32>,
        workspace: &mut Workspace,
        out: &mut ReadBuffer<PixelBufferTag>,
        settings: RenderConfig,
    ) -> Result<Image, SubmitError> {
        self.submit_with_vars(shape, vars, workspace, &settings)?;
        self.gpu.copy(workspace.output(), 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 [`RawDistancePixel`] data) is
    /// available on the GPU in [`buffers.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> {
        workspace.set_image_size(&self.gpu.device, settings.image_size)?;
        let render_size = TileRenderSize::from(workspace.image_size);

        // The WebGPU config type has a mat3x3f, but that type pads each row to
        // 16 bytes, so we'll just use a mat4x4 for simplicity
        let mat =
            settings.world_to_model * workspace.image_size.screen_to_world();
        let mut mat4 = nalgebra::Matrix4x3::<f32>::identity();
        mat4.fixed_view_mut::<3, 3>(0, 0).copy_from(&mat);

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

        {
            // 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();
        }

        // Create a command encoder and dispatch the compute work
        let mut encoder = self.gpu.device.create_command_encoder(
            &wgpu::CommandEncoderDescriptor { label: None },
        );

        // Initial buffer reset pass
        self.reset_ctx.run(&mut 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,
        );
        self.tiles_ctx.run(
            self,
            workspace,
            shape.bytecode.reg_count(),
            &mut compute_pass,
        );
        self.pixels_ctx.run(
            self,
            workspace,
            shape.bytecode.reg_count(),
            &mut compute_pass,
        );

        // Merge filled tiles from large -> small
        self.merge_ctx.run(
            self,
            workspace,
            settings.image_size,
            &mut compute_pass,
        );
        drop(compute_pass);

        // Submit the commands and wait for the GPU to complete
        self.gpu.queue.submit(Some(encoder.finish()));
        Ok(())
    }
}

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

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_ro(0), // tile64_values
                    buffer_ro(1), // tile8_values
                    buffer_rw(2), // pixels
                ],
            });

        // 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,
        render_size: ImageSize,
        compute_pass: &mut wgpu::ComputePass,
    ) {
        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 `count` and `wg_size` members of the tile output buffers
        encoder.clear_buffer(workspace.tile64.tiles.data(), 0, Some(16));
        encoder.clear_buffer(workspace.tile8.tiles.data(), 0, Some(16));
        workspace.tile64.values.clear(encoder);
        workspace.tile8.values.clear(encoder);
        workspace.pixels.clear(encoder);

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

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

#[cfg(test)]
mod test {
    use super::effects::ColorSettings;
    use super::*;
    use crate::color::{ShapeColor, ShapeColorBuffers};

    use fidget_core::{context::Tree, vm::VmShape};

    #[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_pixel_tiles_shader() {
        crate::compile_shader(&pixel_tiles_shader(16), "pixel tiles");
    }

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

    struct RenderOutput {
        distance: Image,
        color: fidget_raster::Image<u32, ImageSize>,
    }

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

        let mut buf = pixel_ctx.workspace();
        let mut merge_buf = effects_ctx.merge_workspace();

        // Render and accumulate each shape
        for (shape, _) in shapes {
            let shape =
                RenderShape::new(&VmShape::from(shape.clone())).unwrap();
            pixel_ctx
                .submit_with_vars(&shape, vars, &mut buf, &render_config)
                .unwrap();
            effects_ctx
                .submit_merge(buf.output(), true, &mut merge_buf)
                .unwrap();
        }
        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 per-pixel colors
        effects_ctx
            .submit_color_with_vars(
                &mut merge_buf,
                ColorSettings {
                    z: 0.0,
                    only_filled: false,
                    world_to_model: render_config.world_to_model,
                },
                &shape_colors,
                &mut color_workspace,
                vars,
            )
            .unwrap();

        let mut distance_out = gpu.read_buffer_for(merge_buf.output_distance());
        let mut color_out =
            gpu.read_buffer_for(merge_buf.output_color().unwrap());
        gpu.copy(merge_buf.output_distance(), &mut distance_out);
        gpu.copy(merge_buf.output_color().unwrap(), &mut color_out);

        RenderOutput {
            color: gpu.map_image(&mut color_out).image(),
            distance: gpu.map_image(&mut distance_out).image(),
        }
    }

    #[test]
    fn pixel_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 (x, y, _z) = Tree::axes();
        let circle = (x.square() + y.square()).sqrt() - Tree::constant(0.5);

        // Test a variety of image sizes for correctness
        for image_size in [
            RenderSize::new(64, 64),
            RenderSize::new(128, 64),
            RenderSize::new(64, 128),
            RenderSize::new(27, 51),
        ] {
            let out = render(
                &[(
                    circle.clone(),
                    ShapeColor::Rgb {
                        r: Tree::x(),
                        g: Tree::y(),
                        b: Tree::constant(0.5),
                    },
                )],
                RenderConfig {
                    image_size,
                    world_to_model: nalgebra::Matrix3::identity(),
                    pixel_perfect: false,
                    z: 0.0,
                },
                &Default::default(),
            );
            assert_eq!(out.color.size(), image_size);
            assert_eq!(out.distance.size(), image_size);

            // Basic circle inside/outside check
            let mat = image_size.screen_to_world();
            for j in 0..image_size.height() {
                for i in 0..image_size.width() {
                    let pos = mat.transform_point(&nalgebra::Point2::new(
                        i as f32, j as f32,
                    ));
                    let p = out.distance[(j as usize, i as usize)];
                    let r = (pos.x.powi(2) + pos.y.powi(2)).sqrt();
                    if r < 0.5 {
                        assert!(
                            p.inside(),
                            "pixel should be inside at pixel ({i}, {j}) \
                            (pos {pos}) with radius {r}"
                        );
                    } else {
                        assert!(
                            !p.inside(),
                            "pixel should be outside at pixel ({i}, {j}) \
                            (pos {pos}) with radius {r}"
                        );
                    }
                }
            }

            for j in 0..image_size.height() {
                for i in 0..image_size.width() {
                    let pos = mat.transform_point(&nalgebra::Point2::new(
                        i as f32, j as f32,
                    ));
                    let p = out.color[(j as usize, i as usize)];
                    let r = (pos.x.powi(2) + pos.y.powi(2)).sqrt();
                    let alpha = if r < 0.5 { 255 } else { 0 };
                    let expected_color = u32::from_ne_bytes([
                        (pos.x.clamp(0.0, 1.0) * 255.0) as u8,
                        (pos.y.clamp(0.0, 1.0) * 255.0) as u8,
                        127,
                        alpha,
                    ]);
                    assert_eq!(
                        p, expected_color,
                        "color mismatch at {i}, {j} ({pos})"
                    );
                }
            }
        }
    }

    #[test]
    fn pixel_multiple_images() {
        // 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 circle_a = ((Tree::x() - 0.5).square() + Tree::y().square()).sqrt()
            - Tree::constant(0.25);
        let circle_b = ((Tree::x() + 0.5).square() + Tree::y().square()).sqrt()
            - Tree::constant(0.25);

        // Test a variety of image sizes for correctness
        let image_size = RenderSize::new(64, 64);
        let out = render(
            &[
                (
                    circle_a,
                    ShapeColor::Rgb {
                        r: Tree::constant(0.0),
                        g: Tree::constant(0.0),
                        b: Tree::constant(1.0),
                    },
                ),
                (
                    circle_b,
                    ShapeColor::Rgb {
                        r: Tree::constant(0.0),
                        g: Tree::constant(1.0),
                        b: Tree::constant(0.0),
                    },
                ),
            ],
            RenderConfig {
                image_size,
                world_to_model: nalgebra::Matrix3::identity(),
                pixel_perfect: false,
                z: 0.0,
            },
            &Default::default(),
        );
        assert_eq!(out.color.size(), image_size);
        assert_eq!(out.distance.size(), image_size);

        let mut pixels = String::new();
        for j in 0..image_size.height() {
            for i in 0..image_size.width() {
                let p = out.color[(j as usize, i as usize)];
                let c = match p.to_ne_bytes() {
                    [0, 0, 255, 0] => "b",
                    [0, 0, 255, 255] => "B",
                    [0, 255, 0, 0] => "g",
                    [0, 255, 0, 255] => "G",
                    _ => panic!("invalid color {p:?}"),
                };
                pixels += c;
            }
            pixels += "\n";
        }

        assert_eq!(
            pixels,
            "\
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
            gggggggggggggGGGGGGGgggggggggggggbbbbbbbbbbbbBBBBBBBbbbbbbbbbbbb
            gggggggggggGGGGGGGGGGGgggggggggggbbbbbbbbbbBBBBBBBBBBBbbbbbbbbbb
            ggggggggggGGGGGGGGGGGGGggggggggggbbbbbbbbbBBBBBBBBBBBBBbbbbbbbbb
            ggggggggggGGGGGGGGGGGGGggggggggggbbbbbbbbbBBBBBBBBBBBBBbbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggbbbbbbbbBBBBBBBBBBBBBBBbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggbbbbbbbbBBBBBBBBBBBBBBBbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggbbbbbbbbBBBBBBBBBBBBBBBbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggbbbbbbbbBBBBBBBBBBBBBBBbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggbbbbbbbbBBBBBBBBBBBBBBBbbbbbbbb
            gggggggggGGGGGGGGGGGGGGGgggggggggggggggbbBBBBBBBBBBBBBBBgggggggg
            gggggggggGGGGGGGGGGGGGGGgggggggggggggggbbBBBBBBBBBBBBBBBbggggggg
            ggggggggggGGGGGGGGGGGGGggggggggggggggggbbbBBBBBBBBBBBBBbbggggggg
            ggggggggggGGGGGGGGGGGGGggggggggggggggggbbbBBBBBBBBBBBBBbbggggggg
            gggggggggggGGGGGGGGGGGgggggggggggggggggbbbbBBBBBBBBBBBbbbggggggg
            gggggggggggggGGGGGGGgggggggggggggggggggbbbbbbBBBBBBBbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggbbbbbbbbbbbbbbbbbbggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            "
            .replace(" ", "")
        );
    }

    #[test]
    fn pixel_vars() {
        // 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 va = fidget_core::var::Var::new();
        let vb = fidget_core::var::Var::new();
        let vc = fidget_core::var::Var::new();
        let vd = fidget_core::var::Var::new();
        let circle_a = ((Tree::x() - 0.5).square() + Tree::y().square()).sqrt()
            - Tree::from(va);
        let circle_b = ((Tree::x() + va).square() + Tree::y().square()).sqrt()
            - Tree::from(vb);

        // Test a variety of image sizes for correctness
        let image_size = RenderSize::new(64, 64);
        let mut vars = fidget_core::shape::ShapeVars::new();
        vars.insert(va.index().unwrap(), 0.25);
        vars.insert(vb.index().unwrap(), 0.55);
        vars.insert(vc.index().unwrap(), 0.75);
        vars.insert(vd.index().unwrap(), 0.5);
        let out = render(
            &[
                (
                    circle_a,
                    ShapeColor::Rgb {
                        r: Tree::constant(0.0),
                        g: Tree::constant(0.0),
                        b: Tree::from(vc),
                    },
                ),
                (
                    circle_b,
                    ShapeColor::Rgb {
                        r: Tree::constant(0.0),
                        g: Tree::from(vd),
                        b: Tree::constant(0.0),
                    },
                ),
            ],
            RenderConfig {
                image_size,
                world_to_model: nalgebra::Matrix3::identity(),
                pixel_perfect: false,
                z: 0.0,
            },
            &vars,
        );
        assert_eq!(out.color.size(), image_size);
        assert_eq!(out.distance.size(), image_size);

        let mut pixels = String::new();
        for j in 0..image_size.height() {
            for i in 0..image_size.width() {
                let p = out.color[(j as usize, i as usize)];
                let c = match p.to_ne_bytes() {
                    [0, 0, 191, 0] => "b",
                    [0, 0, 191, 255] => "B",
                    [0, 127, 0, 0] => "g",
                    [0, 127, 0, 255] => "G",
                    _ => panic!("invalid color {:?}", p.to_ne_bytes()),
                };
                pixels += c;
            }
            pixels += "\n";
        }

        assert_eq!(
            pixels,
            "\
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            ggggggggggggggggggggGGGGGGGGGggggggggggggggggggggggggggggggggggg
            gggggggggggggggggGGGGGGGGGGGGGGGgggggggggggggbbbbbbbbbbbbggggggg
            gggggggggggggggGGGGGGGGGGGGGGGGGGGgggggggggggbbbbbbbbbbbbggggggg
            ggggggggggggggGGGGGGGGGGGGGGGGGGGGGgggggggggbbbbbbbbbbbbbggggggg
            gggggggggggggGGGGGGGGGGGGGGGGGGGGGGGggggggggbbbbbbbbbbbbbggggggg
            ggggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGgggggggbbbbbbbbbbbbbggggggg
            gggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGgggggbbbbbbbbbbbbbbggggggg
            ggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGggggbbbbbbbbbbbbbbggggggg
            gggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgggbbbbbbbbbbbbbbggggggg
            gggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGggbbbbbbbbbbbbbbbbbbbbbb
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgbbbBBBBBBBbbbbbbbbbbbb
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgbBBBBBBBBBBBbbbbbbbbbb
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgBBBBBBBBBBBBBbbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBbbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbbbbbbbb
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBgggggggg
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBBbggggggg
            gggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGBBBBBBBBBBBBBbbggggggg
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgBBBBBBBBBBBBBbbggggggg
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgbBBBBBBBBBBBbbbggggggg
            ggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGgbbbBBBBBBBbbbbbggggggg
            gggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGggbbbbbbbbbbbbbbbggggggg
            gggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGggbbbbbbbbbbbbbbbggggggg
            ggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGGGggggggggggggggggggggggggg
            gggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGGGgggggggggggggggggggggggggg
            ggggggggggggGGGGGGGGGGGGGGGGGGGGGGGGGggggggggggggggggggggggggggg
            gggggggggggggGGGGGGGGGGGGGGGGGGGGGGGgggggggggggggggggggggggggggg
            ggggggggggggggGGGGGGGGGGGGGGGGGGGGGggggggggggggggggggggggggggggg
            gggggggggggggggGGGGGGGGGGGGGGGGGGGgggggggggggggggggggggggggggggg
            gggggggggggggggggGGGGGGGGGGGGGGGgggggggggggggggggggggggggggggggg
            ggggggggggggggggggggGGGGGGGGGggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg
            "
            .replace(" ", "")
        );
    }

    #[test]
    fn pixel_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 circle_a = ((Tree::x() - 0.5).square() + Tree::y().square()).sqrt()
            - Tree::constant(0.25);
        let circle_b = ((Tree::x() + 0.5).square() + Tree::y().square()).sqrt()
            - Tree::constant(0.25);

        // Test a variety of image sizes for correctness
        let image_size = RenderSize::new(64, 64);
        let out = render(
            &[
                (
                    circle_a,
                    ShapeColor::Hsl {
                        h: Tree::constant(0.0),
                        s: Tree::constant(1.0),
                        l: Tree::constant(0.5),
                    },
                ),
                (
                    circle_b,
                    ShapeColor::Hsl {
                        h: Tree::constant(0.5),
                        s: Tree::constant(1.0),
                        l: Tree::constant(0.5),
                    },
                ),
            ],
            RenderConfig {
                image_size,
                world_to_model: nalgebra::Matrix3::identity(),
                pixel_perfect: false,
                z: 0.0,
            },
            &Default::default(),
        );
        assert_eq!(out.color.size(), image_size);
        assert_eq!(out.distance.size(), image_size);

        let mut pixels = String::new();
        for j in 0..image_size.height() {
            for i in 0..image_size.width() {
                let p = out.color[(j as usize, i as usize)];
                let c = match p.to_ne_bytes() {
                    [255, 0, 0, 0] => "r",
                    [255, 0, 0, 255] => "R",
                    [0, 255, 255, 0] => "c",
                    [0, 255, 255, 255] => "C",
                    _ => panic!("invalid color {:?}", p.to_ne_bytes()),
                };
                pixels += c;
            }
            pixels += "\n";
        }

        assert_eq!(
            pixels,
            "\
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr
            cccccccccccccCCCCCCCcccccccccccccrrrrrrrrrrrrRRRRRRRrrrrrrrrrrrr
            cccccccccccCCCCCCCCCCCcccccccccccrrrrrrrrrrRRRRRRRRRRRrrrrrrrrrr
            ccccccccccCCCCCCCCCCCCCccccccccccrrrrrrrrrRRRRRRRRRRRRRrrrrrrrrr
            ccccccccccCCCCCCCCCCCCCccccccccccrrrrrrrrrRRRRRRRRRRRRRrrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccrrrrrrrrRRRRRRRRRRRRRRRrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccrrrrrrrrRRRRRRRRRRRRRRRrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccrrrrrrrrRRRRRRRRRRRRRRRrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccrrrrrrrrRRRRRRRRRRRRRRRrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccrrrrrrrrRRRRRRRRRRRRRRRrrrrrrrr
            cccccccccCCCCCCCCCCCCCCCcccccccccccccccrrRRRRRRRRRRRRRRRcccccccc
            cccccccccCCCCCCCCCCCCCCCcccccccccccccccrrRRRRRRRRRRRRRRRrccccccc
            ccccccccccCCCCCCCCCCCCCccccccccccccccccrrrRRRRRRRRRRRRRrrccccccc
            ccccccccccCCCCCCCCCCCCCccccccccccccccccrrrRRRRRRRRRRRRRrrccccccc
            cccccccccccCCCCCCCCCCCcccccccccccccccccrrrrRRRRRRRRRRRrrrccccccc
            cccccccccccccCCCCCCCcccccccccccccccccccrrrrrrRRRRRRRrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccrrrrrrrrrrrrrrrrrrccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
            "
            .replace(" ", "")
        );
    }

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