Skip to content

Commit

Permalink
Move array subset iterators into separate files
Browse files Browse the repository at this point in the history
  • Loading branch information
@LDeakin
LDeakin committed Feb 12, 2024
1 parent 744d927 commit f7ac30d
Show file tree
Hide file tree
Showing 6 changed files with 417 additions and 375 deletions.
389 changes: 14 additions & 375 deletions src/array_subset/array_subset_iterators.rs
View file
Original file line number Diff line number Diff line change
@@ -1,381 +1,20 @@
use std::{iter::FusedIterator, num::NonZeroU64};

use itertools::izip;

use crate::array::{chunk_shape_to_array_shape, ravel_indices, ArrayIndices};

use super::{ArraySubset, IncompatibleArraySubsetAndShapeError, IncompatibleDimensionalityError};

/// Iterates over element indices in an array subset.
pub struct IndicesIterator {
subset_rev: ArraySubset,
index: u64,
}

impl IndicesIterator {
/// Create a new indices iterator.
#[must_use]
pub fn new(mut subset: ArraySubset) -> Self {
subset.start.reverse();
subset.shape.reverse();
Self {
subset_rev: subset,
index: 0,
}
}
}

impl Iterator for IndicesIterator {
type Item = ArrayIndices;

fn next(&mut self) -> Option<Self::Item> {
let mut current = self.index;
// let mut indices = vec![0u64; self.subset_rev.dimensionality()];
let mut indices = vec![core::mem::MaybeUninit::uninit(); self.subset_rev.dimensionality()];
for (out, &subset_start, &subset_size) in izip!(
indices.iter_mut().rev(),
self.subset_rev.start.iter(),
self.subset_rev.shape.iter(),
) {
out.write(current % subset_size + subset_start);
current /= subset_size;
}
if current == 0 {
self.index += 1;
#[allow(clippy::transmute_undefined_repr)]
Some(unsafe { std::mem::transmute(indices) })
} else {
None
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
let num_elements = self.subset_rev.num_elements_usize();
(num_elements, Some(num_elements))
}
}

impl ExactSizeIterator for IndicesIterator {}

impl FusedIterator for IndicesIterator {}

/// Iterates over linearised element indices of an array subset in an array.
pub struct LinearisedIndicesIterator<'a> {
subset: ArraySubset,
index: u64,
array_shape: &'a [u64],
}

impl<'a> LinearisedIndicesIterator<'a> {
/// Create a new linearised indices iterator.
///
/// # Errors
///
/// Returns [`IncompatibleArraySubsetAndShapeError`] if `array_shape` does not encapsulate `subset`.
pub fn new(
subset: ArraySubset,
array_shape: &'a [u64],
) -> Result<Self, IncompatibleArraySubsetAndShapeError> {
if subset.dimensionality() == array_shape.len()
&& std::iter::zip(subset.end_exc(), array_shape).all(|(end, shape)| end <= *shape)
{
Ok(Self {
subset,
index: 0,
array_shape,
})
} else {
Err(IncompatibleArraySubsetAndShapeError(
subset,
array_shape.to_vec(),
))
}
}

/// Create a new linearised indices iterator.
///
/// # Safety
///
/// `array_shape` must encapsulate `subset`.
#[must_use]
pub unsafe fn new_unchecked(subset: ArraySubset, array_shape: &'a [u64]) -> Self {
debug_assert_eq!(subset.dimensionality(), array_shape.len());
debug_assert!(
std::iter::zip(subset.end_exc(), array_shape).all(|(end, shape)| end <= *shape)
);
Self {
subset,
index: 0,
array_shape,
}
}
}

impl Iterator for LinearisedIndicesIterator<'_> {
type Item = u64;

fn next(&mut self) -> Option<Self::Item> {
let mut current = self.index;
let mut out = 0;
let mut mult = 1;
for (&subset_start, &subset_size, &array_size) in izip!(
self.subset.start.iter().rev(),
self.subset.shape.iter().rev(),
self.array_shape.iter().rev()
) {
let index = current % subset_size + subset_start;
current /= subset_size;
out += index * mult;
mult *= array_size;
}
if current == 0 {
self.index += 1;
Some(out)
} else {
None
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
let num_elements = self.subset.num_elements_usize();
(num_elements, Some(num_elements))
}
}

