sex.rs

use super::assets::{Assets, Sprite};
use std::borrow::Cow;
use std::cell::Cell;
use std::cell::UnsafeCell;
use std::mem;
use std::mem::ManuallyDrop;
use std::mem::MaybeUninit;
use std::num::NonZeroU32;
use std::ptr;
use std::slice;
use thiserror::Error;
use vek::{FrustumPlanes, Mat4, Vec2, Vec3};
use wgpu::util::DeviceExt;
use winit::event_loop::EventLoopWindowTarget;
use winit::window::{Fullscreen, Window as WinitWindow, WindowBuilder};

const MAX_QUADS: u64 = 1000;

pub struct Frame<'a> {
    frame: wgpu::SurfaceFrame,
    window: &'a Window,
}

impl<'a> Frame<'a> {
    pub fn gpu(&self) -> &Gpu {
        self.window.gpu()
    }

    pub fn dimensions(&self) -> Vec2<f32> {
        self.window.dimensions()
    }

    fn view(&self) -> wgpu::TextureView {
        self.frame
            .output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default())
    }
}

pub struct Window {
    window: WinitWindow,
    gpu: Gpu,
    fullscreen: Cell<bool>,
}

impl Window {
    pub(super) async fn new(
        event_loop: &EventLoopWindowTarget<()>,
        name: &str,
    ) -> Result<Self, WindowCreationError> {
        let window = WindowBuilder::new().with_title(name).build(event_loop)?;
        let gpu = Gpu::new(&window).await?;

        Ok(Self {
            window,
            gpu,
            fullscreen: Cell::new(false),
        })
    }

    pub(super) fn new_frame(&self) -> Frame {
        let frame = self.gpu.new_frame();

        Frame {
            frame,
            window: self,
        }
    }

    pub fn gpu(&self) -> &Gpu {
        &self.gpu
    }

    pub fn toggle_fullscreen(&self) {
        let option = if self.fullscreen.get() {
            None
        } else {
            Some(Fullscreen::Borderless(None))
        };

        self.window.set_fullscreen(option);
        self.fullscreen.update(|v| !v);
    }

    pub fn dimensions(&self) -> Vec2<f32> {
        let size = self.window.inner_size();
        Vec2::new(size.width as f32, size.height as f32)
    }

    pub(super) fn id(&self) -> winit::window::WindowId {
        self.window.id()
    }
}

pub struct Gpu {
    surface: wgpu::Surface,
    device: wgpu::Device,
    queue: wgpu::Queue,
}

impl Gpu {
    async fn new(window: &WinitWindow) -> Result<Self, WindowCreationError> {
        let size = window.inner_size();
        let instance = wgpu::Instance::new(wgpu::Backends::PRIMARY);
        let surface = unsafe { instance.create_surface(window) };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
            })
            .await
            .ok_or(WindowCreationError::NoAdapter)?;

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    features: wgpu::Features::empty(),
                    limits: wgpu::Limits::default(),
                    label: None,
                },
                None,
            )
            .await?;

        let gpu = Self {
            surface,
            device,
            queue,
        };

        gpu.reset_error_handler();
        gpu.resize(size);
        Ok(gpu)
    }

    pub(super) fn set_error_handler(&self, handler: impl wgpu::UncapturedErrorHandler) {
        self.device.on_uncaptured_error(handler);
    }

    pub(super) fn reset_error_handler(&self) {
        self.device
            .on_uncaptured_error(|error| panic!("{:?}", error));
    }

    fn new_frame(&self) -> wgpu::SurfaceFrame {
        self.surface.get_current_frame().unwrap()
    }

    pub(super) fn resize(&self, new_size: winit::dpi::PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.surface.configure(
                &self.device,
                &wgpu::SurfaceConfiguration {
                    usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                    format: wgpu::TextureFormat::Bgra8Unorm,
                    width: new_size.width,
                    height: new_size.height,
                    present_mode: wgpu::PresentMode::Fifo,
                },
            );
        }
    }

    fn device(&self) -> &wgpu::Device {
        &self.device
    }

    fn queue(&self) -> &wgpu::Queue {
        &self.queue
    }

    pub fn create_2d_srgb_texture(
        &self,
        width: u32,
        height: u32,
        usage: wgpu::TextureUsages,
        label: Option<&str>,
        data: &[u8],
    ) -> wgpu::Texture {
        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };

        let texture = self.device.create_texture(&wgpu::TextureDescriptor {
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            usage,
            label,
        });

        let data_layout = wgpu::ImageDataLayout {
            offset: 0,
            bytes_per_row: NonZeroU32::new(width * 4),
            rows_per_image: NonZeroU32::new(height),
        };

        self.queue
            .write_texture(texture.as_image_copy(), data, data_layout, size);

        texture
    }

    pub fn create_shader_module(
        &self,
        label: Option<&str>,
        wgsl_source: &str,
    ) -> wgpu::ShaderModule {
        self.device
            .create_shader_module(&wgpu::ShaderModuleDescriptor {
                label,
                source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(wgsl_source)),
            })
    }
}

