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