use crate::internal::uninit;
use crate::linalg::matrix::Matrix;
use num_traits::{ConstOne, ConstZero, Float, One, Zero};
use std::iter::Sum;
use std::ops::{
Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Vector<T, const N: usize> {
data: [T; N],
}
impl<T, const N: usize> Vector<T, N> {
pub const fn new(data: [T; N]) -> Self {
Self { data }
}
pub fn pure(x: T) -> Self
where
T: Clone,
{
let mut data = uninit::new_uninit_array::<T, N>();
for datum in data.iter_mut() {
datum.write(x.clone());
}
let data = unsafe { uninit::array_assume_init(data) };
Self::new(data)
}
pub const fn const_pure(x: T) -> Self
where
T: Copy,
{
Self::new([x; N])
}
pub fn gen<F>(f: F) -> Self
where
F: Fn(usize) -> T,
{
let mut data = uninit::new_uninit_array::<T, N>();
for (i, datum) in data.iter_mut().enumerate() {
datum.write(f(i));
}
let data = unsafe { uninit::array_assume_init(data) };
Self::new(data)
}
pub fn dot(self, rhs: Vector<T, N>) -> T
where
T: Mul<T, Output = T> + Sum,
{
(self * rhs).into_iter().sum()
}
pub fn mag_sq(self) -> T
where
T: Mul<T, Output = T> + Sum + Clone,
{
self.clone().dot(self)
}
pub fn mag(self) -> T
where
T: Float + Sum,
{
self.mag_sq().sqrt()
}
pub fn normalised(self) -> Self
where
T: Float + Sum,
{
self / self.mag()
}
pub fn map<F, U>(self, f: F) -> Vector<U, N>
where
F: Fn(T) -> U,
{
let mut new_data = uninit::new_uninit_array::<U, N>();
for (datum, lhs) in new_data.iter_mut().zip(self.data.into_iter()) {
datum.write(f(lhs));
}
let new_data = unsafe { uninit::array_assume_init(new_data) };
Vector::new(new_data)
}
pub fn map_mut<F>(&mut self, f: F)
where
F: Fn(&mut T),
{
for datum in self.data.iter_mut() {
f(datum);
}
}
pub fn apply<F, U, V>(self, f: F, rhs: Vector<U, N>) -> Vector<V, N>
where
F: Fn(T, U) -> V,
{
let mut new_data = uninit::new_uninit_array::<V, N>();
for (datum, (lhs, rhs)) in new_data
.iter_mut()
.zip(self.data.into_iter().zip(rhs.data.into_iter()))
{
datum.write(f(lhs, rhs));
}
let new_data = unsafe { uninit::array_assume_init(new_data) };
Vector::new(new_data)
}
pub fn apply_mut<F, U>(&mut self, f: F, rhs: Vector<U, N>)
where
F: Fn(&mut T, U),
{
for (datum, rhs) in self.data.iter_mut().zip(rhs.data.into_iter()) {
f(datum, rhs);
}
}
pub(super) fn into_matrix(self) -> Matrix<T, N, 1> {
Matrix::new([self])
}
pub fn into_iter(self) -> <Self as IntoIterator>::IntoIter {
<Self as IntoIterator>::into_iter(self)
}
pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
<&Self as IntoIterator>::into_iter(self)
}
pub fn iter_mut(&mut self) -> <&mut Self as IntoIterator>::IntoIter {
<&mut Self as IntoIterator>::into_iter(self)
}
}
impl<T> Vector<T, 1> {
pub const fn new1(x: T) -> Self {
Vector::new([x])
}
pub const fn x(&self) -> &T {
&self.data[0]
}
}
impl<T> Vector<T, 2> {
pub const fn new2(x: T, y: T) -> Self {
Vector::new([x, y])
}
pub const fn x(&self) -> &T {
&self.data[0]
}
pub const fn y(&self) -> &T {
&self.data[1]
}
}
impl<T> Vector<T, 3> {
pub const fn new3(x: T, y: T, z: T) -> Self {
Vector::new([x, y, z])
}
pub const fn x(&self) -> &T {
&self.data[0]
}
pub const fn y(&self) -> &T {
&self.data[1]
}
pub const fn z(&self) -> &T {
&self.data[2]
}
pub fn cross(self, rhs: Vector<T, 3>) -> Self
where
T: Mul<T, Output = T> + Add<T, Output = T> + Sub<T, Output = T> + Neg<Output = T> + Clone,
{
let x = self.y().clone() * rhs.z().clone() - self.z().clone() * rhs.y().clone();
let y = self.z().clone() * rhs.x().clone() - self.x().clone() * rhs.z().clone();
let z = self.x().clone() * rhs.y().clone() - self.y().clone() * rhs.x().clone();
Vector::new3(x, y, z)
}
}
impl<T> Vector<T, 4> {
pub const fn new4(x: T, y: T, z: T, w: T) -> Self {
Vector::new([x, y, z, w])
}
pub const fn x(&self) -> &T {
&self.data[0]
}
pub const fn y(&self) -> &T {
&self.data[1]
}
pub const fn z(&self) -> &T {
&self.data[2]
}
pub const fn w(&self) -> &T {
&self.data[3]
}
}
impl<T, const N: usize> Neg for Vector<T, N>
where
T: Neg<Output = T>,
{
type Output = Vector<T, N>;
fn neg(self) -> Self::Output {
self.map(T::neg)
}
}
impl<T, const N: usize> Add<Vector<T, N>> for Vector<T, N>
where
T: Add<T, Output = T>,
{
type Output = Vector<T, N>;
fn add(self, rhs: Vector<T, N>) -> Self::Output {
self.apply(T::add, rhs)
}
}
impl<T, const N: usize> AddAssign<Vector<T, N>> for Vector<T, N>
where
T: AddAssign<T>,
{
fn add_assign(&mut self, rhs: Vector<T, N>) {
self.apply_mut(T::add_assign, rhs);
}
}
impl<T, const N: usize> Sub<Vector<T, N>> for Vector<T, N>
where
T: Sub<T, Output = T>,
{
type Output = Vector<T, N>;
fn sub(self, rhs: Vector<T, N>) -> Self::Output {
self.apply(T::sub, rhs)
}
}
impl<T, const N: usize> SubAssign<Vector<T, N>> for Vector<T, N>
where
T: SubAssign<T>,
{
fn sub_assign(&mut self, rhs: Vector<T, N>) {
self.apply_mut(T::sub_assign, rhs);
}
}
impl<T, const N: usize> Mul<Vector<T, N>> for Vector<T, N>
where
T: Mul<T, Output = T>,
{
type Output = Vector<T, N>;
fn mul(self, rhs: Vector<T, N>) -> Self::Output {
self.apply(T::mul, rhs)
}
}
impl<T, const N: usize> Mul<T> for Vector<T, N>
where
T: Mul<T, Output = T> + Clone,
{
type Output = Vector<T, N>;
fn mul(self, rhs: T) -> Self::Output {
self.map(|x| x * rhs.clone())
}
}
impl<T, const N: usize> MulAssign<Vector<T, N>> for Vector<T, N>
where
T: MulAssign,
{
fn mul_assign(&mut self, rhs: Vector<T, N>) {
self.apply_mut(T::mul_assign, rhs);
}
}
impl<T, const N: usize> MulAssign<T> for Vector<T, N>
where
T: MulAssign + Clone,
{
fn mul_assign(&mut self, rhs: T) {
self.map_mut(|x| *x *= rhs.clone());
}
}
impl<T, const N: usize> Div<Vector<T, N>> for Vector<T, N>
where
T: Div<T, Output = T>,
{
type Output = Vector<T, N>;
fn div(self, rhs: Vector<T, N>) -> Self::Output {
self.apply(T::div, rhs)
}
}
impl<T, const N: usize> Div<T> for Vector<T, N>
where
T: Div<T, Output = T> + Clone,
{
type Output = Vector<T, N>;
fn div(self, rhs: T) -> Self::Output {
self.map(|x| x / rhs.clone())
}
}
impl<T, const N: usize> DivAssign<Vector<T, N>> for Vector<T, N>
where
T: DivAssign<T>,
{
fn div_assign(&mut self, rhs: Vector<T, N>) {
self.apply_mut(T::div_assign, rhs);
}
}
impl<T, const N: usize> DivAssign<T> for Vector<T, N>
where
T: DivAssign<T> + Clone,
{
fn div_assign(&mut self, rhs: T) {
self.map_mut(|x| *x /= rhs.clone());
}
}
impl<T, const N: usize> Zero for Vector<T, N>
where
T: Zero + Clone,
{
fn zero() -> Self {
Self::pure(T::zero())
}
fn is_zero(&self) -> bool {
self.data.iter().all(T::is_zero)
}
}
impl<T, const N: usize> ConstZero for Vector<T, N>
where
T: ConstZero + Copy,
{
const ZERO: Self = Self::const_pure(T::ZERO);
}
impl<T, const N: usize> One for Vector<T, N>
where
T: One + Clone,
{
fn one() -> Self {
Vector::pure(T::one())
}
}
impl<T, const N: usize> ConstOne for Vector<T, N>
where
T: ConstOne + Copy,
{
const ONE: Self = Vector::const_pure(T::ONE);
}
impl<T, const N: usize> Sum for Vector<T, N>
where
T: Add<T, Output = T> + Zero + Clone,
{
fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
iter.fold(Vector::zero(), |acc, x| acc + x)
}
}
impl<T, const N: usize> IntoIterator for Vector<T, N> {
type Item = T;
type IntoIter = std::array::IntoIter<T, N>;
fn into_iter(self) -> Self::IntoIter {
self.data.into_iter()
}
}
impl<'a, T, const N: usize> IntoIterator for &'a Vector<T, N> {
type Item = &'a T;
type IntoIter = std::slice::Iter<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.data.iter()
}
}
impl<'a, T, const N: usize> IntoIterator for &'a mut Vector<T, N> {
type Item = &'a mut T;
type IntoIter = std::slice::IterMut<'a, T>;
fn into_iter(self) -> Self::IntoIter {
self.data.iter_mut()
}
}
impl<T, const N: usize> Index<usize> for Vector<T, N> {
type Output = T;
fn index(&self, index: usize) -> &Self::Output {
&self.data[index]
}
}
impl<T, const N: usize> IndexMut<usize> for Vector<T, N> {
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.data[index]
}
}
impl<T, const N: usize> From<Matrix<T, N, 1>> for Vector<T, N> {
fn from(value: Matrix<T, N, 1>) -> Self {
value.into_vector()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn check_add() {
let lhs = Vector::<isize, 4>::new([2, 3, 5, 7]);
let rhs = Vector::<isize, 4>::new([11, 13, 17, 19]);
let expected_sum = Vector::<isize, 4>::new([13, 16, 22, 26]);
let actual_sum = lhs + rhs;
assert_eq!(actual_sum, expected_sum);
}
#[test]
fn check_add_assign() {
let mut lhs = Vector::<isize, 5>::new([23, 29, 31, 37, 41]);
let rhs = Vector::<isize, 5>::new([43, 47, 53, 59, 61]);
let expected_sum = Vector::<isize, 5>::new([66, 76, 84, 96, 102]);
lhs += rhs;
assert_eq!(lhs, expected_sum);
}
#[test]
fn check_sub() {
let lhs = Vector::<isize, 3>::new([67, 71, 73]);
let rhs = Vector::<isize, 3>::new([79, 83, 89]);
let expected_diff = Vector::<isize, 3>::new([-12, -12, -16]);
let actual_diff = lhs - rhs;
assert_eq!(actual_diff, expected_diff);
}
#[test]
fn check_sub_assign() {
let mut lhs = Vector::<isize, 3>::new([97, 101, 103]);
let rhs = Vector::<isize, 3>::new([107, 109, 113]);
let expected_diff = Vector::<isize, 3>::new([-10, -8, -10]);
lhs -= rhs;
assert_eq!(lhs, expected_diff);
}
#[test]
fn check_mul_componentwise() {
let lhs = Vector::<isize, 2>::new([127, 131]);
let rhs = Vector::<isize, 2>::new([137, 139]);
let expected_prod = Vector::<isize, 2>::new([17399, 18209]);
let actual_prod = lhs * rhs;
assert_eq!(actual_prod, expected_prod);
}
#[test]
fn check_mul_scalar() {
let lhs = Vector::<isize, 4>::new([151, 157, 163, 167]);
let rhs = 149isize;
let expected_prod = Vector::<isize, 4>::new([22499, 23393, 24287, 24883]);
let actual_prod = lhs * rhs;
assert_eq!(actual_prod, expected_prod);
}
#[test]
fn check_mul_assign_componentwise() {
let mut lhs = Vector::<isize, 3>::new([173, 179, 181]);
let rhs = Vector::<isize, 3>::new([191, 193, 197]);
let expected_prod = Vector::<isize, 3>::new([33043, 34547, 35657]);
lhs *= rhs;
assert_eq!(lhs, expected_prod);
}
#[test]
fn check_mul_assign_scalar() {
let mut lhs = Vector::<isize, 2>::new([199, 211]);
let rhs = 223isize;
let expected_prod = Vector::<isize, 2>::new([44377, 47053]);
lhs *= rhs;
assert_eq!(lhs, expected_prod);
}
#[test]
fn check_div_componentwise() {
let lhs = Vector::<isize, 3>::new([8, 24, 33]);
let rhs = Vector::<isize, 3>::new([2, 8, 12]);
let expected_quot = Vector::<isize, 3>::new([4, 3, 2]);
let actual_quot = lhs / rhs;
assert_eq!(actual_quot, expected_quot);
}
#[test]
fn check_div_scalar() {
let lhs = Vector::<isize, 4>::new([11, 14, 15, 16]);
let rhs = 3isize;
let expected_quot = Vector::<isize, 4>::new([3, 4, 5, 5]);
let actual_quot = lhs / rhs;
assert_eq!(actual_quot, expected_quot);
}
#[test]
fn check_div_assign_componentwise() {
let mut lhs = Vector::<isize, 3>::new([44, 55, 66]);
let rhs = Vector::<isize, 3>::new([5, 6, 7]);
let expected_quot = Vector::<isize, 3>::new([8, 9, 9]);
lhs /= rhs;
assert_eq!(lhs, expected_quot);
}
#[test]
fn check_div_assign_scalar() {
let mut lhs = Vector::<isize, 2>::new([400, 450]);
let rhs = 50isize;
let expected_quot = Vector::<isize, 2>::new([8, 9]);
lhs /= rhs;
assert_eq!(lhs, expected_quot);
}
}