qt-tracing/src/lib.rs

#![feature(gen_blocks)]

use std::{convert::identity, error::Error, f32::consts::PI};

use glam::{Mat4, UVec2, Vec3, vec3};

use crate::{
	camera::OrbitalCamera,
	render::lines::{LookParams, Mesh, Pipeline, Vertex},
	trace::{Scene, Source, Sphere},
};

mod camera;
mod ray;
mod render;
mod trace;

const OUTPUT_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Bgra8UnormSrgb;

#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(C)]
pub struct SphericalPosition {
	pub yaw: f32,
	pub pitch: f32,
	pub distance: f32,
}

#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct RedrawArgs {
	pub camera_position: SphericalPosition,
	pub light_position: SphericalPosition,
	pub light_radius: f32,
	pub light_spread: f32,
}

pub struct Gpu {
	device: wgpu::Device,
	queue: wgpu::Queue,
	surface: wgpu::Surface<'static>,
}

pub struct Core {
	device: wgpu::Device,
	queue: wgpu::Queue,
	surface: wgpu::Surface<'static>,

	pipeline: Pipeline,
	tripod: Mesh,
}

pub fn new_tripod(device: &wgpu::Device) -> Mesh {
	Mesh::new(
		device,
		&[
			Vertex::new(vec3(0., 0., 0.), vec3(1., 0., 0.)),
			Vertex::new(vec3(1., 0., 0.), vec3(1., 0., 0.)),
			Vertex::new(vec3(0., 0., 0.), vec3(0., 1., 0.)),
			Vertex::new(vec3(0., 1., 0.), vec3(0., 1., 0.)),
			Vertex::new(vec3(0., 0., 0.), vec3(0., 0., 1.)),
			Vertex::new(vec3(0., 0., 1.), vec3(0., 0., 1.)),
		],
	)
}

fn loop_list<T: Clone>(iter: impl IntoIterator<Item = T>) -> impl Iterator<Item = T> {
	loop_list_ex(iter, identity, identity)
}

fn loop_list_ex<T: Clone, U>(
	iter: impl IntoIterator<Item = T>,
	mut fa: impl FnMut(T) -> U,
	mut fb: impl FnMut(T) -> U,
) -> impl Iterator<Item = U> {
	gen move {
		let mut iter = iter.into_iter();
		let Some(first) = iter.next() else { return };
		yield fa(first.clone());
		for item in iter {
			yield fb(item.clone());
			yield fa(item);
		}
		yield fb(first);
	}
}

impl Core {
	pub fn new(gpu: Gpu) -> Self {
		let Gpu {
			device,
			queue,
			surface,
		} = gpu;

		let pipeline = Pipeline::new(&device, OUTPUT_FORMAT);
		let tripod = new_tripod(&device);
		queue.submit([]); // flush buffer updates

		Self {
			device,
			queue,
			surface,
			pipeline,
			tripod,
		}
	}

