codecraft 0.1.1

//! Taking a picture of the frame that was just drawn.
//!
//! The window itself is the source: the swapchain texture is copied into a
//! buffer while the frame is still in hand, then written out as a PNG. That
//! is what makes a screenshot show exactly what the player sees, rather than
//! whatever a desktop screen grab happens to catch.
use std::path::Path;

use crate::gpu::GpuContext;

#[derive(Debug)]
pub enum CaptureError {
    Map(String),
    Write(std::io::Error),
    Encode(png::EncodingError),
}

impl std::fmt::Display for CaptureError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Map(e) => write!(f, "could not read the frame back: {e}"),
            Self::Write(e) => write!(f, "could not write the file: {e}"),
            Self::Encode(e) => write!(f, "could not encode the png: {e}"),
        }
    }
}

impl std::error::Error for CaptureError {}

/// A frame copied out of the swapchain, waiting to be read back.
pub struct Staged {
    buffer: wgpu::Buffer,
    /// Rows are padded to wgpu's copy alignment, so this is usually wider
    /// than the image.
    padded_row: u32,
    width: u32,
    height: u32,
}

/// Rounds a row up to the alignment `copy_texture_to_buffer` demands.
pub fn padded_row(width: u32) -> u32 {
    let unpadded = width * 4;
    let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
    unpadded.div_ceil(align) * align
}

/// Queues a copy of `texture` into a buffer. Must be recorded before the
/// frame is presented, while the texture is still ours.
pub fn stage(
    device: &wgpu::Device,
    encoder: &mut wgpu::CommandEncoder,
    texture: &wgpu::Texture,
    width: u32,
    height: u32,
) -> Staged {
    let padded_row = padded_row(width);
    let buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("screenshot readback"),
        size: (padded_row as u64) * (height as u64),
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
        mapped_at_creation: false,
    });

    encoder.copy_texture_to_buffer(
        wgpu::TexelCopyTextureInfo {
            texture,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        },
        wgpu::TexelCopyBufferInfo {
            buffer: &buffer,
            layout: wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(padded_row),
                rows_per_image: Some(height),
            },
        },
        wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
    );

    Staged {
        buffer,
        padded_row,
        width,
        height,
    }
}

/// Reads the staged copy back and writes it out as a PNG.
///
/// Blocks until the GPU has finished the copy, which is why it runs after the
/// frame is submitted rather than during it.
pub fn write_png(gpu: &GpuContext, staged: Staged, path: &Path) -> Result<(), CaptureError> {
    let slice = staged.buffer.slice(..);
    let (sender, receiver) = std::sync::mpsc::channel();
    slice.map_async(wgpu::MapMode::Read, move |result| {
        let _ = sender.send(result);
    });
    gpu.device
        .poll(wgpu::PollType::wait_indefinitely())
        .map_err(|e| CaptureError::Map(e.to_string()))?;
    receiver
        .recv()
        .map_err(|e| CaptureError::Map(e.to_string()))?
        .map_err(|e| CaptureError::Map(e.to_string()))?;

    let mapped = slice
        .get_mapped_range()
        .map_err(|e| CaptureError::Map(e.to_string()))?;
    let pixels = to_rgba(
        &mapped,
        staged.width,
        staged.height,
        staged.padded_row,
        gpu.format(),
    );
    drop(mapped);
    staged.buffer.unmap();

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).map_err(CaptureError::Write)?;
    }
    let file = std::fs::File::create(path).map_err(CaptureError::Write)?;
    let mut encoder = png::Encoder::new(std::io::BufWriter::new(file), staged.width, staged.height);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut writer = encoder.write_header().map_err(CaptureError::Encode)?;
    writer
        .write_image_data(&pixels)
        .map_err(CaptureError::Encode)?;
    Ok(())
}

/// Drops the row padding, and puts the channels in the order a PNG wants.
///
/// Surfaces are commonly BGRA, and a PNG is RGBA; getting this backwards
/// gives a screenshot with the red and blue swapped, which is subtle enough
/// to miss on a board that is mostly brown.
fn to_rgba(
    mapped: &[u8],
    width: u32,
    height: u32,
    padded_row: u32,
    format: wgpu::TextureFormat,
) -> Vec<u8> {
    let unpadded = (width * 4) as usize;
    let bgra = matches!(
        format,
        wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb
    );

    let mut pixels = Vec::with_capacity(unpadded * height as usize);
    for row in 0..height as usize {
        let start = row * padded_row as usize;
        let row_bytes = &mapped[start..start + unpadded];
        if bgra {
            for pixel in row_bytes.chunks_exact(4) {
                pixels.extend_from_slice(&[pixel[2], pixel[1], pixel[0], pixel[3]]);
            }
        } else {
            pixels.extend_from_slice(row_bytes);
        }
    }
    pixels
}

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

    #[test]
    fn rows_are_padded_up_to_the_copy_alignment() {
        // 256 bytes is what wgpu demands of a row.
        assert_eq!(padded_row(64), 256, "64px is exactly one aligned row");
        assert_eq!(padded_row(65), 512, "65px spills into a second");
        assert_eq!(padded_row(800), 3328, "3200 rounds up to 3328");
        assert_eq!(padded_row(1920), 7680, "already aligned");
        assert!(padded_row(37) >= 37 * 4, "never narrower than the image");
    }

    /// One red and one green pixel per row, with the row padded out.
    fn two_by_two(order: [u8; 4]) -> Vec<u8> {
        let [a, b, c, d] = order;
        let mut rows = Vec::new();
        for _ in 0..2 {
            rows.extend_from_slice(&[a, b, c, d]);
            rows.extend_from_slice(&[d, c, b, a]);
            // Padding out to a 16-byte stride: never part of the image.
            rows.extend_from_slice(&[0xAA; 8]);
        }
        rows
    }

    #[test]
    fn padding_is_dropped_from_every_row() {
        let mapped = two_by_two([10, 20, 30, 255]);
        let pixels = to_rgba(&mapped, 2, 2, 16, wgpu::TextureFormat::Rgba8UnormSrgb);

        assert_eq!(pixels.len(), 2 * 2 * 4, "the padding is not in the image");
        assert!(
            !pixels.contains(&0xAA),
            "padding bytes leaked into the image: {pixels:?}",
        );
    }

    #[test]
    fn an_rgba_surface_is_written_through_unchanged() {
        let mapped = two_by_two([10, 20, 30, 255]);
        let pixels = to_rgba(&mapped, 2, 2, 16, wgpu::TextureFormat::Rgba8UnormSrgb);

        assert_eq!(&pixels[0..4], &[10, 20, 30, 255]);
        assert_eq!(&pixels[4..8], &[255, 30, 20, 10]);
    }

    #[test]
    fn a_bgra_surface_has_its_red_and_blue_swapped() {
        // The surface stored B=10, G=20, R=30; the png wants R, G, B.
        let mapped = two_by_two([10, 20, 30, 255]);
        let pixels = to_rgba(&mapped, 2, 2, 16, wgpu::TextureFormat::Bgra8UnormSrgb);

        assert_eq!(&pixels[0..4], &[30, 20, 10, 255], "red and blue swap");
        assert_eq!(&pixels[4..8], &[20, 30, 255, 10]);
    }
}