minitracing/src/trace.rs

use std::mem;

use bytemuck::{bytes_of, cast_slice, Pod, Zeroable};
use glam::{vec2, Mat3, Vec2, Vec3};
use wgpu::util::DeviceExt;

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct Params {
	pub max_reflections: u32,
	pub min_strength: f32,
	pub sphere_count: u32,
	pub seed: u32,
}

#[derive(Debug, Clone, Copy)]
pub struct Viewport {
	pub corner: Vec3,
}

#[derive(Debug, Clone, Copy)]
pub struct Aperture {
	pub radius: f32,
	pub focal_distance: f32, // from 0 (exclusive) to +∞ (inclusive)
	pub glare_strength: f32,
	pub glare_radius: f32, // at distance 1
}

#[derive(Debug, Clone, Copy)]
pub struct CameraLocation {
	pub eye: Vec3,
	pub view: Mat3,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct ParamsData {
	params: Params,
	camera: CameraData,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct CameraData {
	u: Vec3,
	width: f32,
	v: Vec3,
	height: f32,
	w: Vec3,
	aperture: f32,
	eye: Vec3,
	antifocal: f32,
	glare_strength: f32,
	glare_radius: f32,
}

impl From<(Viewport, CameraLocation, Aperture)> for CameraData {
	fn from(value: (Viewport, CameraLocation, Aperture)) -> Self {
		CameraData {
			u: value.1.view.x_axis,
			v: value.1.view.y_axis,
			w: value.1.view.z_axis,
			eye: value.1.eye,
			width: value.0.corner.x / value.0.corner.z,
			height: value.0.corner.y / value.0.corner.z,
			aperture: value.2.radius,
			antifocal: 1. / value.2.focal_distance,
			glare_strength: value.2.glare_strength,
			glare_radius: value.2.glare_radius,
		}
	}
}

#[derive(Debug, Clone, Copy)]
pub struct Sphere {
	pub center: Vec3,
	pub radius: f32,
	pub emit_color: Vec3,
	pub reflect_color: Vec3,
	pub glossiness: f32,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct Vertex {
	pub screen: Vec2,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct SphereData {
	center: Vec3,
	radius: f32,
	emit_color: Vec3,
	pad1: f32,
	reflect_color: Vec3,
	glossiness: f32,
}

impl From<&Sphere> for SphereData {
	fn from(s: &Sphere) -> Self {
		SphereData {
			center: s.center,
			radius: s.radius,
			emit_color: s.emit_color,
			pad1: 0.0,
			reflect_color: s.reflect_color,
			glossiness: s.glossiness,
		}
	}
}

pub struct Tracer {
	view_buf: wgpu::Buffer,
	pipeline: wgpu::RenderPipeline,
}

pub struct TracerData {
	bindings: wgpu::BindGroup,
}

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

impl Tracer {
	pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
		let screen_coord = [vec2(-1., -1.), vec2(1., -1.), vec2(-1., 1.), vec2(1., 1.)];
		let view_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: None,
			contents: cast_slice(&screen_coord),
			usage: wgpu::BufferUsages::VERTEX,
		});
		let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
			label: None,
			source: wgpu::ShaderSource::Wgsl(SHADER.into()),
		});
		let spheres_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
			label: None,
			entries: &[wgpu::BindGroupLayoutEntry {
				binding: 1,
				visibility: wgpu::ShaderStages::FRAGMENT,
				ty: wgpu::BindingType::Buffer {
					ty: wgpu::BufferBindingType::Storage { read_only: true },
					has_dynamic_offset: false,
					min_binding_size: None,
				},
				count: None,
			}],
		});
		let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
			label: None,
			bind_group_layouts: &[&spheres_bgl],
			push_constant_ranges: &[wgpu::PushConstantRange {
				stages: wgpu::ShaderStages::FRAGMENT,
				range: 0..mem::size_of::<ParamsData>() as u32,
			}],
		});
		let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
			label: None,
			layout: Some(&pipeline_layout),
			vertex: wgpu::VertexState {
				module: &shader,
				entry_point: None,
				compilation_options: wgpu::PipelineCompilationOptions::default(),
				buffers: &[wgpu::VertexBufferLayout {
					array_stride: size_of::<Vertex>() as u64,
					step_mode: wgpu::VertexStepMode::Vertex,
					attributes: &[wgpu::VertexAttribute {
						shader_location: 0,
						offset: 0,
						format: wgpu::VertexFormat::Float32x2,
					}],
				}],
			},
			primitive: wgpu::PrimitiveState {
				topology: wgpu::PrimitiveTopology::TriangleStrip,
				..Default::default()
			},
			depth_stencil: None,
			multisample: wgpu::MultisampleState {
				count: 1,
				mask: !0,
				alpha_to_coverage_enabled: false,
			},
			fragment: Some(wgpu::FragmentState {
				module: &shader,
				entry_point: None,
				compilation_options: wgpu::PipelineCompilationOptions::default(),
				targets: &[Some(wgpu::ColorTargetState {
					format,
					blend: Some(wgpu::BlendState {
						color: wgpu::BlendComponent {
							src_factor: wgpu::BlendFactor::One,
							dst_factor: wgpu::BlendFactor::One,
							operation: wgpu::BlendOperation::Add,
						},
						alpha: wgpu::BlendComponent {
							src_factor: wgpu::BlendFactor::One,
							dst_factor: wgpu::BlendFactor::One,
							operation: wgpu::BlendOperation::Add,
						},
					}),
					write_mask: wgpu::ColorWrites::ALL,
				})],
			}),
			multiview: None,
			cache: None,
		});
		Self { view_buf, pipeline }
	}

	pub fn prepare<'encoder>(
		&self,
		encoder: &'encoder mut wgpu::CommandEncoder,
		target: &wgpu::TextureView,
	) -> wgpu::RenderPass<'encoder> {
		encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			label: None,
			color_attachments: &[Some(wgpu::RenderPassColorAttachment {
				view: &target,
				resolve_target: None,
				ops: wgpu::Operations {
					load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
					store: wgpu::StoreOp::Store,
				},
			})],
			depth_stencil_attachment: None,
			occlusion_query_set: None,
			timestamp_writes: None,
		})
	}

	pub fn render(
		&self,
		pass: &mut wgpu::RenderPass,
		data: &TracerData,
		params: Params,
		viewport: Viewport,
		aperture: Aperture,
		camera: CameraLocation,
	) {
		pass.set_pipeline(&self.pipeline);
		pass.set_push_constants(
			wgpu::ShaderStages::FRAGMENT,
			0,
			bytes_of(&ParamsData {
				params,
				camera: (viewport, camera, aperture).into(),
			}),
		);
		pass.set_vertex_buffer(0, self.view_buf.slice(..));
		pass.set_bind_group(0, &data.bindings, &[]);
		pass.draw(0..4, 0..1);
	}
}

impl TracerData {
	pub fn new(device: &wgpu::Device, tracer: &Tracer, spheres: &[Sphere]) -> Self {
		let spheres: Vec<_> = spheres.iter().map(SphereData::from).collect();
		let spheres_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: None,
			contents: cast_slice(&spheres),
			usage: wgpu::BufferUsages::STORAGE,
		});
		let bindings = device.create_bind_group(&wgpu::BindGroupDescriptor {
			label: None,
			layout: &tracer.pipeline.get_bind_group_layout(0),
			entries: &[wgpu::BindGroupEntry {
				binding: 1,
				resource: spheres_buf.as_entire_binding(),
			}],
		});
		Self { bindings }
	}
}