/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

// The ancestors of the current node are itself, plus the union of all ancestors of all parents.
// Have a Vec of current generation nodes - as they're output, push their parents onto the next
// generation Vec. Once current generation Vec is empty, rotate.

use std::collections::hash_set::IntoIter;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::sync::Arc;

use anyhow::Error;
use changeset_fetcher::ArcChangesetFetcher;
use cloned::cloned;
use context::CoreContext;
use futures::FutureExt;
use futures::TryFutureExt;
use futures_ext::StreamExt;
use futures_old::future::Future;
use futures_old::stream::iter_ok;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use maplit::hashset;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::errors::*;
use crate::BonsaiNodeStream;
use crate::IntersectNodeStream;
use crate::UniqueHeap;

pub struct AncestorsNodeStream {
    ctx: CoreContext,
    changeset_fetcher: ArcChangesetFetcher,
    next_generation: BTreeMap<Generation, HashSet<ChangesetId>>,
    pending_changesets: Box<dyn Stream<Item = (ChangesetId, Generation), Error = Error> + Send>,
    drain: IntoIter<ChangesetId>,

    // max heap of all relevant unique generation numbers
    sorted_unique_generations: UniqueHeap<Generation>,
}

fn make_pending(
    ctx: CoreContext,
    changeset_fetcher: ArcChangesetFetcher,
    hashes: IntoIter<ChangesetId>,
) -> Box<dyn Stream<Item = (ChangesetId, Generation), Error = Error> + Send> {
    let size = hashes.size_hint().0;

    Box::new(
        iter_ok::<_, Error>(hashes)
            .map({
                cloned!(ctx, changeset_fetcher);
                move |hash| {
                    cloned!(ctx, changeset_fetcher);
                    async move { changeset_fetcher.get_parents(&ctx, hash).await }
                        .boxed()
                        .compat()
                        .map(|parents| parents.into_iter())
                        .map_err(|err| err.context(ErrorKind::ParentsFetchFailed))
                }
            })
            .buffered(size)
            .map(iter_ok::<_, Error>)
            .flatten()
            .and_then(move |node_cs| {
                cloned!(ctx, changeset_fetcher);
                async move { changeset_fetcher.get_generation_number(&ctx, node_cs).await }
                    .boxed()
                    .compat()
                    .map(move |gen_id| (node_cs, gen_id))
                    .map_err(|err| err.context(ErrorKind::GenerationFetchFailed))
            }),
    )
}

impl AncestorsNodeStream {
    pub fn new(
        ctx: CoreContext,
        changeset_fetcher: &ArcChangesetFetcher,
        hash: ChangesetId,
    ) -> Self {
        let node_set: HashSet<ChangesetId> = hashset! {hash};
        AncestorsNodeStream {
            ctx: ctx.clone(),
            changeset_fetcher: changeset_fetcher.clone(),
            next_generation: BTreeMap::new(),
            pending_changesets: make_pending(
                ctx,
                changeset_fetcher.clone(),
                node_set.clone().into_iter(),
            ),
            drain: node_set.into_iter(),
            sorted_unique_generations: UniqueHeap::new(),
        }
    }
}

impl Stream for AncestorsNodeStream {
    type Item = ChangesetId;
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // Empty the drain if any - return all items for this generation
        let next_in_drain = self.drain.next();
        if next_in_drain.is_some() {
            return Ok(Async::Ready(next_in_drain));
        }

        // Wait until we've drained pending_changesets - we can't continue until we know about all
        // parents of the just-output generation
        loop {
            match self.pending_changesets.poll()? {
                Async::Ready(Some((hash, generation))) => {
                    self.next_generation
                        .entry(generation)
                        .or_insert_with(HashSet::new)
                        .insert(hash);
                    // insert into our sorted list of generations
                    self.sorted_unique_generations.push(generation);
                }
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => break,
            };
        }

        if self.next_generation.is_empty() {
            // All parents output - nothing more to send
            return Ok(Async::Ready(None));
        }

        let highest_generation = self
            .sorted_unique_generations
            .pop()
            .expect("Expected a non empty heap of generations");
        let current_generation = self
            .next_generation
            .remove(&highest_generation)
            .expect("Highest generation doesn't exist");
        self.pending_changesets = make_pending(
            self.ctx.clone(),
            self.changeset_fetcher.clone(),
            current_generation.clone().into_iter(),
        );
        self.drain = current_generation.into_iter();
        Ok(Async::Ready(Some(self.drain.next().expect(
            "Cannot create a generation without at least one node hash",
        ))))
    }
}

pub fn common_ancestors<I>(
    ctx: CoreContext,
    changeset_fetcher: ArcChangesetFetcher,
    nodes: I,
) -> BonsaiNodeStream
where
    I: IntoIterator<Item = ChangesetId>,
{
    let nodes_iter = nodes.into_iter().map({
        cloned!(ctx, changeset_fetcher);
        move |node| AncestorsNodeStream::new(ctx.clone(), &changeset_fetcher, node).boxify()
    });

    IntersectNodeStream::new(ctx, &Arc::new(changeset_fetcher), nodes_iter).boxify()
}

pub fn greatest_common_ancestor<I>(
    ctx: CoreContext,
    changeset_fetcher: ArcChangesetFetcher,
    nodes: I,
) -> BonsaiNodeStream
where
    I: IntoIterator<Item = ChangesetId>,
{
    common_ancestors(ctx, changeset_fetcher, nodes)
        .take(1)
        .boxify()
}

#[cfg(test)]
mod test {
    use fbinit::FacebookInit;
    use revset_test_helper::assert_changesets_sequence;
    use revset_test_helper::string_to_bonsai;

    use super::*;
    use crate::fixtures::Linear;
    use crate::fixtures::MergeUneven;
    use crate::fixtures::TestRepoFixture;
    use crate::fixtures::UnsharedMergeUneven;
    use crate::tests::TestChangesetFetcher;

