use crate::lattice::PaddedTileLattice;
impl<T> PaddedTileLattice<T> {
pub fn tile_slice(&self, tile_row: usize, tile_col: usize) -> &[T] {
let base = self.layout().tile_base(tile_row, tile_col);
let len = self.geometry().tile_len();
&self.as_storage_slice()[base..base + len]
}
#[inline]
pub fn tile_rows(&self) -> usize {
self.shape().padded_rows / self.geometry().sublanes
}
#[inline]
pub fn tile_cols(&self) -> usize {
self.shape().padded_cols / self.geometry().lanes
}
pub fn iter_tiles(&self) -> TileIter<'_, T> {
TileIter {
lattice: self,
next: 0,
total: self.num_tiles(),
}
}
}
impl<T: Clone> PaddedTileLattice<T> {
pub fn iter_logical(&self) -> impl Iterator<Item = (usize, usize, T)> + '_ {
let rows = self.rows();
let cols = self.cols();
(0..rows).flat_map(move |row| {
(0..cols).map(move |col| (row, col, self.get(row, col).unwrap().clone()))
})
}
}
pub struct TileIter<'a, T> {
lattice: &'a PaddedTileLattice<T>,
next: usize,
total: usize,
}
impl<'a, T> Iterator for TileIter<'a, T> {
type Item = (usize, usize, &'a [T]);
fn next(&mut self) -> Option<Self::Item> {
if self.next >= self.total {
return None;
}
let tiles_per_row = self.lattice.tile_cols();
let tile_row = self.next / tiles_per_row;
let tile_col = self.next % tiles_per_row;
self.next += 1;
Some((
tile_row,
tile_col,
self.lattice.tile_slice(tile_row, tile_col),
))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.total - self.next;
(remaining, Some(remaining))
}
}
impl<'a, T> ExactSizeIterator for TileIter<'a, T> {}