use std::ops::{Range, RangeFrom, RangeFull, RangeTo};
#[derive(Debug, Clone)]
pub enum Selection {
All,
Slice(Vec<SliceInfo>),
}
#[derive(Debug, Clone)]
pub struct SliceInfo {
pub start: u64,
pub end: u64, pub step: u64, }
impl SliceInfo {
pub fn new(start: u64, end: u64) -> Self {
Self { start, end, step: 1 }
}
pub fn count(&self) -> u64 {
if self.end > self.start {
(self.end - self.start + self.step - 1) / self.step
} else {
0
}
}
}
pub trait IntoSelection {
fn into_selection(self, shape: &[u64]) -> Selection;
}
impl IntoSelection for Selection {
fn into_selection(self, _shape: &[u64]) -> Selection {
self
}
}
impl IntoSelection for Range<usize> {
fn into_selection(self, _shape: &[u64]) -> Selection {
Selection::Slice(vec![SliceInfo::new(self.start as u64, self.end as u64)])
}
}
impl IntoSelection for RangeFull {
fn into_selection(self, _shape: &[u64]) -> Selection {
Selection::All
}
}
impl IntoSelection for RangeFrom<usize> {
fn into_selection(self, shape: &[u64]) -> Selection {
let end = if shape.is_empty() { 0 } else { shape[0] };
Selection::Slice(vec![SliceInfo::new(self.start as u64, end)])
}
}
impl IntoSelection for RangeTo<usize> {
fn into_selection(self, _shape: &[u64]) -> Selection {
Selection::Slice(vec![SliceInfo::new(0, self.end as u64)])
}
}
impl IntoSelection for (Range<usize>, Range<usize>) {
fn into_selection(self, _shape: &[u64]) -> Selection {
Selection::Slice(vec![
SliceInfo::new(self.0.start as u64, self.0.end as u64),
SliceInfo::new(self.1.start as u64, self.1.end as u64),
])
}
}
impl Selection {
pub fn output_shape(&self, ds_shape: &[u64]) -> Vec<u64> {
match self {
Selection::All => ds_shape.to_vec(),
Selection::Slice(slices) => {
slices.iter().map(|s| s.count()).collect()
}
}
}
pub fn to_slices(&self, ds_shape: &[u64]) -> Vec<SliceInfo> {
match self {
Selection::All => {
ds_shape.iter().map(|&d| SliceInfo::new(0, d)).collect()
}
Selection::Slice(slices) => slices.clone(),
}
}
}