    #[fbinit::test]
    async fn linear_ancestors(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = AncestorsNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await,
                string_to_bonsai(fb, &repo, "0ed509bf086fadcb8a8a5384dc3b550729b0fc17").await,
                string_to_bonsai(fb, &repo, "eed3a8c0ec67b6a6fe2eb3543334df3f0b4f202b").await,
                string_to_bonsai(fb, &repo, "cb15ca4a43a59acff5388cea9648c162afde8372").await,
                string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await,
                string_to_bonsai(fb, &repo, "607314ef579bd2407752361ba1b0c1729d08b281").await,
                string_to_bonsai(fb, &repo, "3e0e761030db6e479a7fb58b12881883f9f8c63f").await,
                string_to_bonsai(fb, &repo, "2d7d4ba9ce0a6ffd222de7785b249ead9c51c536").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn merge_ancestors_from_merge(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = AncestorsNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            string_to_bonsai(fb, &repo, "d35b1875cdd1ed2c687e86f1604b9d7e989450cb").await,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "d35b1875cdd1ed2c687e86f1604b9d7e989450cb").await,
                string_to_bonsai(fb, &repo, "264f01429683b3dd8042cb3979e8bf37007118bc").await,
                string_to_bonsai(fb, &repo, "5d43888a3c972fe68c224f93d41b30e9f888df7c").await,
                string_to_bonsai(fb, &repo, "fc2cef43395ff3a7b28159007f63d6529d2f41ca").await,
                string_to_bonsai(fb, &repo, "bc7b4d0f858c19e2474b03e442b8495fd7aeef33").await,
                string_to_bonsai(fb, &repo, "795b8133cf375f6d68d27c6c23db24cd5d0cd00f").await,
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "16839021e338500b3cf7c9b871c8a07351697d68").await,
                string_to_bonsai(fb, &repo, "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5").await,
                string_to_bonsai(fb, &repo, "b65231269f651cfe784fd1d97ef02a049a37b8a0").await,
                string_to_bonsai(fb, &repo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await,
                string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await,
                string_to_bonsai(fb, &repo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn merge_ancestors_one_branch(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = AncestorsNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            string_to_bonsai(fb, &repo, "16839021e338500b3cf7c9b871c8a07351697d68").await,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "16839021e338500b3cf7c9b871c8a07351697d68").await,
                string_to_bonsai(fb, &repo, "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5").await,
                string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await,
                string_to_bonsai(fb, &repo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn unshared_merge_all(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        // The UnsharedMergeUneven fixture has a commit after the merge. Pull in everything
        // by starting at the head and working back to the original unshared history commits
        let repo = UnsharedMergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = AncestorsNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            string_to_bonsai(fb, &repo, "dd993aab2bed7276e17c88470286ba8459ba6d94").await,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "dd993aab2bed7276e17c88470286ba8459ba6d94").await,
                string_to_bonsai(fb, &repo, "9c6dd4e2c2f43c89613b094efb426cc42afdee2a").await,
                string_to_bonsai(fb, &repo, "64011f64aaf9c2ad2e674f57c033987da4016f51").await,
                string_to_bonsai(fb, &repo, "c1d5375bf73caab8725d759eaca56037c725c7d1").await,
                string_to_bonsai(fb, &repo, "e819f2dd9a01d3e63d9a93e298968df275e6ad7c").await,
                string_to_bonsai(fb, &repo, "5a3e8d5a475ec07895e64ec1e1b2ec09bfa70e4e").await,
                string_to_bonsai(fb, &repo, "76096af83f52cc9a225ccfd8ddfb05ea18132343").await,
                string_to_bonsai(fb, &repo, "33fb49d8a47b29290f5163e30b294339c89505a2").await,
                string_to_bonsai(fb, &repo, "03b0589d9788870817d03ce7b87516648ed5b33a").await,
                string_to_bonsai(fb, &repo, "2fa8b4ee6803a18db4649a3843a723ef1dfe852b").await,
                string_to_bonsai(fb, &repo, "f01e186c165a2fbe931fd1bf4454235398c591c9").await,
                string_to_bonsai(fb, &repo, "163adc0d0f5d2eb0695ca123addcb92bab202096").await,
                string_to_bonsai(fb, &repo, "0b94a2881dda90f0d64db5fae3ee5695a38e7c8f").await,
                string_to_bonsai(fb, &repo, "eee492dcdeaae18f91822c4359dd516992e0dbcd").await,
                string_to_bonsai(fb, &repo, "f61fdc0ddafd63503dcd8eed8994ec685bfc8941").await,
                string_to_bonsai(fb, &repo, "3775a86c64cceeaf68ffe3f012fc90774c42002b").await,
                string_to_bonsai(fb, &repo, "36ff88dd69c9966c9fad9d6d0457c52153039dde").await,
                string_to_bonsai(fb, &repo, "1700524113b1a3b1806560341009684b4378660b").await,
                string_to_bonsai(fb, &repo, "9d374b7e8180f933e3043ad1ffab0a9f95e2bac6").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn no_common_ancestor(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = UnsharedMergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = greatest_common_ancestor(
            ctx.clone(),
            changeset_fetcher,
            vec![
                string_to_bonsai(fb, &repo, "64011f64aaf9c2ad2e674f57c033987da4016f51").await,
                string_to_bonsai(fb, &repo, "1700524113b1a3b1806560341009684b4378660b").await,
            ],
        );
        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn greatest_common_ancestor_different_branches(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = greatest_common_ancestor(
            ctx.clone(),
            changeset_fetcher,
            vec![
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await,
            ],
        );
        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![string_to_bonsai(fb, &repo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn greatest_common_ancestor_same_branch(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = greatest_common_ancestor(
            ctx.clone(),
            changeset_fetcher,
            vec![
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "264f01429683b3dd8042cb3979e8bf37007118bc").await,
            ],
        );
        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn all_common_ancestors_different_branches(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = common_ancestors(
            ctx.clone(),
            changeset_fetcher,
            vec![
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await,
            ],
        );
        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![string_to_bonsai(fb, &repo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn all_common_ancestors_same_branch(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodestream = common_ancestors(
            ctx.clone(),
            changeset_fetcher,
            vec![
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "264f01429683b3dd8042cb3979e8bf37007118bc").await,
            ],
        );
        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
                string_to_bonsai(fb, &repo, "b65231269f651cfe784fd1d97ef02a049a37b8a0").await,
                string_to_bonsai(fb, &repo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await,
                string_to_bonsai(fb, &repo, "15c40d0abc36d47fb51c8eaec51ac7aad31f669c").await,
            ],
            nodestream,
        )
        .await;
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use mononoke_types::ChangesetId;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("repo error checking for node: {0}")]
    RepoNodeError(HgNodeHash),
    #[error("repo error checking for changeset: {0}")]
    RepoChangesetError(ChangesetId),
    #[error("could not fetch node generation")]
    GenerationFetchFailed,
    #[error("failed to fetch parent nodes")]
    ParentsFetchFailed,
    #[error("Bonsai mapping not found for {0}")]
    BonsaiMappingNotFound(HgChangesetId),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::hash_map::IntoIter;
use std::collections::HashMap;

use anyhow::Error;
use changeset_fetcher::ArcChangesetFetcher;
use context::CoreContext;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::setcommon::*;
use crate::BonsaiNodeStream;

pub struct IntersectNodeStream {
    inputs: Vec<(
        BonsaiInputStream,
        Poll<Option<(ChangesetId, Generation)>, Error>,
    )>,
    current_generation: Option<Generation>,
    accumulator: HashMap<ChangesetId, usize>,
    drain: Option<IntoIter<ChangesetId, usize>>,
}

impl IntersectNodeStream {
    pub fn new<I>(ctx: CoreContext, changeset_fetcher: &ArcChangesetFetcher, inputs: I) -> Self
    where
        I: IntoIterator<Item = BonsaiNodeStream>,
    {
        let csid_and_gen = inputs.into_iter().map({
            move |i| {
                (
                    add_generations_by_bonsai(ctx.clone(), i, changeset_fetcher.clone()),
                    Ok(Async::NotReady),
                )
            }
        });
        Self {
            inputs: csid_and_gen.collect(),
            current_generation: None,
            accumulator: HashMap::new(),
            drain: None,
        }
    }

    fn update_current_generation(&mut self) {
        if all_inputs_ready(&self.inputs) {
            self.current_generation = self
                .inputs
                .iter()
                .filter_map(|(_, state)| match state {
                    Ok(Async::Ready(Some((_, gen_id)))) => Some(*gen_id),
                    Ok(Async::NotReady) => panic!("All states ready, yet some not ready!"),
                    _ => None,
                })
                .min();
        }
    }

    fn accumulate_nodes(&mut self) {
        let mut found_csids = false;
        for &mut (_, ref mut state) in self.inputs.iter_mut() {
            if let Ok(Async::Ready(Some((csid, gen_id)))) = *state {
                if Some(gen_id) == self.current_generation {
                    *self.accumulator.entry(csid).or_insert(0) += 1;
                }
                // Inputs of higher generation than the current one get consumed and dropped
                if Some(gen_id) >= self.current_generation {
                    found_csids = true;
                    *state = Ok(Async::NotReady);
                }
            }
        }
        if !found_csids {
            self.current_generation = None;
        }
    }

    fn any_input_finished(&self) -> bool {
        if self.inputs.is_empty() {
            true
        } else {
            self.inputs
                .iter()
                .map(|(_, state)| match state {
                    &Ok(Async::Ready(None)) => true,
                    _ => false,
                })
                .any(|done| done)
        }
    }
}

impl Stream for IntersectNodeStream {
    type Item = ChangesetId;
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // This feels wrong, but in practice it's fine - it should be quick to hit a return, and
        // the standard futures_old::executor expects you to only return NotReady if blocked on I/O.
        loop {
            // Start by trying to turn as many NotReady as possible into real items
            poll_all_inputs(&mut self.inputs);

            // Empty the drain if any - return all items for this generation
            while self.drain.is_some() {
                let next_in_drain = self.drain.as_mut().and_then(|drain| drain.next());
                if let Some((csid, count)) = next_in_drain {
                    if count == self.inputs.len() {
                        return Ok(Async::Ready(Some(csid)));
                    }
                } else {
                    self.drain = None;
                }
            }

            // Return any errors
            {
                if self.inputs.iter().any(|(_, state)| state.is_err()) {
                    let inputs = std::mem::take(&mut self.inputs);
                    let (_, err) = inputs
                        .into_iter()
                        .find(|(_, state)| state.is_err())
                        .unwrap();
                    return Err(err.unwrap_err());
                }
            }

            // If any input is not ready (we polled above), wait for them all to be ready
            if !all_inputs_ready(&self.inputs) {
                return Ok(Async::NotReady);
            }

            match self.current_generation {
                None => {
                    if self.accumulator.is_empty() {
                        self.update_current_generation();
                    } else {
                        let full_accumulator = std::mem::take(&mut self.accumulator);
                        self.drain = Some(full_accumulator.into_iter());
                    }
                }
                Some(_) => self.accumulate_nodes(),
            }
            // If we cannot ever output another node, we're done.
            if self.drain.is_none() && self.accumulator.is_empty() && self.any_input_finished() {
                return Ok(Async::Ready(None));
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use context::CoreContext;
    use failure_ext::err_downcast;
    use fbinit::FacebookInit;
    use futures::compat::Stream01CompatExt;
    use futures::stream::StreamExt as _;
    use futures_ext::StreamExt;
    use futures_old::executor::spawn;
    use revset_test_helper::assert_changesets_sequence;
    use revset_test_helper::single_changeset_id;
    use revset_test_helper::string_to_bonsai;

    use super::*;
    use crate::errors::ErrorKind;
    use crate::fixtures::Linear;
    use crate::fixtures::TestRepoFixture;
    use crate::fixtures::UnsharedMergeEven;
    use crate::fixtures::UnsharedMergeUneven;
    use crate::setcommon::NotReadyEmptyStream;
    use crate::tests::get_single_bonsai_streams;
    use crate::tests::TestChangesetFetcher;
    use crate::BonsaiNodeStream;
    use crate::UnionNodeStream;

    #[fbinit::test]
    async fn intersect_identical_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let head_csid = string_to_bonsai(fb, &repo, hash).await;

        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), head_csid.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), head_csid.clone(), &repo).boxify(),
        ];

        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![head_csid], nodestream).await;
    }

    #[fbinit::test]
    async fn intersect_three_different_nodes(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_a947 =
            string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await;
        let bcs_3c15 =
            string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await;
        let bcs_d0a = string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await;
        // Note that these are *not* in generation order deliberately.
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), bcs_a947, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_3c15, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
        ];

        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn intersect_three_identical_nodes(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_d0a = string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await;
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
        ];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![bcs_d0a], nodestream).await;
    }

    #[fbinit::test]
    async fn intersect_nesting(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_3c15 =
            string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await;
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), bcs_3c15.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_3c15.clone(), &repo).boxify(),
        ];

        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let inputs: Vec<BonsaiNodeStream> = vec![
            nodestream,
            single_changeset_id(ctx.clone(), bcs_3c15.clone(), &repo).boxify(),
        ];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![bcs_3c15.clone()], nodestream).await;
    }

    #[fbinit::test]
    async fn intersection_of_unions(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash1 = "d0a361e9022d226ae52f689667bd7d212a19cfe0";
        let hash2 = "3c15267ebf11807f3d772eb891272b911ec68759";
        let hash3 = "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157";

        let inputs = get_single_bonsai_streams(ctx.clone(), &repo, &[hash1, hash2]).await;
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // This set has a different node sequence, so that we can demonstrate that we skip nodes
        // when they're not going to contribute.
        let inputs = get_single_bonsai_streams(ctx.clone(), &repo, &[hash3, hash2, hash1]).await;
        let nodestream2 =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let inputs: Vec<BonsaiNodeStream> = vec![nodestream, nodestream2];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await,
                string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn intersect_error_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let changeset = string_to_bonsai(fb, &repo, hash).await;

        let inputs: Vec<BonsaiNodeStream> = vec![
            RepoErrorStream { item: changeset }.boxify(),
            single_changeset_id(ctx.clone(), changeset, &repo).boxify(),
        ];
        let mut nodestream = spawn(
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify(),
        );

        match nodestream.wait_stream() {
            Some(Err(err)) => match err_downcast!(err, err: ErrorKind => err) {
                Ok(ErrorKind::RepoChangesetError(cs)) => assert_eq!(cs, changeset),
                Ok(bad) => panic!("unexpected error {:?}", bad),
                Err(bad) => panic!("unknown error {:?}", bad),
            },
            Some(Ok(bad)) => panic!("unexpected success {:?}", bad),
            None => panic!("no result"),
        };
    }

    #[fbinit::test]
    async fn intersect_nothing(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let inputs: Vec<BonsaiNodeStream> = vec![];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter());
        assert_changesets_sequence(&ctx, &repo, vec![], nodestream.boxify()).await;
    }

    #[fbinit::test]
    async fn slow_ready_intersect_nothing(fb: FacebookInit) {
        // Tests that we handle an input staying at NotReady for a while without panicking
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher = Arc::new(TestChangesetFetcher::new(repo));

        let inputs: Vec<BonsaiNodeStream> = vec![NotReadyEmptyStream::new(10).boxify()];
        let mut nodestream =
            IntersectNodeStream::new(ctx, &changeset_fetcher, inputs.into_iter()).compat();
        assert!(nodestream.next().await.is_none());
    }

    #[fbinit::test]
    async fn intersect_unshared_merge_even(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = UnsharedMergeEven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        // Post-merge, merge, and both unshared branches
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "7fe9947f101acb4acf7d945e69f0d6ce76a81113",
                "d592490c4386cdb3373dd93af04d563de199b2fb",
                "33fb49d8a47b29290f5163e30b294339c89505a2",
                "03b0589d9788870817d03ce7b87516648ed5b33a",
            ],
        )
        .await;

        let left_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // Four commits from one branch
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "03b0589d9788870817d03ce7b87516648ed5b33a",
                "2fa8b4ee6803a18db4649a3843a723ef1dfe852b",
                "0b94a2881dda90f0d64db5fae3ee5695a38e7c8f",
                "f61fdc0ddafd63503dcd8eed8994ec685bfc8941",
            ],
        )
        .await;
        let right_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let inputs: Vec<BonsaiNodeStream> = vec![left_nodestream, right_nodestream];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter());

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![string_to_bonsai(fb, &repo, "03b0589d9788870817d03ce7b87516648ed5b33a").await],
            nodestream.boxify(),
        )
        .await;
    }

    #[fbinit::test]
    async fn intersect_unshared_merge_uneven(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = UnsharedMergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        // Post-merge, merge, and both unshared branches
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "dd993aab2bed7276e17c88470286ba8459ba6d94",
                "9c6dd4e2c2f43c89613b094efb426cc42afdee2a",
                "64011f64aaf9c2ad2e674f57c033987da4016f51",
                "03b0589d9788870817d03ce7b87516648ed5b33a",
            ],
        )
        .await;

        let left_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // Four commits from one branch
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "03b0589d9788870817d03ce7b87516648ed5b33a",
                "2fa8b4ee6803a18db4649a3843a723ef1dfe852b",
                "0b94a2881dda90f0d64db5fae3ee5695a38e7c8f",
                "f61fdc0ddafd63503dcd8eed8994ec685bfc8941",
            ],
        )
        .await;
        let right_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let inputs: Vec<BonsaiNodeStream> = vec![left_nodestream, right_nodestream];
        let nodestream =
            IntersectNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![string_to_bonsai(fb, &repo, "03b0589d9788870817d03ce7b87516648ed5b33a").await],
            nodestream,
        )
        .await;
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use futures_ext::BoxStream;
use mononoke_types::ChangesetId;

