1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};

/// Constructs a new SIMD vector by copying elements from selected elements in other vectors.
///
/// When swizzling one vector, elements are selected like [`Swizzle::swizzle`].
///
/// When swizzling two vectors, elements are selected like [`Swizzle::concat_swizzle`].
///
/// # Examples
///
/// With a single SIMD vector, the const array specifies element indices in that vector:
/// ```
/// # #![feature(portable_simd)]
/// # use core::simd::{u32x2, u32x4, simd_swizzle};
/// let v = u32x4::from_array([10, 11, 12, 13]);
///
/// // Keeping the same size
/// let r: u32x4 = simd_swizzle!(v, [3, 0, 1, 2]);
/// assert_eq!(r.to_array(), [13, 10, 11, 12]);
///
/// // Changing the number of elements
/// let r: u32x2 = simd_swizzle!(v, [3, 1]);
/// assert_eq!(r.to_array(), [13, 11]);
/// ```
///
/// With two input SIMD vectors, the const array specifies element indices in the concatenation of
/// those vectors:
/// ```
/// # #![feature(portable_simd)]
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
/// # use simd::{u32x2, u32x4, simd_swizzle};
/// let a = u32x4::from_array([0, 1, 2, 3]);
/// let b = u32x4::from_array([4, 5, 6, 7]);
///
/// // Keeping the same size
/// let r: u32x4 = simd_swizzle!(a, b, [0, 1, 6, 7]);
/// assert_eq!(r.to_array(), [0, 1, 6, 7]);
///
/// // Changing the number of elements
/// let r: u32x2 = simd_swizzle!(a, b, [0, 4]);
/// assert_eq!(r.to_array(), [0, 4]);
/// ```
#[allow(unused_macros)]
pub macro simd_swizzle {
    (
        $vector:expr, $index:expr $(,)?
    ) => {
        {
            use $crate::simd::Swizzle;
            struct Impl;
            impl Swizzle<{$index.len()}> for Impl {
                const INDEX: [usize; {$index.len()}] = $index;
            }
            Impl::swizzle($vector)
        }
    },
    (
        $first:expr, $second:expr, $index:expr $(,)?
    ) => {
        {
            use $crate::simd::Swizzle;
            struct Impl;
            impl Swizzle<{$index.len()}> for Impl {
                const INDEX: [usize; {$index.len()}] = $index;
            }
            Impl::concat_swizzle($first, $second)
        }
    }
}

/// Create a vector from the elements of another vector.
pub trait Swizzle<const N: usize> {
    /// Map from the elements of the input vector to the output vector.
    const INDEX: [usize; N];

    /// Create a new vector from the elements of `vector`.
    ///
    /// Lane `i` of the output is `vector[Self::INDEX[i]]`.
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    fn swizzle<T, const M: usize>(vector: Simd<T, M>) -> Simd<T, N>
    where
        T: SimdElement,
        LaneCount<N>: SupportedLaneCount,
        LaneCount<M>: SupportedLaneCount,
    {
        // Safety: `vector` is a vector, and the index is a const array of u32.
        unsafe {
            core::intrinsics::simd::simd_shuffle(
                vector,
                vector,
                const {
                    let mut output = [0; N];
                    let mut i = 0;
                    while i < N {
                        let index = Self::INDEX[i];
                        assert!(index as u32 as usize == index);
                        assert!(
                            index < M,
                            "source element index exceeds input vector length"
                        );
                        output[i] = index as u32;
                        i += 1;
                    }
                    output
                },
            )
        }
    }