impl ExactSizeIterator for LinearisedIndicesIterator<'_> {}

impl FusedIterator for LinearisedIndicesIterator<'_> {}

/// Iterates over contiguous element indices in an array subset.
///
/// The iterator item is a tuple: (indices, # contiguous elements).
pub struct ContiguousIndicesIterator {
inner: IndicesIterator,
contiguous_elements: u64,
}

impl ContiguousIndicesIterator {
/// Create a new contiguous indices iterator.
///
/// # Errors
/// Returns [`IncompatibleArraySubsetAndShapeError`] if `array_shape` does not encapsulate `subset`.
pub fn new(
subset: &ArraySubset,
array_shape: &[u64],
) -> Result<Self, IncompatibleArraySubsetAndShapeError> {
if subset.dimensionality() == array_shape.len()
&& std::iter::zip(subset.end_exc(), array_shape).all(|(end, shape)| end <= *shape)
{
Ok(unsafe { Self::new_unchecked(subset, array_shape) })
} else {
Err(IncompatibleArraySubsetAndShapeError(
subset.clone(),
array_shape.to_vec(),
))
}
}

/// Create a new contiguous indices iterator.
///
/// # Safety
/// `array_shape` must encapsulate `subset`.
#[must_use]
#[allow(clippy::missing_panics_doc)]
pub unsafe fn new_unchecked(subset: &ArraySubset, array_shape: &[u64]) -> Self {
debug_assert_eq!(subset.dimensionality(), array_shape.len());
debug_assert!(
std::iter::zip(subset.end_exc(), array_shape).all(|(end, shape)| end <= *shape)
);

let mut contiguous = true;
let mut contiguous_elements = 1;
let mut shape_out = vec![core::mem::MaybeUninit::uninit(); array_shape.len()];
for (&subset_start, &subset_size, &array_size, shape_out_i) in izip!(
subset.start().iter().rev(),
subset.shape().iter().rev(),
array_shape.iter().rev(),
shape_out.iter_mut().rev(),
) {
if contiguous {
contiguous_elements *= subset_size;
shape_out_i.write(1);
contiguous = subset_start == 0 && subset_size == array_size;
} else {
shape_out_i.write(subset_size);
}
}
#[allow(clippy::transmute_undefined_repr)]
let shape_out: Vec<u64> = unsafe { core::mem::transmute(shape_out) };
let subset_contiguous_start =
ArraySubset::new_with_start_shape_unchecked(subset.start().to_vec(), shape_out);
let inner = subset_contiguous_start.iter_indices();
Self {
inner,
contiguous_elements,
}
}

/// Return the number of contiguous elements (fixed on each iteration).
#[must_use]
pub fn contiguous_elements(&self) -> u64 {
self.contiguous_elements
}
}

impl Iterator for ContiguousIndicesIterator {
type Item = (ArrayIndices, u64);

fn next(&mut self) -> Option<Self::Item> {
self.inner
.next()
.map(|indices| (indices, self.contiguous_elements))
}

fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}

impl ExactSizeIterator for ContiguousIndicesIterator {}

impl FusedIterator for ContiguousIndicesIterator {}

/// Iterates over contiguous linearised element indices in an array subset.
///
/// The iterator item is a tuple: (linearised index, # contiguous elements).
pub struct ContiguousLinearisedIndicesIterator<'a> {
inner: ContiguousIndicesIterator,
array_shape: &'a [u64],
}

impl<'a> ContiguousLinearisedIndicesIterator<'a> {
/// Return a new contiguous linearised indices iterator.
///
/// # Errors
///
/// Returns [`IncompatibleArraySubsetAndShapeError`] if `array_shape` does not encapsulate `subset`.
pub fn new(
subset: &ArraySubset,
array_shape: &'a [u64],
) -> Result<Self, IncompatibleArraySubsetAndShapeError> {
let inner = subset.iter_contiguous_indices(array_shape)?;
Ok(Self { inner, array_shape })
}

/// Return a new contiguous linearised indices iterator.
///
/// # Safety
///
/// `array_shape` must encapsulate `subset`.
#[must_use]
pub unsafe fn new_unchecked(subset: &ArraySubset, array_shape: &'a [u64]) -> Self {
let inner = subset.iter_contiguous_indices_unchecked(array_shape);
Self { inner, array_shape }
}

/// Return the number of contiguous elements (fixed on each iteration).
#[must_use]
pub fn contiguous_elements(&self) -> u64 {
self.inner.contiguous_elements()
}
}