#[derive(Error, Debug)]
pub enum WindowCreationError {
    #[error("no graphics adapter found")]
    NoAdapter,

    #[error("failed to request graphics device")]
    RequestDeviceError(#[from] wgpu::RequestDeviceError),

    #[error("failed to create window: {0}")]
    WindowCreationError(#[from] winit::error::OsError),
}

#[allow(dead_code)]
pub struct Camera {
    position: Vec2<f32>,
    zoom: f32,
}

impl Camera {
    pub fn new(position: Vec2<f32>, zoom: f32) -> Self {
        Self { position, zoom }
    }

    #[allow(dead_code)]
    pub fn set_position(&mut self, position: Vec2<f32>) {
        self.position = position;
    }

    #[allow(dead_code)]
    pub fn zoom(&mut self, by: f32) {
        self.zoom += self.zoom * by;
    }

    fn view(&self) -> Mat4<f32> {
        Mat4::identity()
    }
}

pub struct Renderer {
    sprite_sampler: wgpu::Sampler,
    sprite_bind_layout: wgpu::BindGroupLayout,
    sprite_pipeline: wgpu::RenderPipeline,
    sprite_index_buffer: wgpu::Buffer,
    camera: Camera,
    vertex_buffer: wgpu::Buffer,
    matrix_buffer: wgpu::Buffer,
    vertices: UnsafeCell<Vec<[Vertex; 4]>>,
    matrices: UnsafeCell<Vec<[[f32; 4]; 4]>>,
    bind_groups: UnsafeCell<Vec<wgpu::BindGroup>>,
    encoder: UnsafeCell<MaybeUninit<wgpu::CommandEncoder>>,
    output_view: UnsafeCell<MaybeUninit<wgpu::TextureView>>,
}

impl Renderer {
    pub fn new(gpu: &Gpu, assets: &Assets, camera: Camera) -> Self {
        let device = gpu.device();
        let sprite_sampler = sprite_sampler(device);
        let sprite_bind_layout = sprite_bind_group_layout(device);
        let sprite_pipeline = sprite_render_pipeline(device, assets, &sprite_bind_layout);
        let sprite_index_buffer = sprite_index_buffer(device);
        let vertex_buffer = vertex_buffer(device);
        let matrix_buffer = matrix_buffer(device);

        Self {
            sprite_sampler,
            sprite_bind_layout,
            sprite_pipeline,
            sprite_index_buffer,
            camera,
            vertex_buffer,
            matrix_buffer,
            vertices: UnsafeCell::new(Vec::new()),
            matrices: UnsafeCell::new(Vec::new()),
            bind_groups: UnsafeCell::new(Vec::new()),
            encoder: UnsafeCell::new(MaybeUninit::uninit()),
            output_view: UnsafeCell::new(MaybeUninit::uninit()),
        }
    }

    pub fn reload(&mut self, gpu: &Gpu, assets: &Assets) {
        let device = gpu.device();
        self.sprite_pipeline = sprite_render_pipeline(device, assets, &self.sprite_bind_layout);
    }

    #[allow(dead_code)]
    pub fn camera(&mut self) -> &mut Camera {
        &mut self.camera
    }

    pub fn context<'a>(&'a self, frame: Frame<'a>) -> RenderContext<'a> {
        RenderContext::new(self, frame)
    }

    #[allow(clippy::mut_from_ref)]
    unsafe fn encoder(&self) -> &mut MaybeUninit<wgpu::CommandEncoder> {
        &mut *self.encoder.get()
    }
}

pub struct RenderContext<'a> {
    renderer: &'a Renderer,
    render_pass: UnsafeCell<ManuallyDrop<wgpu::RenderPass<'a>>>,
    frame: Frame<'a>,
    view: Mat4<f32>,
    project: Mat4<f32>,
}

impl<'a> RenderContext<'a> {
    fn new(renderer: &'a Renderer, frame: Frame<'a>) -> Self {
        let view = renderer.camera.view();
        let window_size = frame.dimensions();
        let width = window_size.x / 4.0 / frame.window.window.scale_factor() as f32;
        let height = window_size.y / window_size.x * width;
        let frustum = FrustumPlanes {
            left: 0.0,
            right: width,
            bottom: height,
            top: 0.0,
            near: 1.0,
            far: -1.0,
        };

        let project = Mat4::orthographic_rh_no(frustum);
        let output_view = frame.view();
        let output_view_ref = unsafe {
            let view = &mut *renderer.output_view.get();
            view.write(output_view)
        };
        let encoder = unsafe {
            renderer.encoder().write(
                frame
                    .gpu()
                    .device()
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }),
            )
        };

