Skip to main content

ranim_core/traits/
mod.rs

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