Skip to main content

thermite_complex/
vector.rs

1//! Element and vector-trait integration for [`Complex`].
2//!
3//! `Complex<E>` over a scalar float element implements
4//! [`Element`]/[`SignedElement`]/[`FloatElement`], so it can be the element of a
5//! complex vector; `Complex<V>` over a real [`FloatVector`] implements the
6//! [`GenericVector`] -> [`FloatVector`] stack, with `Element = Complex<V::Element>`
7//! and the mask and lanes of `V`.
8//!
9//! The ordering, sign and rounding semantics are in the [crate docs](crate).
10
11use core::marker::PhantomData;
12use core::ops::{Add, Div, Mul, Rem, Sub};
13
14use num_traits::Bounded;
15
16use thermite::Swizzle;
17use thermite::element::{Element, FloatElement, SignedElement};
18use thermite::generic_array::{GenericArray, IntoArrayLength, typenum::Const};
19use thermite::mask::{GenericMask, GenericSelectable};
20use thermite::math::RealMathWithPolicy;
21use thermite::math::algorithms::reduce_in_place;
22use thermite::math::policy::DefaultPolicy;
23use thermite::register::SwizzleIndices;
24use thermite::vector::ops::{AddSubExt, AddSubExtMasked, NegMasked, Square, SquareMasked};
25use thermite::vector::{NewConst, NewVector, SplatConst, SplatVector, VectorValue, const_new, const_splat};
26use thermite::{LargeInt, features, prelude::*};
27
28use crate::{Complex, RealValue};
29
30/// A real [`FloatVector`] usable as the inner storage of a [`Complex`] vector.
31///
32/// The complex primitives (`sqrt`, `abs`, `signum`, `rcp`, ...) are built out of
33/// the inner vector's `hypot`/`reciprocal`, so the policy math library is required
34/// here. `Dual`/`Compensated` split theirs into a separate tier; there is no useful
35/// math-free tier to split out of this one.
36pub trait RealFloatVector:
37    RealValue + FloatVector<Element: RealValue> + CastVector<Self> + RealMathWithPolicy + SwizzleVector
38{
39}
40
41impl<V> RealFloatVector for V where
42    V: RealValue + FloatVector<Element: RealValue> + CastVector<V> + RealMathWithPolicy + SwizzleVector
43{
44}
45
46// Lane swizzles apply to both components: re and im move through the same
47// permutation, so a swizzled complex vector is the complex of the swizzled
48// inputs.
49impl<V: RealFloatVector> Swizzle<V::Lanes> for Complex<V> {
50    #[inline(always)]
51    fn swizzle(self, other: Self, indices: GenericArray<u32, V::Lanes>) -> Self {
52        Self {
53            re: self.re.swizzle(other.re, indices.clone()),
54            im: self.im.swizzle(other.im, indices),
55        }
56    }
57
58    #[inline(always)]
59    fn permute(self, indices: GenericArray<u32, V::Lanes>) -> Self {
60        Self {
61            re: self.re.permute(indices.clone()),
62            im: self.im.permute(indices),
63        }
64    }
65
66    // Forward the `_const` forms per component - the trait defaults route
67    // through the runtime-index methods and lose the immediate-encoded
68    // shuffles.
69    #[inline(always)]
70    fn swizzle_const<I: SwizzleIndices<V::Lanes>>(self, other: Self) -> Self {
71        Self {
72            re: self.re.swizzle_const::<I>(other.re),
73            im: self.im.swizzle_const::<I>(other.im),
74        }
75    }
76
77    #[inline(always)]
78    fn permute_const<I: SwizzleIndices<V::Lanes>>(self) -> Self {
79        Self {
80            re: self.re.permute_const::<I>(),
81            im: self.im.permute_const::<I>(),
82        }
83    }
84}
85
86// --- Element stack: Complex<E> as a scalar element ---
87
88#[rustfmt::skip]
89impl<E: RealValue + Element> Element for Complex<E> {
90    type Signed = <E as Element>::Signed;
91    type Unsigned = <E as Element>::Unsigned;
92
93    const ZERO: Self = Self::ZERO;
94    const ONE: Self = Self::ONE;
95
96    // The (documented, if artificial) order on complex numbers here is
97    // lexicographic (re, im) - see `PartialOrdVector for Complex` - so the
98    // order extremes are extreme in both components, and unordered values
99    // (NaN in either part) exist exactly when the component type has them.
100    const ORDER_MAX: Self = Self { re: E::ORDER_MAX, im: E::ORDER_MAX };
101    const ORDER_MIN: Self = Self { re: E::ORDER_MIN, im: E::ORDER_MIN };
102    const HAS_UNORDERED: bool = E::HAS_UNORDERED;
103    const IS_FLOAT: bool = E::IS_FLOAT;
104
105    #[inline(always)] fn from_i8(value: i8) -> Self { Self::real(E::from_i8(value)) }
106    #[inline(always)] fn from_u8(value: u8) -> Self { Self::real(E::from_u8(value)) }
107    #[inline(always)] fn from_u16(value: u16) -> Self { Self::real(E::from_u16(value)) }
108}
109
110impl<E: RealValue + FloatElement> Complex<E> {
111    /// The modulus `$|z|$` of a complex element, the vector math library not being
112    /// available at the element level.
113    #[inline(always)]
114    fn elem_modulus(self) -> E {
115        E::sqrt(self.re.mul_add(self.re, self.im * self.im))
116    }
117}
118
119impl<E: RealValue + FloatElement> SignedElement for Complex<E> {
120    /// The modulus `$|z|$`, as a real complex number.
121    #[inline(always)]
122    fn abs(self) -> Self {
123        Self::real(self.elem_modulus())
124    }
125
126    /// `$z/|z|$`, the unit complex number along `z`, and zero at the origin.
127    #[inline(always)]
128    fn signum(self) -> Self {
129        let m = self.elem_modulus();
130
131        if m == E::ZERO {
132            return Self::ZERO;
133        }
134
135        Self::new(self.re / m, self.im / m)
136    }
137}
138
139/// Splats a compile-time integer constant as a real `Complex<E>`.
140pub struct ComplexIntConst<E, const VAL: LargeInt>(PhantomData<E>);
141
142/// Splats a compile-time rational constant `N/D` as a real `Complex<E>`.
143pub struct ComplexRatioConst<E, const NUM: LargeInt, const DEN: LargeInt>(PhantomData<E>);
144
145impl<E: RealValue + FloatElement, const VAL: LargeInt> SplatConst<Complex<E>> for ComplexIntConst<E, VAL> {
146    const VALUE: Complex<E> = Complex::real(<E::ConstInt<VAL> as SplatConst<E>>::VALUE);
147}
148
149impl<E: RealValue + FloatElement, const NUM: LargeInt, const DEN: LargeInt> SplatConst<Complex<E>>
150    for ComplexRatioConst<E, NUM, DEN>
151{
152    const VALUE: Complex<E> = Complex::real(<E::ConstRatio<NUM, DEN> as SplatConst<E>>::VALUE);
153}
154
155#[rustfmt::skip]
156impl<E: RealValue + FloatElement> FloatElement for Complex<E> {
157    /// The principal square root, in Kahan's form; see `FloatVector::sqrt` below for
158    /// why the symmetric formula is unusable.
159    #[inline(always)]
160    fn sqrt(this: Self) -> Self {
161        let half = E::from_ratio(1, 2);
162
163        let t = E::sqrt((SignedElement::abs(this.re) + this.elem_modulus()) * half);
164
165        if t == E::ZERO {
166            return Self::ZERO;
167        }
168
169        let half_im = this.im * half;
170
171        if this.re >= E::ZERO {
172            Self::new(t, half_im / t)
173        } else {
174            let i = if this.im < E::ZERO { -t } else { t };
175
176            Self::new(SignedElement::abs(half_im) / t, i)
177        }
178    }
179
180    // Rounding is componentwise; see the crate docs.
181    #[inline(always)] fn floor(this: Self) -> Self { Self::new(E::floor(this.re), E::floor(this.im)) }
182    #[inline(always)] fn ceil(this: Self) -> Self { Self::new(E::ceil(this.re), E::ceil(this.im)) }
183    #[inline(always)] fn round(this: Self) -> Self { Self::new(E::round(this.re), E::round(this.im)) }
184    #[inline(always)] fn trunc(this: Self) -> Self { Self::new(E::trunc(this.re), E::trunc(this.im)) }
185
186    #[inline(always)] fn next_up(this: Self) -> Self { Self::new(E::next_up(this.re), E::next_up(this.im)) }
187    #[inline(always)] fn next_down(this: Self) -> Self { Self::new(E::next_down(this.re), E::next_down(this.im)) }
188
189    #[inline(always)]
190    fn try_from_int(value: LargeInt) -> Option<Self> {
191        E::try_from_int(value).map(Self::real)
192    }
193
194    #[inline(always)]
195    fn try_from_ratio(n: LargeInt, d: LargeInt) -> Option<Self> {
196        E::try_from_ratio(n, d).map(Self::real)
197    }
198
199    const HAS_INFINITY: bool = E::HAS_INFINITY;
200    const HAS_SIGNED_ZERO: bool = E::HAS_SIGNED_ZERO;
201    const HAS_SUBNORMALS: bool = E::HAS_SUBNORMALS;
202
203    type ConstInt<const VAL: LargeInt> = ComplexIntConst<E, VAL>;
204    type ConstRatio<const NUM: LargeInt, const DEN: LargeInt> = ComplexRatioConst<E, NUM, DEN>;
205}
206
207// --- HasIsa / Selectable / Interleave ---
208
209impl<V: thermite::simd::HasIsa> thermite::simd::HasIsa for Complex<V> {
210    type Native = V::Native;
211
212    const ISA: thermite::isa::InstructionSet = V::ISA;
213}
214
215impl<V: RealFloatVector> GenericSelectable for Complex<V> {
216    type SelectableMask = <V as GenericSelectable>::SelectableMask;
217
218    #[inline(always)]
219    fn select<M>(mask: M, t: Self, f: Self) -> Self
220    where
221        Self::SelectableMask: CastMask<M>,
222    {
223        let mask = <Self::SelectableMask as CastMask<M>>::mask_from(mask);
224
225        Self::new(mask.select(t.re, f.re), mask.select(t.im, f.im))
226    }
227}
228
229#[rustfmt::skip]
230impl<V: RealFloatVector> Interleave for Complex<V> {
231    #[inline(always)]
232    fn interleave(self, other: Self) -> (Self, Self) {
233        let (re_lo, re_hi) = self.re.interleave(other.re);
234        let (im_lo, im_hi) = self.im.interleave(other.im);
235
236        (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
237    }
238
239    #[inline(always)]
240    fn deinterleave(self, other: Self) -> (Self, Self) {
241        let (re_lo, re_hi) = self.re.deinterleave(other.re);
242        let (im_lo, im_hi) = self.im.deinterleave(other.im);
243
244        (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
245    }
246}
247
248// --- Compile-time splat / new machinery ---
249
250/// Carriers extracting one component of a `Complex` element constant.
251struct ComplexReSplat<E, V>(PhantomData<(E, V)>);
252struct ComplexImSplat<E, V>(PhantomData<(E, V)>);
253
254impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexReSplat<E, V>
255where
256    E: SplatConst<Complex<V::Element>>,
257{
258    const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.re;
259}
260
261impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexImSplat<E, V>
262where
263    E: SplatConst<Complex<V::Element>>,
264{
265    const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.im;
266}
267
268impl<V: RealFloatVector> SplatVector<Complex<V::Element>> for Complex<V> {
269    type Splat<T: SplatConst<Complex<V::Element>>> = Self;
270}
271
272impl<V: RealFloatVector, E: SplatConst<Complex<V::Element>>> VectorValue<E, Complex<V>> for Complex<V> {
273    const VALUE: Complex<V> = Complex {
274        re: const_splat::<V, ComplexReSplat<E, V>>(),
275        im: const_splat::<V, ComplexImSplat<E, V>>(),
276    };
277}
278
279/// Carriers extracting the per-lane components of a `Complex` element array.
280struct ComplexReNew<C, V>(PhantomData<(C, V)>);
281struct ComplexImNew<C, V>(PhantomData<(C, V)>);
282
283macro_rules! impl_new_const {
284    ($($carrier:ident => $field:ident),* $(,)?) => {$(
285        impl<C, V: RealFloatVector> NewConst<V::Element, V::Lanes> for $carrier<C, V>
286        where
287            C: NewConst<Complex<V::Element>, V::Lanes>,
288        {
289            const VALUES: GenericArray<V::Element, V::Lanes> = const {
290                let c_vals = C::VALUES;
291                let src = c_vals.as_slice();
292                let mut out: GenericArray<V::Element, V::Lanes> = unsafe { core::mem::zeroed() };
293                let dst = out.as_mut_slice();
294                let mut i = 0;
295                while i < V::LANES {
296                    dst[i] = src[i].$field;
297                    i += 1;
298                }
299                core::mem::forget(c_vals);
300                out
301            };
302        }
303    )*};
304}
305
306impl_new_const!(ComplexReNew => re, ComplexImNew => im);
307
308/// `VectorValue` implementor for per-lane (`new`) construction of `Complex` vectors.
309pub struct ComplexNewImpl;
310
311impl<T, V: RealFloatVector> VectorValue<T, Complex<V>> for ComplexNewImpl
312where
313    T: NewConst<Complex<V::Element>, V::Lanes>,
314{
315    const VALUE: Complex<V> = Complex {
316        re: const_new::<V, V::Lanes, ComplexReNew<T, V>>(),
317        im: const_new::<V, V::Lanes, ComplexImNew<T, V>>(),
318    };
319}
320
321impl<V: RealFloatVector> NewVector<Complex<V::Element>, V::Lanes> for Complex<V> {
322    type New<T: NewConst<Complex<V::Element>, V::Lanes>> = ComplexNewImpl;
323}
324
325// --- CastVector ---
326
327impl<FROM, TO> CastVector<Complex<FROM>> for Complex<TO>
328where
329    FROM: RealFloatVector + CastVector<TO>,
330    TO: RealFloatVector + CastVector<FROM>,
331{
332    #[inline(always)]
333    fn cast_into(self) -> Complex<FROM> {
334        Complex::<FROM>::cast_from(self)
335    }
336
337    #[inline(always)]
338    fn cast_from(from: Complex<FROM>) -> Self {
339        Complex::new(TO::cast_from(from.re), TO::cast_from(from.im))
340    }
341}
342
343// --- ComplexVector ---
344
345#[rustfmt::skip]
346impl<V: RealFloatVector> crate::math::specialized::ComplexVector for Complex<V> {
347    type Real = V;
348
349    #[inline(always)] fn re(self) -> V { self.re }
350    #[inline(always)] fn im(self) -> V { self.im }
351    #[inline(always)] fn from_parts(re: V, im: V) -> Self { Complex::new(re, im) }
352
353    #[inline(always)]
354    unsafe fn store_streaming_block(self, ptr: *mut Self) {
355        // Stream each half with the real per-vector NT store (`_mm256_stream_ps` on AVX2; a
356        // plain store where the backend has no NT). `&raw mut (*ptr).re/.im` are the true field
357        // addresses, so no `repr` assumption; a `[Self]` slot is `Self`-aligned, `re` sits at
358        // offset 0 and `im` at `size_of::<V>()`, both aligned for the vector NT store. Preserves
359        // the planar `[re | im]` block layout (NO interleave - see the trait doc).
360        unsafe {
361            self.re.store_streaming((&raw mut (*ptr).re).cast());
362            self.im.store_streaming((&raw mut (*ptr).im).cast());
363        }
364    }
365
366    // Policy-free, so these are inherent on Complex<V> as well, where they also
367    // serve the element-level Complex<f32>. The trait methods forward.
368    #[inline(always)] fn conj(self) -> Self { Complex::conj(self) }
369    #[inline(always)] fn norm_sqr(self) -> V { Complex::norm_sqr(self) }
370    #[inline(always)] fn inv(self) -> Self { Complex::inv(self) }
371
372    #[inline(always)] fn norm_l1(self) -> V { self.re.abs() + self.im.abs() }
373}
374
375// --- GenericVector ---
376
377/// Applies one of the inner vector's radix (de)interleaves to both components.
378///
379/// The four `*_radix*` members differ only in which inner routine they call, so they
380/// share this: split the planar parts out, permute each, weave them back. `f` is a
381/// monomorphized fn item (`V::interleave_radix::<N>` and friends), not a closure over
382/// runtime state, so nothing survives inlining but the permutation itself.
383#[inline(always)]
384fn radix_per_component<V: RealFloatVector, const N: usize>(
385    inputs: [Complex<V>; N],
386    f: impl Fn([V; N]) -> [V; N],
387) -> [Complex<V>; N] {
388    let (mut re, mut im) = ([V::EMPTY; N], [V::EMPTY; N]);
389
390    let mut i = 0;
391    while i < N {
392        re[i] = inputs[i].re;
393        im[i] = inputs[i].im;
394        i += 1;
395    }
396
397    let re = f(re);
398    let im = f(im);
399
400    let mut out = [Complex::<V>::EMPTY; N];
401
402    let mut i = 0;
403    while i < N {
404        out[i] = Complex::new(re[i], im[i]);
405        i += 1;
406    }
407
408    out
409}
410
411impl<V: RealFloatVector> Complex<V> {
412    /// Splat a real and imaginary part across every lane.
413    #[inline(always)]
414    pub fn splat_parts(re: V::Element, im: V::Element) -> Self {
415        Self::new(V::splat(re), V::splat(im))
416    }
417}
418
419impl<V: RealFloatVector> GenericVector for Complex<V> {
420    type Element = Complex<V::Element>;
421
422    const EMPTY: Self = Self::ZERO;
423    const LANES: usize = V::LANES;
424
425    type Lanes = V::Lanes;
426
427    type Unsigned = V::Unsigned;
428    type Signed = V::Signed;
429    type Mask = V::Mask;
430
431    #[inline(always)]
432    fn new<const N: usize>(value: [Self::Element; N]) -> Self
433    where
434        Const<N>: IntoArrayLength<ArrayLength = Self::Lanes>,
435    {
436        let mut re = [<V::Element as Element>::ZERO; N];
437        let mut im = [<V::Element as Element>::ZERO; N];
438
439        let mut i = 0;
440        while i < N {
441            re[i] = value[i].re;
442            im[i] = value[i].im;
443            i += 1;
444        }
445
446        Complex::new(V::new(re), V::new(im))
447    }
448
449    #[inline(always)]
450    fn into_array(self) -> GenericArray<Self::Element, Self::Lanes> {
451        let mut arr = GenericArray::default();
452
453        for i in 0..Self::LANES {
454            arr[i] = Complex::new(self.re.extractv(i), self.im.extractv(i));
455        }
456
457        arr
458    }
459
460    #[inline(always)]
461    fn splat(value: Self::Element) -> Self {
462        Self::splat_parts(value.re, value.im)
463    }
464
465    #[inline(always)]
466    fn single(value: Self::Element) -> Self {
467        Complex::new(V::single(value.re), V::single(value.im))
468    }
469
470    // A Complex element only guarantees the alignment of one V::Element. The aligned
471    // load/store have no fast path to take and forward to the unaligned one.
472    #[inline(always)]
473    unsafe fn load(ptr: *const Self::Element) -> Self {
474        unsafe { Self::load_unaligned(ptr) }
475    }
476
477    /// A `Complex` element is `#[repr(C)]` over two floats, so one element is
478    /// [`load_deinterleaved::<1>`](Self::load_deinterleaved), which routes through the
479    /// inner vector's register engine and not a lane-by-lane loop.
480    #[inline(always)]
481    unsafe fn load_unaligned(ptr: *const Self::Element) -> Self {
482        let [out] = unsafe { Self::load_deinterleaved::<1>(ptr) };
483        out
484    }
485
486    #[inline(always)]
487    unsafe fn load_streaming(ptr: *const Self::Element) -> Self {
488        unsafe { Self::load(ptr) }
489    }
490
491    #[inline(always)]
492    fn interleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
493        let (re_lo, re_hi) = self.re.interleave_by::<GROUP>(other.re);
494        let (im_lo, im_hi) = self.im.interleave_by::<GROUP>(other.im);
495        (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
496    }
497
498    #[inline(always)]
499    fn deinterleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
500        let (re_lo, re_hi) = self.re.deinterleave_by::<GROUP>(other.re);
501        let (im_lo, im_hi) = self.im.deinterleave_by::<GROUP>(other.im);
502        (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
503    }
504
505    #[inline(always)]
506    fn interleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
507        radix_per_component::<V, N>(inputs, V::interleave_radix::<N>)
508    }
509
510    #[inline(always)]
511    fn deinterleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
512        radix_per_component::<V, N>(inputs, V::deinterleave_radix::<N>)
513    }
514
515    #[inline(always)]
516    fn deinterleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
517        radix_per_component::<V, N>(inputs, V::deinterleave_radix_by::<N, GROUP>)
518    }
519
520    #[inline(always)]
521    fn interleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
522        radix_per_component::<V, N>(inputs, V::interleave_radix_by::<N, GROUP>)
523    }
524
525    /// `M` interleaved `Complex` streams are `2 * M` interleaved float streams, i.e.
526    /// a grouped problem with `TAIL = 1` (see [`StreamGroup`]). `M` goes straight to
527    /// the inner vector's [`GenericVector::load_deinterleaved_grouped`] - a NEON
528    /// `LD2`/`LD4`, or a shuffle network on x86 - for any `M`.
529    #[inline(always)]
530    unsafe fn load_deinterleaved<const M: usize>(ptr: *const Self::Element) -> [Self; M] {
531        let groups = unsafe { V::load_deinterleaved_grouped::<M, 1>(ptr as *const V::Element) };
532
533        let mut out = [Self::EMPTY; M];
534        let mut j = 0;
535        while j < M {
536            out[j] = Complex::new(groups[j].head, groups[j].tail[0]);
537            j += 1;
538        }
539        out
540    }
541
542    /// The exact inverse of [`load_deinterleaved`](Self::load_deinterleaved).
543    #[inline(always)]
544    unsafe fn store_interleaved<const M: usize>(ptr: *mut Self::Element, values: [Self; M]) {
545        let mut groups = [StreamGroup {
546            head: V::ZERO,
547            tail: [V::ZERO; 1],
548        }; M];
549
550        let mut j = 0;
551        while j < M {
552            groups[j] = StreamGroup {
553                head: values[j].re,
554                tail: [values[j].im],
555            };
556            j += 1;
557        }
558
559        unsafe { V::store_interleaved_grouped::<M, 1>(ptr as *mut V::Element, groups) }
560    }
561
562    #[inline(always)]
563    unsafe fn load_m(src: Self, mask: Self::Mask, ptr: *const Self::Element) -> Self {
564        let ptr = ptr as *const V::Element;
565
566        // A Complex lane is two consecutive floats, so lane i of the mask covers
567        // positions 2i and 2i+1 of the interleaved layout.
568        let (a_mask, b_mask) = mask.interleave(mask);
569        let (src_a, src_b) = src.re.interleave(src.im);
570
571        let a = unsafe { V::load_m(src_a, a_mask, ptr) };
572        let b = unsafe { V::load_m(src_b, b_mask, ptr.add(V::LANES)) };
573
574        let (re, im) = a.deinterleave(b);
575        Complex::new(re, im)
576    }
577
578    #[inline(always)]
579    unsafe fn load_z(mask: Self::Mask, ptr: *const Self::Element) -> Self {
580        unsafe { Self::load_m(Self::EMPTY, mask, ptr) }
581    }
582
583    // See the note on `load` above: the aligned store forwards to unaligned.
584    #[inline(always)]
585    unsafe fn store(self, ptr: *mut Self::Element) {
586        unsafe { self.store_unaligned(ptr) }
587    }
588
589    #[inline(always)]
590    unsafe fn store_unaligned(self, ptr: *mut Self::Element) {
591        unsafe { Self::store_interleaved::<1>(ptr, [self]) }
592    }
593
594    #[inline(always)]
595    unsafe fn store_streaming(self, ptr: *mut Self::Element) {
596        unsafe { self.store(ptr) }
597    }
598
599    #[inline(always)]
600    unsafe fn store_masked(self, mask: Self::Mask, ptr: *mut Self::Element) {
601        let ptr = ptr as *mut V::Element;
602
603        let (a_mask, b_mask) = mask.interleave(mask);
604        let (a, b) = self.re.interleave(self.im);
605
606        unsafe {
607            a.store_masked(a_mask, ptr);
608            b.store_masked(b_mask, ptr.add(V::LANES));
609        }
610    }
611
612    #[inline(always)]
613    unsafe fn lookup_unchecked(values: &[Self::Element], indices: Self::Unsigned) -> Self {
614        // The table reinterpreted as its interleaved [re, im, re, im, ...] floats:
615        // element i lives at float positions 2i and 2i+1, so each component is one
616        // gather through the inner vector's engine.
617        let floats = unsafe { core::slice::from_raw_parts(values.as_ptr() as *const V::Element, values.len() * 2) };
618
619        let re_idx = indices << 1;
620        let im_idx = re_idx + Self::Unsigned::ONE;
621
622        let re = unsafe { V::lookup_unchecked(floats, re_idx) };
623        let im = unsafe { V::lookup_unchecked(floats, im_idx) };
624
625        Complex::new(re, im)
626    }
627
628    #[inline(always)]
629    fn broadcast<const I: usize>(self) -> Self {
630        Complex::new(V::broadcast::<I>(self.re), V::broadcast::<I>(self.im))
631    }
632
633    #[inline(always)]
634    fn broadcastv(self, idx: usize) -> Self {
635        Complex::new(self.re.broadcastv(idx), self.im.broadcastv(idx))
636    }
637
638    #[inline(always)]
639    fn extract<const I: usize>(self) -> Self::Element {
640        Complex::new(V::extract::<I>(self.re), V::extract::<I>(self.im))
641    }
642
643    #[inline(always)]
644    fn extractv(self, idx: usize) -> Self::Element {
645        Complex::new(self.re.extractv(idx), self.im.extractv(idx))
646    }
647
648    #[inline(always)]
649    fn insert<const I: usize>(self, value: Self::Element) -> Self {
650        Complex::new(V::insert::<I>(self.re, value.re), V::insert::<I>(self.im, value.im))
651    }
652
653    #[inline(always)]
654    fn insertv(self, idx: usize, value: Self::Element) -> Self {
655        Complex::new(self.re.insertv(idx, value.re), self.im.insertv(idx, value.im))
656    }
657
658    #[inline(always)]
659    fn reverse(self) -> Self {
660        Complex::new(self.re.reverse(), self.im.reverse())
661    }
662
663    #[inline(always)]
664    fn swap_bytes(self) -> Self {
665        Complex::new(self.re.swap_bytes(), self.im.swap_bytes())
666    }
667
668    #[inline(always)]
669    fn zz(self, mask: Self::Mask) -> Self {
670        Complex::new(self.re.zz(mask), self.im.zz(mask))
671    }
672
673    #[inline(always)]
674    fn nz(self, mask: Self::Mask) -> Self {
675        Complex::new(self.re.nz(mask), self.im.nz(mask))
676    }
677
678    #[inline(always)]
679    fn compress(self, mask: Self::Mask) -> Self {
680        Complex::new(self.re.compress(mask), self.im.compress(mask))
681    }
682
683    #[inline(always)]
684    fn compress_z(self, mask: Self::Mask) -> Self {
685        Complex::new(self.re.compress_z(mask), self.im.compress_z(mask))
686    }
687
688    // Pure lane movement driven by `mask` alone, so both parts take the same
689    // permutation and no lane ends up with a re/im pair from different sources.
690    // `compress_m` too: its keep-lanes come from the population count of the shared
691    // mask, so they land at the same positions in each part.
692    #[inline(always)]
693    fn compress_m(self, src: Self, mask: Self::Mask) -> Self {
694        Complex::new(self.re.compress_m(src.re, mask), self.im.compress_m(src.im, mask))
695    }
696
697    #[inline(always)]
698    fn expand(self, mask: Self::Mask) -> Self {
699        Complex::new(self.re.expand(mask), self.im.expand(mask))
700    }
701
702    #[inline(always)]
703    fn expand_z(self, mask: Self::Mask) -> Self {
704        Complex::new(self.re.expand_z(mask), self.im.expand_z(mask))
705    }
706
707    #[inline(always)]
708    fn expand_m(self, src: Self, mask: Self::Mask) -> Self {
709        Complex::new(self.re.expand_m(src.re, mask), self.im.expand_m(src.im, mask))
710    }
711
712    #[inline(always)]
713    fn align<const OFFSET: usize>(self, other: Self) -> Self {
714        Complex::new(self.re.align::<OFFSET>(other.re), self.im.align::<OFFSET>(other.im))
715    }
716
717    // Both parts align through `V`, so this is only as native as `V` is.
718    const HAS_NATIVE_ALIGN: bool = V::HAS_NATIVE_ALIGN;
719
720    #[inline(always)]
721    fn map<F>(mut self, f: F) -> Self
722    where
723        F: Fn(Self::Element) -> Self::Element,
724    {
725        for i in 0..Self::LANES {
726            self = self.insertv(i, f(self.extractv(i)));
727        }
728        self
729    }
730
731    #[inline(always)]
732    fn fold<F>(self, mut init: Self::Element, f: F) -> Self::Element
733    where
734        F: Fn(Self::Element, Self::Element) -> Self::Element,
735    {
736        for i in 0..Self::LANES {
737            init = f(init, self.extractv(i));
738        }
739        init
740    }
741
742    #[inline(always)]
743    fn reduce<F>(self, f: F) -> Self::Element
744    where
745        F: Fn(Self::Element, Self::Element) -> Self::Element,
746    {
747        let mut result = self.extractv(0);
748        for i in 1..Self::LANES {
749            result = f(result, self.extractv(i));
750        }
751        result
752    }
753
754    #[rustfmt::skip]
755    #[inline(always)]    fn splat_m(src: Self, mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), src) }
756    #[rustfmt::skip]
757    #[inline(always)]    fn splat_z(mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), Self::EMPTY) }
758    #[rustfmt::skip]
759    #[inline(always)]    fn broadcast_c<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), self) }
760    #[rustfmt::skip]
761    #[inline(always)]    fn broadcast_m<const I: usize>(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), src) }
762    #[rustfmt::skip]
763    #[inline(always)]    fn broadcast_z<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), Self::EMPTY) }
764    #[rustfmt::skip]
765    #[inline(always)]    fn broadcastv_c(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), self) }
766    #[rustfmt::skip]
767    #[inline(always)]    fn broadcastv_m(self, src: Self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), src) }
768    #[rustfmt::skip]
769    #[inline(always)]    fn broadcastv_z(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), Self::EMPTY) }
770    #[rustfmt::skip]
771    #[inline(always)]    fn reverse_c(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), self) }
772    #[rustfmt::skip]
773    #[inline(always)]    fn reverse_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.reverse(), src) }
774    #[rustfmt::skip]
775    #[inline(always)]    fn reverse_z(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), Self::EMPTY) }
776    #[rustfmt::skip]
777    #[inline(always)]    fn swap_bytes_c(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), self) }
778    #[rustfmt::skip]
779    #[inline(always)]    fn swap_bytes_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), src) }
780    #[rustfmt::skip]
781    #[inline(always)]    fn swap_bytes_z(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), Self::EMPTY) }
782}
783
784// --- PartialOrdVector: lexicographic by (re, im) ---
785
786#[rustfmt::skip]
787impl<V: RealFloatVector> PartialOrdVector for Complex<V> {
788    #[inline(always)]
789    fn cmp_eq(self, other: Self) -> Self::Mask {
790        self.re.cmp_eq(other.re) & self.im.cmp_eq(other.im)
791    }
792
793    #[inline(always)]
794    fn cmp_ne(self, other: Self) -> Self::Mask {
795        self.re.cmp_ne(other.re) | self.im.cmp_ne(other.im)
796    }
797
798    // (re < other.re) | (re == other.re & im <op> other.im), in one ternlog.
799    #[inline(always)]
800    fn cmp_lt(self, other: Self) -> Self::Mask {
801        let re_lt = self.re.cmp_lt(other.re);
802        let re_eq = self.re.cmp_eq(other.re);
803        let im_lt = self.im.cmp_lt(other.im);
804
805        GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_lt)
806    }
807
808    #[inline(always)]
809    fn cmp_gt(self, other: Self) -> Self::Mask {
810        let re_gt = self.re.cmp_gt(other.re);
811        let re_eq = self.re.cmp_eq(other.re);
812        let im_gt = self.im.cmp_gt(other.im);
813
814        GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_gt)
815    }
816
817    #[inline(always)]
818    fn cmp_le(self, other: Self) -> Self::Mask {
819        let re_lt = self.re.cmp_lt(other.re);
820        let re_eq = self.re.cmp_eq(other.re);
821        let im_le = self.im.cmp_le(other.im);
822
823        GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_le)
824    }
825
826    #[inline(always)]
827    fn cmp_ge(self, other: Self) -> Self::Mask {
828        let re_gt = self.re.cmp_gt(other.re);
829        let re_eq = self.re.cmp_eq(other.re);
830        let im_ge = self.im.cmp_ge(other.im);
831
832        GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_ge)
833    }
834}
835
836// --- Masked ops ---
837
838macro_rules! impl_masked {
839    (MUL_ADD: $($method:ident),*) => {paste::paste! {
840        impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddExtMasked<V::Mask, A, B> for Complex<V>
841        where
842            Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
843        {
844            $(
845                #[inline(always)]
846                fn [<$method _c>](self, mask: V::Mask, a: A, b: B) -> Self {
847                    mask.select(self.$method(a, b), self)
848                }
849                #[inline(always)]
850                fn [<$method _m>](self, src: Self, mask: V::Mask, a: A, b: B) -> Self {
851                    mask.select(self.$method(a, b), src)
852                }
853                #[inline(always)]
854                fn [<$method _z>](self, mask: V::Mask, a: A, b: B) -> Self {
855                    mask.select(self.$method(a, b), Self::EMPTY)
856                }
857            )*
858        }
859
860        impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddAssignExtMasked<V::Mask, A, B> for Complex<V>
861        where
862            Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
863        {
864            $(
865                #[inline(always)]
866                fn [<$method _assign_c>](&mut self, mask: V::Mask, a: A, b: B) {
867                    *self = mask.select(self.$method(a, b), *self);
868                }
869                #[inline(always)]
870                fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, a: A, b: B) {
871                    *self = mask.select(self.$method(a, b), src);
872                }
873                #[inline(always)]
874                fn [<$method _assign_z>](&mut self, mask: V::Mask, a: A, b: B) {
875                    *self = mask.select(self.$method(a, b), Self::EMPTY);
876                }
877            )*
878        }
879    }};
880
881    ($trait:ident::$method:ident) => {paste::paste! {
882        impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait Masked>]<V::Mask, Rhs> for Complex<V>
883        where
884            Complex<V>: core::ops::$trait<Rhs, Output = Self>,
885        {
886            #[inline(always)]
887            fn [<$method _c>](self, mask: V::Mask, rhs: Rhs) -> Self {
888                mask.select(self.$method(rhs), self)
889            }
890            #[inline(always)]
891            fn [<$method _m>](self, src: Self, mask: V::Mask, rhs: Rhs) -> Self {
892                mask.select(self.$method(rhs), src)
893            }
894            #[inline(always)]
895            fn [<$method _z>](self, mask: V::Mask, rhs: Rhs) -> Self {
896                mask.select(self.$method(rhs), Self::EMPTY)
897            }
898        }
899
900        impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait AssignMasked>]<V::Mask, Rhs> for Complex<V>
901        where
902            Complex<V>: core::ops::$trait<Rhs, Output = Self>,
903        {
904            #[inline(always)]
905            fn [<$method _assign_c>](&mut self, mask: V::Mask, rhs: Rhs) {
906                *self = mask.select(self.$method(rhs), *self);
907            }
908            #[inline(always)]
909            fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, rhs: Rhs) {
910                *self = mask.select(self.$method(rhs), src);
911            }
912            #[inline(always)]
913            fn [<$method _assign_z>](&mut self, mask: V::Mask, rhs: Rhs) {
914                *self = mask.select(self.$method(rhs), Self::EMPTY);
915            }
916        }
917    }};
918}
919
920impl_masked!(MUL_ADD: mul_add, mul_sub, nmul_add, nmul_sub, mul_adde, mul_sube, nmul_adde, nmul_sube);
921impl_masked!(Add::add);
922impl_masked!(Sub::sub);
923impl_masked!(Mul::mul);
924impl_masked!(Div::div);
925impl_masked!(Rem::rem);
926
927// =====================================================================================
928// Lane-alternating add/sub (`AddSubExt`), over the *inner vector's* lanes - so the
929// even/odd parity applies per complex number. `Add`/`Sub` are component-wise on
930// `re`/`im`, so this is exact: `neg_even` flips the even-lane signs of both parts,
931// then:
932//   addsub(a, b)      = a + neg_even(b)
933//   fmaddsub(a, b, c) = a*b + neg_even(c)   (via the complex product-rule mul_adde)
934//   fmsubadd(a, b, c) = a*b - neg_even(c)
935// =====================================================================================
936
937#[inline(always)]
938fn neg_even_complex<V: RealFloatVector>(x: Complex<V>) -> Complex<V> {
939    // `addsub(0, w) = [-w0, w1, -w2, ...]` flips the even lanes exactly.
940    Complex::new(V::ZERO.addsub(x.re), V::ZERO.addsub(x.im))
941}
942
943impl<V: RealFloatVector> AddSubExt for Complex<V> {
944    type Output = Self;
945
946    #[inline(always)]
947    fn addsub(self, b: Self) -> Self {
948        self + neg_even_complex(b)
949    }
950    #[inline(always)]
951    fn fmaddsub(self, b: Self, c: Self) -> Self {
952        self.mul_adde(b, neg_even_complex(c))
953    }
954    #[inline(always)]
955    fn fmsubadd(self, b: Self, c: Self) -> Self {
956        self.mul_sube(b, neg_even_complex(c))
957    }
958}
959
960impl<V: RealFloatVector> AddSubExtMasked<V::Mask> for Complex<V> {
961    #[inline(always)]
962    fn addsub_c(self, mask: V::Mask, b: Self) -> Self {
963        mask.select(self.addsub(b), self)
964    }
965    #[inline(always)]
966    fn addsub_m(self, src: Self, mask: V::Mask, b: Self) -> Self {
967        mask.select(self.addsub(b), src)
968    }
969    #[inline(always)]
970    fn addsub_z(self, mask: V::Mask, b: Self) -> Self {
971        mask.select(self.addsub(b), Self::EMPTY)
972    }
973
974    #[inline(always)]
975    fn fmaddsub_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
976        mask.select(self.fmaddsub(b, c), self)
977    }
978    #[inline(always)]
979    fn fmaddsub_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
980        mask.select(self.fmaddsub(b, c), src)
981    }
982    #[inline(always)]
983    fn fmaddsub_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
984        mask.select(self.fmaddsub(b, c), Self::EMPTY)
985    }
986
987    #[inline(always)]
988    fn fmsubadd_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
989        mask.select(self.fmsubadd(b, c), self)
990    }
991    #[inline(always)]
992    fn fmsubadd_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
993        mask.select(self.fmsubadd(b, c), src)
994    }
995    #[inline(always)]
996    fn fmsubadd_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
997        mask.select(self.fmsubadd(b, c), Self::EMPTY)
998    }
999}
1000
1001impl<V: RealFloatVector> SquareMasked<V::Mask> for Complex<V> {
1002    #[inline(always)]
1003    fn square_c(self, mask: V::Mask) -> Self::Output {
1004        mask.select(self.square(), self)
1005    }
1006
1007    #[inline(always)]
1008    fn square_m(self, src: Self, mask: V::Mask) -> Self::Output {
1009        mask.select(self.square(), src)
1010    }
1011
1012    #[inline(always)]
1013    fn square_z(self, mask: V::Mask) -> Self::Output {
1014        mask.select(self.square(), Self::ZERO)
1015    }
1016}
1017
1018// The _c/_m/_z variants of the unary (fn m(self) -> Self) and binary
1019// (fn m(self, Self) -> Self) ops, as plain blends, like impl_masked! above.
1020macro_rules! complex_masked {
1021    (unary: $($m:ident),* $(,)?) => { paste::paste! {
1022        $(
1023            #[inline(always)] fn [<$m _c>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), self) }
1024            #[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask) -> Self { mask.select(self.$m(), src) }
1025            #[inline(always)] fn [<$m _z>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), Self::ZERO) }
1026        )*
1027    }};
1028    (binary: $($m:ident),* $(,)?) => { paste::paste! {
1029        $(
1030            #[inline(always)] fn [<$m _c>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), self) }
1031            #[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), src) }
1032            #[inline(always)] fn [<$m _z>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), Self::ZERO) }
1033        )*
1034    }};
1035}
1036
1037// --- NumericVector ---
1038
1039#[rustfmt::skip]
1040impl<V: RealFloatVector> Bounded for Complex<V> {
1041    #[inline(always)] fn min_value() -> Self { Complex::new(V::MIN, V::MIN) }
1042    #[inline(always)] fn max_value() -> Self { Complex::new(V::MAX, V::MAX) }
1043}
1044
1045#[rustfmt::skip]
1046/// The lane-sort key: strictly-before under the lexicographic (re, im)
1047/// order, i.e. `cmp_lt`. See `thermite::sort::SortKey` for why this is a
1048/// static trait method and not a closure.
1049impl<V: RealFloatVector> thermite::sort::SortKey<Self> for Complex<V> {
1050    #[inline(always)]
1051    fn key_lt(a: Self, b: Self) -> V::Mask {
1052        a.cmp_lt(b)
1053    }
1054}
1055
1056/// Scalar insertion walk over whole lanes, for widths past the network ladder.
1057/// Quadratic, like core's `sort_any`; compares composite elements through
1058/// `PartialOrd` (lexicographic, matching the vector comparisons).
1059#[inline(always)]
1060fn sort_lanes_scalar<V: NumericVector, O: thermite::sort::SortOrder>(v: V) -> V
1061where
1062    V::Element: PartialOrd,
1063{
1064    let mut out = v;
1065    let mut i = 1;
1066    while i < V::LANES {
1067        let key = out.extractv(i);
1068        let mut j = i;
1069        while j > 0 {
1070            let prev = out.extractv(j - 1);
1071            let misplaced = if O::IS_ASCENDING { prev > key } else { prev < key };
1072            if !misplaced {
1073                break;
1074            }
1075            out = out.insertv(j, prev);
1076            j -= 1;
1077        }
1078        out = out.insertv(j, key);
1079        i += 1;
1080    }
1081    out
1082}
1083
1084impl<V: RealFloatVector> NumericVector for Complex<V> {
1085    // The integer conversions are real/value-only in both directions: an integer has no
1086    // derivative, no imaginary part and no error term, so converting one in yields a
1087    // constant, and converting out is the value part alone.
1088    #[inline(always)]
1089    fn to_signed_integer(self) -> Self::Signed {
1090        self.re.to_signed_integer()
1091    }
1092
1093    #[inline(always)]
1094    fn from_signed_integer(v: Self::Signed) -> Self {
1095        Self::real(V::from_signed_integer(v))
1096    }
1097
1098    #[inline(always)]
1099    fn to_unsigned_integer(self) -> Self::Unsigned {
1100        self.re.to_unsigned_integer()
1101    }
1102
1103    #[inline(always)]
1104    fn from_unsigned_integer(v: Self::Unsigned) -> Self {
1105        Self::real(V::from_unsigned_integer(v))
1106    }
1107
1108    const ZERO: Self = Complex::new(V::ZERO, V::ZERO);
1109    const ONE: Self = Complex::new(V::ONE, V::ZERO);
1110    const TWO: Self = Self::real(V::TWO);
1111
1112    // The lexicographic extremes, consistent with the ordering above.
1113    const MIN: Self = Complex::new(V::MIN, V::MIN);
1114    const MAX: Self = Complex::new(V::MAX, V::MAX);
1115
1116    #[inline(always)]
1117    fn is_zero(self) -> Self::Mask {
1118        self.re.is_zero() & self.im.is_zero()
1119    }
1120    #[inline(always)]
1121    fn is_all_zero(self) -> bool {
1122        self.re.is_all_zero() && self.im.is_all_zero()
1123    }
1124
1125    // Lane sorts are keyed on the lexicographic (re, im) order - which is exactly
1126    // `cmp_lt` here, so the key IS the comparison. Each compare-exchange derives one
1127    // routing mask from it and moves both components through the same permutation
1128    // and select (`thermite::sort::sort_lanes_by_key`).
1129    #[inline(always)]
1130    fn sort_by<O: thermite::sort::SortOrder>(self) -> Self {
1131        if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
1132            thermite::sort::sort_lanes_by_key::<Self, O, Self>(self)
1133        } else {
1134            sort_lanes_scalar::<Self, O>(self)
1135        }
1136    }
1137
1138    #[inline(always)]
1139    fn bitonic_clean_by<O: thermite::sort::SortOrder>(self) -> Self {
1140        if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
1141            thermite::sort::bitonic_clean_lanes_by_key::<Self, O, Self>(self)
1142        } else {
1143            // A full sort trivially cleans a bitonic input.
1144            sort_lanes_scalar::<Self, O>(self)
1145        }
1146    }
1147
1148    #[inline(always)]
1149    fn min(self, other: Self) -> Self {
1150        self.cmp_lt(other).select(self, other)
1151    }
1152    #[inline(always)]
1153    fn max(self, other: Self) -> Self {
1154        self.cmp_gt(other).select(self, other)
1155    }
1156
1157    // Compare against both bounds once, then blend per component.
1158    #[inline(always)]
1159    fn clamp(self, min: Self, max: Self) -> Self {
1160        let is_lt = self.cmp_lt(min);
1161        let is_gt = self.cmp_gt(max);
1162
1163        Complex::new(
1164            is_lt.select(min.re, is_gt.select(max.re, self.re)),
1165            is_lt.select(min.im, is_gt.select(max.im, self.im)),
1166        )
1167    }
1168
1169    // Lexicographic order spans both components: there is no single inner vector to
1170    // hand to arg_minmax. Reduce the extracted elements in log2(LANES) steps with the
1171    // derived PartialOrd.
1172    #[inline(always)]
1173    fn min_element(self) -> Self::Element {
1174        let mut arr = self.into_array();
1175        reduce_in_place(&mut arr, |a, b| if b < a { b } else { a });
1176        arr[0]
1177    }
1178
1179    #[inline(always)]
1180    fn max_element(self) -> Self::Element {
1181        let mut arr = self.into_array();
1182        reduce_in_place(&mut arr, |a, b| if b > a { b } else { a });
1183        arr[0]
1184    }
1185
1186    #[inline(always)]
1187    fn min_max_element(self) -> (Self::Element, Self::Element) {
1188        (self.min_element(), self.max_element())
1189    }
1190
1191    #[inline(always)]
1192    fn arg_minmax(self) -> (usize, usize) {
1193        let arr = self.into_array();
1194
1195        let (mut lo, mut hi) = (0, 0);
1196
1197        for i in 1..Self::LANES {
1198            if arr[i] < arr[lo] {
1199                lo = i;
1200            }
1201            if arr[i] > arr[hi] {
1202                hi = i;
1203            }
1204        }
1205
1206        (lo, hi)
1207    }
1208
1209    // Sum is linear and commutes with the re/im split: each component reduces through
1210    // the inner vector's horizontal sum. No extracting and folding LANES scalar
1211    // complex numbers.
1212    #[inline(always)]
1213    fn sum_elements(self) -> Self::Element {
1214        Complex::new(self.re.sum_elements(), self.im.sum_elements())
1215    }
1216
1217    // Complex addition is componentwise, so the scan is too - the same argument as
1218    // `sum_elements`, one step at a time instead of all the way down.
1219    #[inline(always)]
1220    fn prefix_sum(self) -> Self {
1221        Complex::new(self.re.prefix_sum(), self.im.prefix_sum())
1222    }
1223
1224    #[inline(always)]
1225    fn reverse_prefix_sum(self) -> Self {
1226        Complex::new(self.re.reverse_prefix_sum(), self.im.reverse_prefix_sum())
1227    }
1228
1229    // min/max are lexicographic over both parts (see the ordering above), so there is
1230    // no per-component scan to delegate to: scanning `re` and `im` separately would
1231    // pair a real part from one lane with an imaginary part from another. The ladder
1232    // runs on whole complex values through `Self::min`/`Self::max`.
1233    #[inline(always)]
1234    fn prefix_min(self) -> Self {
1235        thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::min)
1236    }
1237
1238    #[inline(always)]
1239    fn prefix_max(self) -> Self {
1240        thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::max)
1241    }
1242
1243    #[inline(always)]
1244    fn reverse_prefix_min(self) -> Self {
1245        thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::min)
1246    }
1247
1248    #[inline(always)]
1249    fn reverse_prefix_max(self) -> Self {
1250        thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::max)
1251    }
1252
1253    // Product is not componentwise (the parts cross-multiply) and needs complex
1254    // multiplies across lanes. The tree reduction keeps the dependency chain at log
1255    // depth.
1256    #[inline(always)]
1257    fn prod_elements(self) -> Self::Element {
1258        let mut arr = self.into_array();
1259        reduce_in_place(&mut arr, |a, b| a * b);
1260        arr[0]
1261    }
1262
1263    #[inline(always)]
1264    fn offset() -> Self {
1265        Self::real(V::offset())
1266    }
1267    #[inline(always)]
1268    fn indexed() -> Self {
1269        Self::real(V::indexed())
1270    }
1271
1272    #[inline(always)]
1273    fn scale(self, factor: Self::Element) -> Self {
1274        self * Self::splat(factor)
1275    }
1276
1277    #[inline(always)]
1278    fn scale_c(self, mask: Self::Mask, factor: Self::Element) -> Self {
1279        mask.select(<Self as NumericVector>::scale(self, factor), self)
1280    }
1281    #[inline(always)]
1282    fn scale_m(self, src: Self, mask: Self::Mask, factor: Self::Element) -> Self {
1283        mask.select(<Self as NumericVector>::scale(self, factor), src)
1284    }
1285    #[inline(always)]
1286    fn scale_z(self, mask: Self::Mask, factor: Self::Element) -> Self {
1287        mask.select(<Self as NumericVector>::scale(self, factor), Self::ZERO)
1288    }
1289    complex_masked!(binary: min, max);
1290
1291    // pairwise_sum is a lane rearrange-and-add, and a lane's components sit at the
1292    // same index in re and im, so it applies componentwise.
1293    #[inline(always)]
1294    fn pairwise_sum(lo: Self, hi: Self) -> Self {
1295        Complex::new(V::pairwise_sum(lo.re, hi.re), V::pairwise_sum(lo.im, hi.im))
1296    }
1297
1298    #[inline(always)]
1299    fn relaxed_pairwise_sum(lo: Self, hi: Self) -> Self {
1300        Complex::new(
1301            V::relaxed_pairwise_sum(lo.re, hi.re),
1302            V::relaxed_pairwise_sum(lo.im, hi.im),
1303        )
1304    }
1305}
1306
1307// --- SignedVector ---
1308
1309impl<V: RealFloatVector> NegMasked<V::Mask> for Complex<V> {
1310    #[inline(always)]
1311    fn neg_c(self, mask: V::Mask) -> Self {
1312        Complex::new(self.re.neg_c(mask), self.im.neg_c(mask))
1313    }
1314
1315    #[inline(always)]
1316    fn neg_m(self, src: Self, mask: V::Mask) -> Self {
1317        Complex::new(self.re.neg_m(src.re, mask), self.im.neg_m(src.im, mask))
1318    }
1319
1320    #[inline(always)]
1321    fn neg_z(self, mask: V::Mask) -> Self {
1322        Complex::new(self.re.neg_z(mask), self.im.neg_z(mask))
1323    }
1324}
1325
1326impl<V: RealFloatVector> Complex<V> {
1327    /// The modulus `$|z|$` under the default policy, for the vector-trait methods
1328    /// that take no policy. The policy-aware form is
1329    /// [`norm_p`](crate::math::ComplexMathWithPolicy::norm_p).
1330    #[inline(always)]
1331    pub(crate) fn modulus(self) -> V {
1332        self.re.hypot_p::<DefaultPolicy>(self.im)
1333    }
1334}
1335
1336#[rustfmt::skip]
1337impl<V: RealFloatVector> SignedVector for Complex<V> {
1338    const NEG_ONE: Self = Self::real(V::NEG_ONE);
1339    const MIN_POSITIVE: Self = Self::real(V::MIN_POSITIVE);
1340
1341    /// The modulus `$|z|$`, as a real complex number.
1342    #[inline(always)]
1343    fn abs(self) -> Self {
1344        Self::real(self.modulus())
1345    }
1346
1347    /// `$z/|z|$`, the unit complex number along `z`, and zero at the origin.
1348    ///
1349    /// Preserves the real identity `abs(z) * signum(z) == z`.
1350    #[inline(always)]
1351    fn signum(self) -> Self {
1352        let m = self.modulus();
1353        let is_zero = m.is_zero();
1354        let inv = m.reciprocal_p::<DefaultPolicy>();
1355
1356        // nz() zeroes where the mask is set, pinning signum(0) to 0; the division
1357        // there gives 0 * inf = NaN.
1358        Complex::new((self.re * inv).nz(is_zero), (self.im * inv).nz(is_zero))
1359    }
1360
1361    // Sign-bit ops are componentwise. A mask has one bit per lane, and the sign
1362    // predicates report the sign of the real part. See the crate docs.
1363    #[inline(always)] fn is_positive(self) -> Self::Mask { self.re.is_positive() }
1364    #[inline(always)] fn is_negative(self) -> Self::Mask { self.re.is_negative() }
1365    #[inline(always)] fn select_negative(self, if_neg: Self, if_pos: Self) -> Self { self.is_negative().select(if_neg, if_pos) }
1366
1367    #[inline(always)]
1368    fn copysign(self, sign: Self) -> Self {
1369        Complex::new(self.re.copysign(sign.re), self.im.copysign(sign.im))
1370    }
1371
1372    complex_masked!(unary: abs);
1373    complex_masked!(binary: copysign);
1374}
1375
1376// --- FloatVector ---
1377
1378/// Kahan's principal square root, plus a mask of the lanes that degenerated.
1379///
1380/// Split out so those lanes can re-enter it on a rescaled argument; see
1381/// [`sqrt_rescaled`]. Both `t` and the modulus have to be tested, and which one fires
1382/// depends on how the build handles denormals - see the mask's own comment.
1383#[inline(always)]
1384fn sqrt_kahan<V: RealFloatVector>(z: Complex<V>) -> (Complex<V>, V::Mask) {
1385    let half = <V as FloatVector>::HALF;
1386
1387    let m = z.modulus();
1388
1389    // Halving each term before the sum, not after: `|re| + |z|` overflows for a `re`
1390    // near the top of the range even though `sqrt(z)` is comfortably representable
1391    // there (`sqrt(1.7e308)` gave `inf` for a true 1.3e154). Scaling by a power of two
1392    // is exact, so this is the same value at no extra instruction - the multiply folds
1393    // into the FMA.
1394    let t = z.re.abs().mul_adde(half, m * half).sqrt(); // sqrt((|re| + |z|)/2)
1395    let half_im = z.im * half;
1396
1397    // One quotient serves both branches: t >= 0, so |im/2|/t is |im/2t| exactly.
1398    let q = half_im / t;
1399
1400    // re >= 0: (t, im/2t).   re < 0: (|im|/2t, sign(im)*t).
1401    let re_pos = z.re.is_positive();
1402
1403    // *Which* of the two goes to zero depends on how the build handles denormals, and
1404    // neither implies the other - so the flag picks the test, at one instruction either
1405    // way. Testing only one unconditionally returns silent garbage in the other
1406    // configuration, which is exactly how this was found.
1407    //
1408    // - Flushing (the default): `hypot` drops a subnormal `|z|` to zero while the
1409    //   `|re|` beside it survives, so `m == 0` while `t == sqrt(|re|/2)` is non-zero -
1410    //   and a factor of sqrt(2) below the answer. A normal `|z|` keeps `t >= sqrt(m/2)`
1411    //   positive, so `m` alone decides.
1412    // - Preserving or ignoring: `m` keeps the subnormal, so it is zero only for a true
1413    //   zero - but halving the *smallest* subnormals underflows, leaving `t == 0` with
1414    //   `m != 0`. A true zero sets `t` too, so `t` alone decides.
1415    let degenerate = if const { features::PRESERVE_DENORMALS || features::IGNORE_DENORMALS } {
1416        t.is_zero()
1417    } else {
1418        m.is_zero()
1419    };
1420
1421    (
1422        Complex::new(re_pos.select(t, q.abs()), re_pos.select(q, t.mul_sign(z.im))),
1423        degenerate,
1424    )
1425}
1426
1427/// The degenerate lanes of [`sqrt_kahan`].
1428///
1429/// That means one of two things. Either `z` is genuinely zero, where the root is zero
1430/// and both quotients were `0/0`; or `|z|` is subnormal, and either `hypot` flushed it
1431/// even though the *answer* is an entirely ordinary number - `sqrt(5e-324)` is 2.2e-162,
1432/// a perfectly normal double. Denormal flushing exists to avoid hardware stalls on
1433/// denormal intermediates and results; here there are none, only a denormal input, so
1434/// preserving it costs nothing at runtime and buys back 160 orders of magnitude.
1435///
1436/// Scaling into the normal range recovers it. `1/MIN_POSITIVE` is an exact power of two
1437/// whose exponent is even in both binary formats (2^1022 for f64, 2^126 for f32), so
1438/// `sqrt(MIN_POSITIVE)` is exact as well and undoing the scale introduces no rounding of
1439/// its own.
1440#[inline(always)]
1441fn sqrt_rescaled<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
1442    let up = V::ONE / <V as SignedVector>::MIN_POSITIVE;
1443    let down = <V as SignedVector>::MIN_POSITIVE.sqrt();
1444
1445    let (r, _) = sqrt_kahan(Complex::new(z.re * up, z.im * up));
1446
1447    let is_zero = z.is_zero();
1448
1449    Complex::new((r.re * down).nz(is_zero), (r.im * down).nz(is_zero))
1450}
1451
1452#[rustfmt::skip]
1453impl<V: RealFloatVector> FloatVector for Complex<V> {
1454    const HALF: Self = Self::real(<V as FloatVector>::HALF);
1455    const NEG_ZERO: Self = Self::real(<V as FloatVector>::NEG_ZERO);
1456    const EPSILON: Self = Self::real(<V as FloatVector>::EPSILON);
1457
1458    // Directed along the real axis. A NaN in either component makes the whole
1459    // number NaN.
1460    const INFINITY: Self = Self::real(<V as FloatVector>::INFINITY);
1461    const NEG_INFINITY: Self = Self::real(<V as FloatVector>::NEG_INFINITY);
1462    const NAN: Self = Complex::new(<V as FloatVector>::NAN, <V as FloatVector>::NAN);
1463
1464    type ExtendedPrecision = Self;
1465
1466    // rcp/rsqrt go through the inner vector's reciprocal, and are approximate exactly
1467    // when it is.
1468    const HAS_APPROX_RCP: bool = V::HAS_APPROX_RCP;
1469    const HAS_APPROX_RSQRT: bool = V::HAS_APPROX_RCP;
1470
1471    #[inline(always)] fn is_nan(self) -> Self::Mask { self.re.is_nan() | self.im.is_nan() }
1472    #[inline(always)] fn is_infinite(self) -> Self::Mask { self.re.is_infinite() | self.im.is_infinite() }
1473    #[inline(always)] fn is_finite(self) -> Self::Mask { self.re.is_finite() & self.im.is_finite() }
1474
1475    // |z| is negligible only if both components are, and z is normal if either
1476    // component is and the other is finite.
1477    #[inline(always)] fn is_zero_or_subnormal(self) -> Self::Mask { self.re.is_zero_or_subnormal() & self.im.is_zero_or_subnormal() }
1478    #[inline(always)] fn is_normal(self) -> Self::Mask { (self.re.is_normal() | self.im.is_normal()) & self.is_finite() }
1479    #[inline(always)] fn is_subnormal(self) -> Self::Mask { (self.re.is_subnormal() | self.im.is_subnormal()) & self.is_zero_or_subnormal() }
1480
1481    /// The principal square root, with a branch cut on the negative real axis.
1482    #[inline(always)]
1483    fn sqrt(self) -> Self {
1484        // Kahan's formulation. The symmetric form
1485        //
1486        //     sqrt((|z| + re)/2) + i*sign(im)*sqrt((|z| - re)/2)
1487        //
1488        // is only safe on one side: for a nearly-real z the smaller of |z| +- re is a
1489        // difference of nearly equal numbers, so it cancels to noise and the small
1490        // component has no correct digits left (asin/acos/asinh feed in arguments of
1491        // exactly that shape). Take the large component from the modulus, always
1492        // adding, and recover the small one from re*im_out = im/2. One division, no
1493        // cancellation.
1494        let (res, degenerate) = sqrt_kahan(self);
1495
1496        // Branching rather than blending costs the common path nothing: it trades the two
1497        // unconditional `nz` this used to end with for a test the predictor will call
1498        // correctly essentially always, and takes them off the dependency chain.
1499        if thermite::unlikely(degenerate.any()) {
1500            return degenerate.select(sqrt_rescaled(self), res);
1501        }
1502
1503        res
1504    }
1505
1506    /// `$1/z = \bar{z}/|z|^2$`
1507    #[inline(always)]
1508    fn rcp(self) -> Self {
1509        let inv = self.norm_sqr().rcp();
1510
1511        Complex::new(self.re * inv, -(self.im * inv))
1512    }
1513
1514    /// `$1/\sqrt{z} = \overline{\sqrt{z}}/|z|$`, since `$|\sqrt{z}|^2 = |z|$`.
1515    #[inline(always)]
1516    fn rsqrt(self) -> Self {
1517        let s = self.sqrt();
1518        let inv = self.modulus().rcp();
1519
1520        Complex::new(s.re * inv, -(s.im * inv))
1521    }
1522
1523    // Rounding is componentwise; see the crate docs.
1524    #[inline(always)] fn floor(self) -> Self { Complex::new(self.re.floor(), self.im.floor()) }
1525    #[inline(always)] fn ceil(self) -> Self { Complex::new(self.re.ceil(), self.im.ceil()) }
1526    #[inline(always)] fn round(self) -> Self { Complex::new(self.re.round(), self.im.round()) }
1527    #[inline(always)] fn trunc(self) -> Self { Complex::new(self.re.trunc(), self.im.trunc()) }
1528    #[inline(always)] fn fract(self) -> Self { Complex::new(self.re.fract(), self.im.fract()) }
1529
1530    #[inline(always)] fn mul_sign(self, sign: Self) -> Self { Complex::new(self.re.mul_sign(sign.re), self.im.mul_sign(sign.im)) }
1531    #[inline(always)] fn signed_zero(self) -> Self { Complex::new(self.re.signed_zero(), self.im.signed_zero()) }
1532
1533    #[inline(always)] fn next_up(self) -> Self { Complex::new(self.re.next_up(), self.im.next_up()) }
1534    #[inline(always)] fn next_down(self) -> Self { Complex::new(self.re.next_down(), self.im.next_down()) }
1535
1536    #[inline(always)]
1537    unsafe fn block_autovectorization(&mut self) {
1538        unsafe {
1539            self.re.block_autovectorization();
1540            self.im.block_autovectorization();
1541        }
1542    }
1543
1544    // mix(t) = a*(1 - t) + b*t = (b - a)*t + a, one complex FMA (four inner FMAs)
1545    // rather than a complex multiply and a complex add.
1546    #[inline(always)] fn mix(self, a: Self, b: Self) -> Self { (b - a).mul_adde(self, a) }
1547
1548    complex_masked!(unary: sqrt, rsqrt, rcp, floor, ceil, round, trunc, fract, signed_zero, next_up, next_down);
1549    complex_masked!(binary: mul_sign);
1550}