gear-teeth/src/lib.rs

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

use glam::{Affine3A, FloatExt, Mat4, Quat, UVec2, Vec3, Vec4, vec3, vec4};

use crate::render::mesh::GpuParams;

mod bevel_gear;
mod camera;
mod crown_gear;
mod profile;
mod render;

const BACKGROUND: Vec4 = vec4(0.00, 0.39, 0.52, 1.00);

#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct ResetArgs {
	pub tooth: u32,
	pub cone_halfangle: f32,
}

impl Default for ResetArgs {
	fn default() -> Self {
		Self {
			tooth: 24,
			cone_halfangle: std::f32::consts::FRAC_PI_4,
		}
	}
}

#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct RedrawArgs {
	pub cam_yaw: f32,
	pub cam_pitch: f32,
	pub cam_view_y: f32,
	pub cam_view_angle: f32,
}

impl Default for RedrawArgs {
	fn default() -> Self {
		Self {
			cam_yaw: 3. * std::f32::consts::FRAC_PI_4,
			cam_pitch: std::f32::consts::FRAC_PI_6,
			cam_view_y: 80.,
			cam_view_angle: std::f32::consts::FRAC_PI_2,
		}
	}
}

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>,
	depth: wgpu::Texture,

	mesh_pipeline: render::mesh::Pipeline,

	mesh_cutter: render::mesh::GpuMeshIndexed,

	params_gear: render::mesh::GpuParams,
	params_cutter: render::mesh::GpuParams,

	bevel_gear_desc: bevel_gear::BevelGearDesc,
	bevel_gear_shape: bevel_gear::BevelGearShape,
	crown_gear_desc: crown_gear::CrownGearDesc,
	t: f32,
}

fn srgb_to_linear(color: Vec3) -> Vec3 {
	color.map(|x| {
		if x > 0.04045 {
			((x + 0.055) / 1.055).powf(2.4)
		} else {
			x / 12.92
		}
	})
}

fn srgba_to_linear(color: Vec4) -> Vec4 {
	let alpha = color.w;
	let color = srgb_to_linear(vec3(color.x, color.y, color.z));
	vec4(color.x, color.y, color.z, alpha)
}