        let render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: None,
            depth_stencil_attachment: None,
            color_attachments: &[wgpu::RenderPassColorAttachment {
                view: output_view_ref,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color {
                        r: 0.1,
                        g: 0.2,
                        b: 0.3,
                        a: 1.0,
                    }),
                    store: true,
                },
            }],
        });

        Self {
            renderer,
            render_pass: UnsafeCell::new(ManuallyDrop::new(render_pass)),
            frame,
            view,
            project,
        }
    }

    fn gpu(&self) -> &Gpu {
        self.frame.gpu()
    }

    pub fn pass<'b>(&'b self) -> SpritePass<'a, 'b> {
        SpritePass::new(self)
    }
}

impl<'a> Drop for RenderContext<'a> {
    fn drop(&mut self) {
        unsafe {
            ptr::drop_in_place(&mut **self.render_pass.get());
            ptr::drop_in_place(self.renderer.output_view.get());
            let bind_groups = &mut *self.renderer.bind_groups.get();
            bind_groups.clear();
            let encoder = self.renderer.encoder().assume_init_read();
            let commands = Some(encoder.finish());
            self.frame.gpu().queue().submit(commands);
        }
    }
}

pub struct SpritePass<'a, 'b> {
    context: &'b RenderContext<'a>,
}

impl<'a, 'b> SpritePass<'a, 'b> {
    fn new(context: &'b RenderContext<'a>) -> Self {
        Self { context }
    }

    pub fn target<'c>(&'c self) -> Target<'a, 'b, 'c> {
        Target::new(self)
    }
}

