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}