itertools icon indicating copy to clipboard operation
itertools copied to clipboard

multi fold?

Open cynecx opened this issue 3 years ago • 0 comments

I am currently looking for a fold-like iterator adaptor but instead of producing one folded value it produces many partial ones.

Example code:

trait IteratorEx: Iterator {
    // `Result` might not be the best choice here.
    // (B, B) = (next init state, state to be yielded)
    fn multi_fold<B, F: FnMut(B, Self::Item) -> Result<(B, B), B>>(
        self,
        init: B,
        folder: F,
    ) -> MultiFold<B, Self, F>
    where
        Self: Sized,
    {
        MultiFold {
            iter: self,
            state: Some(init),
            folder,
        }
    }
}

impl<T: Iterator> IteratorEx for T {}

struct MultiFold<B, I, F> {
    iter: I,
    state: Option<B>,
    folder: F,
}

impl<I: Iterator, B, F: FnMut(B, I::Item) -> Result<(B, B), B>> Iterator
    for MultiFold<B, I, F>
{
    type Item = B;

    fn next(&mut self) -> Option<Self::Item> {
        let mut state = if let Some(state) = self.state.take() {
            state
        } else {
            return None;
        };
        while let Some(x) = self.iter.next() {
            match (self.folder)(state, x) {
                Ok((init, state)) => {
                    self.state = Some(init);
                    return Some(state);
                }
                Err(next_state) => {
                    state = next_state;
                }
            }
        }
        self.state = None;
        Some(state)
    }
}

fn main() {
    let l = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
    let coll = l
        .iter()
        .multi_fold(1, |prev, x| {
            if prev > 6 {
                Ok((*x, prev))
            } else {
                Err(prev + x)
            }
        })
        .collect::<Vec<_>>();
    println!("{:?}", coll);
}

Am I perhaps missing something in itertools that allows me to do this? If not, is there interest in adding such utility to itertools?

cynecx avatar Dec 26 '21 23:12 cynecx