impl<'a, 'b> Drop for SpritePass<'a, 'b> {
    fn drop(&mut self) {
        let (vertices, matrices, bind_groups) = unsafe {
            (
                &mut *self.context.renderer.vertices.get(),
                &mut *self.context.renderer.matrices.get(),
                &*self.context.renderer.bind_groups.get(),
            )
        };

        let queue = self.context.gpu().queue();
        let raw_vertices = unsafe {
            slice::from_raw_parts(
                &**vertices as *const _ as *const u8,
                mem::size_of_val(&**vertices),
            )
        };

        let raw_matrices = unsafe {
            slice::from_raw_parts(
                &**matrices as *const _ as *const u8,
                mem::size_of_val(&**matrices),
            )
        };

        queue.write_buffer(&self.context.renderer.vertex_buffer, 0, raw_vertices);
        queue.write_buffer(&self.context.renderer.matrix_buffer, 0, raw_matrices);
        vertices.clear();
        matrices.clear();

        let render_pass = unsafe { &mut **self.context.render_pass.get() };
        render_pass.set_pipeline(&self.context.renderer.sprite_pipeline);
        render_pass.set_vertex_buffer(0, self.context.renderer.vertex_buffer.slice(..));
        render_pass.set_index_buffer(
            self.context.renderer.sprite_index_buffer.slice(..),
            wgpu::IndexFormat::Uint16,
        );

        for (instance, bind_group) in bind_groups.iter().enumerate() {
            let instances = (instance as u32)..(instance as u32 + 1);
            render_pass.set_bind_group(0, bind_group, &[]);
            render_pass.draw_indexed(0..6, 0, instances);
        }
    }
}

#[derive(Clone, Copy)]
pub struct Target<'a, 'b, 'c> {
    pass: &'c SpritePass<'a, 'b>,
    transform: Mat4<f32>,
}

impl<'a, 'b, 'c> Target<'a, 'b, 'c> {
    fn new(pass: &'c SpritePass<'a, 'b>) -> Self {
        Self {
            pass,
            transform: Mat4::identity(),
        }
    }

    fn object_transform(&self) -> Mat4<f32> {
        self.transform * self.pass.context.view * self.pass.context.project
    }

    #[allow(dead_code)]
    pub fn translate(&self, by: Vec2<f32>) -> Self {
        Self {
            pass: self.pass,
            transform: self.transform.translated_2d(by),
        }
    }

    pub fn scale(&self, by: f32) -> Self {
        Self {
            pass: self.pass,
            transform: self.transform.scaled_3d(Vec3::new(by, by, 0.0)),
        }
    }

    #[allow(dead_code)]
    pub fn rotate(&self, by: f32) -> Self {
        Self {
            pass: self.pass,
            transform: self.transform.rotated_z(by),
        }
    }

    pub fn draw(&self, sprite: &Sprite) {
        let (vertices, matrices) = unsafe {
            (
                &mut *self.pass.context.renderer.vertices.get(),
                &mut *self.pass.context.renderer.matrices.get(),
            )
        };

        let matrix = self.object_transform();
        let tl = [0.0, 0.0];
        let tr = [sprite.size.x, 0.0];
        let bl = [0.0, sprite.size.y];
        let br = [sprite.size.x, sprite.size.y];

        vertices.push([
            Vertex {
                position: tl,
                tex_coords: [0.0, 0.0],
            },
            Vertex {
                position: tr,
                tex_coords: [1.0, 0.0],
            },
            Vertex {
                position: bl,
                tex_coords: [0.0, 1.0],
            },
            Vertex {
                position: br,
                tex_coords: [1.0, 1.0],
            },
        ]);
        matrices.push(matrix.into_col_arrays());

        let bind_group = bind_group_for_texture(
            self.pass.context.gpu().device(),
            &self.pass.context.renderer.sprite_bind_layout,
            &self.pass.context.renderer.sprite_sampler,
            &sprite.texture,
            self.pass
                .context
                .renderer
                .matrix_buffer
                .as_entire_buffer_binding(),
        );

        unsafe {
            let binds = &mut *self.pass.context.renderer.bind_groups.get();
            binds.push(bind_group);
        }
    }
}

