Skip to content

Commit 675401b

Browse files
committed
Add extend special swizzle fn, and implement special swizzle fns for masks
1 parent 7cd6f95 commit 675401b

File tree

1 file changed

+179
-1
lines changed

1 file changed

+179
-1
lines changed

crates/core_simd/src/swizzle.rs

Lines changed: 179 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,9 @@ where
312312
///
313313
/// ```
314314
/// # #![feature(portable_simd)]
315-
/// # use core::simd::Simd;
315+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
316+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
317+
/// # use simd::Simd;
316318
/// let a = Simd::from_array([0, 4, 1, 5]);
317319
/// let b = Simd::from_array([2, 6, 3, 7]);
318320
/// let (x, y) = a.deinterleave(b);
@@ -383,4 +385,180 @@ where
383385
}
384386
Resize::<N>::concat_swizzle(self, Simd::splat(value))
385387
}
388+
389+
/// Extract a vector from another vector.
390+
///
391+
/// ```
392+
/// # #![feature(portable_simd)]
393+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
394+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
395+
/// # use simd::u32x4;
396+
/// let x = u32x4::from_array([0, 1, 2, 3]);
397+
/// assert_eq!(x.extract::<1, 2>().to_array(), [1, 2]);
398+
/// ```
399+
#[inline]
400+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
401+
pub fn extract<const START: usize, const LEN: usize>(self) -> Simd<T, LEN>
402+
where
403+
LaneCount<LEN>: SupportedLaneCount,
404+
{
405+
struct Extract<const N: usize, const START: usize>;
406+
impl<const N: usize, const START: usize, const LEN: usize> Swizzle<LEN> for Extract<N, START> {
407+
const INDEX: [usize; LEN] = const {
408+
assert!(START + LEN <= N, "index out of bounds");
409+
let mut index = [0; LEN];
410+
let mut i = 0;
411+
while i < LEN {
412+
index[i] = START + i;
413+
i += 1;
414+
}
415+
index
416+
};
417+
}
418+
Extract::<N, START>::swizzle(self)
419+
}
420+
}
421+
422+
impl<T, const N: usize> Mask<T, N>
423+
where
424+
T: MaskElement,
425+
LaneCount<N>: SupportedLaneCount,
426+
{
427+
/// Reverse the order of the elements in the mask.
428+
#[inline]
429+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
430+
pub fn reverse(self) -> Self {
431+
// Safety: swizzles are safe for masks
432+
unsafe { Self::from_int_unchecked(self.to_int().reverse()) }
433+
}
434+
435+
/// Rotates the mask such that the first `OFFSET` elements of the slice move to the end
436+
/// while the last `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`,
437+
/// the element previously at index `OFFSET` will become the first element in the slice.
438+
#[inline]
439+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
440+
pub fn rotate_elements_left<const OFFSET: usize>(self) -> Self {
441+
// Safety: swizzles are safe for masks
442+
unsafe { Self::from_int_unchecked(self.to_int().rotate_elements_left::<OFFSET>()) }
443+
}
444+
445+
/// Rotates the mask such that the first `self.len() - OFFSET` elements of the mask move to
446+
/// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`,
447+
/// the element previously at index `self.len() - OFFSET` will become the first element in the slice.
448+
#[inline]
449+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
450+
pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self {
451+
// Safety: swizzles are safe for masks
452+
unsafe { Self::from_int_unchecked(self.to_int().rotate_elements_right::<OFFSET>()) }
453+
}
454+
455+
/// Interleave two masks.
456+
///
457+
/// The resulting masks contain elements taken alternatively from `self` and `other`, first
458+
/// filling the first result, and then the second.
459+
///
460+
/// The reverse of this operation is [`Mask::deinterleave`].
461+
///
462+
/// ```
463+
/// # #![feature(portable_simd)]
464+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
465+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
466+
/// # use simd::mask32x4;
467+
/// let a = mask32x4::from_array([false, true, false, true]);
468+
/// let b = mask32x4::from_array([false, false, true, true]);
469+
/// let (x, y) = a.interleave(b);
470+
/// assert_eq!(x.to_array(), [false, false, true, false]);
471+
/// assert_eq!(y.to_array(), [false, true, true, true]);
472+
/// ```
473+
#[inline]
474+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
475+
pub fn interleave(self, other: Self) -> (Self, Self) {
476+
// Safety: swizzles are safe for masks
477+
let (lo, hi) = self.to_int().interleave(other.to_int());
478+
unsafe { (Self::from_int_unchecked(lo), Self::from_int_unchecked(hi)) }
479+
}
480+
481+
/// Deinterleave two masks.
482+
///
483+
/// The first result takes every other element of `self` and then `other`, starting with
484+
/// the first element.
485+
///
486+
/// The second result takes every other element of `self` and then `other`, starting with
487+
/// the second element.
488+
///
489+
/// The reverse of this operation is [`Mask::interleave`].
490+
///
491+
/// ```
492+
/// # #![feature(portable_simd)]
493+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
494+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
495+
/// # use simd::mask32x4;
496+
/// let a = mask32x4::from_array([false, true, false, true]);
497+
/// let b = mask32x4::from_array([false, false, true, true]);
498+
/// let (x, y) = a.deinterleave(b);
499+
/// assert_eq!(x.to_array(), [false, false, false, true]);
500+
/// assert_eq!(y.to_array(), [true, true, false, true]);
501+
/// ```
502+
#[inline]
503+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
504+
pub fn deinterleave(self, other: Self) -> (Self, Self) {
505+
// Safety: swizzles are safe for masks
506+
let (even, odd) = self.to_int().deinterleave(other.to_int());
507+
unsafe {
508+
(
509+
Self::from_int_unchecked(even),
510+
Self::from_int_unchecked(odd),
511+
)
512+
}
513+
}
514+
515+
/// Resize a mask.
516+
///
517+
/// If `M` > `N`, extends the length of a mask, setting the new elements to `value`.
518+
/// If `M` < `N`, truncates the mask to the first `M` elements.
519+
///
520+
/// ```
521+
/// # #![feature(portable_simd)]
522+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
523+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
524+
/// # use simd::mask32x4;
525+
/// let x = mask32x4::from_array([false, true, true, false]);
526+
/// assert_eq!(x.resize::<8>(true).to_array(), [false, true, true, false, true, true, true, true]);
527+
/// assert_eq!(x.resize::<2>(true).to_array(), [false, true]);
528+
/// ```
529+
#[inline]
530+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
531+
pub fn resize<const M: usize>(self, value: bool) -> Mask<T, M>
532+
where
533+
LaneCount<M>: SupportedLaneCount,
534+
{
535+
// Safety: swizzles are safe for masks
536+
unsafe {
537+
Mask::<T, M>::from_int_unchecked(self.to_int().resize::<M>(if value {
538+
T::TRUE
539+
} else {
540+
T::FALSE
541+
}))
542+
}
543+
}
544+
545+
/// Extract a vector from another vector.
546+
///
547+
/// ```
548+
/// # #![feature(portable_simd)]
549+
/// # #[cfg(feature = "as_crate")] use core_simd::simd;
550+
/// # #[cfg(not(feature = "as_crate"))] use core::simd;
551+
/// # use simd::mask32x4;
552+
/// let x = mask32x4::from_array([false, true, true, false]);
553+
/// assert_eq!(x.extract::<1, 2>().to_array(), [true, true]);
554+
/// ```
555+
#[inline]
556+
#[must_use = "method returns a new vector and does not mutate the original inputs"]
557+
pub fn extract<const START: usize, const LEN: usize>(self) -> Mask<T, LEN>
558+
where
559+
LaneCount<LEN>: SupportedLaneCount,
560+
{
561+
// Safety: swizzles are safe for masks
562+
unsafe { Mask::<T, LEN>::from_int_unchecked(self.to_int().extract::<START, LEN>()) }
563+
}
386564
}

0 commit comments

Comments
 (0)