use std::{iter::FusedIterator, ops::ControlFlow};
use crate::collector::{Collector, CollectorBase, Fuse as CollectorFuse};
#[derive(Debug)]
pub struct Driver<'a, I, C> {
iter: I,
collector: CollectorFuse<&'a mut C>,
}
impl<'a, I, C> Driver<'a, I, C>
where
I: Iterator,
C: for<'i> Collector<&'i mut I::Item>,
{
pub(in crate::iter) fn new(iter: I, collector: &'a mut C) -> Self {
Self {
iter,
collector: <&mut C>::fuse(collector),
}
}
}
impl<I, C> Driver<'_, I, C>
where
I: Iterator,
C: for<'i> Collector<&'i mut I::Item>,
{
fn fold_then_forward_once<B>(
mut self,
init: B,
mut forwardable_fold: impl ForwardableFold<B, I>,
) -> B {
match self.iter.try_fold(init, {
let forwardable_fold = &mut forwardable_fold;
move |accum, mut item| {
let cf = self.collector.collect(&mut item);
let accum = forwardable_fold.fold(accum, item);
if cf.is_continue() {
ControlFlow::Continue(accum)
} else {
ControlFlow::Break(accum)
}
}
}) {
ControlFlow::Continue(accum) => accum,
ControlFlow::Break(accum) => forwardable_fold.forward(accum, self.iter),
}
}
}
impl<I, C> Iterator for Driver<'_, I, C>
where
I: Iterator,
C: for<'i> Collector<&'i mut I::Item>,
{
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
let mut item = self.iter.next()?;
let _ = self.collector.collect(&mut item);
Some(item)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
fn count(self) -> usize
where
Self: Sized,
{
struct ForwardableCount;
impl<I: Iterator> ForwardableFold<usize, I> for ForwardableCount {
#[inline]
fn fold(&mut self, accum: usize, _item: <I as Iterator>::Item) -> usize {
accum + 1
}
#[inline]
fn forward(self, accum: usize, items: I) -> usize
where
Self: Sized,
{
accum + items.count()
}
}
self.fold_then_forward_once(0, ForwardableCount)
}
fn last(self) -> Option<Self::Item>
where
Self: Sized,
{
struct ForwardableLast;
impl<T, I: Iterator<Item = T>> ForwardableFold<Option<T>, I> for ForwardableLast {
#[inline]
fn fold(&mut self, _accum: Option<T>, item: <I as Iterator>::Item) -> Option<T> {
Some(item)
}
#[inline]
fn forward(self, accum: Option<T>, items: I) -> Option<T>
where
Self: Sized,
{
items.last().or(accum)
}
}
self.fold_then_forward_once(None, ForwardableLast)
}
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let mut item = self.iter.next()?;
for n in (0..n).rev() {
if self.collector.collect(&mut item).is_break() {
return self.iter.nth(n);
}
item = self.iter.next()?;
}
Some(item)
}
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.fold_then_forward_once(init, f)
}
}
impl<I, C> ExactSizeIterator for Driver<'_, I, C>
where
I: ExactSizeIterator,
C: for<'i> Collector<&'i mut I::Item>,
{
#[inline]
fn len(&self) -> usize {
self.iter.len()
}
}
impl<I, C> FusedIterator for Driver<'_, I, C>
where
I: FusedIterator,
C: for<'i> Collector<&'i mut I::Item>,
{
}
trait ForwardableFold<A, I: Iterator> {
fn fold(&mut self, accum: A, item: I::Item) -> A;
#[inline]
fn forward(self, accum: A, items: I) -> A
where
Self: Sized,
{
let mut this = self;
items.fold(accum, move |accum, item| this.fold(accum, item))
}
}
impl<A, I, F> ForwardableFold<A, I> for F
where
I: Iterator,
F: FnMut(A, I::Item) -> A,
{
#[inline]
fn fold(&mut self, accum: A, item: <I as Iterator>::Item) -> A {
self(accum, item)
}
}