use crate::tensor_ops::TryStack;
pub struct Stacker<I> {
iter: I,
}
impl<I: Iterator> Iterator for Stacker<I>
where
I::Item: TryStack,
{
type Item = <I::Item as TryStack>::Stacked;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next().map(|i| i.stack())
}
}
impl<I: ExactSizeIterator> ExactSizeIterator for Stacker<I>
where
Self: Iterator,
{
fn len(&self) -> usize {
self.iter.len()
}
}
pub trait IteratorStackExt: Iterator {
fn stack(self) -> Stacker<Self>
where
Self: Sized,
{
Stacker { iter: self }
}
}
impl<I: Iterator> IteratorStackExt for I {}