mod setcommon;
pub use setcommon::add_generations_by_bonsai;

mod intersectnodestream;
pub use crate::intersectnodestream::IntersectNodeStream;

mod unionnodestream;
pub use crate::unionnodestream::UnionNodeStream;

mod setdifferencenodestream;
pub use crate::setdifferencenodestream::SetDifferenceNodeStream;

pub mod errors;
pub use crate::errors::ErrorKind;

pub type BonsaiNodeStream = BoxStream<ChangesetId, Error>;

mod validation;
pub use crate::validation::ValidateNodeStream;

mod ancestors;
use uniqueheap::UniqueHeap;

pub use crate::ancestors::common_ancestors;
pub use crate::ancestors::greatest_common_ancestor;
pub use crate::ancestors::AncestorsNodeStream;
pub use crate::test::*;

#[cfg(test)]
mod test {
    pub use ascii;
    pub use fixtures;
    pub use quickcheck;
}
#[cfg(not(test))]
mod test {}
#[cfg(test)]
mod quickchecks;
#[cfg(test)]
mod tests;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use crate::tests::*;

#[cfg(test)]
mod test {
    use std::collections::HashSet;
    use std::sync::Arc;

    use blobrepo::BlobRepo;
    use bookmarks::BookmarksMaybeStaleExt;
    use bookmarks::BookmarksRef;
    use changeset_fetcher::ArcChangesetFetcher;
    use changeset_fetcher::ChangesetFetcherArc;
    use cloned::cloned;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use futures::compat::Stream01CompatExt;
    use futures::stream::StreamExt as _;
    use futures::TryStreamExt;
    use futures_ext::StreamExt;
    use futures_old::Stream;
    use mononoke_types::ChangesetId;
    use quickcheck::quickcheck;
    use quickcheck::Arbitrary;
    use quickcheck::Gen;
    use rand::seq::SliceRandom;
    use rand::thread_rng;
    use rand::Rng;
    use revset_test_helper::single_changeset_id;

