holly 0.5.3

use crate::stream::Closable;
use futures::{
    future::{self, Either::{A, B}, Loop},
    prelude::*,
    stream::{futures_unordered::FuturesUnordered, StreamFuture},
    sync::mpsc,
    Stream
};
use void::{Void, unreachable};

/// The result of each folding step.
pub enum State<S, T, A> {
    /// Ready for the next step.
    Ready(S),

    /// Ready for the next step.
    ///
    /// The given [`Stream`] should be added to the existing stream of items.
    Stream(S, Box<dyn Stream<Item = T, Error = Void> + Send>),

    /// The main stream should be closed.
    ///
    /// The function being applied to the stream is responsible for closing
    /// any auxiliary streams added via [`State::Stream`]. [`State::Close`]
    /// will only close the [`mpsc::Receiver`] given to [`dynamic_fold`].
    Close(S),

    /// The fold should stop and return the given value of type `A`.
    Done(A)
}

/// Fold a function over a stream of items.
///
/// This is very much like a regular stream fold over an [`mpsc::Receiver`].
/// It is "dynamic" because at each step, the function being applied can
/// augment the stream of items with additional streams.
pub fn dynamic_fold<S, F, R, T, E, A>
    (init: S, main: mpsc::Receiver<T>, fun: F) -> impl Future<Item = Option<A>, Error = E>
where
    F: FnMut(S, Option<T>) -> R,
    R: Future<Item = State<S, T, A>, Error = E>,
    T: Send + 'static
{
    // Attach a kill cord to the `mpsc::Receiver` which will close
    // the stream when pulled. We pull if the function gives us
    // a `State::Close` back. Closing means that no more items can
    // be sent to this channel but we will process the ones already
    // enqueued.
    let (killcord, streams) = {
        let (killcord, main) = Closable::new(main);
        (killcord, Streams::new(Box::new(main)))
    };

    let init = (Some(killcord), streams, init, fun);

    future::loop_fn(init, move |(killcord, streams, state, mut fun)| {
        streams.into_future()
            .and_then(move |(item, mut streams)| {
                match item {
                    Some(m) => A(fun(state, Some(m)).then(move |result| {
                        match result {
                            Ok(State::Ready(a)) =>
                                A(future::ok(Loop::Continue((killcord, streams, a, fun)))),
                            Ok(State::Stream(a, s)) => {
                                let s = s.map_err(|void| unreachable(void));
                                streams.add(Box::new(s));
                                A(future::ok(Loop::Continue((killcord, streams, a, fun))))
                            }
                            Ok(State::Close(a)) =>
                                A(future::ok(Loop::Continue((None, streams, a, fun)))),
                            Ok(State::Done(x)) =>
                                B(future::ok(Loop::Break(Ok(Some(x))))),
                            Err(e) =>
                                B(future::ok(Loop::Break(Err(e))))
                        }
                    })),
                    // The stream has ended. We pass the terminal `None` to
                    // the function so it knows too and can perform any resource
                    // cleanup if necessary.
                    None => B(fun(state, None).then(|result| {
                        match result {
                            Ok(State::Done(x)) => Ok(Loop::Break(Ok(Some(x)))),
                            Ok(_) => Ok(Loop::Break(Ok(None))),
                            Err(e) => Ok(Loop::Break(Err(e)))
                        }
                    }))
                }
            })
            // We only need to handle errors here because `mpsc::Receiver`'s `Stream`
            // implementation defines an error type `()` which it never produces.
            // `Stream`s that have been added via `State::Stream` are required to
            // use `Void` as their error type so we should never enter this `or_else`.
            .or_else(|((), _)| Ok(Loop::Break(Ok(None))))
    })
    .then(|r: Result<_, ()>| {
        match r {
            Ok(Ok(x)) => Ok(x),
            Ok(Err(e)) => Err(e),
            Err(()) => Ok(None)
        }
    })
}

type UnorderedStreams<T> =
    FuturesUnordered<StreamFuture<Box<dyn Stream<Item = T, Error = Void> + Send>>>;

struct Streams<T> {
    main: Box<dyn Stream<Item = T, Error = ()> + Send>,
    unordered: Option<UnorderedStreams<T>>,
    done: bool // true iff `main` has ended
}

impl<T> Streams<T> {
    fn new(main: Box<dyn Stream<Item = T, Error = ()> + Send>) -> Self {
        Streams { main, unordered: None, done: false }
    }

    fn add(&mut self, s: Box<dyn Stream<Item = T, Error = Void> + Send>) {
        if let Some(ref mut t) = self.unordered {
            t.push(s.into_future())
        } else {
            let mut t = FuturesUnordered::new();
            t.push(s.into_future());
            self.unordered = Some(t)
        }
    }
}

impl<T> Stream for Streams<T> {
    type Item = T;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.main.poll()? {
            Async::Ready(Some(x)) => return Ok(Async::Ready(Some(x))),
            Async::Ready(None) => { self.done = true }
            Async::NotReady => {}
        }
        loop {
            match self.unordered {
                None => if self.done {
                    return Ok(Async::Ready(None))
                } else {
                    return Ok(Async::NotReady)
                }
                Some(ref mut set) => match set.poll() {
                    Ok(Async::Ready(Some((Some(x), stream)))) => {
                        set.push(stream.into_future());
                        return Ok(Async::Ready(Some(x)))
                    }
                    Ok(Async::Ready(Some((None, _)))) => continue,
                    Ok(Async::Ready(None)) => if self.done {
                        return Ok(Async::Ready(None))
                    } else {
                        return Ok(Async::NotReady)
                    }
                    Ok(Async::NotReady) => return Ok(Async::NotReady),
                    Err((void, _)) => unreachable(void)
                }
            }
        }
    }
}