use crate::AsyncIterator;
#[derive(Clone, Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Take<I> {
iter: I,
n: usize,
}
impl<I> Take<I> {
pub fn new(iter: I, n: usize) -> Take<I> {
Take { iter, n }
}
}
impl<I: AsyncIterator> AsyncIterator for Take<I> {
type Item = I::Item;
#[inline]
async fn next(&mut self) -> Option<I::Item> {
if self.n != 0 {
self.n -= 1;
self.iter.next().await
} else {
None
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
if self.n == 0 {
return (0, Some(0));
}
let (lower, upper) = self.iter.size_hint();
let lower = std::cmp::min(lower, self.n);
let upper = match upper {
Some(x) if x < self.n => Some(x),
_ => Some(self.n),
};
(lower, upper)
}
}