    use super::*;
    use crate::ancestors::AncestorsNodeStream;
    use crate::fixtures::BranchEven;
    use crate::fixtures::BranchUneven;
    use crate::fixtures::BranchWide;
    use crate::fixtures::Linear;
    use crate::fixtures::MergeEven;
    use crate::fixtures::MergeUneven;
    use crate::fixtures::TestRepoFixture;
    use crate::fixtures::UnsharedMergeEven;
    use crate::fixtures::UnsharedMergeUneven;
    use crate::intersectnodestream::IntersectNodeStream;
    use crate::setdifferencenodestream::SetDifferenceNodeStream;
    use crate::unionnodestream::UnionNodeStream;
    use crate::validation::ValidateNodeStream;
    use crate::BonsaiNodeStream;

    #[derive(Clone, Copy, Debug)]
    enum RevsetEntry {
        SingleNode(Option<ChangesetId>),
        SetDifference,
        Intersect(usize),
        Union(usize),
    }

    #[derive(Clone, Debug)]
    pub struct RevsetSpec {
        rp_entries: Vec<RevsetEntry>,
    }

    async fn get_changesets_from_repo(ctx: CoreContext, repo: &BlobRepo) -> Vec<ChangesetId> {
        let mut all_changesets_stream = repo
            .bookmarks()
            .get_heads_maybe_stale(ctx.clone())
            .compat() // conversion is needed as AncestorsNodeStream is an OldStream
            .map({
                cloned!(ctx);
                move |head| {
                    AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), head)
                }
            })
            .flatten()
            .compat();

        let mut all_changesets: Vec<ChangesetId> = Vec::new();
        loop {
            all_changesets.push(match all_changesets_stream.next().await {
                None => break,
                Some(changeset) => changeset.expect("Failed to get changesets from repo"),
            });
        }

        assert!(!all_changesets.is_empty(), "Repo has no changesets");
        all_changesets
    }

    impl RevsetSpec {
        pub async fn add_hashes<G>(&mut self, ctx: CoreContext, repo: &BlobRepo, random: &mut G)
        where
            G: Rng,
        {
            let all_changesets = get_changesets_from_repo(ctx, repo).await;
            for elem in self.rp_entries.iter_mut() {
                if let &mut RevsetEntry::SingleNode(None) = elem {
                    *elem =
                        RevsetEntry::SingleNode(all_changesets.as_slice().choose(random).cloned());
                }
            }
        }

        pub fn as_hashes(&self) -> HashSet<ChangesetId> {
            let mut output: Vec<HashSet<ChangesetId>> = Vec::new();
            for entry in self.rp_entries.iter() {
                match *entry {
                    RevsetEntry::SingleNode(None) => panic!("You need to add_hashes first!"),
                    RevsetEntry::SingleNode(Some(hash)) => {
                        let mut item = HashSet::new();
                        item.insert(hash);
                        output.push(item)
                    }
                    RevsetEntry::SetDifference => {
                        let keep = output.pop().expect("No keep for setdifference");
                        let remove = output.pop().expect("No remove for setdifference");
                        output.push(keep.difference(&remove).copied().collect())
                    }
                    RevsetEntry::Union(size) => {
                        let idx = output.len() - size;
                        let mut inputs = output.split_off(idx).into_iter();
                        let first = inputs.next().expect("No first element");
                        output.push(inputs.fold(first, |a, b| a.union(&b).copied().collect()))
                    }
                    RevsetEntry::Intersect(size) => {
                        let idx = output.len() - size;
                        let mut inputs = output.split_off(idx).into_iter();
                        let first = inputs.next().expect("No first element");
                        output
                            .push(inputs.fold(first, |a, b| a.intersection(&b).copied().collect()))
                    }
                }
            }
            assert!(
                output.len() == 1,
                "output should have been length 1, was {}",
                output.len()
            );
            output.pop().expect("No revisions").into_iter().collect()
        }

        pub fn as_revset(&self, ctx: CoreContext, repo: BlobRepo) -> BonsaiNodeStream {
            let mut output: Vec<BonsaiNodeStream> = Vec::with_capacity(self.rp_entries.len());
            let changeset_fetcher: ArcChangesetFetcher =
                Arc::new(TestChangesetFetcher::new(repo.clone()));
            for entry in self.rp_entries.iter() {
                let next_node = ValidateNodeStream::new(
                    ctx.clone(),
                    match *entry {
                        RevsetEntry::SingleNode(None) => panic!("You need to add_hashes first!"),
                        RevsetEntry::SingleNode(Some(hash)) => {
                            single_changeset_id(ctx.clone(), hash, &repo).boxify()
                        }
                        RevsetEntry::SetDifference => {
                            let keep = output.pop().expect("No keep for setdifference");
                            let remove = output.pop().expect("No remove for setdifference");
                            SetDifferenceNodeStream::new(
                                ctx.clone(),
                                &changeset_fetcher,
                                keep,
                                remove,
                            )
                            .boxify()
                        }
                        RevsetEntry::Union(size) => {
                            let idx = output.len() - size;
                            let inputs = output.split_off(idx);

                            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs).boxify()
                        }
                        RevsetEntry::Intersect(size) => {
                            let idx = output.len() - size;
                            let inputs = output.split_off(idx);
                            IntersectNodeStream::new(
                                ctx.clone(),
                                &repo.changeset_fetcher_arc(),
                                inputs,
                            )
                            .boxify()
                        }
                    },
                    &repo.changeset_fetcher_arc(),
                )
                .boxify();
                output.push(next_node);
            }
            assert!(
                output.len() == 1,
                "output should have been length 1, was {}",
                output.len()
            );
            output.pop().expect("No revset entries")
        }
    }

    impl Arbitrary for RevsetSpec {
        fn arbitrary(g: &mut Gen) -> Self {
            let mut revset: Vec<RevsetEntry> = Vec::with_capacity(g.size());
            let mut revspecs_in_set: usize = 0;

            for _ in 0..g.size() {
                if revspecs_in_set == 0 {
                    // Can't add a set operator if we have don't have at least one node
                    revset.push(RevsetEntry::SingleNode(None));
                } else {
                    let input_count = (usize::arbitrary(g) % revspecs_in_set) + 1;
                    revset.push(
                        // Bias towards SingleNode if we only have 1 rev
                        match g.choose(&[0, 1, 2, 3]).unwrap() {
                            0 => RevsetEntry::SingleNode(None),
                            1 => {
                                if revspecs_in_set >= 2 {
                                    revspecs_in_set -= 2;
                                    RevsetEntry::SetDifference
                                } else {
                                    RevsetEntry::SingleNode(None)
                                }
                            }
                            2 => {
                                revspecs_in_set -= input_count;
                                RevsetEntry::Intersect(input_count)
                            }
                            3 => {
                                revspecs_in_set -= input_count;
                                RevsetEntry::Union(input_count)
                            }
                            _ => panic!("Range returned too wide a variation"),
                        },
                    );
                }
                revspecs_in_set += 1;
            }
            assert!(revspecs_in_set > 0, "Did not produce enough revs");

            if revspecs_in_set > 1 {
                revset.push(match bool::arbitrary(g) {
                    true => RevsetEntry::Intersect(revspecs_in_set),
                    false => RevsetEntry::Union(revspecs_in_set),
                });
            }

            RevsetSpec { rp_entries: revset }
        }

        // TODO(simonfar) We should implement shrink(), but we face the issue of ensuring that the
        // resulting revset only contains one final item.
        // Rough sketch: Take the last element of the Vec, so that we're using the same final reduction
        // type. Vector shrink the rest of the Vec using the standard shrinker. Re-add the final
        // reduction type. Note that we then need to handle the case where the final reduction type
        // is a SetDifference by pure chance.
    }

    async fn match_hashset_to_revset(
        ctx: CoreContext,
        repo: BlobRepo,
        mut set: RevsetSpec,
    ) -> bool {
        set.add_hashes(ctx.clone(), &repo, &mut thread_rng()).await;
        let mut hashes = set.as_hashes();
        let mut nodestream = set.as_revset(ctx, repo).compat();

        while !hashes.is_empty() {
            let hash = nodestream
                .next()
                .await
                .expect("Unexpected end of stream")
                .expect("Unexpected error");
            if !hashes.remove(&hash) {
                return false;
            }
        }
        nodestream.next().await.is_none() && hashes.is_empty()
    }

    // This is slightly icky. I would like to construct $test_name as setops_$repo, but concat_idents!
    // does not work the way I'd like it to. For now, make the user of this macro pass in both idents
    macro_rules! quickcheck_setops {
        ($test_name:ident, $repo:ident) => {
            #[test]
            fn $test_name() {
                #[tokio::main(flavor = "current_thread")]
                async fn prop(fb: FacebookInit, set: RevsetSpec) -> bool {
                    let ctx = CoreContext::test_mock(fb);
                    let repo = $repo::getrepo(fb).await;
                    match_hashset_to_revset(ctx, repo, set).await
                }

                quickcheck(prop as fn(FacebookInit, RevsetSpec) -> bool)
            }
        };
    }

    quickcheck_setops!(setops_branch_even, BranchEven);
    quickcheck_setops!(setops_branch_uneven, BranchUneven);
    quickcheck_setops!(setops_branch_wide, BranchWide);
    quickcheck_setops!(setops_linear, Linear);
    quickcheck_setops!(setops_merge_even, MergeEven);
    quickcheck_setops!(setops_merge_uneven, MergeUneven);
    quickcheck_setops!(setops_unshared_merge_even, UnsharedMergeEven);
    quickcheck_setops!(setops_unshared_merge_uneven, UnsharedMergeUneven);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use changeset_fetcher::ArcChangesetFetcher;
