strided_traits/
element_op.rs

1//! Element-wise operations applied lazily to strided views.
2//!
3//! Julia's StridedViews.jl supports four element operations that form a group:
4//! - `identity`: No transformation
5//! - `conj`: Complex conjugate
6//! - `transpose`: Element-wise transpose (for matrix elements)
7//! - `adjoint`: Element-wise adjoint (conj + transpose)
8//!
9//! These operations are composed at the type level to avoid runtime dispatch.
10//!
11//! # Key Design: `ElementOp<T>` is Generic Over T
12//!
13//! `Identity` implements `ElementOp<T>` for any `T: Copy`, requiring no
14//! additional bounds. `Conj`, `Transpose`, and `Adjoint` require
15//! `T: ElementOpApply`. This allows custom scalar types (e.g., tropical
16//! semiring types) to use `Identity` views without implementing
17//! `ElementOpApply`.
18//!
19//! Composition associated types (Inverse, ComposeConj, etc.) are separated
20//! into `ComposableElementOp<T>`, only available when `T: ElementOpApply`.
21
22use num_complex::Complex;
23use num_traits::Num;
24
25// ---------------------------------------------------------------------------
26// ElementOpApply: trait for types that support conj/transpose/adjoint
27// ---------------------------------------------------------------------------
28
29/// Trait for types that support element operations (conj, transpose, adjoint).
30///
31/// Default implementations return `self` unchanged, so real-valued types
32/// (and custom types that don't need complex operations) can simply write:
33/// ```ignore
34/// impl ElementOpApply for MyType {}
35/// ```
36pub trait ElementOpApply: Copy {
37    #[inline(always)]
38    fn conj(self) -> Self {
39        self
40    }
41    #[inline(always)]
42    fn transpose(self) -> Self {
43        self
44    }
45    #[inline(always)]
46    fn adjoint(self) -> Self {
47        self
48    }
49}
50
51// Real types: use default identity implementations
52macro_rules! impl_element_op_apply_real {
53    ($($t:ty),*) => {
54        $(impl ElementOpApply for $t {})*
55    };
56}
57
58impl_element_op_apply_real!(
59    f32, f64, i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize
60);
61
62// Complex types: override with actual conjugation
63impl<T: Num + Copy + Clone + std::ops::Neg<Output = T>> ElementOpApply for Complex<T> {
64    #[inline(always)]
65    fn conj(self) -> Self {
66        Complex::conj(&self)
67    }
68
69    #[inline(always)]
70    fn transpose(self) -> Self {
71        self
72    }
73
74    #[inline(always)]
75    fn adjoint(self) -> Self {
76        Complex::conj(&self)
77    }
78}
79
80// ---------------------------------------------------------------------------
81// Marker types
82// ---------------------------------------------------------------------------
83
84/// Identity operation: f(x) = x
85#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
86pub struct Identity;
87
88/// Complex conjugate operation: f(x) = conj(x)
89#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
90pub struct Conj;
91
92/// Transpose operation: f(x) = transpose(x)
93/// For scalar numbers, this is identity.
94/// For matrix elements, this would transpose each element.
95#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
96pub struct Transpose;
97
98/// Adjoint operation: f(x) = adjoint(x) = conj(transpose(x))
99/// For scalar numbers, this is conj.
100#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
101pub struct Adjoint;
102
103// ---------------------------------------------------------------------------
104// ElementOp<T>: generic over element type
105// ---------------------------------------------------------------------------
106
107/// Trait for element-wise operations applied to strided views.
108///
109/// Generic over the element type `T`. `Identity` implements this for any
110/// `T: Copy`, while `Conj`, `Transpose`, and `Adjoint` require
111/// `T: ElementOpApply`.
112///
113/// Operations form a group under composition:
114/// ```text
115///   compose | Id   | Conj | Trans | Adj
116/// ---------|------|------|-------|------
117///   Id     | Id   | Conj | Trans | Adj
118///   Conj   | Conj | Id   | Adj   | Trans
119///   Trans  | Trans| Adj  | Id    | Conj
120///   Adj    | Adj  | Trans| Conj  | Id
121/// ```
122pub trait ElementOp<T>: Copy + Default + 'static {
123    /// Whether this operation is the identity (no-op).
124    const IS_IDENTITY: bool = false;
125
126    /// Apply the operation to a value.
127    fn apply(value: T) -> T;
128}
129
130// Identity: works with ANY Copy type (no ElementOpApply needed)
131impl<T: Copy> ElementOp<T> for Identity {
132    const IS_IDENTITY: bool = true;
133
134    #[inline(always)]
135    fn apply(value: T) -> T {
136        value
137    }
138}
139
140// Conj, Transpose, Adjoint: only work with ElementOpApply types
141impl<T: ElementOpApply> ElementOp<T> for Conj {
142    #[inline(always)]
143    fn apply(value: T) -> T {
144        value.conj()
145    }
146}
147
148impl<T: ElementOpApply> ElementOp<T> for Transpose {
149    #[inline(always)]
150    fn apply(value: T) -> T {
151        value.transpose()
152    }
153}
154
155impl<T: ElementOpApply> ElementOp<T> for Adjoint {
156    #[inline(always)]
157    fn apply(value: T) -> T {
158        value.adjoint()
159    }
160}
161
162// ---------------------------------------------------------------------------
163// ComposableElementOp<T>: composition associated types
164// ---------------------------------------------------------------------------
165
166/// Trait for element operations that support type-level composition.
167///
168/// Only available when `T: ElementOpApply`, since composition with
169/// `Conj`/`Transpose`/`Adjoint` requires the element type to support
170/// those operations.
171pub trait ComposableElementOp<T: ElementOpApply>: ElementOp<T> {
172    /// The inverse operation (for this group, each element is its own inverse).
173    type Inverse: ComposableElementOp<T>;
174
175    /// Compose with Conj: Self then Conj
176    type ComposeConj: ComposableElementOp<T>;
177
178    /// Compose with Transpose: Self then Transpose
179    type ComposeTranspose: ComposableElementOp<T>;
180
181    /// Compose with Adjoint: Self then Adjoint
182    type ComposeAdjoint: ComposableElementOp<T>;
183}
184
185impl<T: ElementOpApply> ComposableElementOp<T> for Identity {
186    type Inverse = Identity;
187    type ComposeConj = Conj;
188    type ComposeTranspose = Transpose;
189    type ComposeAdjoint = Adjoint;
190}
191
192impl<T: ElementOpApply> ComposableElementOp<T> for Conj {
193    type Inverse = Conj;
194    type ComposeConj = Identity;
195    type ComposeTranspose = Adjoint;
196    type ComposeAdjoint = Transpose;
197}
198
199impl<T: ElementOpApply> ComposableElementOp<T> for Transpose {
200    type Inverse = Transpose;
201    type ComposeConj = Adjoint;
202    type ComposeTranspose = Identity;
203    type ComposeAdjoint = Conj;
204}
205
206impl<T: ElementOpApply> ComposableElementOp<T> for Adjoint {
207    type Inverse = Adjoint;
208    type ComposeConj = Transpose;
209    type ComposeTranspose = Conj;
210    type ComposeAdjoint = Identity;
211}
212
213// ---------------------------------------------------------------------------
214// Compose<Other>: helper trait for composing two ElementOp types
215// ---------------------------------------------------------------------------
216
217/// Helper trait for composing two ElementOp types.
218///
219/// Only available when `T: ElementOpApply`.
220pub trait Compose<T: ElementOpApply, Other: ComposableElementOp<T>>:
221    ComposableElementOp<T>
222{
223    type Result: ComposableElementOp<T>;
224}
225
226impl<T: ElementOpApply, Op: ComposableElementOp<T>> Compose<T, Identity> for Op {
227    type Result = Op;
228}
229
230impl<T: ElementOpApply> Compose<T, Conj> for Identity {
231    type Result = Conj;
232}
233
234impl<T: ElementOpApply> Compose<T, Conj> for Conj {
235    type Result = Identity;
236}
237
238impl<T: ElementOpApply> Compose<T, Conj> for Transpose {
239    type Result = Adjoint;
240}
241
242impl<T: ElementOpApply> Compose<T, Conj> for Adjoint {
243    type Result = Transpose;
244}
245
246impl<T: ElementOpApply> Compose<T, Transpose> for Identity {
247    type Result = Transpose;
248}
249
250impl<T: ElementOpApply> Compose<T, Transpose> for Conj {
251    type Result = Adjoint;
252}
253
254impl<T: ElementOpApply> Compose<T, Transpose> for Transpose {
255    type Result = Identity;
256}
257
258impl<T: ElementOpApply> Compose<T, Transpose> for Adjoint {
259    type Result = Conj;
260}
261
262impl<T: ElementOpApply> Compose<T, Adjoint> for Identity {
263    type Result = Adjoint;
264}
265
266impl<T: ElementOpApply> Compose<T, Adjoint> for Conj {
267    type Result = Transpose;
268}
269
270impl<T: ElementOpApply> Compose<T, Adjoint> for Transpose {
271    type Result = Conj;
272}
273
274impl<T: ElementOpApply> Compose<T, Adjoint> for Adjoint {
275    type Result = Identity;
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use num_complex::Complex64;
282
283    #[test]
284    fn test_identity() {
285        let x = Complex64::new(3.0, 4.0);
286        assert_eq!(<Identity as ElementOp<Complex64>>::apply(x), x);
287    }
288
289    #[test]
290    fn test_identity_custom_type() {
291        // Custom type that is Copy but does NOT implement ElementOpApply
292        #[derive(Debug, Clone, Copy, PartialEq)]
293        struct MyCustom(f64);
294
295        let x = MyCustom(42.0);
296        assert_eq!(<Identity as ElementOp<MyCustom>>::apply(x), x);
297    }
298
299    #[test]
300    fn test_conj() {
301        let x = Complex64::new(3.0, 4.0);
302        assert_eq!(
303            <Conj as ElementOp<Complex64>>::apply(x),
304            Complex64::new(3.0, -4.0)
305        );
306    }
307
308    #[test]
309    fn test_conj_real() {
310        let x = 3.0f64;
311        assert_eq!(<Conj as ElementOp<f64>>::apply(x), 3.0);
312    }
313
314    #[test]
315    fn test_adjoint_complex() {
316        let x = Complex64::new(3.0, 4.0);
317        assert_eq!(
318            <Adjoint as ElementOp<Complex64>>::apply(x),
319            Complex64::new(3.0, -4.0)
320        );
321    }
322
323    #[test]
324    fn test_composition_conj_conj() {
325        let x = Complex64::new(3.0, 4.0);
326        let result =
327            <Conj as ElementOp<Complex64>>::apply(<Conj as ElementOp<Complex64>>::apply(x));
328        assert_eq!(result, x);
329    }
330
331    #[test]
332    fn test_composable_types() {
333        fn assert_same<A: 'static, B: 'static>() {
334            assert_eq!(
335                std::any::TypeId::of::<A>(),
336                std::any::TypeId::of::<B>(),
337                "types should be the same"
338            );
339        }
340
341        // Identity composed with Conj = Conj
342        assert_same::<<Identity as Compose<f64, Conj>>::Result, Conj>();
343
344        // Conj composed with Conj = Identity
345        assert_same::<<Conj as Compose<f64, Conj>>::Result, Identity>();
346
347        // Transpose composed with Conj = Adjoint
348        assert_same::<<Transpose as Compose<f64, Conj>>::Result, Adjoint>();
349
350        // Adjoint composed with Adjoint = Identity
351        assert_same::<<Adjoint as Compose<f64, Adjoint>>::Result, Identity>();
352    }
353
354    #[test]
355    fn test_element_op_apply_defaults() {
356        // A type using default identity implementations
357        #[derive(Debug, Clone, Copy, PartialEq)]
358        struct Real(f64);
359        impl ElementOpApply for Real {}
360
361        let x = Real(3.0);
362        assert_eq!(x.conj(), x);
363        assert_eq!(x.transpose(), x);
364        assert_eq!(x.adjoint(), x);
365
366        // Can use with all ops
367        assert_eq!(<Conj as ElementOp<Real>>::apply(x), x);
368        assert_eq!(<Transpose as ElementOp<Real>>::apply(x), x);
369        assert_eq!(<Adjoint as ElementOp<Real>>::apply(x), x);
370    }
371}