minitracing/src/bin/envmap/main.rs

use std::error::Error;

use glam::{mat3, vec2, vec3, Mat3, Vec3};
use image::buffer::ConvertBuffer;
use raytracing3::perlin::{self, Pipeline, Vertex};

const EXTENT: u32 = 1024;

fn make_viewport(m: Mat3) -> [Vertex; 4] {
	let screen_coord = [vec2(-1., -1.), vec2(1., -1.), vec2(-1., 1.), vec2(1., 1.)];
	let world_coord = screen_coord.map(|s| m * vec3(s.x, s.y, 1.));
	[0, 1, 2, 3].map(|k| Vertex {
		dir: world_coord[k],
		screen: screen_coord[k],
	})
}

fn main() {
	let (device, queue) = pollster::block_on(init_gpu()).unwrap();

	let format = wgpu::TextureFormat::Rgba8UnormSrgb;
	let mut noiser = Pipeline::new(&device, format);
	noiser.set_params(
		&queue,
		perlin::Params {
			origin: vec3(0., 0., -30.),
			radius: 30.,
		},
	);
	let faces = [
		mat3(-Vec3::Z, -Vec3::Y, Vec3::X),
		mat3(Vec3::Z, -Vec3::Y, -Vec3::X),
		mat3(Vec3::X, -Vec3::Z, -Vec3::Y),
		mat3(Vec3::X, Vec3::Z, Vec3::Y),
		mat3(Vec3::X, -Vec3::Y, Vec3::Z),
		mat3(-Vec3::X, -Vec3::Y, -Vec3::Z),
	];
	let output = device.create_texture(&wgpu::TextureDescriptor {
		label: None,
		size: wgpu::Extent3d {
			width: EXTENT,
			height: EXTENT,
			depth_or_array_layers: 6,
		},
		mip_level_count: 1,
		sample_count: 1,
		dimension: wgpu::TextureDimension::D2,
		format,
		usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
		view_formats: &[],
	});
	let buffers = faces.map(|_| {
		device.create_buffer(&wgpu::BufferDescriptor {
			label: None,
			size: (4 * EXTENT * EXTENT) as u64,
			usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
			mapped_at_creation: false,
		})
	});
	for (face, m) in faces.iter().enumerate() {
		let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
		let buffer = &buffers[face];
		let view = output.create_view(&wgpu::TextureViewDescriptor {
			dimension: Some(wgpu::TextureViewDimension::D2),
			base_array_layer: face as u32,
			..Default::default()
		});
		let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			label: None,
			color_attachments: &[Some(wgpu::RenderPassColorAttachment {
				view: &view,
				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,
		});
		noiser.set_view(&queue, &make_viewport(*m));
		noiser.render(&mut render_pass);
		drop(render_pass);
		encoder.copy_texture_to_buffer(
			wgpu::ImageCopyTexture {
				texture: &output,
				mip_level: 0,
				origin: wgpu::Origin3d {
					x: 0,
					y: 0,
					z: face as u32,
				},
				aspect: wgpu::TextureAspect::All,
			},
			wgpu::ImageCopyBuffer {
				buffer: &buffer,
				layout: wgpu::ImageDataLayout {
					offset: 0,
					bytes_per_row: Some(4 * EXTENT),
					rows_per_image: Some(EXTENT),
				},
			},
			wgpu::Extent3d {
				width: EXTENT,
				height: EXTENT,
				depth_or_array_layers: 1,
			},
		);
		queue.submit([encoder.finish()]);
	}
	for buf in &buffers {
		buf.slice(..).map_async(wgpu::MapMode::Read, |res| res.unwrap());
	}
	device.poll(wgpu::Maintain::Wait);
	std::thread::scope(|s| {
		for (face, buf) in buffers.iter().enumerate() {
			s.spawn(move || {
				let img =
					image::RgbaImage::from_raw(EXTENT, EXTENT, buf.slice(..).get_mapped_range().to_vec()).unwrap();
				let img: image::RgbImage = img.convert();
				img.save(format!("textures/env{face}.jpeg")).unwrap();
			});
		}
	})
}

async fn init_gpu() -> Result<(wgpu::Device, wgpu::Queue), Box<dyn Error>> {
	let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
		backends: wgpu::Backends::PRIMARY,
		..Default::default()
	});
	let adapter = instance
		.request_adapter(&wgpu::RequestAdapterOptions {
			power_preference: wgpu::PowerPreference::default(),
			compatible_surface: None,
			force_fallback_adapter: false,
		})
		.await
		.unwrap();
	let (device, queue) = adapter
		.request_device(
			&wgpu::DeviceDescriptor {
				label: None,
				required_features: wgpu::Features::empty(),
				required_limits: wgpu::Limits::default(),
				memory_hints: Default::default(),
			},
			None,
		)
		.await
		.unwrap();
	Ok((device, queue))
}