use cloned::cloned;
use context::CoreContext;
use futures::FutureExt;
use futures::TryFutureExt;
use futures_ext::BoxStream;
use futures_ext::StreamExt;
use futures_old::future::Future;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::errors::*;
use crate::BonsaiNodeStream;

type GenericStream<T> = BoxStream<(T, Generation), Error>;
pub type BonsaiInputStream = GenericStream<ChangesetId>;

pub fn add_generations_by_bonsai(
    ctx: CoreContext,
    stream: BonsaiNodeStream,
    changeset_fetcher: ArcChangesetFetcher,
) -> BonsaiInputStream {
    stream
        .map(move |changesetid| {
            cloned!(ctx, changeset_fetcher);
            async move {
                changeset_fetcher
                    .get_generation_number(&ctx, changesetid)
                    .await
            }
            .boxed()
            .compat()
            .map(move |gen_id| (changesetid, gen_id))
            .map_err(|err| err.context(ErrorKind::GenerationFetchFailed))
        })
        .buffered(100)
        .boxify()
}

pub fn all_inputs_ready<T>(
    inputs: &[(GenericStream<T>, Poll<Option<(T, Generation)>, Error>)],
) -> bool {
    inputs
        .iter()
        .map(|(_, state)| match state {
            Err(_) => false,
            Ok(p) => p.is_ready(),
        })
        .all(|ready| ready)
}

pub fn poll_all_inputs<T>(inputs: &mut [(GenericStream<T>, Poll<Option<(T, Generation)>, Error>)]) {
    for (input, state) in inputs.iter_mut() {
        if let Ok(Async::NotReady) = *state {
            *state = input.poll();
        }
    }
}

#[cfg(test)]
mod test_utils {
    use std::marker::PhantomData;

    use anyhow::bail;
    use futures_old::task;
    use mercurial_types::HgNodeHash;

    use super::*;

    pub struct NotReadyEmptyStream<T> {
        pub poll_count: usize,
        __phantom: PhantomData<T>,
    }

    impl<T> NotReadyEmptyStream<T> {
        pub fn new(poll_count: usize) -> Self {
            Self {
                poll_count,
                __phantom: PhantomData,
            }
        }
    }

    impl<T> Stream for NotReadyEmptyStream<T> {
        type Item = T;
        type Error = Error;

        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
            if self.poll_count == 0 {
                Ok(Async::Ready(None))
            } else {
                self.poll_count -= 1;
                task::current().notify();
                Ok(Async::NotReady)
            }
        }
    }

    pub struct RepoErrorStream<T> {
        pub item: T,
    }

    impl Stream for RepoErrorStream<HgNodeHash> {
        type Item = HgNodeHash;
        type Error = Error;

        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
            bail!(ErrorKind::RepoNodeError(self.item));
        }
    }

    impl Stream for RepoErrorStream<ChangesetId> {
        type Item = ChangesetId;
        type Error = Error;

        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
            bail!(ErrorKind::RepoChangesetError(self.item));
        }
    }
}

#[cfg(test)]
pub use test_utils::NotReadyEmptyStream;
#[cfg(test)]
pub use test_utils::RepoErrorStream;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;

use anyhow::Error;
use anyhow::Result;
use changeset_fetcher::ArcChangesetFetcher;
use context::CoreContext;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::setcommon::*;
use crate::BonsaiNodeStream;

pub struct SetDifferenceNodeStream {
    keep_input: BonsaiInputStream,
    next_keep: Async<Option<(ChangesetId, Generation)>>,

    remove_input: BonsaiInputStream,
    next_remove: Async<Option<(ChangesetId, Generation)>>,

    remove_nodes: HashSet<ChangesetId>,
    remove_generation: Option<Generation>,
}

impl SetDifferenceNodeStream {
    pub fn new(
        ctx: CoreContext,
        changeset_fetcher: &ArcChangesetFetcher,
        keep_input: BonsaiNodeStream,
        remove_input: BonsaiNodeStream,
    ) -> SetDifferenceNodeStream {
        SetDifferenceNodeStream {
            keep_input: add_generations_by_bonsai(
                ctx.clone(),
                keep_input,
                changeset_fetcher.clone(),
            ),
            next_keep: Async::NotReady,
            remove_input: add_generations_by_bonsai(ctx, remove_input, changeset_fetcher.clone()),
            next_remove: Async::NotReady,
            remove_nodes: HashSet::new(),
            remove_generation: None,
        }
    }

    fn next_keep(&mut self) -> Result<&Async<Option<(ChangesetId, Generation)>>> {
        if self.next_keep.is_not_ready() {
            self.next_keep = self.keep_input.poll()?;
        }
        Ok(&self.next_keep)
    }

    fn next_remove(&mut self) -> Result<&Async<Option<(ChangesetId, Generation)>>> {
        if self.next_remove.is_not_ready() {
            self.next_remove = self.remove_input.poll()?;
        }
        Ok(&self.next_remove)
    }
}

