refraction/src/mathx.rs

use glam::{FloatExt, Mat2, Mat3, Vec2, Vec3};

mod bounds {
	pub trait Pair<T> {}

	impl<T> Pair<T> for (T, T) {}
}

pub trait FloatExt2<T>: bounds::Pair<T> {
	fn lerp(self, t: T) -> T;
	fn inverse_lerp(self, y: T) -> T;
}

impl<F: FloatExt> FloatExt2<F> for (F, F) {
	fn lerp(self, t: F) -> F {
		F::lerp(self.0, self.1, t)
	}
	fn inverse_lerp(self, y: F) -> F {
		F::inverse_lerp(self.0, self.1, y)
	}
}

pub trait MatExt {
	fn orthonormalize(&self) -> Self;
}

impl MatExt for Mat2 {
	fn orthonormalize(&self) -> Self {
		let fx = self.x_axis.normalize();
		let fy = (self.y_axis - self.y_axis.project_onto_normalized(fx)).normalize();
		Self::from_cols(fx, fy)
	}
}

impl MatExt for Mat3 {
	fn orthonormalize(&self) -> Self {
		let fx = self.x_axis.normalize();
		let fy = (self.y_axis - self.y_axis.project_onto_normalized(fx)).normalize();
		let fz = (self.z_axis
			- self.z_axis.project_onto_normalized(fx)
			- self.z_axis.project_onto_normalized(fy))
		.normalize();
		Self::from_cols(fx, fy, fz)
	}
}

/// Represents a 2×2 matrix decomposed as O^T D O, where O is orthogonal and D is diagonal.
///
/// Not every matrix can be decomposed like this, only that of a symmetric bilinear function.
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct Decomp2 {
	/// The orthogonal part.
	///
	/// Using a non-orthogonal matrix will yield to incorrect results (but no UB).
	pub ortho: Mat2,

	/// The diagonal part.
	pub diag: Vec2,
}

impl Decomp2 {
	/// Computes the square of this, more efficiently than doing that with a matrix.
	pub fn square(&self) -> Self {
		Self {
			ortho: self.ortho,
			diag: self.diag * self.diag,
		}
	}

	/// Computes the inverse of this, more efficiently than doing that with a matrix.
	pub fn inverse(&self) -> Self {
		Self {
			ortho: self.ortho,
			diag: Vec2::splat(1.0) / self.diag,
		}
	}
}

impl From<Decomp2> for Mat2 {
	fn from(value: Decomp2) -> Self {
		value.ortho.transpose() * Mat2::from_diagonal(value.diag) * value.ortho
	}
}

impl<T> std::ops::Mul<T> for Decomp2
where
	Mat2: std::ops::Mul<T>,
{
	type Output = <Mat2 as std::ops::Mul<T>>::Output;

	fn mul(self, rhs: T) -> Self::Output {
		Mat2::from(self) * rhs
	}
}

/// Represents a 3×3 matrix decomposed as O^T D O, where O is orthogonal and D is diagonal.
///
/// Not every matrix can be decomposed like this, only that of a symmetric bilinear function.
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct Decomp3 {
	/// The orthogonal part.
	///
	/// Using a non-orthogonal matrix will yield to incorrect results (but no UB).
	pub ortho: Mat3,

	/// The diagonal part.
	pub diag: Vec3,
}

impl Decomp3 {
	/// Computes the square of this, more efficiently than doing that with a matrix.
	pub fn square(&self) -> Self {
		Self {
			ortho: self.ortho,
			diag: self.diag * self.diag,
		}
	}

	/// Computes the inverse of this, more efficiently than doing that with a matrix.
	pub fn inverse(&self) -> Self {
		Self {
			ortho: self.ortho,
			diag: Vec3::splat(1.0) / self.diag,
		}
	}
}

impl From<Decomp3> for Mat3 {
	fn from(value: Decomp3) -> Self {
		value.ortho.transpose() * Mat3::from_diagonal(value.diag) * value.ortho
	}
}

