Skip to main content

ear_algae/
vector.rs

1use std::{array, iter::Sum, ops::*};
2
3#[cfg(feature = "bytemuck")]
4use bytemuck::{Pod, Zeroable};
5#[cfg(feature = "serde")]
6use serde::{Deserialize, Serialize};
7
8use crate::{
9    Nrml,
10    ops::{Dot, ProjRej, Refl},
11    traits::{Field, Ring},
12};
13
14#[repr(transparent)]
15#[derive(Copy, Clone, PartialEq, Eq, Hash)]
16#[cfg_attr(
17    feature = "serde",
18    derive(Serialize, Deserialize),
19    serde(
20        bound(
21            serialize = "[S; N]:  Serialize",
22            deserialize = "[S; N]: Deserialize<'de>",
23        ),
24        transparent,
25    )
26)]
27pub struct Vect<const N: usize, S>(pub [S; N]);
28
29#[cfg(feature = "bytemuck")]
30unsafe impl<const N: usize, S: Copy> Zeroable for Vect<N, S> where [S; N]: Zeroable {}
31#[cfg(feature = "bytemuck")]
32unsafe impl<const N: usize, S: Copy> Pod for Vect<N, S> where [S; N]: Pod {}
33
34impl<const N: usize, S: Ring> Default for Vect<N, S> {
35    fn default() -> Self {
36        Self::ZERO
37    }
38}
39
40impl<S: Ring, const N: usize> Vect<N, S> {
41    pub const ZERO: Self = Vect([S::ZERO; N]);
42
43    pub const fn axis(i: usize, value: S) -> Self {
44        let mut v = [S::ZERO; N];
45        v[i] = value;
46        Self(v)
47    }
48
49    pub const fn splat(s: S) -> Self {
50        Self([s; N])
51    }
52
53    pub fn from_fn<F: Fn(usize) -> S>(f: F) -> Self {
54        Vect(array::from_fn(f))
55    }
56
57    pub fn swiz<const M: usize>(self, f: impl Fn([S; N]) -> [S; M]) -> Vect<M, S> {
58        Vect(f(self.0))
59    }
60
61    pub fn map<T: Ring>(self, f: impl Fn(S) -> T) -> Vect<N, T> {
62        Vect(self.0.map(f))
63    }
64
65    pub fn scale(self, other: Self) -> Self {
66        Vect::from_fn(|i| self[i].mul(other[i]))
67    }
68
69    pub fn zero_extend<const M: usize>(self, i: usize) -> Vect<M, S>
70    where
71        S: Sized,
72    {
73        const { assert!(M > N) }
74        Vect::<M, S>::from_fn(|j| {
75            if j < i {
76                self[i]
77            } else if j > i + M - N {
78                self[j - M - N]
79            } else {
80                S::ZERO
81            }
82        })
83    }
84
85    pub fn is_nan(self) -> bool {
86        self.0.into_iter().any(S::is_nan)
87    }
88
89    pub fn is_finite(self) -> bool {
90        self.0.into_iter().all(S::is_finite)
91    }
92}
93
94impl<const N: usize> Vect<N, f32> {
95    pub fn to_f64(self) -> Vect<N, f64> {
96        self.map(|x| x as _)
97    }
98}
99
100impl<const N: usize> Vect<N, f64> {
101    pub fn to_f32(self) -> Vect<N, f32> {
102        self.map(|x| x as _)
103    }
104}
105
106impl<S: Ring, const N: usize> Index<usize> for Vect<N, S> {
107    type Output = S;
108
109    fn index(&self, i: usize) -> &Self::Output {
110        &self.0[i]
111    }
112}
113
114impl<S: Ring, const N: usize> IndexMut<usize> for Vect<N, S> {
115    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
116        &mut self.0[i]
117    }
118}
119
120impl<S: Ring, const N: usize> Index<RangeFull> for Vect<N, S> {
121    type Output = [S];
122
123    fn index(&self, index: RangeFull) -> &Self::Output {
124        &self.0[index]
125    }
126}
127
128impl<S: Ring, const N: usize> Add for Vect<N, S> {
129    type Output = Self;
130
131    fn add(self, rhs: Self) -> Self::Output {
132        Self::from_fn(|i| self[i].add(rhs[i]))
133    }
134}
135
136impl<S: Ring, const N: usize> Sub for Vect<N, S> {
137    type Output = Self;
138
139    fn sub(self, rhs: Self) -> Self::Output {
140        Self::from_fn(|i| self[i].sub(rhs[i]))
141    }
142}
143
144impl<S: Ring, const N: usize> Neg for Vect<N, S> {
145    type Output = Self;
146
147    fn neg(self) -> Self::Output {
148        Self::from_fn(|i| self[i].neg())
149    }
150}
151
152impl<S: Ring, const N: usize> Mul<S> for Vect<N, S> {
153    type Output = Self;
154
155    fn mul(self, rhs: S) -> Self::Output {
156        Self::from_fn(|i| self[i].mul(rhs))
157    }
158}
159
160impl<S: Ring, const N: usize> Div<S> for Vect<N, S> {
161    type Output = Self;
162
163    fn div(self, rhs: S) -> Self::Output {
164        Self::from_fn(|i| self[i].div(rhs))
165    }
166}
167
168impl<S: Ring, const N: usize> Dot<Self> for Vect<N, S> {
169    type Output = S;
170
171    fn dot(self, other: Self) -> Self::Output {
172        (0..N)
173            .map(|i| self[i].mul(other[i]))
174            .fold(S::ZERO, |c, n| c.add(n))
175    }
176}
177
178impl<S: Ring, const N: usize> Vect<N, S> {
179    pub fn sqr_magn(self) -> S {
180        self.dot(self)
181    }
182}
183
184impl<S: Field, const N: usize> Vect<N, S> {
185    pub fn magn(self) -> S {
186        self.sqr_magn().sqrt()
187    }
188
189    pub fn normal(self) -> Option<Nrml<N, S>> {
190        Some(self.magn_normal()?.1)
191    }
192
193    pub fn normal_or_zero(self) -> Vect<N, S> {
194        if let Some(normal) = self.normal() {
195            normal.into()
196        } else {
197            Vect::ZERO
198        }
199    }
200
201    pub fn magn_normal(self) -> Option<(S, Nrml<N, S>)> {
202        if !self.is_finite() {
203            if let Some(value) = self.divide_by_infinity() {
204                return Some((S::INFINITY, unsafe { Nrml::new_unchecked(value.0) }));
205            }
206            return None;
207        }
208        let magn = self.magn();
209        if magn.is_zero() {
210            None
211        } else {
212            Some((magn, unsafe { Nrml::new_unchecked((self / magn).0) }))
213        }
214    }
215
216    pub fn magn_normal_or_zero(self) -> (S, Vect<N, S>) {
217        if let Some((magn, normal)) = self.magn_normal() {
218            return (magn, normal.into());
219        }
220        (S::ZERO, Vect::ZERO)
221    }
222
223    pub fn divide_by_infinity(self) -> Option<Vect<N, S>> {
224        if self.is_nan() {
225            return None;
226        }
227        let mut array = self.0;
228        let mut magn = S::ZERO;
229        for e in &mut array {
230            if e.is_finite() {
231                *e = S::ZERO;
232            } else {
233                *e = e.sign();
234                magn.add_assign(S::ONE);
235            }
236        }
237        magn = magn.sqrt();
238
239        let vect = Vect(array) / magn;
240
241        if vect.is_finite() {
242            return Some(vect);
243        }
244        None
245    }
246}
247
248impl<S: Ring, const N: usize> Sum for Vect<N, S> {
249    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
250        iter.fold(Vect::ZERO, |c, n| c + n)
251    }
252}
253
254impl<S: Ring, const N: usize> AddAssign for Vect<N, S> {
255    fn add_assign(&mut self, rhs: Self) {
256        for i in 0..N {
257            self[i].add_assign(rhs[i])
258        }
259    }
260}
261
262impl<S: Ring, const N: usize> SubAssign for Vect<N, S> {
263    fn sub_assign(&mut self, rhs: Self) {
264        for i in 0..N {
265            self[i].sub_assign(rhs[i])
266        }
267    }
268}
269
270impl<S: Ring, const N: usize> MulAssign<S> for Vect<N, S> {
271    fn mul_assign(&mut self, rhs: S) {
272        for i in 0..N {
273            self[i].mul_assign(rhs)
274        }
275    }
276}
277
278impl<S: Ring, const N: usize> DivAssign<S> for Vect<N, S> {
279    fn div_assign(&mut self, rhs: S) {
280        for i in 0..N {
281            self[i].div_assign(rhs)
282        }
283    }
284}
285
286impl<S: Field, const N: usize> ProjRej<Nrml<N, S>> for Vect<N, S> {
287    type Output = Self;
288
289    fn proj(self, axis: Nrml<N, S>) -> Self {
290        axis * self.dot(axis)
291    }
292
293    fn rej(self, axis: Nrml<N, S>) -> Self {
294        self - self.proj(axis)
295    }
296
297    fn proj_rej(self, axis: Nrml<N, S>) -> (Self, Self) {
298        let proj = self.proj(axis);
299        (proj, self - proj)
300    }
301}
302
303impl<S: Field, const N: usize> ProjRej<Option<Nrml<N, S>>> for Vect<N, S> {
304    type Output = Self;
305
306    fn proj(self, axis: Option<Nrml<N, S>>) -> Self {
307        match axis {
308            Some(axis) => self.proj(axis),
309            None => Vect::ZERO,
310        }
311    }
312
313    fn rej(self, axis: Option<Nrml<N, S>>) -> Self {
314        self - self.proj(axis)
315    }
316
317    fn proj_rej(self, axis: Option<Nrml<N, S>>) -> (Self, Self) {
318        let proj = self.proj(axis);
319        (proj, self - proj)
320    }
321}
322
323impl<S: Field, const N: usize> ProjRej<Vect<N, S>> for Vect<N, S> {
324    type Output = Self;
325
326    fn proj(self, axis: Vect<N, S>) -> Self::Output {
327        let magn2 = axis.sqr_magn();
328        if magn2.is_zero() {
329            Vect::ZERO
330        } else {
331            axis * self.dot(axis).div(magn2)
332        }
333    }
334
335    fn rej(self, axis: Vect<N, S>) -> Self::Output {
336        self - self.proj(axis)
337    }
338
339    fn proj_rej(self, axis: Vect<N, S>) -> (Self::Output, Self::Output) {
340        let proj = self.proj(axis);
341        (proj, self - proj)
342    }
343}
344
345impl<S: Field, const N: usize> Refl<Nrml<N, S>> for Vect<N, S> {
346    type Output = Self;
347
348    fn refl(self, axis: Nrml<N, S>) -> Self::Output {
349        self - self.proj(axis) * S::TWO
350    }
351}
352
353impl<S: Field, const N: usize> Refl<Option<Nrml<N, S>>> for Vect<N, S> {
354    type Output = Self;
355
356    fn refl(self, axis: Option<Nrml<N, S>>) -> Self::Output {
357        self - self.proj(axis) * S::TWO
358    }
359}
360
361impl<S: Field, const N: usize> Refl<Vect<N, S>> for Vect<N, S> {
362    type Output = Self;
363
364    fn refl(self, axis: Vect<N, S>) -> Self::Output {
365        self - self.proj(axis) * S::TWO
366    }
367}
368
369impl<S: Field> Vect<1, S> {
370    pub fn x(self) -> S {
371        self[0]
372    }
373}
374
375impl<S: Field> Vect<2, S> {
376    pub fn x(self) -> S {
377        self[0]
378    }
379
380    pub fn y(self) -> S {
381        self[1]
382    }
383}
384
385impl<S: Field> Vect<3, S> {
386    pub fn x(self) -> S {
387        self[0]
388    }
389
390    pub fn y(self) -> S {
391        self[1]
392    }
393
394    pub fn z(self) -> S {
395        self[2]
396    }
397}
398
399impl<S: Field> Vect<4, S> {
400    pub fn x(self) -> S {
401        self[0]
402    }
403
404    pub fn y(self) -> S {
405        self[1]
406    }
407
408    pub fn z(self) -> S {
409        self[2]
410    }
411
412    pub fn w(self) -> S {
413        self[3]
414    }
415}