    /// Create a new vector from the elements of `first` and `second`.
    ///
    /// Lane `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
    /// `first` and `second`.
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    fn concat_swizzle<T, const M: usize>(first: Simd<T, M>, second: Simd<T, M>) -> Simd<T, N>
    where
        T: SimdElement,
        LaneCount<N>: SupportedLaneCount,
        LaneCount<M>: SupportedLaneCount,
    {
        // Safety: `first` and `second` are vectors, and the index is a const array of u32.
        unsafe {
            core::intrinsics::simd::simd_shuffle(
                first,
                second,
                const {
                    let mut output = [0; N];
                    let mut i = 0;
                    while i < N {
                        let index = Self::INDEX[i];
                        assert!(index as u32 as usize == index);
                        assert!(
                            index < 2 * M,
                            "source element index exceeds input vector length"
                        );
                        output[i] = index as u32;
                        i += 1;
                    }
                    output
                },
            )
        }
    }

    /// Create a new mask from the elements of `mask`.
    ///
    /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
    /// `first` and `second`.
    #[inline]
    #[must_use = "method returns a new mask and does not mutate the original inputs"]
    fn swizzle_mask<T, const M: usize>(mask: Mask<T, M>) -> Mask<T, N>
    where
        T: MaskElement,
        LaneCount<N>: SupportedLaneCount,
        LaneCount<M>: SupportedLaneCount,
    {
        // SAFETY: all elements of this mask come from another mask
        unsafe { Mask::from_int_unchecked(Self::swizzle(mask.to_int())) }
    }

    /// Create a new mask from the elements of `first` and `second`.
    ///
    /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
    /// `first` and `second`.
    #[inline]
    #[must_use = "method returns a new mask and does not mutate the original inputs"]
    fn concat_swizzle_mask<T, const M: usize>(first: Mask<T, M>, second: Mask<T, M>) -> Mask<T, N>
    where
        T: MaskElement,
        LaneCount<N>: SupportedLaneCount,
        LaneCount<M>: SupportedLaneCount,
    {
        // SAFETY: all elements of this mask come from another mask
        unsafe { Mask::from_int_unchecked(Self::concat_swizzle(first.to_int(), second.to_int())) }
    }
}

