Extract image savers to function

src/bin/rec.rs

@@ -1,8 +1,9 @@
 use std::env::args;
 use std::error::Error;
-use std::fs;
-use std::path::PathBuf;
+use std::io::{stdout, Write};
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::{fs, io};
 
 use glam::{uvec2, UVec2};
 use image::buffer::ConvertBuffer;
@@ -11,25 +12,61 @@ use raytracing3::scene::{load_envmap, Renderer, SceneParams};
 
 const SIZE: UVec2 = uvec2(1920, 1080);
 const FRAME_RATE: u32 = 60;
-const DURATION_SECONDS: u32 = 120;
-const N_FRAMES: u32 = DURATION_SECONDS * FRAME_RATE;
 const N_SPHERES: u32 = 100;
 const RAYS_PER_PIXEL: u32 = 1024;
 
 fn main() {
-	let args: Vec<_> = args().collect();
-	let [_, path] = args.as_slice() else {
+	let path: PathBuf;
+	let start_frame;
+	let stop_frame;
+	match args().collect::<Vec<_>>().as_slice() {
+		[_, apath] => {
+			path = apath.into();
+			start_frame = 0;
+			stop_frame = 5;
+		}
+		[_, apath, aframes] => {
+			path = apath.into();
+			start_frame = 0;
+			stop_frame = aframes.parse().unwrap();
+		}
+		[_, apath, astart_frame, astop_frame] => {
+			path = apath.into();
+			start_frame = astart_frame.parse().unwrap();
+			stop_frame = astop_frame.parse().unwrap();
+		}
+		_ => {
 			panic!("invalid arguments");
-	};
-	let path: PathBuf = path.into();
-	fs::create_dir(&path).expect("failed to create the output directory");
+		}
+	}
+	println!(
+		"Rendering frames {start_frame} to {stop_frame} ({} total)",
+		stop_frame - start_frame
+	);
+	match fs::create_dir(&path) {
+		Ok(_) => (),
+		Err(err) if err.kind() == io::ErrorKind::AlreadyExists => (),
+		Err(err) => panic!("failed to create the output directory: {err}"),
+	}
 
+	let n_threads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(2);
+	let (img_sender, img_receiver) = async_channel::bounded::<Frame>(2 * n_threads);
 	std::thread::scope(|s| {
-		let (img_sender, img_receiver) = async_channel::bounded::<(u32, Arc<wgpu::Buffer>)>(50);
-		for _ in 0..16 {
-			let img_receiver = img_receiver.clone();
-			let path = &path;
-			s.spawn(move || {
+		for _ in 0..n_threads {
+			start_saver(s, &path, &img_receiver);
+		}
+		do_work(img_sender, start_frame, stop_frame);
+	});
+}
+
+type Frame = (u32, Arc<wgpu::Buffer>);
+
+fn start_saver<'a>(
+	scope: &'a std::thread::Scope<'a, '_>,
+	path: &'a Path,
+	img_receiver: &'a async_channel::Receiver<Frame>,
+) {
+	scope.spawn(move || {
 		while let Ok((frame, buffer)) = img_receiver.recv_blocking() {
 			let img = image::RgbaImage::from_raw(SIZE.x, SIZE.y, buffer.slice(..).get_mapped_range().to_vec())
 				.expect("read failure!");
@@ -38,12 +75,10 @@ fn main() {
 				.expect("save failure!");
 		}
 	});
-		}
-		do_work(img_sender);
-	});
 }
 
-fn do_work(img_sender: async_channel::Sender<(u32, Arc<wgpu::Buffer>)>) {
+fn do_work(img_sender: async_channel::Sender<Frame>, start_frame: u32, stop_frame: u32) {
+	println!("Loading...");
 	let img_sender = Arc::new(img_sender);
 
 	let (device, queue) = pollster::block_on(init_gpu()).unwrap();
@@ -63,7 +98,11 @@ fn do_work(img_sender: async_channel::Sender<(u32, Arc<wgpu::Buffer>)>) {
 	let renderer = Renderer::new(&device, envmap);
 	let presenter = Presenter::new(&device, output_format);
 
-	for frame in 0..N_FRAMES {
+	println!("Rendering...");
+	for frame in start_frame..stop_frame {
+		if frame % FRAME_RATE == 0 {
+			println!("{frame}");
+		}
 		let output = device.create_texture(&wgpu::TextureDescriptor {
 			label: None,
 			size: texsize,
@@ -136,8 +175,12 @@ fn do_work(img_sender: async_channel::Sender<(u32, Arc<wgpu::Buffer>)>) {
 				res.unwrap();
 				img_sender.send_blocking((frame, buffer)).unwrap();
 			});
+		print!(".");
+		stdout().flush().unwrap();
 	}
+	println!("{stop_frame}");
 	device.poll(wgpu::Maintain::Wait);
+	println!("Done!");
 }
 
 async fn init_gpu() -> Result<(wgpu::Device, wgpu::Queue), Box<dyn Error>> {