use std::iter;
use std::marker::PhantomData;
use crate::util::PeekableIterator;
pub struct BatchesWithHeader<II, I, F> {
input: Input<II, I, F>,
}
struct Input<II, I, F> {
unfiltered: I,
batch_starting: F,
marker: PhantomData<fn() -> II>,
}
pub struct Batches<II, I, F> {
input: Input<II, I, F>,
no_drain: Option<NoDrainToken>,
yield_one: Option<EvenYieldOneBatchStarting>,
}
struct NoDrainToken;
struct EvenYieldOneBatchStarting;
pub struct Batch<'p, II, I, F> {
parent: &'p mut Batches<II, I, F>,
}
impl<II, I, F> Input<II, I, F>
where
I: Iterator<Item = II> + PeekableIterator,
F: FnMut(&II) -> bool,
{
fn next_non_starting(&mut self) -> Option<II> {
let item = self.unfiltered.peek()?;
if (self.batch_starting)(item) {
return None;
};
self.unfiltered.next()
}
}
impl<II, I, F> Iterator for BatchesWithHeader<II, I, F>
where
I: Iterator<Item = II> + PeekableIterator,
F: FnMut(&II) -> bool,
{
type Item = II;
fn next(&mut self) -> Option<II> {
self.input.next_non_starting()
}
}
impl<II, I, F> BatchesWithHeader<II, I, F>
where
I: Iterator<Item = II> + PeekableIterator,
F: FnMut(&II) -> bool,
{
pub fn subsequent(mut self) -> Batches<II, I, F> {
let _ = self.by_ref().count();
Batches {
input: self.input,
yield_one: None,
no_drain: None,
}
}
}
impl<II, I, F> Iterator for Batch<'_, II, I, F>
where
I: Iterator<Item = II> + PeekableIterator,
F: FnMut(&II) -> bool,
{
type Item = II;
fn next(&mut self) -> Option<II> {
if self.parent.yield_one.take().is_some() {
self.parent.input.unfiltered.next()
} else {
self.parent.input.next_non_starting()
}
}
}
impl<II, I: Iterator<Item = II> + PeekableIterator, F: FnMut(&II) -> bool> Batches<II, I, F> {
pub fn next_batch(&mut self) -> Option<Batch<'_, II, I, F>> {
if self.no_drain.take().is_none() {
let _ = Batch { parent: self }.count();
}
let _: &II = self.input.unfiltered.peek()?;
self.yield_one = Some(EvenYieldOneBatchStarting);
Some(Batch { parent: self })
}
pub fn map<T>(
mut self,
mut f: impl FnMut(Batch<'_, II, I, F>) -> T,
) -> impl Iterator<Item = T> {
iter::from_fn(move || {
let batch = self.next_batch()?;
Some(f(batch))
})
}
}
pub trait IteratorExt: Iterator + Sized {
fn batching_split_before_with_header<F>(
self,
batch_starting: F,
) -> BatchesWithHeader<Self::Item, Self, F>
where
F: FnMut(&Self::Item) -> bool,
{
let input = Input {
unfiltered: self,
batch_starting,
marker: PhantomData,
};
BatchesWithHeader { input }
}
fn batching_split_before_loose<F>(self, batch_starting: F) -> Batches<Self::Item, Self, F>
where
F: FnMut(&Self::Item) -> bool,
{
let input = Input {
unfiltered: self,
batch_starting,
marker: PhantomData,
};
Batches {
input,
no_drain: Some(NoDrainToken),
yield_one: None,
}
}
}
impl<I: Iterator> IteratorExt for I {}
#[cfg(test)]
mod tests {
#![allow(clippy::bool_assert_comparison)]
#![allow(clippy::clone_on_copy)]
#![allow(clippy::dbg_macro)]
#![allow(clippy::mixed_attributes_style)]
#![allow(clippy::print_stderr)]
#![allow(clippy::print_stdout)]
#![allow(clippy::single_char_pattern)]
#![allow(clippy::unwrap_used)]
#![allow(clippy::unchecked_time_subtraction)]
#![allow(clippy::useless_vec)]
#![allow(clippy::needless_pass_by_value)]
use super::*;
use crate::util::*;
use itertools::chain;
use std::fmt::Debug;
use std::iter;
struct TrackingPeekable<I: Iterator>(Peekable<I>);
impl<I: Iterator> Iterator for TrackingPeekable<I>
where
I::Item: Debug,
{
type Item = I::Item;
fn next(&mut self) -> Option<I::Item> {
let v = self.0.next();
eprintln!(" iter yielded {v:?}");
v
}
}
impl<I: Iterator> PeekableIterator for TrackingPeekable<I>
where
I::Item: Debug,
{
fn peek(&mut self) -> Option<&I::Item> {
let v = self.0.peek();
eprintln!(" iter peeked {v:?}");
v
}
}
#[test]
fn test_batching_split_before() {
fn chk_exp(mut iter: impl Iterator<Item = u32>, exp: &[u32]) {
eprintln!(" exp {exp:?}");
for exp in exp.iter().cloned() {
assert_eq!(iter.next(), Some(exp));
}
assert_eq!(iter.next(), None);
assert_eq!(iter.next(), None);
assert_eq!(iter.next(), None);
}
let chk_breakdown = |input: &[u32], iexp: &[u32], sexp: &[&[u32]]| {
let chk_batches = |mut subseq: Batches<_, _, _>, sexp: &mut dyn Iterator<Item = _>| {
loop {
match (subseq.next_batch(), sexp.next()) {
(Some(batch), Some(sexp)) => chk_exp(batch, sexp),
(None, None) => break,
(b, e) => panic!("({:?}, {e:?}", b.map(|_| ())),
}
}
assert!(subseq.next_batch().is_none());
assert!(subseq.next_batch().is_none());
};
eprintln!("input {input:?}");
let input = || TrackingPeekable(input.iter().cloned().peekable());
let is_starting = |v: &u32| *v >= 10;
{
let mut header = input().batching_split_before_with_header(is_starting);
chk_exp(&mut header, iexp);
eprintln!(" subsequent...");
let subseq = header.subsequent();
let mut sexp = sexp.iter().cloned();
chk_batches(subseq, &mut sexp);
}
{
let batches = input().batching_split_before_loose(is_starting);
let mut sexp =
chain!(iter::once(iexp), sexp.iter().cloned(),).filter(|s| !s.is_empty());
chk_batches(batches, &mut sexp);
}
};
chk_breakdown(&[], &[], &[]);
chk_breakdown(&[10], &[], &[&[10]]);
chk_breakdown(
&[1, 2, 30, 4, 5, 60, 7, 8],
&[1, 2],
&[&[30, 4, 5], &[60, 7, 8]],
);
}
}