	fn render(&self, output: &wgpu::Texture, args: &RedrawArgs) {
		let camera = OrbitalCamera {
			position_yaw: args.camera_position.yaw,
			position_pitch: args.camera_position.pitch,
			distance: args.camera_position.distance,
		};
		let aspect = {
			let size = output.size();
			let w = size.width as f32;
			let h = size.height as f32;
			w / h
		};
		let perspective = Mat4::perspective_lh(PI / 3., aspect, 1e-2, 1e2);
		self.pipeline.set_look(
			&self.queue,
			LookParams {
				m: perspective * camera.transform(),
			},
		);
		self.queue.submit([]); // flush buffer updates

		let view = output.create_view(&wgpu::TextureViewDescriptor::default());
		let mut encoder = self
			.device
			.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
		let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			color_attachments: &[Some(wgpu::RenderPassColorAttachment {
				view: &view,
				depth_slice: None,
				resolve_target: None,
				ops: wgpu::Operations {
					load: wgpu::LoadOp::Clear(wgpu::Color {
						r: 0.1,
						g: 0.2,
						b: 0.8,
						a: 1.0,
					}),
					store: wgpu::StoreOp::Store,
				},
			})],
			depth_stencil_attachment: None,
			..Default::default()
		});
		self.pipeline.render(&mut pass, [&self.tripod]);

		let source = Source {
			position_yaw: args.light_position.yaw,
			position_pitch: args.light_position.pitch,
			distance: args.light_position.distance,
			radius: args.light_radius,
			spread: args.light_spread,
		};

		let contour: Vec<Vertex> = loop_list(source.contour(17))
			.map(|pos| Vertex {
				pos,
				color: vec3(1., 1., 1.),
			})
			.collect();
		self.pipeline
			.render(&mut pass, [&Mesh::new(&self.device, &contour)]);

		const BASE_R: f32 = 2.;
		const BASE_POS: Vec3 = vec3(0., 0., -BASE_R);
		const BASE: Sphere = Sphere {
			position: vec3(0., 0., -BASE_R),
			radius: BASE_R,
		};
		fn sphere(pos: Vec3) -> Sphere {
			Sphere {
				position: pos,
				radius: BASE_POS.distance(pos) - BASE_R,
			}
		}
		let scene = Scene {
			objects: vec![
				BASE,
				sphere(vec3(0., 0., 0.1)),
				sphere(vec3(0.3, 0., 0.1)),
				sphere(vec3(0.1, 0.3, 0.1)),
			],
		};

		let mut prng = rand_pcg::Pcg64::new(42, 0);
		let rays: Vec<Vertex> = (0..10000)
			.flat_map(|_| {
				let ray = source.make_ray(&mut prng);
				if let Some(ray) = scene.trace_ray(ray) {
					[
						Vertex {
							pos: ray.base - 0.02 * ray.dir,
							color: vec3(1., 1., 1.),
						},
						Vertex {
							pos: ray.base,
							color: vec3(0., 1., 0.),
						},
					]
				} else {
					[
						Vertex {
							pos: ray.base,
							color: vec3(1., 1., 1.),
						},
						Vertex {
							pos: ray.base + 0.1 * ray.dir,
							color: vec3(1., 0., 0.),
						},
					]
				}
			})
			.collect();
		self.pipeline
			.render(&mut pass, [&Mesh::new(&self.device, &rays)]);

		drop(pass);
		self.queue.submit(std::iter::once(encoder.finish()));
	}

	/// Configures the renderer for a given target size.
	pub fn configure(&mut self, pixel_size: UVec2) {
		self.surface.configure(
			&self.device,
			&wgpu::SurfaceConfiguration {
				usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST,
				format: OUTPUT_FORMAT,
				width: pixel_size.x,
				height: pixel_size.y,
				present_mode: wgpu::PresentMode::Fifo,
				alpha_mode: wgpu::CompositeAlphaMode::Auto,
				view_formats: vec![],
				desired_maximum_frame_latency: 2,
			},
		);
	}

	/// Redraws the entire surface.
	///
	/// [`Self::configure`] must be called at least once before this.
	pub fn redraw(&mut self, args: &RedrawArgs) {
		let output = self.surface.get_current_texture().unwrap();
		self.render(&output.texture, args);
		output.present();
	}
}

pub async fn init_gpu_inner<E: Error + 'static>(
	make_surface: impl FnOnce(&wgpu::Instance) -> Result<wgpu::Surface<'static>, E>,
) -> Result<Gpu, Box<dyn Error>> {
	let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
		backends: wgpu::Backends::PRIMARY,
		..Default::default()
	});
	let surface = make_surface(&instance)?;
	let adapter = instance
		.request_adapter(&wgpu::RequestAdapterOptions {
			power_preference: wgpu::PowerPreference::default(),
			compatible_surface: Some(&surface),
			force_fallback_adapter: false,
		})
		.await
		.unwrap();
	let (device, queue) = adapter
		.request_device(&wgpu::DeviceDescriptor::default())
		.await
		.unwrap();
	Ok(Gpu {
		device,
		queue,
		surface,
	})
}