impl<T> std::ops::Mul<T> for Decomp3
where
	Mat3: std::ops::Mul<T>,
{
	type Output = <Mat3 as std::ops::Mul<T>>::Output;

	fn mul(self, rhs: T) -> Self::Output {
		Mat3::from(self) * rhs
	}
}

pub fn make_vec2(f: impl Fn(usize) -> f32) -> Vec2 {
	Vec2::from_array(std::array::from_fn(|i| f(i)))
}

pub fn make_mat2(f: impl Fn(usize, usize) -> f32) -> Mat2 {
	Mat2::from_cols_array_2d(&std::array::from_fn(|i| std::array::from_fn(|j| f(i, j))))
}

pub fn make_vec3(f: impl Fn(usize) -> f32) -> Vec3 {
	Vec3::from_array(std::array::from_fn(|i| f(i)))
}

pub fn make_mat3(f: impl Fn(usize, usize) -> f32) -> Mat3 {
	Mat3::from_cols_array_2d(&std::array::from_fn(|i| std::array::from_fn(|j| f(i, j))))
}

#[cfg(test)]
mod tests {
	use super::*;

	use approx::assert_abs_diff_eq;
	use glam::{mat2, mat3, vec2, vec3};

	#[test]
	fn test_lerp() {
		assert_eq!((3., 7.).lerp(-0.5), 1.);
		assert_eq!((3., 7.).lerp(0.0), 3.);
		assert_eq!((3., 7.).lerp(0.5), 5.);
		assert_eq!((3., 7.).lerp(1.0), 7.);
		assert_eq!((3., 7.).lerp(1.5), 9.);
		assert_eq!((3., 7.).inverse_lerp(1.), -0.5);
		assert_eq!((3., 7.).inverse_lerp(3.), 0.0);
		assert_eq!((3., 7.).inverse_lerp(5.), 0.5);
		assert_eq!((3., 7.).inverse_lerp(7.), 1.0);
		assert_eq!((3., 7.).inverse_lerp(9.), 1.5);
	}

	#[test]
	fn test_orthonormalize_2d() {
		assert_abs_diff_eq!(
			mat2(vec2(1., 0.), vec2(0., 1.)).orthonormalize(),
			mat2(vec2(1., 0.), vec2(0., 1.)),
		);
		assert_abs_diff_eq!(
			mat2(vec2(2., 0.), vec2(3., 5.)).orthonormalize(),
			mat2(vec2(1., 0.), vec2(0., 1.)),
		);
		assert_abs_diff_eq!(
			mat2(vec2(0., -3.), vec2(5., 1.)).orthonormalize(),
			mat2(vec2(0., -1.), vec2(1., 0.)),
		);
		assert_abs_diff_eq!(
			mat2(vec2(3., 4.), vec2(5., 1.)).orthonormalize(),
			mat2(vec2(0.6, 0.8), vec2(0.8, -0.6)),
		);
		assert_abs_diff_eq!(
			mat2(vec2(3., 4.), vec2(1., 5.)).orthonormalize(),
			mat2(vec2(0.6, 0.8), vec2(-0.8, 0.6)),
		);
	}

	#[test]
	fn test_orthonormalize_3d() {
		assert_abs_diff_eq!(
			mat3(vec3(1., 0., 0.), vec3(0., 1., 0.), vec3(0., 0., 1.)).orthonormalize(),
			mat3(vec3(1., 0., 0.), vec3(0., 1., 0.), vec3(0., 0., 1.)),
		);
		assert_abs_diff_eq!(
			mat3(vec3(2., 0., 0.), vec3(3., 4., 0.), vec3(5., 6., 7.)).orthonormalize(),
			mat3(vec3(1., 0., 0.), vec3(0., 1., 0.), vec3(0., 0., 1.)),
		);
		assert_abs_diff_eq!(
			mat3(vec3(0., 5., 0.), vec3(0., 7., 6.), vec3(9., 2., 3.)).orthonormalize(),
			mat3(vec3(0., 1., 0.), vec3(0., 0., 1.), vec3(1., 0., 0.)),
		);
	}
}