fn color_to_wgpu(color: Vec4) -> wgpu::Color {
	wgpu::Color {
		r: color.x.into(),
		g: color.y.into(),
		b: color.z.into(),
		a: color.w.into(),
	}
}

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

		let depth = Self::create_depth_buffer(&device, pixel_size);
		let mesh_pipeline = render::mesh::Pipeline::new(&device);

		let bevel_gear_desc = bevel_gear::BevelGearDesc {
			profile: Default::default(),
			teeth: args.tooth,
			cone_halfangle: args.cone_halfangle,
		};
		let bevel_gear_shape =
			bevel_gear::BevelGearShape::new_blank(&bevel_gear_desc, 16 * args.tooth as usize);

		let crown_gear_desc = bevel_gear_desc.cutter();
		let crown_gear_mesh = crown_gear::make_mesh_simple(
			&crown_gear_desc,
			&crown_gear::CrownGearVisual {
				inner_radius: crown_gear_desc.pitch_radius / 2.,
				outer_radius: crown_gear_desc.pitch_radius
					+ 2. * bevel_gear_desc.cone_halfangle.tan(),
				base_thickness: 2.0,
			},
			(16. * crown_gear_desc.teeth()).floor() as usize,
		)
		.to_mesh()
		.to_standard_mesh()
		.into_indexed();

		let mesh_cutter = render::mesh::GpuMeshIndexed::new(&device, &crown_gear_mesh);

		let params_gear = GpuParams::new(&device, &mesh_pipeline);
		let params_cutter = GpuParams::new(&device, &mesh_pipeline);

		Self::configure_surface(&surface, &device, pixel_size);

		Self {
			device,
			queue,
			surface,
			depth,
			mesh_pipeline,
			mesh_cutter,
			params_gear,
			params_cutter,
			bevel_gear_desc,
			bevel_gear_shape,
			crown_gear_desc,
			t: 0.,
		}
	}

	fn gear_tfm(&self) -> Affine3A {
		let turns = -self.t / self.bevel_gear_desc.teeth as f32;
		Affine3A::from_rotation_translation(Quat::from_rotation_z(2. * PI * turns), Vec3::ZERO)
	}

	fn cutter_tfm(&self) -> Affine3A {
		let skip_teeth = self.crown_gear_desc.pitch_radius.floor();
		let turns = (self.t.fract_gl() + skip_teeth) / self.crown_gear_desc.teeth();
		Affine3A::from_rotation_translation(
			Quat::from_rotation_y(FRAC_PI_2 + self.bevel_gear_desc.cone_halfangle)
				* Quat::from_rotation_z(2. * PI * turns),
			vec3(0., 0., self.bevel_gear_desc.pitch_distance()),
		)
	}

	fn render(&self, output: &wgpu::Texture, args: &RedrawArgs) {
		let y_frac = {
			let size = output.size();
			let w = size.width as f32;
			let h = size.height as f32;
			h / w
		};
		let proj = camera::Camera {
			y_frac,
			y_angle: args.cam_view_angle,
			y_range: args.cam_view_y,
			z_min: -0.5 * args.cam_view_y,
			z_max: 0.5 * args.cam_view_y,
		}
		.clamp_z_min(0.01 * args.cam_view_y)
		.to_mat4();
		let view = Mat4::from_mat3(camera::rotation(-args.cam_yaw, -args.cam_pitch));
		let light = vec3(0.7, 0.5, 1.0).normalize();

		self.mesh_pipeline.set_params(
			&self.queue,
			render::mesh::SceneParams {
				projection: proj * view,
				light,
				_zero: 0.,
			},
		);

		self.params_gear.set(
			&self.queue,
			render::mesh::ObjectParams {
				transform: self.gear_tfm().into(),
				color: srgba_to_linear(vec4(0.2, 0.8, 0.1, 1.0)),
			},
		);
		self.params_cutter.set(
			&self.queue,
			render::mesh::ObjectParams {
				transform: self.cutter_tfm().into(),
				color: srgba_to_linear(vec4(0.8, 0.2, 0.1, 1.0)),
			},
		);

		let bevel_gear_mesh = bevel_gear::make_mesh(
			&self.bevel_gear_desc,
			&bevel_gear::BevelGearVisual {
				pad_in: 4.,
				pad_out: 0.,
			},
			&self.bevel_gear_shape,
		)
		.to_mesh()
		.to_standard_mesh()
		.into_indexed();
		let mesh_gear = render::mesh::GpuMeshIndexed::new(&self.device, &bevel_gear_mesh);

		let view = output.create_view(&wgpu::TextureViewDescriptor::default());
		let depth_view = self
			.depth
			.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(color_to_wgpu(srgba_to_linear(BACKGROUND))),
					store: wgpu::StoreOp::Store,
				},
			})],
			depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
				view: &depth_view,
				depth_ops: Some(wgpu::Operations {
					load: wgpu::LoadOp::Clear(1.),
					store: wgpu::StoreOp::Discard,
				}),
				stencil_ops: None,
			}),
			..Default::default()
		});
		self.mesh_cutter
			.render(&mut pass, &self.mesh_pipeline, &self.params_cutter);
		mesh_gear.render(&mut pass, &self.mesh_pipeline, &self.params_gear);
		drop(pass);
		self.queue.submit(std::iter::once(encoder.finish()));
	}

	fn create_depth_buffer(device: &wgpu::Device, pixel_size: UVec2) -> wgpu::Texture {
		device.create_texture(&wgpu::TextureDescriptor {
			label: Some("depth buffer"),
			size: wgpu::Extent3d {
				width: pixel_size.x,
				height: pixel_size.y,
				depth_or_array_layers: 1,
			},
			mip_level_count: 1,
			sample_count: 1,
			dimension: wgpu::TextureDimension::D2,
			format: render::DEPTH_FORMAT,
			usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
			view_formats: &[],
		})
	}

	fn configure_surface(surface: &wgpu::Surface, device: &wgpu::Device, pixel_size: UVec2) {
		surface.configure(
			device,
			&wgpu::SurfaceConfiguration {
				usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST,
				format: render::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,
			},
		);
	}

	/// Configures the renderer for a given target size.
	pub fn configure(&mut self, pixel_size: UVec2) {
		Self::configure_surface(&self.surface, &self.device, pixel_size);
		self.depth = Self::create_depth_buffer(&self.device, pixel_size);
	}

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

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,
	})
}