impl Stream for SetDifferenceNodeStream {
    type Item = ChangesetId;
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // This feels wrong, but in practice it's fine - it should be quick to hit a return, and
        // the standard futures_old::executor expects you to only return NotReady if blocked on I/O.
        loop {
            let (keep_hash, keep_gen) = match self.next_keep()? {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => return Ok(Async::Ready(None)),
                Async::Ready(Some((hash, gen))) => (*hash, *gen),
            };

            // Clear nodes that won't affect future results
            if self.remove_generation != Some(keep_gen) {
                self.remove_nodes.clear();
                self.remove_generation = Some(keep_gen);
            }

            // Gather the current generation's remove hashes
            loop {
                let remove_hash = match self.next_remove()? {
                    Async::NotReady => return Ok(Async::NotReady),
                    Async::Ready(Some((hash, gen))) if *gen == keep_gen => *hash,
                    Async::Ready(Some((_, gen))) if *gen > keep_gen => {
                        // Refers to a generation that's already past (probably nothing on keep
                        // side of this generation). Skip it.
                        self.next_remove = Async::NotReady;
                        continue;
                    }
                    _ => break, // Either no more or gen < keep_gen
                };
                self.remove_nodes.insert(remove_hash);
                self.next_remove = Async::NotReady; // will cause polling of remove_input
            }

            self.next_keep = Async::NotReady; // will cause polling of keep_input

            if !self.remove_nodes.contains(&keep_hash) {
                return Ok(Async::Ready(Some(keep_hash)));
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use changeset_fetcher::ArcChangesetFetcher;
    use context::CoreContext;
    use failure_ext::err_downcast;
    use fbinit::FacebookInit;
    use futures::compat::Stream01CompatExt;
    use futures::stream::StreamExt as _;
    use futures_ext::StreamExt;
    use futures_old::executor::spawn;
    use revset_test_helper::assert_changesets_sequence;
    use revset_test_helper::single_changeset_id;
    use revset_test_helper::string_to_bonsai;

    use super::*;
    use crate::errors::ErrorKind;
    use crate::fixtures::Linear;
    use crate::fixtures::MergeEven;
    use crate::fixtures::MergeUneven;
    use crate::fixtures::TestRepoFixture;
    use crate::setcommon::NotReadyEmptyStream;
    use crate::tests::get_single_bonsai_streams;
    use crate::tests::TestChangesetFetcher;
    use crate::UnionNodeStream;

    #[fbinit::test]
    async fn difference_identical_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let changeset = string_to_bonsai(fb, &repo, hash).await;
        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            single_changeset_id(ctx.clone(), changeset.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), changeset.clone(), &repo).boxify(),
        )
        .boxify();
        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn difference_node_and_empty(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let changeset = string_to_bonsai(fb, &repo, hash).await;
        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            single_changeset_id(ctx.clone(), changeset.clone(), &repo).boxify(),
            NotReadyEmptyStream::new(0).boxify(),
        )
        .boxify();
        assert_changesets_sequence(&ctx, &repo, vec![changeset], nodestream).await;
    }

    #[fbinit::test]
    async fn difference_empty_and_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_id = string_to_bonsai(fb, &repo, "a5ffa77602a066db7d5cfb9fb5823a0895717c5a").await;

        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            NotReadyEmptyStream::new(0).boxify(),
            single_changeset_id(ctx.clone(), bcs_id, &repo).boxify(),
        )
        .boxify();

        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn difference_two_nodes(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_id_1 = string_to_bonsai(
            fb,
            &repo.clone(),
            "d0a361e9022d226ae52f689667bd7d212a19cfe0",
        )
        .await;
        let bcs_id_2 = string_to_bonsai(
            fb,
            &repo.clone(),
            "3c15267ebf11807f3d772eb891272b911ec68759",
        )
        .await;
        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            single_changeset_id(ctx.clone(), bcs_id_1.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_id_2, &repo).boxify(),
        )
        .boxify();

        assert_changesets_sequence(&ctx, &repo, vec![bcs_id_1], nodestream).await;
    }

    #[fbinit::test]
    async fn difference_error_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let changeset = string_to_bonsai(fb, &repo, hash).await;
        let mut nodestream = spawn(
            SetDifferenceNodeStream::new(
                ctx.clone(),
                &changeset_fetcher,
                RepoErrorStream {
                    item: changeset.clone(),
                }
                .boxify(),
                single_changeset_id(ctx.clone(), changeset, &repo).boxify(),
            )
            .boxify(),
        );

        match nodestream.wait_stream() {
            Some(Err(err)) => match err_downcast!(err, err: ErrorKind => err) {
                Ok(ErrorKind::RepoChangesetError(cs)) => assert_eq!(cs, changeset),
                Ok(bad) => panic!("unexpected error {:?}", bad),
                Err(bad) => panic!("unknown error {:?}", bad),
            },
            Some(Ok(bad)) => panic!("unexpected success {:?}", bad),
            None => panic!("no result"),
        };
    }

    #[fbinit::test]
    async fn slow_ready_difference_nothing(fb: FacebookInit) {
        // Tests that we handle an input staying at NotReady for a while without panicking
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher = Arc::new(TestChangesetFetcher::new(repo));

        let mut nodestream = SetDifferenceNodeStream::new(
            ctx,
            &changeset_fetcher,
            NotReadyEmptyStream::new(10).boxify(),
            NotReadyEmptyStream::new(10).boxify(),
        )
        .compat();

        assert!(nodestream.next().await.is_none());
    }

    #[fbinit::test]
    async fn difference_union_with_single_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "3c15267ebf11807f3d772eb891272b911ec68759",
                "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157",
                "d0a361e9022d226ae52f689667bd7d212a19cfe0",
            ],
        )
        .await;

        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let bcs_id = string_to_bonsai(
            fb,
            &repo.clone(),
            "3c15267ebf11807f3d772eb891272b911ec68759",
        )
        .await;
        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            nodestream,
            single_changeset_id(ctx.clone(), bcs_id, &repo).boxify(),
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await,
                string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn difference_single_node_with_union(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "3c15267ebf11807f3d772eb891272b911ec68759",
                "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157",
                "d0a361e9022d226ae52f689667bd7d212a19cfe0",
            ],
        )
        .await;
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let bcs_id = string_to_bonsai(
            fb,
            &repo.clone(),
            "3c15267ebf11807f3d772eb891272b911ec68759",
        )
        .await;
        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            single_changeset_id(ctx.clone(), bcs_id, &repo).boxify(),
            nodestream,
        )
        .boxify();

        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn difference_merge_even(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeEven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        // Top three commits in my hg log -G -r 'all()' output
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "1f6bc010883e397abeca773192f3370558ee1320",
                "4f7f3fd428bec1a48f9314414b063c706d9c1aed",
                "16839021e338500b3cf7c9b871c8a07351697d68",
            ],
        )
        .await;

        let left_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // Everything from base to just before merge on one side
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "4f7f3fd428bec1a48f9314414b063c706d9c1aed",
                "b65231269f651cfe784fd1d97ef02a049a37b8a0",
                "d7542c9db7f4c77dab4b315edd328edf1514952f",
                "15c40d0abc36d47fb51c8eaec51ac7aad31f669c",
            ],
        )
        .await;
        let right_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            left_nodestream,
            right_nodestream,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "1f6bc010883e397abeca773192f3370558ee1320").await,
                string_to_bonsai(fb, &repo, "16839021e338500b3cf7c9b871c8a07351697d68").await,
            ],
            nodestream,
        )
        .await;
    }

    #[fbinit::test]
    async fn difference_merge_uneven(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = MergeUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        // Merge commit, and one from each branch
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "d35b1875cdd1ed2c687e86f1604b9d7e989450cb",
                "4f7f3fd428bec1a48f9314414b063c706d9c1aed",
                "16839021e338500b3cf7c9b871c8a07351697d68",
            ],
        )
        .await;
        let left_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // Everything from base to just before merge on one side
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "16839021e338500b3cf7c9b871c8a07351697d68",
                "1d8a907f7b4bf50c6a09c16361e2205047ecc5e5",
                "3cda5c78aa35f0f5b09780d971197b51cad4613a",
                "15c40d0abc36d47fb51c8eaec51ac7aad31f669c",
            ],
        )
        .await;
        let right_nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let nodestream = SetDifferenceNodeStream::new(
            ctx.clone(),
            &changeset_fetcher,
            left_nodestream,
            right_nodestream,
        )
        .boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "d35b1875cdd1ed2c687e86f1604b9d7e989450cb").await,
                string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
            ],
            nodestream,
        )
        .await;
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::any::Any;
use std::collections::HashMap;

use anyhow::Error;
use async_trait::async_trait;
use blobrepo::BlobRepo;
use changeset_fetcher::ChangesetFetcher;
use changeset_fetcher::ChangesetFetcherRef;
use context::CoreContext;
use futures_ext::StreamExt;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use revset_test_helper::single_changeset_id;
use revset_test_helper::string_to_bonsai;

use crate::BonsaiNodeStream;

pub struct TestChangesetFetcher {
    repo: BlobRepo,
}

impl TestChangesetFetcher {
    pub fn new(repo: BlobRepo) -> Self {
        Self { repo }
    }
}

