use crate::array::iterator::{distributed_iterator::*, IterLockFuture};
#[derive(Clone, Debug)]
pub struct Take<I> {
iter: I,
count: usize,
cur_index: usize,
}
impl<I: InnerIter> InnerIter for Take<I> {
fn lock_if_needed(&self, _s: Sealed) -> Option<IterLockFuture> {
self.iter.lock_if_needed(_s)
}
fn iter_clone(&self, _s: Sealed) -> Self {
Take {
iter: self.iter.iter_clone(Sealed),
count: self.count,
cur_index: self.cur_index,
}
}
}
impl<I> Take<I>
where
I: IndexedDistributedIterator,
{
pub(crate) fn new(iter: I, count: usize, start_i: usize) -> Take<I> {
Take {
iter,
count,
cur_index: start_i,
}
}
}
impl<I> DistributedIterator for Take<I>
where
I: IndexedDistributedIterator,
{
type Item = <I as DistributedIterator>::Item;
type Array = <I as DistributedIterator>::Array;
fn init(&self, start_i: usize, len: usize, _s: Sealed) -> Take<I> {
let val = Take::new(self.iter.init(start_i, len, _s), self.count, start_i);
val
}
fn array(&self) -> Self::Array {
self.iter.array()
}
fn next(&mut self) -> Option<Self::Item> {
if self.iterator_index(self.cur_index)? < self.count {
self.cur_index += 1;
let val = self.iter.next();
val
} else {
None
}
}
fn elems(&self, in_elems: usize) -> usize {
let in_elems = self.iter.elems(in_elems);
std::cmp::max(0, in_elems)
}
fn advance_index(&mut self, count: usize) {
self.iter.advance_index(count);
self.cur_index += count;
}
}
impl<I> IndexedDistributedIterator for Take<I>
where
I: IndexedDistributedIterator,
{
fn iterator_index(&self, index: usize) -> Option<usize> {
let Some(g_index) = self.iter.iterator_index(index) else {
return None;
};
if g_index < self.count {
Some(g_index)
} else {
None
}
}
}