qt-tracing/src/lib.rs

400 lines
10 KiB
Rust

#![feature(gen_blocks)]
use std::{convert::identity, error::Error, f32::consts::PI};
use glam::{Mat4, UVec2, Vec2, Vec3, vec3};
use rand_distr::Distribution as _;
use crate::{
camera::OrbitalCamera,
ray::Ray,
render::lines::{LookParams, Mesh, Pipeline, Vertex},
trace::{Hit, Lambertian, Reflector, 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 accum_sigma: f32,
pub accum_scale: f32,
pub reflections: u32,
pub show_axes: bool,
pub show_shapes: bool,
pub show_hit_emission: bool,
pub show_miss_emission: bool,
pub show_direct_hit: bool,
pub show_indirect_hit: bool,
pub show_light: bool,
}
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 OrbitalCamera {
fn make_ray(&self, rng: &mut impl rand::Rng) -> Ray {
let off: f32 = rand_distr::StandardUniform.sample(rng);
let side: Vec2 = rand_distr::UnitCircle.sample(rng).into();
let m = self.transform().inverse();
let fwd = 1. - 0.1 * off;
let side_scale = (1. - fwd.powi(2)).sqrt();
let dir = Vec3::from((side_scale * side, fwd));
Ray {
base: self.position(),
dir: m.transform_vector3(dir),
}
}
}
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::BLACK),
store: wgpu::StoreOp::Store,
},
})],
depth_stencil_attachment: None,
..Default::default()
});
if args.show_axes {
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,
};
if args.show_shapes {
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 source_rays: Vec<Ray> = (0..10240).map(|_| source.make_ray(&mut prng)).collect();
let camera_rays: Vec<Ray> = (0..10240).map(|_| camera.make_ray(&mut prng)).collect();
let mut source_ray_display: Vec<Vertex> = Vec::with_capacity(source_rays.len());
let mut hits: Vec<Hit> = Vec::with_capacity(source_rays.len());
for ray in source_rays {
if let Some(hit) = scene.trace_ray(ray) {
hits.push(hit);
if args.show_hit_emission {
source_ray_display.extend([
Vertex {
pos: ray.base,
color: vec3(1., 1., 1.),
},
Vertex {
pos: ray.base + 0.1 * ray.dir,
color: vec3(0., 1., 0.),
},
]);
}
if args.show_direct_hit {
source_ray_display.extend([
Vertex {
pos: hit.incident.base - 0.02 * hit.incident.dir,
color: vec3(0., 0., 1.),
},
Vertex {
pos: hit.incident.base,
color: vec3(1., 1., 1.),
},
]);
}
} else {
if args.show_miss_emission {
source_ray_display.extend([
Vertex {
pos: ray.base,
color: vec3(1., 1., 1.),
},
Vertex {
pos: ray.base + 0.1 * ray.dir,
color: vec3(1., 0., 0.),
},
]);
}
}
}
if args.reflections > 0 {
let mut hits1 = hits.clone();
for _ in 0..args.reflections {
let mut hits2: Vec<Hit> = Vec::with_capacity(hits1.len());
for hit in &hits1 {
let reflector = Lambertian;
let reflected = reflector.reflect(&mut prng, hit.normal, hit.incident.dir);
let ray = Ray::new(hit.incident.base, reflected);
let Some(hit2) = scene.trace_ray(ray) else {
continue;
};
hits2.push(hit2);
if args.show_indirect_hit {
source_ray_display.extend([
Vertex {
pos: hit2.incident.base - 0.02 * hit2.incident.dir,
color: vec3(1., 0., 1.),
},
Vertex {
pos: hit2.incident.base,
color: vec3(1., 1., 1.),
},
]);
}
}
hits.extend(&hits2);
hits1 = hits2;
}
}
let mut camera_ray_display: Vec<Vertex> = Vec::with_capacity(camera_rays.len());
if args.show_light {
let sigma2 = args.accum_sigma.powi(2);
let accum_normalizator = (2. * PI * sigma2).sqrt().recip();
for ray in camera_rays {
let Some(hit) = scene.trace_ray(ray) else {
continue;
};
let mut total_cd = 0.0f32;
for light_hit in &hits {
let d2 = hit.incident.base.distance_squared(light_hit.incident.base);
if d2 > 9. * sigma2 {
continue;
}
assert!(hit.normal.is_normalized());
assert!(hit.incident.dir.is_normalized());
let reflector = Lambertian;
let in_lm = 1.0;
let out_cd = in_lm * reflector.brdf(hit.normal, hit.incident.dir, -ray.dir);
let weight = accum_normalizator * (-0.5 * d2 / sigma2).exp();
total_cd += weight * out_cd;
}
let brightness = 3. * (1. - (1. + total_cd * args.accum_scale).recip());
let r = args.accum_sigma;
let color = vec3(brightness, brightness - 1., brightness - 2.)
.clamp(Vec3::splat(0.), Vec3::splat(1.));
let vertex = |off: Vec3| Vertex {
pos: hit.incident.base + r * off,
color,
};
camera_ray_display.extend([
vertex(-Vec3::X),
vertex(Vec3::X),
vertex(-Vec3::Y),
vertex(Vec3::Y),
vertex(-Vec3::Z),
vertex(Vec3::Z),
]);
}
}
if !source_ray_display.is_empty() {
self.pipeline
.render(&mut pass, [&Mesh::new(&self.device, &source_ray_display)]);
}
if !camera_ray_display.is_empty() {
self.pipeline
.render(&mut pass, [&Mesh::new(&self.device, &camera_ray_display)]);
}
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,
})
}