1use crate::simd::intrinsics;
2use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};
3
4/// Constructs a new SIMD vector by copying elements from selected elements in other vectors.
5///
6/// When swizzling one vector, elements are selected like [`Swizzle::swizzle`].
7///
8/// When swizzling two vectors, elements are selected like [`Swizzle::concat_swizzle`].
9///
10/// # Examples
11///
12/// With a single SIMD vector, the const array specifies element indices in that vector:
13/// ```
14/// # #![feature(portable_simd)]
15/// # use core::simd::{u32x2, u32x4, simd_swizzle};
16/// let v = u32x4::from_array([10, 11, 12, 13]);
17///
18/// // Keeping the same size
19/// let r: u32x4 = simd_swizzle!(v, [3, 0, 1, 2]);
20/// assert_eq!(r.to_array(), [13, 10, 11, 12]);
21///
22/// // Changing the number of elements
23/// let r: u32x2 = simd_swizzle!(v, [3, 1]);
24/// assert_eq!(r.to_array(), [13, 11]);
25/// ```
26///
27/// With two input SIMD vectors, the const array specifies element indices in the concatenation of
28/// those vectors:
29/// ```
30/// # #![feature(portable_simd)]
31/// # #[cfg(feature = "as_crate")] use core_simd::simd;
32/// # #[cfg(not(feature = "as_crate"))] use core::simd;
33/// # use simd::{u32x2, u32x4, simd_swizzle};
34/// let a = u32x4::from_array([0, 1, 2, 3]);
35/// let b = u32x4::from_array([4, 5, 6, 7]);
36///
37/// // Keeping the same size
38/// let r: u32x4 = simd_swizzle!(a, b, [0, 1, 6, 7]);
39/// assert_eq!(r.to_array(), [0, 1, 6, 7]);
40///
41/// // Changing the number of elements
42/// let r: u32x2 = simd_swizzle!(a, b, [0, 4]);
43/// assert_eq!(r.to_array(), [0, 4]);
44/// ```
45#[allow(unused_macros)]
46pub macro simd_swizzle {
47 (
48 $vector:expr, $index:expr $(,)?
49 ) => {
50 {
51 use $crate::simd::Swizzle;
52 struct Impl;
53 impl Swizzle<{$index.len()}> for Impl {
54 const INDEX: [usize; {$index.len()}] = $index;
55 }
56 Impl::swizzle($vector)
57 }
58 },
59 (
60 $first:expr, $second:expr, $index:expr $(,)?
61 ) => {
62 {
63 use $crate::simd::Swizzle;
64 struct Impl;
65 impl Swizzle<{$index.len()}> for Impl {
66 const INDEX: [usize; {$index.len()}] = $index;
67 }
68 Impl::concat_swizzle($first, $second)
69 }
70 }
71}
72
73/// Create a vector from the elements of another vector.
74pub trait Swizzle<const N: usize> {
75 /// Map from the elements of the input vector to the output vector.
76 const INDEX: [usize; N];
77
78 /// Create a new vector from the elements of `vector`.
79 ///
80 /// Lane `i` of the output is `vector[Self::INDEX[i]]`.
81 #[inline]
82 #[must_use = "method returns a new vector and does not mutate the original inputs"]
83 fn swizzle<T, const M: usize>(vector: Simd<T, M>) -> Simd<T, N>
84 where
85 T: SimdElement,
86 LaneCount<N>: SupportedLaneCount,
87 LaneCount<M>: SupportedLaneCount,
88 {
89 // Safety: `vector` is a vector, and the index is a const array of u32.
90 unsafe {
91 intrinsics::simd_shuffle(
92 vector,
93 vector,
94 const {
95 let mut output = [0; N];
96 let mut i = 0;
97 while i < N {
98 let index = Self::INDEX[i];
99 assert!(index as u32 as usize == index);
100 assert!(
101 index < M,
102 "source element index exceeds input vector length"
103 );
104 output[i] = index as u32;
105 i += 1;
106 }
107 output
108 },
109 )
110 }
111 }
112
113 /// Create a new vector from the elements of `first` and `second`.
114 ///
115 /// Lane `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
116 /// `first` and `second`.
117 #[inline]
118 #[must_use = "method returns a new vector and does not mutate the original inputs"]
119 fn concat_swizzle<T, const M: usize>(first: Simd<T, M>, second: Simd<T, M>) -> Simd<T, N>
120 where
121 T: SimdElement,
122 LaneCount<N>: SupportedLaneCount,
123 LaneCount<M>: SupportedLaneCount,
124 {
125 // Safety: `first` and `second` are vectors, and the index is a const array of u32.
126 unsafe {
127 intrinsics::simd_shuffle(
128 first,
129 second,
130 const {
131 let mut output = [0; N];
132 let mut i = 0;
133 while i < N {
134 let index = Self::INDEX[i];
135 assert!(index as u32 as usize == index);
136 assert!(
137 index < 2 * M,
138 "source element index exceeds input vector length"
139 );
140 output[i] = index as u32;
141 i += 1;
142 }
143 output
144 },
145 )
146 }
147 }
148
149 /// Create a new mask from the elements of `mask`.
150 ///
151 /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
152 /// `first` and `second`.
153 #[inline]
154 #[must_use = "method returns a new mask and does not mutate the original inputs"]
155 fn swizzle_mask<T, const M: usize>(mask: Mask<T, M>) -> Mask<T, N>
156 where
157 T: MaskElement,
158 LaneCount<N>: SupportedLaneCount,
159 LaneCount<M>: SupportedLaneCount,
160 {
161 // SAFETY: all elements of this mask come from another mask
162 unsafe { Mask::from_int_unchecked(Self::swizzle(mask.to_int())) }
163 }
164
165 /// Create a new mask from the elements of `first` and `second`.
166 ///
167 /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
168 /// `first` and `second`.
169 #[inline]
170 #[must_use = "method returns a new mask and does not mutate the original inputs"]
171 fn concat_swizzle_mask<T, const M: usize>(first: Mask<T, M>, second: Mask<T, M>) -> Mask<T, N>
172 where
173 T: MaskElement,
174 LaneCount<N>: SupportedLaneCount,
175 LaneCount<M>: SupportedLaneCount,
176 {
177 // SAFETY: all elements of this mask come from another mask
178 unsafe { Mask::from_int_unchecked(Self::concat_swizzle(first.to_int(), second.to_int())) }
179 }
180}
181
182impl<T, const N: usize> Simd<T, N>
183where
184 T: SimdElement,
185 LaneCount<N>: SupportedLaneCount,
186{
187 /// Reverse the order of the elements in the vector.
188 #[inline]
189 #[must_use = "method returns a new vector and does not mutate the original inputs"]
190 pub fn reverse(self) -> Self {
191 struct Reverse;
192
193 impl<const N: usize> Swizzle<N> for Reverse {
194 const INDEX: [usize; N] = const {
195 let mut index = [0; N];
196 let mut i = 0;
197 while i < N {
198 index[i] = N - i - 1;
199 i += 1;
200 }
201 index
202 };
203 }
204
205 Reverse::swizzle(self)
206 }
207
208 /// Rotates the vector such that the first `OFFSET` elements of the slice move to the end
209 /// while the last `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`,
210 /// the element previously at index `OFFSET` will become the first element in the slice.
211 #[inline]
212 #[must_use = "method returns a new vector and does not mutate the original inputs"]
213 pub fn rotate_elements_left<const OFFSET: usize>(self) -> Self {
214 struct Rotate<const OFFSET: usize>;
215
216 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
217 const INDEX: [usize; N] = const {
218 let offset = OFFSET % N;
219 let mut index = [0; N];
220 let mut i = 0;
221 while i < N {
222 index[i] = (i + offset) % N;
223 i += 1;
224 }
225 index
226 };
227 }
228
229 Rotate::<OFFSET>::swizzle(self)
230 }
231
232 /// Rotates the vector such that the first `self.len() - OFFSET` elements of the vector move to
233 /// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`,
234 /// the element previously at index `self.len() - OFFSET` will become the first element in the slice.
235 #[inline]
236 #[must_use = "method returns a new vector and does not mutate the original inputs"]
237 pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self {
238 struct Rotate<const OFFSET: usize>;
239
240 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
241 const INDEX: [usize; N] = const {
242 let offset = N - OFFSET % N;
243 let mut index = [0; N];
244 let mut i = 0;
245 while i < N {
246 index[i] = (i + offset) % N;
247 i += 1;
248 }
249 index
250 };
251 }
252
253 Rotate::<OFFSET>::swizzle(self)
254 }
255
256 /// Interleave two vectors.
257 ///
258 /// The resulting vectors contain elements taken alternatively from `self` and `other`, first
259 /// filling the first result, and then the second.
260 ///
261 /// The reverse of this operation is [`Simd::deinterleave`].
262 ///
263 /// ```
264 /// # #![feature(portable_simd)]
265 /// # use core::simd::Simd;
266 /// let a = Simd::from_array([0, 1, 2, 3]);
267 /// let b = Simd::from_array([4, 5, 6, 7]);
268 /// let (x, y) = a.interleave(b);
269 /// assert_eq!(x.to_array(), [0, 4, 1, 5]);
270 /// assert_eq!(y.to_array(), [2, 6, 3, 7]);
271 /// ```
272 #[inline]
273 #[must_use = "method returns a new vector and does not mutate the original inputs"]
274 pub fn interleave(self, other: Self) -> (Self, Self) {
275 const fn interleave<const N: usize>(high: bool) -> [usize; N] {
276 let mut idx = [0; N];
277 let mut i = 0;
278 while i < N {
279 let dst_index = if high { i + N } else { i };
280 let src_index = dst_index / 2 + (dst_index % 2) * N;
281 idx[i] = src_index;
282 i += 1;
283 }
284 idx
285 }
286
287 struct Lo;
288 struct Hi;
289
290 impl<const N: usize> Swizzle<N> for Lo {
291 const INDEX: [usize; N] = interleave::<N>(false);
292 }
293
294 impl<const N: usize> Swizzle<N> for Hi {
295 const INDEX: [usize; N] = interleave::<N>(true);
296 }
297
298 (
299 Lo::concat_swizzle(self, other),
300 Hi::concat_swizzle(self, other),
301 )
302 }
303
304 /// Deinterleave two vectors.
305 ///
306 /// The first result takes every other element of `self` and then `other`, starting with
307 /// the first element.
308 ///
309 /// The second result takes every other element of `self` and then `other`, starting with
310 /// the second element.
311 ///
312 /// The reverse of this operation is [`Simd::interleave`].
313 ///
314 /// ```
315 /// # #![feature(portable_simd)]
316 /// # use core::simd::Simd;
317 /// let a = Simd::from_array([0, 4, 1, 5]);
318 /// let b = Simd::from_array([2, 6, 3, 7]);
319 /// let (x, y) = a.deinterleave(b);
320 /// assert_eq!(x.to_array(), [0, 1, 2, 3]);
321 /// assert_eq!(y.to_array(), [4, 5, 6, 7]);
322 /// ```
323 #[inline]
324 #[must_use = "method returns a new vector and does not mutate the original inputs"]
325 pub fn deinterleave(self, other: Self) -> (Self, Self) {
326 const fn deinterleave<const N: usize>(second: bool) -> [usize; N] {
327 let mut idx = [0; N];
328 let mut i = 0;
329 while i < N {
330 idx[i] = i * 2 + second as usize;
331 i += 1;
332 }
333 idx
334 }
335
336 struct Even;
337 struct Odd;
338
339 impl<const N: usize> Swizzle<N> for Even {
340 const INDEX: [usize; N] = deinterleave::<N>(false);
341 }
342
343 impl<const N: usize> Swizzle<N> for Odd {
344 const INDEX: [usize; N] = deinterleave::<N>(true);
345 }
346
347 (
348 Even::concat_swizzle(self, other),
349 Odd::concat_swizzle(self, other),
350 )
351 }
352
353 /// Resize a vector.
354 ///
355 /// If `M` > `N`, extends the length of a vector, setting the new elements to `value`.
356 /// If `M` < `N`, truncates the vector to the first `M` elements.
357 ///
358 /// ```
359 /// # #![feature(portable_simd)]
360 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
361 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
362 /// # use simd::u32x4;
363 /// let x = u32x4::from_array([0, 1, 2, 3]);
364 /// assert_eq!(x.resize::<8>(9).to_array(), [0, 1, 2, 3, 9, 9, 9, 9]);
365 /// assert_eq!(x.resize::<2>(9).to_array(), [0, 1]);
366 /// ```
367 #[inline]
368 #[must_use = "method returns a new vector and does not mutate the original inputs"]
369 pub fn resize<const M: usize>(self, value: T) -> Simd<T, M>
370 where
371 LaneCount<M>: SupportedLaneCount,
372 {
373 struct Resize<const N: usize>;
374 impl<const N: usize, const M: usize> Swizzle<M> for Resize<N> {
375 const INDEX: [usize; M] = const {
376 let mut index = [0; M];
377 let mut i = 0;
378 while i < M {
379 index[i] = if i < N { i } else { N };
380 i += 1;
381 }
382 index
383 };
384 }
385 Resize::<N>::concat_swizzle(self, Simd::splat(value))
386 }
387}
388