use crate::{structs::utility::set_some, Generator, GeneratorResult, IntoGenerator, ValueResult};
pub struct Flatten<Src>
where
Src: Generator,
Src::Output: IntoGenerator,
{
source: Src,
current_generator: Option<<Src::Output as IntoGenerator>::IntoGen>,
}
impl<Src> Flatten<Src>
where
Src: Generator,
Src::Output: IntoGenerator,
{
#[inline]
pub(crate) fn new(source: Src) -> Self {
Self {
source,
current_generator: None,
}
}
}
impl<Src> Clone for Flatten<Src>
where
Src: Generator + Clone,
Src::Output: IntoGenerator,
<Src::Output as IntoGenerator>::IntoGen: Clone,
{
#[inline]
fn clone(&self) -> Self {
Self {
source: self.source.clone(),
current_generator: self.current_generator.clone(),
}
}
}
impl<Src> Generator for Flatten<Src>
where
Src: Generator,
Src::Output: IntoGenerator,
{
type Output = <<Src as Generator>::Output as IntoGenerator>::Output;
#[inline]
fn run(&mut self, mut output: impl FnMut(Self::Output) -> ValueResult) -> GeneratorResult {
if let Some(current) = self.current_generator.as_mut() {
if current.run(|x| output(x)) == GeneratorResult::Stopped {
return GeneratorResult::Stopped;
}
}
let current_generator = &mut self.current_generator;
self.source.run(|x| {
match set_some(current_generator, x.into_gen()).run(|value| output(value)) {
GeneratorResult::Stopped => ValueResult::Stop,
GeneratorResult::Complete => ValueResult::MoreValues,
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test::StoppingGen;
use crate::{GeneratorExt, SliceGenerator};
#[test]
fn vector_flatten() {
let data = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9], vec![10]];
let mut output: Vec<i32> = Vec::new();
let result = SliceGenerator::new(data.as_slice())
.map(|x| SliceGenerator::new(x.as_slice()))
.flatten()
.for_each(|x| output.push(*x));
assert_eq!(output, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
assert_eq!(result, GeneratorResult::Complete);
}
#[test]
fn slice_flatten() {
let data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]];
let mut output = Vec::new();
let result = SliceGenerator::new(&data)
.map(|x| SliceGenerator::new(x))
.flatten()
.for_each(|x| output.push(*x));
assert_eq!(result, GeneratorResult::Complete);
assert_eq!(output, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
}
#[test]
fn stopping_generator() {
let data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]];
let expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
for x in 0..3 {
let mut gen = crate::test::StoppingGen::new(x, &data)
.map(|x| SliceGenerator::new(x))
.flatten();
let mut output = Vec::new();
let mut num_stops = 0;
while gen.for_each(|x| output.push(*x)) == GeneratorResult::Stopped {
num_stops += 1;
}
assert_eq!(num_stops, 1);
assert_eq!(output, expected);
}
}
#[test]
fn stopping_nested_generator() {
let data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]];
let expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
for i in 0..4 {
let mut gen = SliceGenerator::new(&data)
.map(|x| StoppingGen::new(i, x))
.flatten();
let mut output = Vec::new();
let mut num_stops = 0;
while gen.for_each(|x| output.push(*x)) == GeneratorResult::Stopped {
num_stops += 1;
}
assert_eq!(num_stops, 3);
assert_eq!(output, expected);
}
}
#[test]
fn only_stopping_generators() {
let data = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]];
let expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
for outer in 0..3 {
for i in 0..4 {
let mut gen = StoppingGen::new(outer, &data)
.map(|x| StoppingGen::new(i, x))
.flatten();
let mut output = Vec::new();
let mut num_stops = 0;
while gen.for_each(|x| output.push(*x)) == GeneratorResult::Stopped {
num_stops += 1;
}
assert_eq!(num_stops, 4);
assert_eq!(output, expected);
}
}
}
}