#[cfg(feature = "rayon")]
use rayon::prelude::*;
use super::*;
use std::convert::{AsMut, AsRef};
use std::ops::Range;
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct Select<S, I = Vec<usize>> {
pub indices: I,
pub target: S,
}
pub type SelectView<'a, S> = Select<S, &'a [usize]>;
pub trait AsIndexSlice: AsRef<[usize]> {}
pub trait AsIndexSliceMut: AsMut<[usize]> {}
impl AsIndexSlice for [usize] {}
impl AsIndexSliceMut for [usize] {}
impl AsIndexSlice for &[usize] {}
impl AsIndexSlice for &mut [usize] {}
impl AsIndexSliceMut for &mut [usize] {}
impl AsIndexSlice for Vec<usize> {}
impl AsIndexSliceMut for Vec<usize> {}
impl<S, I> Select<S, I>
where
S: Set,
I: AsIndexSlice,
{
#[inline]
pub fn new(indices: I, target: S) -> Self {
Self::validate(Select { indices, target })
}
#[inline]
fn validate(self) -> Self {
if !self.indices.as_ref().iter().all(|&i| i < self.target.len()) {
panic!("Select index out of bounds.");
}
self
}
}
impl<S: Set> Select<S, Range<usize>> {
#[inline]
pub fn from_range(indices: Range<usize>, target: S) -> Self {
Self::validate(Select { indices, target })
}
#[inline]
fn validate(self) -> Self {
assert!(
self.indices.end <= self.target.len(),
"Select index out of bounds."
);
self
}
}
impl<'a, S, I> Select<S, I>
where
S: Set + IntoOwned + Get<'a, usize>,
<S as IntoOwned>::Owned: std::iter::FromIterator<<S as Set>::Elem>,
<S as Get<'a, usize>>::Output: IntoOwned<Owned = <S as Set>::Elem>,
I: AsIndexSlice,
{
#[inline]
pub fn collapse(self) -> S::Owned {
self.indices
.as_ref()
.iter()
.map(|&i| self.target.at(i).into_owned())
.collect()
}
}
impl<S: Set, I: AsIndexSlice> Set for Select<S, I> {
type Elem = S::Elem;
type Atom = S::Atom;
#[inline]
fn len(&self) -> usize {
self.indices.as_ref().len()
}
}
impl<S: Set> Set for Select<S, Range<usize>> {
type Elem = S::Elem;
type Atom = S::Atom;
#[inline]
fn len(&self) -> usize {
Set::len(&self.indices)
}
}
impl<'a, S, I> View<'a> for Select<S, I>
where
S: View<'a>,
I: AsIndexSlice,
{
type Type = Select<S::Type, &'a [usize]>;
#[inline]
fn view(&'a self) -> Self::Type {
Select {
indices: self.indices.as_ref(),
target: self.target.view(),
}
}
}
impl<'a, S> View<'a> for Select<S, Range<usize>>
where
S: View<'a>,
{
type Type = Select<S::Type, Range<usize>>;
#[inline]
fn view(&'a self) -> Self::Type {
Select {
indices: self.indices.clone(),
target: self.target.view(),
}
}
}
impl<'a, S, I> ViewMut<'a> for Select<S, I>
where
S: Set + ViewMut<'a>,
I: AsIndexSlice,
{
type Type = Select<S::Type, &'a [usize]>;
#[inline]
fn view_mut(&'a mut self) -> Self::Type {
Select {
indices: self.indices.as_ref(),
target: self.target.view_mut(),
}
}
}
impl<'a, S> ViewMut<'a> for Select<S, Range<usize>>
where
S: Set + ViewMut<'a>,
{
type Type = Select<S::Type, Range<usize>>;
#[inline]
fn view_mut(&'a mut self) -> Self::Type {
Select {
indices: self.indices.clone(),
target: self.target.view_mut(),
}
}
}
impl<V, I> SplitAt for Select<V, I>
where
V: Set + Clone,
I: SplitAt,
{
#[inline]
fn split_at(self, mid: usize) -> (Self, Self) {
let Select { target, indices } = self;
let (indices_l, indices_r) = indices.split_at(mid);
(
Select {
indices: indices_l,
target: target.clone(),
},
Select {
indices: indices_r,
target,
},
)
}
}
impl<S, I: RemovePrefix> RemovePrefix for Select<S, I> {
#[inline]
fn remove_prefix(&mut self, n: usize) {
self.indices.remove_prefix(n);
}
}
impl<'a, S, I> Select<S, I>
where
S: Set + Get<'a, usize, Output = &'a <S as Set>::Elem> + View<'a>,
I: AsIndexSlice,
<S as View<'a>>::Type: IntoIterator<Item = S::Output>,
<S as Set>::Elem: 'a,
{
pub fn clone_values_into<V>(&'a self, other: &'a mut V)
where
V: ViewMut<'a> + ?Sized,
<V as ViewMut<'a>>::Type: Set + IntoIterator<Item = &'a mut S::Elem>,
<S as Set>::Elem: Clone,
{
let other_view = other.view_mut();
assert_eq!(other_view.len(), self.len());
for (theirs, mine) in other_view.into_iter().zip(self.iter()) {
theirs.clone_from(mine.1);
}
}
}
impl<'a, S> Select<S, Range<usize>>
where
S: Set + Get<'a, usize, Output = &'a <S as Set>::Elem> + View<'a>,
<S as View<'a>>::Type: IntoIterator<Item = S::Output>,
<S as Set>::Elem: 'a,
{
pub fn clone_values_into<V>(&'a self, other: &'a mut V)
where
V: ViewMut<'a> + ?Sized,
<V as ViewMut<'a>>::Type: Set + IntoIterator<Item = &'a mut S::Elem>,
<S as Set>::Elem: Clone,
{
let other_view = other.view_mut();
assert_eq!(other_view.len(), self.len());
for (theirs, mine) in other_view.into_iter().zip(self.iter()) {
theirs.clone_from(mine.1);
}
}
}
impl<'a, S, I> GetIndex<'a, Select<S, I>> for usize
where
I: AsIndexSlice,
S: Get<'a, usize>,
{
type Output = (usize, <S as Get<'a, usize>>::Output);
#[inline]
fn get(self, selection: &Select<S, I>) -> Option<Self::Output> {
selection
.indices
.as_ref()
.get(self)
.and_then(|&idx| selection.target.get(idx).map(|val| (idx, val)))
}
}
impl<'a, S> GetIndex<'a, Select<S, Range<usize>>> for usize
where
S: Get<'a, usize>,
{
type Output = (usize, <S as Get<'a, usize>>::Output);
#[inline]
fn get(self, selection: &Select<S, Range<usize>>) -> Option<Self::Output> {
selection
.indices
.clone()
.nth(self)
.and_then(|idx| selection.target.get(idx).map(|val| (idx, val)))
}
}
impl<'a, S, I> GetIndex<'a, Select<S, I>> for &usize
where
I: AsIndexSlice,
S: Get<'a, usize>,
{
type Output = (usize, <S as Get<'a, usize>>::Output);
#[inline]
fn get(self, selection: &Select<S, I>) -> Option<Self::Output> {
GetIndex::get(*self, selection)
}
}
impl<'a, S> GetIndex<'a, Select<S, Range<usize>>> for &usize
where
S: Get<'a, usize>,
{
type Output = (usize, <S as Get<'a, usize>>::Output);
#[inline]
fn get(self, selection: &Select<S, Range<usize>>) -> Option<Self::Output> {
GetIndex::get(*self, selection)
}
}
impl<S, I> IsolateIndex<Select<S, I>> for usize
where
I: Isolate<usize>,
<I as Isolate<usize>>::Output: std::borrow::Borrow<usize>,
S: Isolate<usize>,
{
type Output = (I::Output, S::Output);
#[inline]
unsafe fn isolate_unchecked(self, selection: Select<S, I>) -> Self::Output {
use std::borrow::Borrow;
let Select { indices, target } = selection;
let idx = indices.isolate_unchecked(self);
let val = target.isolate_unchecked(*idx.borrow());
(idx, val)
}
#[inline]
fn try_isolate(self, selection: Select<S, I>) -> Option<Self::Output> {
use std::borrow::Borrow;
let Select { indices, target } = selection;
let idx = indices.try_isolate(self)?;
target.try_isolate(*idx.borrow()).map(|val| (idx, val))
}
}
impl<S, I> IsolateIndex<Select<S, I>> for std::ops::Range<usize>
where
I: Isolate<std::ops::Range<usize>>,
{
type Output = Select<S, I::Output>;
#[inline]
unsafe fn isolate_unchecked(self, selection: Select<S, I>) -> Self::Output {
let Select { indices, target } = selection;
let indices = indices.isolate_unchecked(self);
Select { indices, target }
}
#[inline]
fn try_isolate(self, selection: Select<S, I>) -> Option<Self::Output> {
let Select { indices, target } = selection;
Some(Select {
indices: indices.try_isolate(self)?,
target,
})
}
}
impl<S, I> std::ops::Index<usize> for Select<S, I>
where
S: std::ops::Index<usize> + Set + ValueType,
I: AsIndexSlice,
{
type Output = S::Output;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target.index(self.indices.as_ref()[idx])
}
}
impl<S> std::ops::Index<usize> for Select<S, Range<usize>>
where
S: std::ops::Index<usize> + Set + ValueType,
{
type Output = S::Output;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target
.index(self.indices.clone().nth(idx).expect("Index out of bounds"))
}
}
impl<S, I> std::ops::IndexMut<usize> for Select<S, I>
where
S: std::ops::IndexMut<usize> + Set + ValueType,
I: AsIndexSlice,
{
#[inline]
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
self.target.index_mut(self.indices.as_ref()[idx])
}
}
impl<S> std::ops::IndexMut<usize> for Select<S, Range<usize>>
where
S: std::ops::IndexMut<usize> + Set + ValueType,
{
#[inline]
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
self.target
.index_mut(self.indices.clone().nth(idx).expect("Index out of bounds"))
}
}
impl<'a, T, I> std::ops::Index<usize> for Select<&'a [T], I>
where
I: AsIndexSlice,
{
type Output = T;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target.index(self.indices.as_ref()[idx])
}
}
impl<'a, T> std::ops::Index<usize> for Select<&'a [T], Range<usize>> {
type Output = T;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target
.index(self.indices.clone().nth(idx).expect("Index out of bounds"))
}
}
impl<'a, T, I> std::ops::Index<usize> for Select<&'a mut [T], I>
where
I: AsIndexSlice,
{
type Output = T;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target.index(self.indices.as_ref()[idx])
}
}
impl<'a, T> std::ops::Index<usize> for Select<&'a mut [T], Range<usize>> {
type Output = T;
#[inline]
fn index(&self, idx: usize) -> &Self::Output {
self.target
.index(self.indices.clone().nth(idx).expect("Index out of bounds"))
}
}
impl<'a, T, I> std::ops::IndexMut<usize> for Select<&'a mut [T], I>
where
I: AsIndexSlice,
{
#[inline]
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
self.target.index_mut(self.indices.as_ref()[idx])
}
}
impl<'a, T> std::ops::IndexMut<usize> for Select<&'a mut [T], Range<usize>> {
#[inline]
fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
self.target
.index_mut(self.indices.clone().nth(idx).expect("Index out of bounds"))
}
}
impl<'a, S, I> Select<S, I>
where
S: Set + Get<'a, usize> + View<'a>,
I: AsIndexSlice,
{
#[inline]
pub fn iter(&'a self) -> impl Iterator<Item = (usize, <S as Get<'a, usize>>::Output)> + Clone {
self.indices
.as_ref()
.iter()
.cloned()
.filter_map(move |idx| self.target.get(idx).map(|val| (idx, val)))
}
}
impl<'a, S, I> Select<S, I>
where
S: Set + Get<'a, usize> + View<'a> + Sync,
<S as Get<'a, usize>>::Output: Send,
I: AsIndexSlice + Sync,
{
#[cfg(feature = "rayon")]
#[inline]
pub fn par_iter(
&'a self,
) -> impl IndexedParallelIterator<Item = (usize, <S as Get<'a, usize>>::Output)> + Clone {
self.indices
.as_ref()
.par_iter()
.cloned()
.map(move |idx| (idx, unsafe { self.target.at_unchecked(idx) }))
}
}
impl<'a, S> Select<S, Range<usize>>
where
S: Set + Get<'a, usize> + View<'a>,
{
#[inline]
pub fn iter(&'a self) -> impl Iterator<Item = (usize, <S as Get<'a, usize>>::Output)> + Clone {
self.indices
.clone()
.filter_map(move |idx| self.target.get(idx).map(|val| (idx, val)))
}
}
impl<'a, S> Select<S, Range<usize>>
where
S: Set + Get<'a, usize> + View<'a> + Sync,
<S as Get<'a, usize>>::Output: Send,
{
#[cfg(feature = "rayon")]
#[inline]
pub fn par_iter(
&'a self,
) -> impl IndexedParallelIterator<Item = (usize, <S as Get<'a, usize>>::Output)> + Clone {
self.indices
.clone()
.into_par_iter()
.map(move |idx| (idx, unsafe { self.target.at_unchecked(idx) }))
}
}
impl<S, I> Select<S, I>
where
I: AsIndexSlice,
{
#[inline]
pub fn index_iter(&self) -> std::slice::Iter<'_, usize> {
self.indices.as_ref().iter()
}
#[cfg(feature = "rayon")]
#[inline]
pub fn index_par_iter(&self) -> rayon::slice::Iter<'_, usize> {
self.indices.as_ref().par_iter()
}
}
impl<S, I> Select<S, I>
where
I: AsMut<[usize]>,
{
#[inline]
pub fn index_iter_mut(&mut self) -> std::slice::IterMut<'_, usize> {
self.indices.as_mut().iter_mut()
}
#[cfg(feature = "rayon")]
#[inline]
pub fn index_par_iter_mut(&mut self) -> rayon::slice::IterMut<'_, usize> {
self.indices.as_mut().par_iter_mut()
}
}
impl<S: Dummy, I: Dummy> Dummy for Select<S, I> {
#[inline]
unsafe fn dummy() -> Self {
Select {
indices: Dummy::dummy(),
target: Dummy::dummy(),
}
}
}
impl<S, I: Truncate> Truncate for Select<S, I> {
#[inline]
fn truncate(&mut self, new_len: usize) {
self.indices.truncate(new_len);
}
}
impl<S, I: Clear> Clear for Select<S, I> {
#[inline]
fn clear(&mut self) {
self.indices.clear();
}
}
impl<S: StorageInto<T>, I, T> StorageInto<T> for Select<S, I> {
type Output = Select<S::Output, I>;
#[inline]
fn storage_into(self) -> Self::Output {
Select {
target: self.target.storage_into(),
indices: self.indices,
}
}
}
impl<S: MapStorage<Out>, I, Out> MapStorage<Out> for Select<S, I> {
type Input = S::Input;
type Output = Select<S::Output, I>;
#[inline]
fn map_storage<F: FnOnce(Self::Input) -> Out>(self, f: F) -> Self::Output {
Select {
target: self.target.map_storage(f),
indices: self.indices,
}
}
}
impl<'a, S: StorageView<'a>, I> StorageView<'a> for Select<S, I> {
type StorageView = S::StorageView;
#[inline]
fn storage_view(&'a self) -> Self::StorageView {
self.target.storage_view()
}
}
impl<S: Storage, I> Storage for Select<S, I> {
type Storage = S::Storage;
#[inline]
fn storage(&self) -> &Self::Storage {
self.target.storage()
}
}
impl<S: StorageMut, I> StorageMut for Select<S, I> {
#[inline]
fn storage_mut(&mut self) -> &mut Self::Storage {
self.target.storage_mut()
}
}
impl<S: ChunkSize, I> ChunkSize for Select<S, I> {
#[inline]
fn chunk_size(&self) -> usize {
self.target.chunk_size()
}
}
impl<S: IntoOwned, I: IntoOwned> IntoOwned for Select<S, I> {
type Owned = Select<S::Owned, I::Owned>;
#[inline]
fn into_owned(self) -> Self::Owned {
Select {
indices: self.indices.into_owned(),
target: self.target.into_owned(),
}
}
}
impl<S, I> IntoOwnedData for Select<S, I>
where
S: IntoOwnedData,
{
type OwnedData = Select<S::OwnedData, I>;
#[inline]
fn into_owned_data(self) -> Self::OwnedData {
Select {
indices: self.indices,
target: self.target.into_owned_data(),
}
}
}
impl<S, I: Reserve> Reserve for Select<S, I> {
#[inline]
fn reserve_with_storage(&mut self, n: usize, storage_n: usize) {
self.indices.reserve_with_storage(n, storage_n);
}
}
impl<S, I, M> UniChunkable<M> for Select<S, I> {
type Chunk = Select<S, I>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn into_owned() {
let indices = vec![1, 2, 3];
let target = 0..4;
let select = Select {
indices: indices.as_slice(),
target: target.clone(),
};
assert_eq!(select.into_owned(), Select { indices, target });
}
}