refraction/src/bin/wireframe/meshes.rs

use std::mem;

use bytemuck::{bytes_of, cast_slice, Pod, Zeroable};
use glam::{Vec3, Vec4};
use wgpu::util::DeviceExt as _;

#[repr(C)]
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
struct Vertex {
	pub position: [f32; 3],
	pub normal: [f32; 3],
}

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct PushConsts {
	pub color: [f32; 4],
}

#[derive(Copy, Clone)]
pub struct Attrs {
	pub color: Vec4,
}

impl Attrs {
	fn consts(&self) -> PushConsts {
		PushConsts {
			color: self.color.to_array(),
		}
	}
}

pub struct Mesh {
	consts: PushConsts,
	npoints: u32,
	buf: wgpu::Buffer,
}

impl Mesh {
	pub fn new_list(device: &wgpu::Device, attrs: Attrs, tris: Vec<[Vec3; 3]>) -> Self {
		let data: Vec<Vertex> = tris
			.into_iter()
			.flat_map(|[a, b, c]| {
				let n = (b - a).cross(c - a).normalize();
				[
					Vertex {
						position: a.into(),
						normal: n.into(),
					},
					Vertex {
						position: b.into(),
						normal: n.into(),
					},
					Vertex {
						position: c.into(),
						normal: n.into(),
					},
				]
			})
			.collect();
		let buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: Some("Mesh Vertex Buffer"),
			contents: cast_slice(&data),
			usage: wgpu::BufferUsages::VERTEX,
		});
		Mesh {
			consts: attrs.consts(),
			npoints: data.len() as u32,
			buf,
		}
	}
}

pub struct Renderer {
	pipeline: wgpu::RenderPipeline,
}

static SHADER: &'static str = include_str!("mesh.wgsl");

impl Renderer {
	pub fn new(
		device: &wgpu::Device,
		cam_layout: &wgpu::BindGroupLayout,
		target_format: wgpu::TextureFormat,
		depth_stencil: Option<wgpu::DepthStencilState>,
		multisample: wgpu::MultisampleState,
	) -> Renderer {
		let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
			label: Some("Mesh Shader"),
			source: wgpu::ShaderSource::Wgsl(SHADER.into()),
		});

		let consts_range = wgpu::PushConstantRange {
			stages: wgpu::ShaderStages::VERTEX,
			range: 0..mem::size_of::<PushConsts>() as u32,
		};

		let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
			label: Some("Mesh RenderPipelineLayout"),
			bind_group_layouts: &[cam_layout],
			push_constant_ranges: &[consts_range],
		});

		let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
			label: Some("Mesh RenderPipeline"),
			layout: Some(&layout),
			vertex: wgpu::VertexState {
				module: &shader,
				entry_point: "vs_main",
				buffers: &[wgpu::VertexBufferLayout {
					array_stride: mem::size_of::<Vertex>() as wgpu::BufferAddress,
					step_mode: wgpu::VertexStepMode::Vertex,
					attributes: &[
						wgpu::VertexAttribute {
							offset: mem::offset_of!(Vertex, position) as u64,
							shader_location: 0,
							format: wgpu::VertexFormat::Float32x3,
						},
						wgpu::VertexAttribute {
							offset: mem::offset_of!(Vertex, normal) as u64,
							shader_location: 1,
							format: wgpu::VertexFormat::Float32x3,
						},
					],
				}],
				compilation_options: Default::default(),
			},
			fragment: Some(wgpu::FragmentState {
				module: &shader,
				entry_point: "fs_main",
				targets: &[Some(wgpu::ColorTargetState {
					format: target_format,
					blend: Some(wgpu::BlendState {
						color: wgpu::BlendComponent::OVER,
						alpha: wgpu::BlendComponent::OVER,
					}),
					write_mask: wgpu::ColorWrites::ALL,
				})],
				compilation_options: Default::default(),
			}),
			primitive: wgpu::PrimitiveState {
				topology: wgpu::PrimitiveTopology::TriangleList,
				..Default::default()
			},
			depth_stencil,
			multisample,
			multiview: None,
			cache: None,
		});

		Renderer { pipeline }
	}

	pub fn render<'a>(
		&self,
		pass: &mut wgpu::RenderPass,
		cam_bind: &wgpu::BindGroup,
		meshes: impl Iterator<Item = &'a Mesh>,
	) {
		pass.set_pipeline(&self.pipeline);
		pass.set_bind_group(0, cam_bind, &[]);
		for mesh in meshes {
			pass.set_push_constants(wgpu::ShaderStages::VERTEX, 0, bytes_of(&mesh.consts));
			pass.set_vertex_buffer(0, mesh.buf.slice(..));
			pass.draw(0..mesh.npoints, 0..1);
		}
	}
}