Skip to main content

ranim_core/traits/
mod.rs

1/// Transform related things
2pub mod transform;
3
4pub use crate::anchor::{Aabb, AabbPoint, Locate};
5pub use transform::{
6    ApplyTransform, Diag, NotSimilarity, Rigid, RotateTransform, ScaleHint, ScaleTransform,
7    ScaleTransformExt, ScaleTransformStrokeExt, ShiftTransform, ShiftTransformExt, Similarity,
8    TransformGroup, Translation, UniformScaleTransform,
9};
10
11use std::ops::Range;
12
13use color::{AlphaColor, ColorSpace, OpaqueColor, Srgb};
14use glam::{
15    DAffine2, DAffine3, DMat3, DMat4, DQuat, DVec2, DVec3, Mat4, USizeVec3, Vec3, Vec3Swizzles,
16    Vec4, dvec3,
17};
18use num::complex::Complex64;
19
20use crate::{components::width::Width, utils::resize_preserving_order_with_repeated_indices};
21
22// MARK: With
23/// A trait for mutating a value in place.
24///
25/// This trait is automatically implemented for `T`.
26///
27/// # Example
28/// ```ignore
29/// use ranim::prelude::*;
30///
31/// let mut a = 1;
32/// a = a.with(|x| *x = 2);
33/// assert_eq!(a, 2);
34/// ```
35pub trait With {
36    /// Mutating a value in place
37    fn with(mut self, f: impl Fn(&mut Self)) -> Self
38    where
39        Self: Sized,
40    {
41        f(&mut self);
42        self
43    }
44}
45
46impl<T> With for T {}
47
48/// A trait for discarding a value.
49///
50/// It is useful when you want a short closure:
51/// ```ignore
52/// let x = Square::new(1.0).with(|x| {
53///     x.set_color(manim::BLUE_C);
54/// });
55/// let x = Square::new(1.0).with(|x|
56///     x.set_color(manim::BLUE_C).discard()
57/// );
58/// ```
59pub trait Discard {
60    /// Simply returns `()`
61    fn discard(&self) {}
62}
63
64impl<T> Discard for T {}
65
66// MARK: Interpolatable
67/// A trait for interpolating to values
68///
69/// It uses the reference of two values and produce an owned interpolated value.
70pub trait Interpolatable {
71    /// Lerping between values
72    fn lerp(&self, target: &Self, t: f64) -> Self;
73}
74
75macro_rules! impl_interpolatable_for_int {
76    ($($t:ty),*) => {
77        $(
78            impl Interpolatable for $t {
79                fn lerp(&self, target: &Self, t: f64) -> Self {
80                    (*self as f64).lerp(&(*target as f64), t) as $t
81                }
82            }
83        )*
84    };
85}
86
87impl_interpolatable_for_int!(
88    u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
89);
90
91impl Interpolatable for f32 {
92    fn lerp(&self, target: &Self, t: f64) -> Self {
93        self + (target - self) * t as f32
94    }
95}
96
97impl Interpolatable for f64 {
98    fn lerp(&self, target: &Self, t: f64) -> Self {
99        self + (target - self) * t
100    }
101}
102
103impl Interpolatable for DVec3 {
104    fn lerp(&self, target: &Self, t: f64) -> Self {
105        self + (target - self) * t
106    }
107}
108
109impl Interpolatable for Vec3 {
110    fn lerp(&self, target: &Self, t: f64) -> Self {
111        self + (target - self) * t as f32
112    }
113}
114
115impl Interpolatable for Vec4 {
116    fn lerp(&self, target: &Self, t: f64) -> Self {
117        self + (target - self) * t as f32
118    }
119}
120
121impl Interpolatable for DVec2 {
122    fn lerp(&self, target: &Self, t: f64) -> Self {
123        self + (target - self) * t
124    }
125}
126
127impl Interpolatable for DQuat {
128    fn lerp(&self, target: &Self, t: f64) -> Self {
129        self.slerp(*target, t)
130    }
131}
132
133impl<CS: ColorSpace> Interpolatable for AlphaColor<CS> {
134    fn lerp(&self, target: &Self, t: f64) -> Self {
135        // TODO: figure out to use `lerp_rect` or `lerp`
136        AlphaColor::lerp_rect(*self, *target, t as f32)
137    }
138}
139
140impl<CS: ColorSpace> Interpolatable for OpaqueColor<CS> {
141    fn lerp(&self, target: &Self, t: f64) -> Self {
142        // TODO: figure out to use `lerp_rect` or `lerp`
143        OpaqueColor::lerp_rect(*self, *target, t as f32)
144    }
145}
146
147impl Interpolatable for DMat4 {
148    fn lerp(&self, target: &Self, t: f64) -> Self {
149        let mut result = DMat4::ZERO;
150        for i in 0..4 {
151            for j in 0..4 {
152                result.col_mut(i)[j] = self.col(i)[j].lerp(&target.col(i)[j], t);
153            }
154        }
155        result
156    }
157}
158
159impl Interpolatable for DAffine3 {
160    fn lerp(&self, target: &Self, t: f64) -> Self {
161        // Component-wise lerp of `matrix3` + `translation` is trivially closed
162        // within affine transforms (no homogeneous bottom row to preserve).
163        Self {
164            matrix3: DMat3::from_cols(
165                self.matrix3.x_axis.lerp(target.matrix3.x_axis, t),
166                self.matrix3.y_axis.lerp(target.matrix3.y_axis, t),
167                self.matrix3.z_axis.lerp(target.matrix3.z_axis, t),
168            ),
169            translation: self.translation.lerp(target.translation, t),
170        }
171    }
172}
173
174impl Interpolatable for Mat4 {
175    fn lerp(&self, other: &Self, t: f64) -> Self {
176        let t = t as f32;
177        let mut result = Mat4::ZERO;
178        for i in 0..4 {
179            for j in 0..4 {
180                result.col_mut(i)[j] = self.col(i)[j] + (other.col(i)[j] - self.col(i)[j]) * t;
181            }
182        }
183        result
184    }
185}
186
187impl<T: Interpolatable> Interpolatable for Vec<T> {
188    fn lerp(&self, target: &Self, t: f64) -> Self {
189        self.iter().zip(target).map(|(a, b)| a.lerp(b, t)).collect()
190    }
191}
192
193impl<T: Interpolatable, const N: usize> Interpolatable for [T; N] {
194    fn lerp(&self, target: &Self, t: f64) -> Self {
195        core::array::from_fn(|i| self[i].lerp(&target[i], t))
196    }
197}
198
199macro_rules! impl_interpolatable_tuple {
200    ($(($T:ident, $s:ident)),*) => {
201        impl<$($T: Interpolatable),*> Interpolatable for ($($T,)*) {
202            #[allow(non_snake_case)]
203            fn lerp(&self, target: &Self, t: f64) -> Self {
204                let ($($s,)*) = self;
205                let ($($T,)*) = target;
206                ($($s.lerp($T, t),)*)
207            }
208        }
209    }
210}
211variadics_please::all_tuples!(impl_interpolatable_tuple, 1, 12, T, S);
212
213impl<T: Opacity + Alignable + Clone> Alignable for Vec<T> {
214    fn is_aligned(&self, other: &Self) -> bool {
215        self.len() == other.len() && self.iter().zip(other).all(|(a, b)| a.is_aligned(b))
216    }
217    fn align_with(&mut self, other: &mut Self) {
218        let len = self.len().max(other.len());
219
220        let transparent_repeated = |items: &mut Vec<T>, repeat_idxs: Vec<usize>| {
221            for idx in repeat_idxs {
222                items[idx].set_opacity(0.0);
223            }
224        };
225        if self.len() != len {
226            let (mut items, idxs) = resize_preserving_order_with_repeated_indices(self, len);
227            transparent_repeated(&mut items, idxs);
228            *self = items;
229        }
230        if other.len() != len {
231            let (mut items, idxs) = resize_preserving_order_with_repeated_indices(other, len);
232            transparent_repeated(&mut items, idxs);
233            *other = items;
234        }
235        self.iter_mut()
236            .zip(other)
237            .for_each(|(a, b)| a.align_with(b));
238    }
239}
240
241// MARK: Alignable
242/// A trait for aligning two items
243///
244/// Alignment is actually the meaning of preparation for interpolation.
245///
246/// For example, if we want to interpolate two VItems, we need to
247/// align all their inner components like `ComponentVec<VPoint>` to the same length.
248pub trait Alignable: Clone {
249    /// Checking if two items are aligned
250    fn is_aligned(&self, other: &Self) -> bool;
251    /// Aligning two items
252    fn align_with(&mut self, other: &mut Self);
253}
254
255impl Alignable for DVec3 {
256    fn align_with(&mut self, _other: &mut Self) {}
257    fn is_aligned(&self, _other: &Self) -> bool {
258        true
259    }
260}
261
262// MARK: Opacity
263/// A trait for items with opacity
264pub trait Opacity {
265    /// Setting opacity of an item
266    fn set_opacity(&mut self, opacity: f32) -> &mut Self;
267}
268
269impl<T: Opacity, I> Opacity for I
270where
271    for<'a> &'a mut I: IntoIterator<Item = &'a mut T>,
272{
273    fn set_opacity(&mut self, opacity: f32) -> &mut Self {
274        self.into_iter().for_each(|x: &mut T| {
275            x.set_opacity(opacity);
276        });
277        self
278    }
279}
280
281// MARK: Partial
282/// A trait for items that can be displayed partially
283pub trait Partial {
284    /// Getting a partial item
285    fn get_partial(&self, range: Range<f64>) -> Self;
286    /// Getting a partial item closed
287    fn get_partial_closed(&self, range: Range<f64>) -> Self;
288}
289
290// MARK: Empty
291/// A trait for items that can be empty
292pub trait Empty {
293    /// Getting an empty item
294    fn empty() -> Self;
295}
296
297// MARK: FillColor
298/// A trait for items that have fill color
299pub trait FillColor {
300    /// Getting fill color of an item
301    fn fill_color(&self) -> AlphaColor<Srgb>;
302    /// Setting fill opacity of an item
303    fn set_fill_opacity(&mut self, opacity: f32) -> &mut Self;
304    /// Setting fill color(rgba) of an item
305    fn set_fill_color(&mut self, color: AlphaColor<Srgb>) -> &mut Self;
306}
307
308impl<T: FillColor> FillColor for [T] {
309    fn fill_color(&self) -> color::AlphaColor<color::Srgb> {
310        self[0].fill_color()
311    }
312    fn set_fill_color(&mut self, color: color::AlphaColor<color::Srgb>) -> &mut Self {
313        self.iter_mut()
314            .for_each(|x| x.set_fill_color(color).discard());
315        self
316    }
317    fn set_fill_opacity(&mut self, opacity: f32) -> &mut Self {
318        self.iter_mut()
319            .for_each(|x| x.set_fill_opacity(opacity).discard());
320        self
321    }
322}
323
324// MARK: StrokeColor
325/// A trait for items that have stroke color
326pub trait StrokeColor {
327    /// Getting stroke color of an item
328    fn stroke_color(&self) -> AlphaColor<Srgb>;
329    /// Setting stroke opacity of an item
330    fn set_stroke_opacity(&mut self, opacity: f32) -> &mut Self;
331    /// Setting stroke color(rgba) of an item
332    fn set_stroke_color(&mut self, color: AlphaColor<Srgb>) -> &mut Self;
333}
334
335impl<T: StrokeColor> StrokeColor for [T] {
336    fn stroke_color(&self) -> AlphaColor<Srgb> {
337        self[0].stroke_color()
338    }
339    fn set_stroke_color(&mut self, color: color::AlphaColor<color::Srgb>) -> &mut Self {
340        self.iter_mut().for_each(|x| {
341            x.set_stroke_color(color);
342        });
343        self
344    }
345    fn set_stroke_opacity(&mut self, opacity: f32) -> &mut Self {
346        self.iter_mut().for_each(|x| {
347            x.set_stroke_opacity(opacity);
348        });
349        self
350    }
351}
352
353// MARK: StrokeWidth
354/// A trait for items have stroke width
355pub trait StrokeWidth {
356    // TODO: Make this better
357    /// Get the stroke width
358    fn stroke_width(&self) -> f32;
359    /// Applying stroke width function to an item
360    fn apply_stroke_func(&mut self, f: impl for<'a> Fn(&'a mut [Width])) -> &mut Self;
361    /// Setting stroke width of an item
362    fn set_stroke_width(&mut self, width: f32) -> &mut Self {
363        self.apply_stroke_func(|widths| widths.fill(width.into()))
364    }
365}
366
367impl<T: StrokeWidth> StrokeWidth for [T] {
368    fn stroke_width(&self) -> f32 {
369        self[0].stroke_width()
370    }
371    fn apply_stroke_func(
372        &mut self,
373        f: impl for<'a> Fn(&'a mut [crate::components::width::Width]),
374    ) -> &mut Self {
375        self.iter_mut().for_each(|x| {
376            x.apply_stroke_func(&f);
377        });
378        self
379    }
380}
381
382// MARK: Color
383/// A trait for items that have both fill color and stroke color
384///
385/// This trait is auto implemented for items that implement [`FillColor`] and [`StrokeColor`].
386pub trait Color: FillColor + StrokeColor {
387    /// Setting color(rgba) of an item
388    fn set_color(&mut self, color: AlphaColor<Srgb>) -> &mut Self {
389        self.set_fill_color(color);
390        self.set_stroke_color(color);
391        self
392    }
393}
394
395impl<T: FillColor + StrokeColor + ?Sized> Color for T {}
396
397// MARK: PointsFunc
398/// A trait for items that can apply points function.
399pub trait PointsFunc {
400    /// Applying points function to an item
401    fn apply_points_func(&mut self, f: impl for<'a> Fn(&'a mut [DVec3])) -> &mut Self;
402
403    /// Applying affine transform in xy plane to an item
404    fn apply_affine2(&mut self, affine: DAffine2) -> &mut Self {
405        self.apply_point_func(|p| {
406            let transformed = affine.transform_point2(p.xy());
407            p.x = transformed.x;
408            p.y = transformed.y;
409        });
410        self
411    }
412
413    /// Applying affine transform to an item
414    fn apply_affine3(&mut self, affine: DAffine3) -> &mut Self {
415        self.apply_point_func(|p| *p = affine.transform_point3(*p));
416        self
417    }
418
419    /// Applying point function to an item
420    fn apply_point_func(&mut self, f: impl Fn(&mut DVec3)) -> &mut Self {
421        self.apply_points_func(|points| {
422            points.iter_mut().for_each(&f);
423        });
424        self
425    }
426    /// Applying point function to an item
427    fn apply_point_map(&mut self, f: impl Fn(DVec3) -> DVec3) -> &mut Self {
428        self.apply_points_func(|points| {
429            points.iter_mut().for_each(|p| *p = f(*p));
430        });
431        self
432    }
433
434    /// Applying complex function to an item.
435    ///
436    /// The point's x and y coordinates will be used as real and imaginary parts of a complex number.
437    fn apply_complex_func(&mut self, f: impl Fn(&mut Complex64)) -> &mut Self {
438        self.apply_point_func(|p| {
439            let mut c = Complex64::new(p.x, p.y);
440            f(&mut c);
441            p.x = c.re;
442            p.y = c.im;
443        });
444        self
445    }
446    /// Applying complex function to an item.
447    ///
448    /// The point's x and y coordinates will be used as real and imaginary parts of a complex number.
449    fn apply_complex_map(&mut self, f: impl Fn(Complex64) -> Complex64) -> &mut Self {
450        self.apply_complex_func(|p| {
451            *p = f(*p);
452        });
453        self
454    }
455}
456
457impl PointsFunc for DVec3 {
458    fn apply_points_func(&mut self, f: impl for<'a> Fn(&'a mut [DVec3])) -> &mut Self {
459        f(std::slice::from_mut(self));
460        self
461    }
462}
463
464impl<T: PointsFunc> PointsFunc for [T] {
465    fn apply_points_func(&mut self, f: impl for<'a> Fn(&'a mut [DVec3])) -> &mut Self {
466        self.iter_mut()
467            .for_each(|x| x.apply_points_func(&f).discard());
468        self
469    }
470}
471
472// MARK: Align
473/// Align a slice of items
474pub trait AlignSlice<T: ShiftTransformExt>: AsMut<[T]> {
475    /// Align items' anchors in a given axis, based on the first item.
476    fn align_anchor<A>(&mut self, axis: DVec3, anchor: A) -> &mut Self
477    where
478        A: Locate<T> + Clone,
479    {
480        let Some(dir) = axis.try_normalize() else {
481            return self;
482        };
483        let Some(point) = self.as_mut().first().map(|x| anchor.locate(x)) else {
484            return self;
485        };
486
487        self.as_mut().iter_mut().for_each(|x| {
488            let p = anchor.locate(x);
489
490            let v = p - point;
491            let proj = dir * v.dot(dir);
492            let closest = point + proj;
493            let displacement = closest - p;
494            x.shift(displacement);
495        });
496        self
497    }
498    /// Align items' centers in a given axis, based on the first item.
499    fn align(&mut self, axis: DVec3) -> &mut Self
500    where
501        T: Aabb,
502    {
503        self.align_anchor(axis, AabbPoint::CENTER)
504    }
505}
506
507// MARK: Arrange
508/// A trait for arranging operations.
509pub trait ArrangeSlice<T: ShiftTransformExt>: AsMut<[T]> {
510    /// Arrange the items by a given function.
511    ///
512    /// The `pos_func` takes index as input and output the center position.
513    fn arrange_with(&mut self, pos_func: impl Fn(usize) -> DVec3)
514    where
515        AabbPoint: Locate<T>,
516    {
517        self.as_mut().iter_mut().enumerate().for_each(|(i, x)| {
518            x.move_to(pos_func(i));
519        });
520    }
521    /// Arrange the items in a col
522    fn arrange_in_y(&mut self, gap: f64)
523    where
524        T: Aabb,
525        AabbPoint: Locate<T>,
526    {
527        let Some(mut bbox) = self.as_mut().first().map(|x| x.aabb()) else {
528            return;
529        };
530
531        self.as_mut().iter_mut().for_each(|x| {
532            x.move_next_to_padded(bbox.as_slice(), AabbPoint(DVec3::Y), gap);
533            bbox = x.aabb();
534        });
535    }
536    /// Arrange the items in a grid.
537    fn arrange_in_grid(&mut self, cell_cnt: USizeVec3, cell_size: DVec3, gap: DVec3) -> &mut Self
538    where
539        AabbPoint: Locate<T>,
540    {
541        // x -> y -> z
542        let pos_func = |idx: usize| {
543            let x = idx % cell_cnt.x;
544            let temp = idx / cell_cnt.x;
545
546            let y = temp % cell_cnt.y;
547            let z = temp / cell_cnt.y;
548            dvec3(x as f64, y as f64, z as f64) * cell_size
549                + gap * dvec3(x as f64, y as f64, z as f64)
550        };
551        self.arrange_with(pos_func);
552        self
553    }
554    /// Arrange the items in a grid with given number of columns.
555    ///
556    /// The `pos_func` takes row and column index as input and output the center position.
557    fn arrange_in_cols_with(&mut self, ncols: usize, pos_func: impl Fn(usize, usize) -> DVec3)
558    where
559        AabbPoint: Locate<T>,
560    {
561        let pos_func = |idx: usize| {
562            let row = idx / ncols;
563            let col = idx % ncols;
564            pos_func(row, col)
565        };
566        self.arrange_with(pos_func);
567    }
568    /// Arrange the items in a grid with given number of rows.
569    ///
570    /// The `pos_func` takes row and column index as input and output the center position.
571    fn arrange_in_rows_with(&mut self, nrows: usize, pos_func: impl Fn(usize, usize) -> DVec3)
572    where
573        AabbPoint: Locate<T>,
574    {
575        let ncols = self.as_mut().len().div_ceil(nrows);
576        self.arrange_in_cols_with(ncols, pos_func);
577    }
578}
579
580impl<T: ShiftTransformExt, E: AsMut<[T]>> ArrangeSlice<T> for E {}