canary-rs/renderers/wgpu/src/lib.rs

use std::sync::Arc;

use bytemuck::{Pod, Zeroable};
use canary::{DrawCommand, Vec2};
use wgpu::{util::*, *};

pub use wgpu;

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct Vertex {
    pub position: [f32; 2],
    pub color: u32,
}

impl Vertex {
    pub fn from_canary(size: Vec2, v: &canary::MeshVertex) -> Self {
        // TODO do this in the vertex shader with a size uniform
        Self {
            position: [
                (v.position.x / size.x) * 2.0 - 1.0,
                (v.position.y / size.y) * -2.0 + 1.0,
            ],
            color: v.color.0,
        }
    }

    pub const BUFFER_LAYOUT: VertexBufferLayout<'static> = VertexBufferLayout {
        array_stride: std::mem::size_of::<Self>() as BufferAddress,
        step_mode: VertexStepMode::Vertex,
        attributes: &[
            VertexAttribute {
                format: VertexFormat::Float32x2,
                offset: 0,
                shader_location: 0,
            },
            VertexAttribute {
                format: VertexFormat::Unorm8x4,
                offset: std::mem::size_of::<[f32; 2]>() as BufferAddress,
                shader_location: 1,
            },
        ],
    };
}

pub struct Renderer {
    device: Arc<Device>,
    queue: Arc<Queue>,
    pipeline_layout: PipelineLayout,
    pipeline: RenderPipeline,
}

impl Renderer {
    pub fn new(device: Arc<Device>, queue: Arc<Queue>) -> Self {
        let shader = device.create_shader_module(include_wgsl!("shader.wgsl"));

        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: Some("Canary Pipeline Layout"),
            bind_group_layouts: &[],
            push_constant_ranges: &[],
        });

        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: Some("Canary Render Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[Vertex::BUFFER_LAYOUT],
            },
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: FrontFace::Ccw,
                cull_mode: None,
                unclipped_depth: false,
                polygon_mode: PolygonMode::Fill,
                conservative: false,
            },
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[Some(ColorTargetState {
                    format: TextureFormat::Bgra8Unorm,
                    blend: Some(BlendState::ALPHA_BLENDING),
                    write_mask: ColorWrites::ALL,
                })],
            }),
            multiview: None,
        });

        Self {
            device,
            queue,
            pipeline_layout,
            pipeline,
        }
    }

    pub fn render(&self, target: DrawTarget<'_>, cmds: &[DrawCommand]) {
        let size = target.size;

        let mut joined_vs = Vec::new();
        let mut joined_is = Vec::new();

        for cmd in cmds.iter() {
            match cmd {
                canary::DrawCommand::Mesh { vertices, indices } => {
                    let voff = joined_vs.len() as canary::MeshIndex;
                    joined_vs.extend(vertices.iter().map(|v| Vertex::from_canary(size, v)));
                    joined_is.extend(indices.iter().map(|i| i + voff));
                }
                _ => unimplemented!(),
            }
        }

        let vertex_buffer = self.device.create_buffer_init(&BufferInitDescriptor {
            label: Some("Canary Vertex Buffer"),
            contents: bytemuck::cast_slice(joined_vs.as_slice()),
            usage: BufferUsages::VERTEX,
        });

        let index_buffer = self.device.create_buffer_init(&BufferInitDescriptor {
            label: Some("Canary Index Buffer"),
            contents: bytemuck::cast_slice(joined_is.as_slice()),
            usage: BufferUsages::INDEX,
        });

        let mut cmd_encoder = self.device.create_command_encoder(&Default::default());

        let mut rp = cmd_encoder.begin_render_pass(&RenderPassDescriptor {
            label: Some("Canary Render Pass"),
            color_attachments: &[Some(RenderPassColorAttachment {
                view: target.texture,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Clear(Color::TRANSPARENT),
                    store: true,
                },
            })],
            depth_stencil_attachment: None,
        });

        let indices = 0..(joined_is.len() as u32);
        rp.set_vertex_buffer(0, vertex_buffer.slice(..));
        rp.set_index_buffer(index_buffer.slice(..), IndexFormat::Uint32);
        rp.set_pipeline(&self.pipeline);
        rp.draw_indexed(indices, 0, 0..1);

        drop(rp);

        let cmd_buf = cmd_encoder.finish();
        self.queue.submit(std::iter::once(cmd_buf));
    }
}

pub struct DrawTarget<'a> {
    pub texture: &'a TextureView,
    pub size: Vec2,
}