#[async_trait]
impl ChangesetFetcher for TestChangesetFetcher {
    async fn get_generation_number(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Generation, Error> {
        self.repo
            .changeset_fetcher()
            .get_generation_number(ctx, cs_id)
            .await
    }

    async fn get_parents(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>, Error> {
        self.repo.changeset_fetcher().get_parents(ctx, cs_id).await
    }

    fn get_stats(&self) -> HashMap<String, Box<dyn Any>> {
        HashMap::new()
    }
}

pub async fn get_single_bonsai_streams(
    ctx: CoreContext,
    repo: &BlobRepo,
    hashes: &[&str],
) -> Vec<BonsaiNodeStream> {
    let mut ret = vec![];

    for hash in hashes {
        let bonsai = string_to_bonsai(ctx.fb, repo, hash).await;
        let stream = single_changeset_id(ctx.clone(), bonsai, repo).boxify();
        ret.push(stream)
    }

    ret
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::hash_set::IntoIter;
use std::collections::HashSet;

use anyhow::Error;
use changeset_fetcher::ArcChangesetFetcher;
use context::CoreContext;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::setcommon::*;
use crate::BonsaiNodeStream;

pub struct UnionNodeStream {
    inputs: Vec<(
        BonsaiInputStream,
        Poll<Option<(ChangesetId, Generation)>, Error>,
    )>,
    current_generation: Option<Generation>,
    accumulator: HashSet<ChangesetId>,
    drain: Option<IntoIter<ChangesetId>>,
}

impl UnionNodeStream {
    pub fn new<I>(ctx: CoreContext, changeset_fetcher: &ArcChangesetFetcher, inputs: I) -> Self
    where
        I: IntoIterator<Item = BonsaiNodeStream>,
    {
        let csid_and_gen = inputs.into_iter().map(move |i| {
            (
                add_generations_by_bonsai(ctx.clone(), i, changeset_fetcher.clone()),
                Ok(Async::NotReady),
            )
        });
        Self {
            inputs: csid_and_gen.collect(),
            current_generation: None,
            accumulator: HashSet::new(),
            drain: None,
        }
    }

    fn gc_finished_inputs(&mut self) {
        self.inputs.retain(|(_, state)| {
            if let Ok(Async::Ready(None)) = *state {
                false
            } else {
                true
            }
        });
    }

    fn update_current_generation(&mut self) {
        if all_inputs_ready(&self.inputs) {
            self.current_generation = self
                .inputs
                .iter()
                .filter_map(|(_, state)| match state {
                    Ok(Async::Ready(Some((_, gen_id)))) => Some(*gen_id),
                    Ok(Async::NotReady) => panic!("All states ready, yet some not ready!"),
                    _ => None,
                })
                .max();
        }
    }

    fn accumulate_nodes(&mut self) {
        let mut found_csids = false;
        for &mut (_, ref mut state) in self.inputs.iter_mut() {
            if let Ok(Async::Ready(Some((csid, gen_id)))) = *state {
                if Some(gen_id) == self.current_generation {
                    found_csids = true;
                    self.accumulator.insert(csid);
                    *state = Ok(Async::NotReady);
                }
            }
        }
        if !found_csids {
            self.current_generation = None;
        }
    }
}

impl Stream for UnionNodeStream {
    type Item = ChangesetId;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // This feels wrong, but in practice it's fine - it should be quick to hit a return, and
        // the standard futures_old::executor expects you to only return NotReady if blocked on I/O.
        loop {
            // Start by trying to turn as many NotReady as possible into real items
            poll_all_inputs(&mut self.inputs);

            // Empty the drain if any - return all items for this generation
            let next_in_drain = self.drain.as_mut().and_then(|drain| drain.next());
            if next_in_drain.is_some() {
                return Ok(Async::Ready(next_in_drain));
            } else {
                self.drain = None;
            }

            // Return any errors
            {
                if self.inputs.iter().any(|(_, state)| state.is_err()) {
                    let inputs = std::mem::take(&mut self.inputs);
                    let err = inputs
                        .into_iter()
                        .find_map(|(_, state)| state.err())
                        .unwrap();
                    return Err(err);
                }
            }

            self.gc_finished_inputs();

            // If any input is not ready (we polled above), wait for them all to be ready
            if !all_inputs_ready(&self.inputs) {
                return Ok(Async::NotReady);
            }

            match self.current_generation {
                None => {
                    if self.accumulator.is_empty() {
                        self.update_current_generation();
                    } else {
                        let full_accumulator = std::mem::take(&mut self.accumulator);
                        self.drain = Some(full_accumulator.into_iter());
                    }
                }
                Some(_) => self.accumulate_nodes(),
            }
            // If we cannot ever output another node, we're done.
            if self.inputs.is_empty() && self.drain.is_none() && self.accumulator.is_empty() {
                return Ok(Async::Ready(None));
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use context::CoreContext;
    use failure_ext::err_downcast;
    use fbinit::FacebookInit;
    use futures::compat::Stream01CompatExt;
    use futures::stream::StreamExt as _;
    use futures_ext::StreamExt;
    use futures_old::executor::spawn;
    use revset_test_helper::assert_changesets_sequence;
    use revset_test_helper::single_changeset_id;
    use revset_test_helper::string_to_bonsai;

    use super::*;
    use crate::errors::ErrorKind;
    use crate::fixtures::BranchEven;
    use crate::fixtures::BranchUneven;
    use crate::fixtures::BranchWide;
    use crate::fixtures::Linear;
    use crate::fixtures::TestRepoFixture;
    use crate::setcommon::NotReadyEmptyStream;
    use crate::setcommon::RepoErrorStream;
    use crate::tests::get_single_bonsai_streams;
    use crate::tests::TestChangesetFetcher;
    use crate::BonsaiNodeStream;

    #[fbinit::test]
    async fn union_identical_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let head_csid = string_to_bonsai(fb, &repo, hash).await;

        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), head_csid.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), head_csid.clone(), &repo).boxify(),
        ];
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![head_csid.clone()], nodestream).await;
    }

    #[fbinit::test]
    async fn union_error_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let hash = "a5ffa77602a066db7d5cfb9fb5823a0895717c5a";
        let expected_csid = string_to_bonsai(fb, &repo, hash).await;

        let inputs: Vec<BonsaiNodeStream> = vec![
            RepoErrorStream {
                item: expected_csid,
            }
            .boxify(),
            single_changeset_id(ctx.clone(), expected_csid.clone(), &repo).boxify(),
        ];
        let mut nodestream = spawn(
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify(),
        );

        match nodestream.wait_stream() {
            Some(Err(err)) => match err_downcast!(err, err: ErrorKind => err) {
                Ok(ErrorKind::RepoChangesetError(cs)) => assert_eq!(cs, expected_csid),
                Ok(bad) => panic!("unexpected error {:?}", bad),
                Err(bad) => panic!("unknown error {:?}", bad),
            },
            Some(Ok(bad)) => panic!("unexpected success {:?}", bad),
            None => panic!("no result"),
        };
    }

    #[fbinit::test]
    async fn union_three_nodes(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_d0a = string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await;
        let bcs_3c1 = string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await;
        let bcs_a947 =
            string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await;
        // Note that these are *not* in generation order deliberately.
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), bcs_a947, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_3c1, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
        ];
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        // But, once I hit the asserts, I expect them in generation order.
        assert_changesets_sequence(&ctx, &repo, vec![bcs_3c1, bcs_a947, bcs_d0a], nodestream).await;
    }

    #[fbinit::test]
    async fn union_nothing(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let inputs: Vec<BonsaiNodeStream> = vec![];
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();
        assert_changesets_sequence(&ctx, &repo, vec![], nodestream).await;
    }

    #[fbinit::test]
    async fn union_nesting(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let bcs_d0a = string_to_bonsai(fb, &repo, "d0a361e9022d226ae52f689667bd7d212a19cfe0").await;
        let bcs_3c1 = string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await;
        // Note that these are *not* in generation order deliberately.
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), bcs_d0a, &repo).boxify(),
            single_changeset_id(ctx.clone(), bcs_3c1, &repo).boxify(),
        ];

        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        let bcs_a947 =
            string_to_bonsai(fb, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await;
        let inputs: Vec<BonsaiNodeStream> = vec![
            nodestream,
            single_changeset_id(ctx.clone(), bcs_a947, &repo).boxify(),
        ];
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![bcs_3c1, bcs_a947, bcs_d0a], nodestream).await;
    }

    #[fbinit::test]
    async fn slow_ready_union_nothing(fb: FacebookInit) {
        // Tests that we handle an input staying at NotReady for a while without panicking
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher = Arc::new(TestChangesetFetcher::new(repo));

        let inputs: Vec<BonsaiNodeStream> = vec![NotReadyEmptyStream::new(10).boxify()];
        let mut nodestream =
            UnionNodeStream::new(ctx, &changeset_fetcher, inputs.into_iter()).compat();

        assert!(nodestream.next().await.is_none());
    }

    #[fbinit::test]
    async fn union_branch_even_repo(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = BranchEven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let nodes = vec![
            string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await,
            string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await,
            string_to_bonsai(fb, &repo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await,
        ];

        // Two nodes should share the same generation number
        let inputs: Vec<BonsaiNodeStream> = nodes
            .clone()
            .into_iter()
            .map(|cs| single_changeset_id(ctx.clone(), cs, &repo).boxify())
            .collect();
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();
        assert_changesets_sequence(&ctx, &repo, nodes, nodestream).await;
    }

    #[fbinit::test]
    async fn union_branch_uneven_repo(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = BranchUneven::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let cs_1 = string_to_bonsai(fb, &repo, "3cda5c78aa35f0f5b09780d971197b51cad4613a").await;
        let cs_2 = string_to_bonsai(fb, &repo, "d7542c9db7f4c77dab4b315edd328edf1514952f").await;
        let cs_3 = string_to_bonsai(fb, &repo, "4f7f3fd428bec1a48f9314414b063c706d9c1aed").await;
        let cs_4 = string_to_bonsai(fb, &repo, "bc7b4d0f858c19e2474b03e442b8495fd7aeef33").await;
        let cs_5 = string_to_bonsai(fb, &repo, "264f01429683b3dd8042cb3979e8bf37007118bc").await;
        // Two nodes should share the same generation number
        let inputs: Vec<BonsaiNodeStream> = vec![
            single_changeset_id(ctx.clone(), cs_1.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), cs_2.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), cs_3.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), cs_4.clone(), &repo).boxify(),
            single_changeset_id(ctx.clone(), cs_5.clone(), &repo).boxify(),
        ];
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(&ctx, &repo, vec![cs_5, cs_4, cs_3, cs_1, cs_2], nodestream)
            .await;
    }

    #[fbinit::test]
    async fn union_branch_wide_repo(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = BranchWide::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        // Two nodes should share the same generation number
        let inputs = get_single_bonsai_streams(
            ctx.clone(),
            &repo,
            &[
                "49f53ab171171b3180e125b918bd1cf0af7e5449",
                "4685e9e62e4885d477ead6964a7600c750e39b03",
                "c27ef5b7f15e9930e5b93b1f32cc2108a2aabe12",
                "9e8521affb7f9d10e9551a99c526e69909042b20",
            ],
        )
        .await;
        let nodestream =
            UnionNodeStream::new(ctx.clone(), &changeset_fetcher, inputs.into_iter()).boxify();

        assert_changesets_sequence(
            &ctx,
            &repo,
            vec![
                string_to_bonsai(fb, &repo, "49f53ab171171b3180e125b918bd1cf0af7e5449").await,
                string_to_bonsai(fb, &repo, "c27ef5b7f15e9930e5b93b1f32cc2108a2aabe12").await,
                string_to_bonsai(fb, &repo, "4685e9e62e4885d477ead6964a7600c750e39b03").await,
                string_to_bonsai(fb, &repo, "9e8521affb7f9d10e9551a99c526e69909042b20").await,
            ],
            nodestream,
        )
        .await;
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;

use anyhow::Error;
use changeset_fetcher::ArcChangesetFetcher;
use context::CoreContext;
use futures_ext::StreamExt;
use futures_old::stream::Stream;
use futures_old::Async;
use futures_old::Poll;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

use crate::setcommon::add_generations_by_bonsai;
use crate::setcommon::BonsaiInputStream;
use crate::BonsaiNodeStream;

/// A wrapper around a NodeStream that asserts that the two revset invariants hold:
/// 1. The generation number never increases
/// 2. No hash is seen twice
/// This uses memory proportional to the number of hashes in the revset.
pub struct ValidateNodeStream {
    wrapped: BonsaiInputStream,
    last_generation: Option<Generation>,
    seen_hashes: HashSet<ChangesetId>,
}

impl ValidateNodeStream {
    pub fn new(
        ctx: CoreContext,
        wrapped: BonsaiNodeStream,
        changeset_fetcher: &ArcChangesetFetcher,
    ) -> ValidateNodeStream {
        ValidateNodeStream {
            wrapped: add_generations_by_bonsai(ctx, wrapped, changeset_fetcher.clone()).boxify(),
            last_generation: None,
            seen_hashes: HashSet::new(),
        }
    }
}

impl Stream for ValidateNodeStream {
    type Item = ChangesetId;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let next = self.wrapped.poll()?;

        let (hash, gen) = match next {
            Async::NotReady => return Ok(Async::NotReady),
            Async::Ready(None) => return Ok(Async::Ready(None)),
            Async::Ready(Some((hash, gen))) => (hash, gen),
        };

        assert!(self.seen_hashes.insert(hash), "Hash {} seen twice", hash);

        assert!(
            self.last_generation.is_none() || self.last_generation >= Some(gen),
            "Generation number increased unexpectedly"
        );

        self.last_generation = Some(gen);

        Ok(Async::Ready(Some(hash)))
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use fbinit::FacebookInit;
    use futures::compat::Stream01CompatExt;
    use futures::stream::StreamExt as _;
    use futures_ext::StreamExt;
    use revset_test_helper::assert_changesets_sequence;
    use revset_test_helper::single_changeset_id;
    use revset_test_helper::string_to_bonsai;

    use super::*;
    use crate::fixtures::Linear;
    use crate::fixtures::TestRepoFixture;
    use crate::setcommon::NotReadyEmptyStream;
    use crate::tests::TestChangesetFetcher;

    #[fbinit::test]
    async fn validate_accepts_single_node(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new(repo.clone()));
        let repo = Arc::new(repo);

        let head_csid =
            string_to_bonsai(fb, &repo, "a5ffa77602a066db7d5cfb9fb5823a0895717c5a").await;

        let nodestream = single_changeset_id(ctx.clone(), head_csid.clone(), &repo).boxify();

        let nodestream =
            ValidateNodeStream::new(ctx.clone(), nodestream, &changeset_fetcher).boxify();
        assert_changesets_sequence(&ctx, &repo, vec![head_csid], nodestream).await;
    }

    #[fbinit::test]
    async fn slow_ready_validates(fb: FacebookInit) {
        // Tests that we handle an input staying at NotReady for a while without panicking
        let ctx = CoreContext::test_mock(fb);
        let repo = Linear::getrepo(fb).await;
        let changeset_fetcher: ArcChangesetFetcher = Arc::new(TestChangesetFetcher::new(repo));

        let mut nodestream = ValidateNodeStream::new(
            ctx,
            NotReadyEmptyStream::new(10).boxify(),
            &changeset_fetcher,
        )
        .compat();

        assert!(nodestream.next().await.is_none());
    }

    #[fbinit::test]
    #[should_panic]
    async fn repeat_hash_panics(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Arc::new(Linear::getrepo(fb).await);

        let head_csid =
            string_to_bonsai(fb, &repo, "a5ffa77602a066db7d5cfb9fb5823a0895717c5a").await;
        let nodestream = single_changeset_id(ctx.clone(), head_csid.clone(), &repo)
            .chain(single_changeset_id(ctx.clone(), head_csid.clone(), &repo));

        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new((*repo).clone()));

        let mut nodestream =
            ValidateNodeStream::new(ctx, nodestream.boxify(), &changeset_fetcher).boxify();

        loop {
            match nodestream.poll() {
                Ok(Async::Ready(None)) => return,
                _ => {}
            }
        }
    }

    #[fbinit::test]
    #[should_panic]
    async fn wrong_order_panics(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let repo = Arc::new(Linear::getrepo(fb).await);

        let nodestream = single_changeset_id(
            ctx.clone(),
            string_to_bonsai(fb, &repo, "cb15ca4a43a59acff5388cea9648c162afde8372")
                .await
                .clone(),
            &repo,
        )
        .chain(single_changeset_id(
            ctx.clone(),
            string_to_bonsai(fb, &repo, "3c15267ebf11807f3d772eb891272b911ec68759").await,
            &repo,
        ));

        let changeset_fetcher: ArcChangesetFetcher =
            Arc::new(TestChangesetFetcher::new((*repo).clone()));

        let mut nodestream =
            ValidateNodeStream::new(ctx, nodestream.boxify(), &changeset_fetcher).boxify();

        loop {
            match nodestream.poll() {
                Ok(Async::Ready(None)) => return,
                _ => {}
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use crate::ScscApp;

base_app::subcommands! {
    type App = ScscApp;
    mod cat;
    mod blame;
    mod common_base;
    mod create_bookmark if "SCSC_WRITES_ENABLED";
    mod delete_bookmark if "SCSC_WRITES_ENABLED";
    mod diff;
    mod export;
    mod find_files;
    mod info;
    mod is_ancestor;
    mod land_stack if "SCSC_WRITES_ENABLED";
    mod list_bookmarks;
    mod log;
    mod lookup;
    mod ls;
    mod move_bookmark if "SCSC_WRITES_ENABLED";
    mod prepare_commits if "SCSC_WRITES_ENABLED";
    mod pushrebase_history;
    mod repos;
    mod run_hooks;
    mod sparse_profile_delta;
    mod sparse_profile_size;
    mod xrepo_lookup;
}