impl<T, const N: usize> Simd<T, N>
where
    T: SimdElement,
    LaneCount<N>: SupportedLaneCount,
{
    /// Reverse the order of the elements in the vector.
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn reverse(self) -> Self {
        struct Reverse;

        impl<const N: usize> Swizzle<N> for Reverse {
            const INDEX: [usize; N] = const {
                let mut index = [0; N];
                let mut i = 0;
                while i < N {
                    index[i] = N - i - 1;
                    i += 1;
                }
                index
            };
        }

        Reverse::swizzle(self)
    }

    /// Rotates the vector such that the first `OFFSET` elements of the slice move to the end
    /// while the last `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`,
    /// the element previously at index `OFFSET` will become the first element in the slice.
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn rotate_elements_left<const OFFSET: usize>(self) -> Self {
        struct Rotate<const OFFSET: usize>;

        impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
            const INDEX: [usize; N] = const {
                let offset = OFFSET % N;
                let mut index = [0; N];
                let mut i = 0;
                while i < N {
                    index[i] = (i + offset) % N;
                    i += 1;
                }
                index
            };
        }

        Rotate::<OFFSET>::swizzle(self)
    }

    /// Rotates the vector such that the first `self.len() - OFFSET` elements of the vector move to
    /// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`,
    /// the element previously at index `self.len() - OFFSET` will become the first element in the slice.
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self {
        struct Rotate<const OFFSET: usize>;

        impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
            const INDEX: [usize; N] = const {
                let offset = N - OFFSET % N;
                let mut index = [0; N];
                let mut i = 0;
                while i < N {
                    index[i] = (i + offset) % N;
                    i += 1;
                }
                index
            };
        }

        Rotate::<OFFSET>::swizzle(self)
    }

    /// Interleave two vectors.
    ///
    /// The resulting vectors contain elements taken alternatively from `self` and `other`, first
    /// filling the first result, and then the second.
    ///
    /// The reverse of this operation is [`Simd::deinterleave`].
    ///
    /// ```
    /// # #![feature(portable_simd)]
    /// # use core::simd::Simd;
    /// let a = Simd::from_array([0, 1, 2, 3]);
    /// let b = Simd::from_array([4, 5, 6, 7]);
    /// let (x, y) = a.interleave(b);
    /// assert_eq!(x.to_array(), [0, 4, 1, 5]);
    /// assert_eq!(y.to_array(), [2, 6, 3, 7]);
    /// ```
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn interleave(self, other: Self) -> (Self, Self) {
        const fn interleave<const N: usize>(high: bool) -> [usize; N] {
            let mut idx = [0; N];
            let mut i = 0;
            while i < N {
                let dst_index = if high { i + N } else { i };
                let src_index = dst_index / 2 + (dst_index % 2) * N;
                idx[i] = src_index;
                i += 1;
            }
            idx
        }

        struct Lo;
        struct Hi;

        impl<const N: usize> Swizzle<N> for Lo {
            const INDEX: [usize; N] = interleave::<N>(false);
        }

        impl<const N: usize> Swizzle<N> for Hi {
            const INDEX: [usize; N] = interleave::<N>(true);
        }

        (
            Lo::concat_swizzle(self, other),
            Hi::concat_swizzle(self, other),
        )
    }

    /// Deinterleave two vectors.
    ///
    /// The first result takes every other element of `self` and then `other`, starting with
    /// the first element.
    ///
    /// The second result takes every other element of `self` and then `other`, starting with
    /// the second element.
    ///
    /// The reverse of this operation is [`Simd::interleave`].
    ///
    /// ```
    /// # #![feature(portable_simd)]
    /// # use core::simd::Simd;
    /// let a = Simd::from_array([0, 4, 1, 5]);
    /// let b = Simd::from_array([2, 6, 3, 7]);
    /// let (x, y) = a.deinterleave(b);
    /// assert_eq!(x.to_array(), [0, 1, 2, 3]);
    /// assert_eq!(y.to_array(), [4, 5, 6, 7]);
    /// ```
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn deinterleave(self, other: Self) -> (Self, Self) {
        const fn deinterleave<const N: usize>(second: bool) -> [usize; N] {
            let mut idx = [0; N];
            let mut i = 0;
            while i < N {
                idx[i] = i * 2 + second as usize;
                i += 1;
            }
            idx
        }

        struct Even;
        struct Odd;

        impl<const N: usize> Swizzle<N> for Even {
            const INDEX: [usize; N] = deinterleave::<N>(false);
        }

        impl<const N: usize> Swizzle<N> for Odd {
            const INDEX: [usize; N] = deinterleave::<N>(true);
        }

        (
            Even::concat_swizzle(self, other),
            Odd::concat_swizzle(self, other),
        )
    }

    /// Resize a vector.
    ///
    /// If `M` > `N`, extends the length of a vector, setting the new elements to `value`.
    /// If `M` < `N`, truncates the vector to the first `M` elements.
    ///
    /// ```
    /// # #![feature(portable_simd)]
    /// # #[cfg(feature = "as_crate")] use core_simd::simd;
    /// # #[cfg(not(feature = "as_crate"))] use core::simd;
    /// # use simd::u32x4;
    /// let x = u32x4::from_array([0, 1, 2, 3]);
    /// assert_eq!(x.resize::<8>(9).to_array(), [0, 1, 2, 3, 9, 9, 9, 9]);
    /// assert_eq!(x.resize::<2>(9).to_array(), [0, 1]);
    /// ```
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn resize<const M: usize>(self, value: T) -> Simd<T, M>
    where
        LaneCount<M>: SupportedLaneCount,
    {
        struct Resize<const N: usize>;
        impl<const N: usize, const M: usize> Swizzle<M> for Resize<N> {
            const INDEX: [usize; M] = const {
                let mut index = [0; M];
                let mut i = 0;
                while i < M {
                    index[i] = if i < N { i } else { N };
                    i += 1;
                }
                index
            };
        }
        Resize::<N>::concat_swizzle(self, Simd::splat(value))
    }
}