fn bind_group_for_texture(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    sampler: &wgpu::Sampler,
    texture: &wgpu::Texture,
    matrix: wgpu::BufferBinding,
) -> wgpu::BindGroup {
    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());

    device.create_bind_group(&wgpu::BindGroupDescriptor {
        layout,
        entries: &[
            wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&view),
            },
            wgpu::BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::Sampler(sampler),
            },
            wgpu::BindGroupEntry {
                binding: 2,
                resource: wgpu::BindingResource::Buffer(matrix),
            },
        ],
        label: None,
    })
}

fn sprite_sampler(device: &wgpu::Device) -> wgpu::Sampler {
    device.create_sampler(&wgpu::SamplerDescriptor {
        address_mode_u: wgpu::AddressMode::ClampToEdge,
        address_mode_v: wgpu::AddressMode::ClampToEdge,
        address_mode_w: wgpu::AddressMode::ClampToEdge,
        mag_filter: wgpu::FilterMode::Nearest,
        min_filter: wgpu::FilterMode::Nearest,
        mipmap_filter: wgpu::FilterMode::Nearest,
        ..Default::default()
    })
}

fn sprite_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    multisampled: false,
                    view_dimension: wgpu::TextureViewDimension::D2,
                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler {
                    comparison: false,
                    filtering: true,
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 2,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only: true },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            },
        ],
        label: None,
    })
}

fn sprite_index_buffer(device: &wgpu::Device) -> wgpu::Buffer {
    let indices: &[u16] = &[0, 2, 3, 3, 1, 0, 0, 0];
    let indices_raw = unsafe {
        slice::from_raw_parts(indices as *const _ as *const u8, mem::size_of_val(indices))
    };

    device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None,
        contents: indices_raw,
        usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
    })
}

fn vertex_buffer(device: &wgpu::Device) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: mem::size_of::<[Vertex; 4]>() as u64 * MAX_QUADS,
        mapped_at_creation: false,
        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
    })
}

fn matrix_buffer(device: &wgpu::Device) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: mem::size_of::<Mat4<f32>>() as u64 * MAX_QUADS,
        mapped_at_creation: false,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    })
}

fn sprite_render_pipeline(
    device: &wgpu::Device,
    assets: &Assets,
    sprite_bind_layout: &wgpu::BindGroupLayout,
) -> wgpu::RenderPipeline {
    let shader = assets.get_shader("shader/sprite").unwrap();

    let sprite_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &[sprite_bind_layout],
        push_constant_ranges: &[],
    });

    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: None,
        layout: Some(&sprite_pipeline_layout),
        vertex: wgpu::VertexState {
            module: shader,
            entry_point: "main",
            buffers: &[Vertex::desc()],
        },
        fragment: Some(wgpu::FragmentState {
            module: shader,
            entry_point: "main",
            targets: &[wgpu::ColorTargetState {
                format: wgpu::TextureFormat::Bgra8Unorm,
                blend: Some(wgpu::BlendState::REPLACE),
                write_mask: wgpu::ColorWrites::ALL,
            }],
        }),
        primitive: wgpu::PrimitiveState {
            topology: wgpu::PrimitiveTopology::TriangleList,
            strip_index_format: None,
            front_face: wgpu::FrontFace::Ccw,
            cull_mode: Some(wgpu::Face::Back),
            polygon_mode: wgpu::PolygonMode::Fill,
            clamp_depth: false,
            conservative: false,
        },
        depth_stencil: None,
        multisample: wgpu::MultisampleState {
            count: 1,
            mask: !0,
            alpha_to_coverage_enabled: false,
        },
    })
}

#[repr(C)]
#[derive(Copy, Clone, Debug)]
struct Vertex {
    position: [f32; 2],
    tex_coords: [f32; 2],
}

impl Vertex {
    fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x2,
                },
                wgpu::VertexAttribute {
                    offset: mem::size_of::<[f32; 2]>() as wgpu::BufferAddress,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x2,
                },
            ],
        }
    }
}