impl Iterator for ContiguousLinearisedIndicesIterator<'_> {
type Item = (u64, u64);

fn next(&mut self) -> Option<Self::Item> {
self.inner
.next()
.map(|(indices, elements)| (ravel_indices(&indices, self.array_shape), elements))
}

fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}

impl ExactSizeIterator for ContiguousLinearisedIndicesIterator<'_> {}

impl FusedIterator for ContiguousLinearisedIndicesIterator<'_> {}

/// Iterates over the regular sized chunks overlapping this array subset.
/// All chunks have the same size, and may extend over the bounds of the array subset.
///
/// The iterator item is a ([`ArrayIndices`], [`ArraySubset`]) tuple corresponding to the chunk indices and array subset.
pub struct ChunksIterator {
inner: IndicesIterator,
chunk_shape: Vec<u64>,
}

impl ChunksIterator {
/// Create a new chunks iterator.
///
/// # Errors
///
/// Returns [`IncompatibleDimensionalityError`] if `chunk_shape` does not match the dimensionality of `subset`.
pub fn new(
subset: &ArraySubset,
chunk_shape: &[NonZeroU64],
) -> Result<Self, IncompatibleDimensionalityError> {
if subset.dimensionality() == chunk_shape.len() {
Ok(unsafe { Self::new_unchecked(subset, chunk_shape) })
} else {
Err(IncompatibleDimensionalityError(
chunk_shape.len(),
subset.dimensionality(),
))
}
}

/// Create a new chunks iterator.
///
/// # Safety
///
/// The dimensionality of `chunk_shape` must match the dimensionality of `subset`.
#[must_use]
pub unsafe fn new_unchecked(subset: &ArraySubset, chunk_shape: &[NonZeroU64]) -> Self {
debug_assert_eq!(subset.dimensionality(), chunk_shape.len());
let chunk_shape = chunk_shape_to_array_shape(chunk_shape);
let chunk_start: ArrayIndices = std::iter::zip(subset.start(), &chunk_shape)
.map(|(s, c)| s / c)
.collect();
let chunk_end_inc: ArrayIndices = std::iter::zip(subset.end_inc(), &chunk_shape)
.map(|(e, c)| e / c)
.collect();
let subset_chunks =
unsafe { ArraySubset::new_with_start_end_inc_unchecked(chunk_start, chunk_end_inc) };
let inner = IndicesIterator::new(subset_chunks);
Self { inner, chunk_shape }
}
}

impl Iterator for ChunksIterator {
type Item = (ArrayIndices, ArraySubset);

fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|chunk_indices| {
let start = std::iter::zip(&chunk_indices, &self.chunk_shape)
.map(|(i, c)| i * c)
.collect();
let chunk_subset = unsafe {
ArraySubset::new_with_start_shape_unchecked(start, self.chunk_shape.clone())
};
(chunk_indices, chunk_subset)
})
}

fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}

impl ExactSizeIterator for ChunksIterator {}

impl FusedIterator for ChunksIterator {}
mod chunks_iterator;
mod contiguous_indices_iterator;
mod contiguous_linearised_indices_iterator;
mod indices_iterator;
mod linearised_indices_iterator;

pub use chunks_iterator::ChunksIterator;
pub use contiguous_indices_iterator::ContiguousIndicesIterator;
pub use contiguous_linearised_indices_iterator::ContiguousLinearisedIndicesIterator;
pub use indices_iterator::IndicesIterator;
pub use linearised_indices_iterator::LinearisedIndicesIterator;

#[cfg(test)]
mod tests {
use super::*;
use std::num::NonZeroU64;

use crate::array_subset::ArraySubset;

#[test]
fn array_subset_iter_indices() {
Expand Down
Loading

0 comments on commit f7ac30d

Please sign in to comment.