/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! # nameset
//!
//! See [`NameSet`] for the main structure.

use std::any::Any;
use std::cmp;
use std::fmt;
use std::fmt::Debug;
use std::ops::BitAnd;
use std::ops::BitOr;
use std::ops::Deref;
use std::ops::Sub;
use std::pin::Pin;
use std::sync::Arc;

use futures::future::BoxFuture;
use futures::Stream;
use futures::StreamExt;
use nonblocking::non_blocking;

use crate::default_impl;
use crate::ops::DagAlgorithm;
use crate::ops::IdConvert;
use crate::ops::IdMapSnapshot;
use crate::ops::Parents;
use crate::Id;
use crate::IdSet;
use crate::Result;
use crate::VertexName;

pub mod difference;
pub mod hints;
pub mod id_lazy;
pub mod id_static;
pub mod intersection;
pub mod lazy;
#[cfg(any(test, feature = "indexedlog-backend"))]
pub mod legacy;
pub mod meta;
pub mod slice;
pub mod r#static;
pub mod union;

use self::hints::Flags;
use self::hints::Hints;
use self::id_static::IdStaticSet;
use self::meta::MetaSet;
use self::r#static::StaticSet;

/// A [`NameSet`] contains an immutable list of names.
///
/// It provides order-preserving iteration and set operations,
/// and is cheaply clonable.
#[derive(Clone)]
pub struct NameSet(Arc<dyn AsyncNameSetQuery>);

impl NameSet {
    pub(crate) fn from_query(query: impl AsyncNameSetQuery) -> Self {
        Self(Arc::new(query))
    }

    /// Creates an empty set.
    pub fn empty() -> Self {
        Self::from_query(r#static::StaticSet::empty())
    }

    /// Creates from a (short) list of known names.
    pub fn from_static_names(names: impl IntoIterator<Item = VertexName>) -> NameSet {
        Self::from_query(r#static::StaticSet::from_names(names))
    }

    /// Creates from a (lazy) iterator of names.
    pub fn from_iter<I>(iter: I, hints: Hints) -> NameSet
    where
        I: IntoIterator<Item = Result<VertexName>> + 'static,
        <I as IntoIterator>::IntoIter: Send + Sync,
    {
        Self::from_query(lazy::LazySet::from_iter(iter, hints))
    }

    /// Creates from a (lazy) stream of names with hints.
    pub fn from_stream(stream: BoxVertexStream, hints: Hints) -> NameSet {
        Self::from_query(lazy::LazySet::from_stream(stream, hints))
    }

    /// Creates from a (lazy) iterator of Ids, an IdMap, and a Dag.
    pub fn from_id_iter_idmap_dag<I>(
        iter: I,
        map: Arc<dyn IdConvert + Send + Sync>,
        dag: Arc<dyn DagAlgorithm + Send + Sync>,
    ) -> NameSet
    where
        I: IntoIterator<Item = Result<Id>> + 'static,
        <I as IntoIterator>::IntoIter: Send + Sync,
    {
        Self::from_query(id_lazy::IdLazySet::from_iter_idmap_dag(iter, map, dag))
    }

    /// Creates from a (lazy) iterator of Ids and a struct with snapshot abilities.
    pub fn from_id_iter_dag<I>(
        iter: I,
        dag: &(impl DagAlgorithm + IdMapSnapshot),
    ) -> Result<NameSet>
    where
        I: IntoIterator<Item = Result<Id>> + 'static,
        <I as IntoIterator>::IntoIter: Send + Sync,
    {
        let map = dag.id_map_snapshot()?;
        let dag = dag.dag_snapshot()?;
        Ok(Self::from_id_iter_idmap_dag(iter, map, dag))
    }

    /// Creates from [`IdSet`], [`IdMap`] and [`DagAlgorithm`].
    pub fn from_spans_idmap_dag(
        spans: IdSet,
        map: Arc<dyn IdConvert + Send + Sync>,
        dag: Arc<dyn DagAlgorithm + Send + Sync>,
    ) -> NameSet {
        Self::from_query(IdStaticSet::from_spans_idmap_dag(spans, map, dag))
    }

    /// Creates from [`IdSet`] and a struct with snapshot abilities.
    pub fn from_spans_dag(spans: IdSet, dag: &(impl DagAlgorithm + IdMapSnapshot)) -> Result<Self> {
        let map = dag.id_map_snapshot()?;
        let dag = dag.dag_snapshot()?;
        Ok(Self::from_spans_idmap_dag(spans, map, dag))
    }

    /// Creates from a function that evaluates to a [`NameSet`], and a
    /// `contains` fast path.
    pub fn from_evaluate_contains<C>(
        evaluate: impl Fn() -> Result<NameSet> + Send + Sync + 'static,
        contains: C,
        hints: Hints,
    ) -> NameSet
    where
        C: for<'a> Fn(&'a MetaSet, &'a VertexName) -> Result<bool>,
        C: Send + Sync + 'static,
    {
        let evaluate = move || -> BoxFuture<_> {
            let result = evaluate();
            Box::pin(async move { result })
        };
        let contains = Arc::new(contains);
        Self::from_async_evaluate_contains(
            Box::new(evaluate),
            Box::new(move |m, v| {
                let contains = contains.clone();
                Box::pin(async move { contains(m, v) })
            }),
            hints,
        )
    }

    /// Creates from an async function that evaluates to a [`NameSet`], and a
    /// async `contains` fast path.
    pub fn from_async_evaluate_contains(
        evaluate: Box<dyn Fn() -> BoxFuture<'static, Result<NameSet>> + Send + Sync>,
        contains: Box<
            dyn for<'a> Fn(&'a MetaSet, &'a VertexName) -> BoxFuture<'a, Result<bool>>
                + Send
                + Sync,
        >,
        hints: Hints,
    ) -> NameSet {
        Self::from_query(MetaSet::from_evaluate_hints(evaluate, hints).with_contains(contains))
    }

    /// Calculates the subset that is only in self, not in other.
    pub fn difference(&self, other: &NameSet) -> NameSet {
        if other.hints().contains(Flags::FULL)
            && other.hints().dag_version() >= self.hints().dag_version()
            && self.hints().dag_version() > None
        {
            tracing::debug!(
                "difference(x={:.6?}, y={:.6?}) = () (fast path 1)",
                self,
                other
            );
            return Self::empty();
        }
        if self.hints().contains(Flags::EMPTY) || other.hints().contains(Flags::EMPTY) {
            tracing::debug!(
                "difference(x={:.6?}, y={:.6?}) = x (fast path 2)",
                self,
                other
            );
            return self.clone();
        }
        if let (Some(this), Some(other)) = (
            self.as_any().downcast_ref::<IdStaticSet>(),
            other.as_any().downcast_ref::<IdStaticSet>(),
        ) {
            let order = this.map.map_version().partial_cmp(other.map.map_version());
            if order.is_some() {
                // Fast path for IdStaticSet
                let result = Self::from_spans_idmap_dag(
                    this.spans.difference(&other.spans),
                    this.map.clone(),
                    this.dag.clone(),
                );
                tracing::debug!(
                    "difference(x={:.6?}, y={:.6?}) = {:.6?} (fast path 3)",
                    self,
                    other,
                    &result
                );
                return result;
            }
        }
        tracing::debug!("difference(x={:.6?}, y={:.6?}) (slow path)", self, other);
        Self::from_query(difference::DifferenceSet::new(self.clone(), other.clone()))
    }

    /// Calculates the intersection of two sets.
    pub fn intersection(&self, other: &NameSet) -> NameSet {
        if self.hints().contains(Flags::FULL)
            && self.hints().dag_version() >= other.hints().dag_version()
            && other.hints().dag_version() > None
        {
            tracing::debug!(
                "intersection(x={:.6?}, y={:.6?}) = y (fast path 1)",
                self,
                other
            );
            return other.clone();
        }
        if other.hints().contains(Flags::FULL)
            && other.hints().dag_version() >= self.hints().dag_version()
            && self.hints().dag_version() > None
        {
            tracing::debug!(
                "intersection(x={:.6?}, y={:.6?}) = x (fast path 2)",
                self,
                other
            );
            return self.clone();
        }
        if self.hints().contains(Flags::EMPTY) || other.hints().contains(Flags::EMPTY) {
            tracing::debug!(
                "intersection(x={:.6?}, y={:.6?}) = () (fast path 3)",
                self,
                other
            );
            return Self::empty();
        }
        if let (Some(this), Some(other)) = (
            self.as_any().downcast_ref::<IdStaticSet>(),
            other.as_any().downcast_ref::<IdStaticSet>(),
        ) {
            let order = this.map.map_version().partial_cmp(other.map.map_version());
            if let Some(order) = order {
                // Fast path for IdStaticSet
                let result = Self::from_spans_idmap_dag(
                    this.spans.intersection(&other.spans),
                    pick(order, &this.map, &other.map).clone(),
                    pick(order, &this.dag, &other.dag).clone(),
                );
                tracing::debug!(
                    "intersection(x={:.6?}, y={:.6?}) = {:?} (IdStatic fast path)",
                    self,
                    other,
                    &result
                );
                return result;
            }
        }
        tracing::debug!("intersection(x={:.6?}, y={:.6?}) (slow path)", self, other,);
        Self::from_query(intersection::IntersectionSet::new(
            self.clone(),
            other.clone(),
        ))
    }

    /// Calculates the union of two sets.
    pub fn union(&self, other: &NameSet) -> NameSet {
        if (self.hints().contains(Flags::FULL)
            && self.hints().dag_version() >= other.hints().dag_version()
            && other.hints().dag_version() > None)
            || other.hints().contains(Flags::EMPTY)
        {
            tracing::debug!("union(x={:.6?}, y={:.6?}) = x (fast path 1)", self, other);
            return self.clone();
        }
        if self.hints().contains(Flags::EMPTY)
            || (other.hints().contains(Flags::FULL)
                && other.hints().dag_version() >= self.hints().dag_version()
                && self.hints().dag_version() > None)
        {
            tracing::debug!("union(x={:.6?}, y={:.6?}) = y (fast path 2)", self, other);
            return other.clone();
        }
        if let (Some(this), Some(other)) = (
            self.as_any().downcast_ref::<IdStaticSet>(),
            other.as_any().downcast_ref::<IdStaticSet>(),
        ) {
            let order = this.map.map_version().partial_cmp(other.map.map_version());
            if let Some(order) = order {
                // Fast path for IdStaticSet
                let result = Self::from_spans_idmap_dag(
                    this.spans.union(&other.spans),
                    pick(order, &this.map, &other.map).clone(),
                    pick(order, &this.dag, &other.dag).clone(),
                );
                tracing::debug!(
                    "union(x={:.6?}, y={:.6?}) = {:.6?} (fast path 3)",
                    self,
                    other,
                    &result
                );
                return result;
            }
        }
        tracing::debug!("union(x={:.6?}, y={:.6?}) (slow path)", self, other);
        Self::from_query(union::UnionSet::new(self.clone(), other.clone()))
    }

    /// Filter using the given async function. If `filter_func` returns `true`
    /// for a vertex, then the vertex will be taken, other it will be skipped.
    pub fn filter(
        &self,
        filter_func: Box<dyn Fn(&VertexName) -> BoxFuture<Result<bool>> + Send + Sync + 'static>,
    ) -> Self {
        let filter_func = Arc::new(filter_func);
        let this = self.clone();
        let hints = {
            // Drop ANCESTORS | FULL and add FILTER.
            let hints = self.hints().clone();
            hints.update_flags_with(|f| (f - Flags::ANCESTORS - Flags::FULL) | Flags::FILTER);
            hints
        };
        let result = Self::from_async_evaluate_contains(
            Box::new({
                let filter_func = filter_func.clone();
                let this = this.clone();
                let hints = hints.clone();
                move || {
                    let filter_func = filter_func.clone();
                    let this = this.clone();
                    let hints = hints.clone();
                    Box::pin(async move {
                        let stream = this.0.iter().await?;
                        let stream = stream.filter_map(move |v| {
                            let filter_func = filter_func.clone();
                            async move {
                                match v {
                                    Ok(v) => match (&filter_func)(&v).await {
                                        Ok(true) => Some(Ok(v)),
                                        Ok(false) => None,
                                        Err(e) => Some(Err(e)),
                                    },
                                    Err(e) => Some(Err(e)),
                                }
                            }
                        });
                        Ok(Self::from_stream(Box::pin(stream), hints))
                    })
                }
            }),
            Box::new(move |_, v| {
                let filter_func = filter_func.clone();
                let this = this.clone();
                Box::pin(async move { Ok(this.0.contains(v).await? && (&filter_func)(v).await?) })
            }),
            hints,
        );
        result.hints().add_flags(Flags::FILTER);
        result
    }

    /// Convert the set to a graph containing only the vertexes in the set. This can be slow on
    /// larger sets.
    pub async fn to_parents(&self) -> Result<Option<impl Parents>> {
        default_impl::set_to_parents(self).await
    }

    /// Obtain the attached dag if available.
    pub fn dag(&self) -> Option<Arc<dyn DagAlgorithm + Send + Sync>> {
        self.hints().dag()
    }

    /// Obtain the attached IdMap if available.
    pub fn id_map(&self) -> Option<Arc<dyn IdConvert + Send + Sync>> {
        self.hints().id_map()
    }

    /// Convert the current set into a flat static set so it can be used in some
    /// fast paths. This is useful for some common sets like `obsolete()` that
    /// might be represented by a complex expression.
    pub async fn flatten(&self) -> Result<NameSet> {
        match (self.id_map(), self.dag()) {
            (Some(id_map), Some(dag)) => {
                // Convert to IdStaticSet
                self.flatten_id(id_map, dag).await
            }
            _ => {
                // Convert to StaticSet
                self.flatten_names().await
            }
        }
    }

    /// Convert this set to a static id set.
    pub async fn flatten_id(
        &self,
        id_map: Arc<dyn IdConvert + Send + Sync>,
        dag: Arc<dyn DagAlgorithm + Send + Sync>,
    ) -> Result<NameSet> {
        if self.as_any().is::<IdStaticSet>() {
            return Ok(self.clone());
        }
        let mut ids = Vec::with_capacity(self.count()?);
        for vertex in self.iter()? {
            let id = id_map.vertex_id(vertex?).await?;
            ids.push(id);
        }
        ids.sort_unstable_by_key(|i| u64::MAX - i.0);
        let spans = IdSet::from_sorted_spans(ids);
        let flat_set = NameSet::from_spans_idmap_dag(spans, id_map, dag);
        flat_set.hints().inherit_flags_min_max_id(self.hints());
        Ok(flat_set)
    }

    /// Convert this set to a static name set.
    pub async fn flatten_names(&self) -> Result<NameSet> {
        if self.as_any().is::<StaticSet>() {
            return Ok(self.clone());
        }
        let names = self.iter()?.collect::<Result<Vec<_>>>()?;
        let flat_set = Self::from_static_names(names);
        flat_set.hints().inherit_flags_min_max_id(self.hints());
        Ok(flat_set)
    }

    /// Take the first `n` items.
    pub fn take(&self, n: u64) -> NameSet {
        if let Some(set) = self.as_any().downcast_ref::<IdStaticSet>() {
            tracing::debug!("take(x={:.6?}, {}) (fast path)", self, n);
            Self::from_spans_idmap_dag(set.spans.take(n), set.map.clone(), set.dag.clone())
        } else {
            tracing::debug!("take(x={:.6?}, {}) (slow path)", self, n);
            let set = slice::SliceSet::new(self.clone(), 0, Some(n));
            Self::from_query(set)
        }
    }

    /// Skip the first `n` items.
    pub fn skip(&self, n: u64) -> NameSet {
        if n == 0 {
            return self.clone();
        }
        if let Some(set) = self.as_any().downcast_ref::<IdStaticSet>() {
            tracing::debug!("skip(x={:.6?}, {}) (fast path)", self, n);
            Self::from_spans_idmap_dag(set.spans.skip(n), set.map.clone(), set.dag.clone())
        } else {
            tracing::debug!("skip(x={:.6?}, {}) (slow path)", self, n);
            let set = slice::SliceSet::new(self.clone(), n, None);
            Self::from_query(set)
        }
    }

    /// Converts to `(IdSet, IdConvert)` pair in O(1). If the underlying set
    /// cannot provide such information in O(1), return `None`.
    ///
    /// Useful if the callsite wants to have random access (ex.  bisect) and
    /// control how to resolve in batches.
    pub fn to_id_set_and_id_map_in_o1(&self) -> Option<(IdSet, Arc<dyn IdConvert + Send + Sync>)> {
        let id_map = self.id_map()?;
        let id_set = self.as_any().downcast_ref::<IdStaticSet>()?.spans.clone();
        Some((id_set, id_map))
    }
}

impl BitAnd for NameSet {
    type Output = Self;

    fn bitand(self, other: Self) -> Self {
        self.intersection(&other)
    }
}

impl BitOr for NameSet {
    type Output = Self;

    fn bitor(self, other: Self) -> Self {
        self.union(&other)
    }
}

impl Sub for NameSet {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        self.difference(&other)
    }
}

impl Deref for NameSet {
    type Target = dyn AsyncNameSetQuery;

    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

impl fmt::Debug for NameSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

/// Read-only queries required by [`NameSet`]: Iteration, length and contains.
///
/// Types implementating this trait should rewrite methods to use fast paths
/// when possible.
#[async_trait::async_trait]
pub trait AsyncNameSetQuery: Any + Debug + Send + Sync {
    /// Iterate through the set in defined order.
    async fn iter(&self) -> Result<BoxVertexStream>;

    /// Iterate through the set in the reversed order.
    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let mut iter = self.iter().await?;
        let mut items = Vec::new();
        while let Some(item) = iter.next().await {
            items.push(item);
        }
        Ok(Box::pin(futures::stream::iter(items.into_iter().rev())))
    }

    /// Number of names in this set.
    async fn count(&self) -> Result<usize> {
        let mut iter = self.iter().await?;
        let mut count = 0;
        while let Some(item) = iter.next().await {
            item?;
            count += 1;
        }
        Ok(count)
    }

    /// The first name in the set.
    async fn first(&self) -> Result<Option<VertexName>> {
        self.iter().await?.next().await.transpose()
    }

    /// The last name in the set.
    async fn last(&self) -> Result<Option<VertexName>> {
        self.iter_rev().await?.next().await.transpose()
    }

    /// Test if this set is empty.
    async fn is_empty(&self) -> Result<bool> {
        self.first().await.map(|n| n.is_none())
    }

    /// Test if this set contains a given name.
    async fn contains(&self, name: &VertexName) -> Result<bool> {
        let mut iter = self.iter().await?;
        while let Some(item) = iter.next().await {
            if &item? == name {
                return Ok(true);
            }
        }
        Ok(false)
    }

    /// Test contains in less than O(N) time.
    /// Returns None if cannot achieve in less than O(N) time.
    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        let _ = name;
        Ok(None)
    }

    /// For downcasting.
    fn as_any(&self) -> &dyn Any;

    /// Get or set optimization hints.
    fn hints(&self) -> &Hints;

    /// Get an optional IdConvert interface to check hints.
    fn id_convert(&self) -> Option<&dyn IdConvert> {
        None
    }
}

/// Sync version of `AsyncNameSetQuery`.
pub trait SyncNameSetQuery {
    /// Iterate through the set in defined order.
    fn iter(&self) -> Result<Box<dyn NameIter>>;

    /// Iterate through the set in the reversed order.
    fn iter_rev(&self) -> Result<Box<dyn NameIter>>;

    /// Number of names in this set.
    fn count(&self) -> Result<usize>;

    /// The first name in the set.
    fn first(&self) -> Result<Option<VertexName>>;

    /// The last name in the set.
    fn last(&self) -> Result<Option<VertexName>>;

    /// Test if this set is empty.
    fn is_empty(&self) -> Result<bool>;

    /// Test if this set contains a given name.
    fn contains(&self, name: &VertexName) -> Result<bool>;

    /// For downcasting.
    fn as_any(&self) -> &dyn Any;

    /// Get or set optimization hints.
    fn hints(&self) -> &Hints;

    /// Get an optional IdConvert interface to check hints.
    fn id_convert(&self) -> Option<&dyn IdConvert>;
}

impl<T: AsyncNameSetQuery> SyncNameSetQuery for T {
    fn iter(&self) -> Result<Box<dyn NameIter>> {
        non_blocking(AsyncNameSetQuery::iter(self))?.map(to_iter)
    }

    fn iter_rev(&self) -> Result<Box<dyn NameIter>> {
        non_blocking(AsyncNameSetQuery::iter_rev(self))?.map(to_iter)
    }

    fn count(&self) -> Result<usize> {
        non_blocking(AsyncNameSetQuery::count(self))?
    }

    fn first(&self) -> Result<Option<VertexName>> {
        non_blocking(AsyncNameSetQuery::first(self))?
    }

    fn last(&self) -> Result<Option<VertexName>> {
        non_blocking(AsyncNameSetQuery::last(self))?
    }

    fn is_empty(&self) -> Result<bool> {
        non_blocking(AsyncNameSetQuery::is_empty(self))?
    }

    fn contains(&self, name: &VertexName) -> Result<bool> {
        non_blocking(AsyncNameSetQuery::contains(self, name))?
    }

    fn as_any(&self) -> &dyn Any {
        AsyncNameSetQuery::as_any(self)
    }

    fn hints(&self) -> &Hints {
        AsyncNameSetQuery::hints(self)
    }

    fn id_convert(&self) -> Option<&dyn IdConvert> {
        AsyncNameSetQuery::id_convert(self)
    }
}

impl SyncNameSetQuery for NameSet {
    fn iter(&self) -> Result<Box<dyn NameIter>> {
        non_blocking(AsyncNameSetQuery::iter(self.0.deref()))?.map(to_iter)
    }

    fn iter_rev(&self) -> Result<Box<dyn NameIter>> {
        non_blocking(AsyncNameSetQuery::iter_rev(self.0.deref()))?.map(to_iter)
    }

    fn count(&self) -> Result<usize> {
        non_blocking(AsyncNameSetQuery::count(self.0.deref()))?
    }

    fn first(&self) -> Result<Option<VertexName>> {
        non_blocking(AsyncNameSetQuery::first(self.0.deref()))?
    }

    fn last(&self) -> Result<Option<VertexName>> {
        non_blocking(AsyncNameSetQuery::last(self.0.deref()))?
    }

    fn is_empty(&self) -> Result<bool> {
        non_blocking(AsyncNameSetQuery::is_empty(self.0.deref()))?
    }

    fn contains(&self, name: &VertexName) -> Result<bool> {
        non_blocking(AsyncNameSetQuery::contains(self.0.deref(), name))?
    }

    fn as_any(&self) -> &dyn Any {
        AsyncNameSetQuery::as_any(self.0.deref())
    }

    fn hints(&self) -> &Hints {
        AsyncNameSetQuery::hints(self.0.deref())
    }

    fn id_convert(&self) -> Option<&dyn IdConvert> {
        AsyncNameSetQuery::id_convert(self.0.deref())
    }
}

/// Iterator of [`NameSet`].
/// Types implementing this should consider replacing `iter_rev` with a fast
/// path if possible.
pub trait NameIter: Iterator<Item = Result<VertexName>> + Send {}
impl<T> NameIter for T where T: Iterator<Item = Result<VertexName>> + Send {}

/// Abstract async iterator that yields `Vertex`es.
pub trait VertexStream: Stream<Item = Result<VertexName>> + Send {}
impl<T> VertexStream for T where T: Stream<Item = Result<VertexName>> + Send {}

/// Boxed async iterator that yields `Vertex`es.
pub type BoxVertexStream = Pin<Box<dyn VertexStream>>;

/// A wrapper that converts `VertexStream` to `NameIter`.
struct NonblockingNameIter(BoxVertexStream);

impl Iterator for NonblockingNameIter {
    type Item = Result<VertexName>;

    fn next(&mut self) -> Option<Self::Item> {
        match non_blocking(self.0.next()) {
            Err(e) => Some(Err(e.into())),
            Ok(v) => v,
        }
    }
}

fn to_iter(stream: BoxVertexStream) -> Box<dyn NameIter> {
    Box::new(NonblockingNameIter(stream))
}

impl From<VertexName> for NameSet {
    fn from(name: VertexName) -> NameSet {
        NameSet::from_static_names(std::iter::once(name))
    }
}

impl From<&VertexName> for NameSet {
    fn from(name: &VertexName) -> NameSet {
        NameSet::from_static_names(std::iter::once(name.clone()))
    }
}

/// Pick `left` if `order` is "greater or equal".
/// Pick `right` otherwise.
fn pick<T>(order: cmp::Ordering, left: T, right: T) -> T {
    match order {
        cmp::Ordering::Greater | cmp::Ordering::Equal => left,
        cmp::Ordering::Less => right,
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use nonblocking::non_blocking_result as r;

    use super::*;
    use crate::Id;

    pub(crate) fn nb<F, R>(future: F) -> R
    where
        F: std::future::Future<Output = R>,
    {
        non_blocking(future).unwrap()
    }

    // Converts async Stream to Iterator.
    pub(crate) fn ni<F>(future: F) -> Result<Box<dyn NameIter>>
    where
        F: std::future::Future<Output = Result<BoxVertexStream>>,
    {
        nb(future).map(to_iter)
    }

    // For easier testing.
    impl From<&str> for NameSet {
        fn from(name: &str) -> NameSet {
            NameSet::from_static_names(
                name.split_whitespace()
                    .map(|n| VertexName::copy_from(n.as_bytes())),
            )
        }
    }

    impl NameSet {
        pub(crate) fn assert_eq(&self, other: NameSet) {
            assert!(
                other.count().unwrap() == self.count().unwrap()
                    && (other.clone() & self.clone()).count().unwrap() == self.count().unwrap(),
                "set {:?} ({:?}) != {:?} ({:?})",
                self,
                self.iter().unwrap().map(|i| i.unwrap()).collect::<Vec<_>>(),
                &other,
                other
                    .iter()
                    .unwrap()
                    .map(|i| i.unwrap())
                    .collect::<Vec<_>>(),
            );
        }
    }

    #[derive(Default, Debug)]
    pub(crate) struct VecQuery(Vec<VertexName>, Hints);

    #[async_trait::async_trait]
    impl AsyncNameSetQuery for VecQuery {
        async fn iter(&self) -> Result<BoxVertexStream> {
            let iter = self.0.clone().into_iter().map(Ok);
            Ok(Box::pin(futures::stream::iter(iter)))
        }

        fn as_any(&self) -> &dyn Any {
            self
        }

        fn hints(&self) -> &Hints {
            &self.1
        }
    }

    impl VecQuery {
        /// Quickly create [`VecQuery`] that contains `len(bytes)` items.
        pub(crate) fn from_bytes(bytes: &[u8]) -> Self {
            let mut used = [false; 256];
            Self(
                bytes
                    .iter()
                    .filter_map(|&b| {
                        if used[b as usize] {
                            None
                        } else {
                            used[b as usize] = true;
                            Some(to_name(b))
                        }
                    })
                    .collect(),
                Hints::default(),
            )
        }
    }

    /// Create a [`VertexName`] from `u8` by repeating them.
    pub(crate) fn to_name(value: u8) -> VertexName {
        VertexName::from(vec![value; 2])
    }

    /// Shorten a [`VertexName`] result.
    pub(crate) fn shorten_name(name: VertexName) -> String {
        name.to_hex()[..2].to_string()
    }

    /// Shorten a [`NameIter`] result.
    pub(crate) fn shorten_iter(iter: Result<Box<dyn NameIter>>) -> Vec<String> {
        iter.unwrap()
            .map(|v| shorten_name(v.unwrap()))
            .collect::<Vec<_>>()
    }

    #[test]
    fn test_empty_query() -> Result<()> {
        let query = VecQuery::default();
        check_invariants(&query)?;
        assert_eq!(SyncNameSetQuery::iter(&query)?.count(), 0);
        assert_eq!(SyncNameSetQuery::iter_rev(&query)?.count(), 0);
        assert_eq!(SyncNameSetQuery::first(&query)?, None);
        assert_eq!(SyncNameSetQuery::last(&query)?, None);
        assert_eq!(SyncNameSetQuery::count(&query)?, 0);
        assert!(SyncNameSetQuery::is_empty(&query)?);
        assert!(!SyncNameSetQuery::contains(&query, &to_name(0))?);
        Ok(())
    }

    #[test]
    fn test_vec_query() -> Result<()> {
        let query = VecQuery::from_bytes(b"\xab\xef\xcd");
        check_invariants(&query)?;
        assert_eq!(
            shorten_iter(SyncNameSetQuery::iter(&query)),
            ["ab", "ef", "cd"]
        );
        assert_eq!(
            shorten_iter(SyncNameSetQuery::iter_rev(&query)),
            ["cd", "ef", "ab"]
        );
        assert_eq!(
            shorten_name(SyncNameSetQuery::first(&query)?.unwrap()),
            "ab"
        );
        assert_eq!(shorten_name(SyncNameSetQuery::last(&query)?.unwrap()), "cd");
        assert!(!SyncNameSetQuery::is_empty(&query)?);
        assert!(SyncNameSetQuery::contains(&query, &to_name(0xef))?);
        assert!(!SyncNameSetQuery::contains(&query, &to_name(0))?);
        Ok(())
    }

    #[test]
    fn test_debug() {
        let set = NameSet::from_static_names(vec![to_name(2)])
            .union(&NameSet::from_static_names(vec![to_name(1)]))
            .difference(
                &NameSet::from_static_names(vec![to_name(3)])
                    .intersection(&NameSet::from_static_names(vec![to_name(2), to_name(3)])),
            );
        assert_eq!(
            format!("{:?}", &set),
            "<diff <or <static [0202]> <static [0101]>> <and <static [0303]> <static [0202, 0303]>>>"
        );
        assert_eq!(
            format!("\n{:#?}", &set),
            r#"
<diff
  <or
    <static [
        0202,
    ]>
    <static [
        0101,
    ]>>
  <and
    <static [
        0303,
    ]>
    <static [
        0202,
        0303,
    ]>>>"#
        );
    }

    #[test]
    fn test_flatten() {
        let set = NameSet::from_static_names(vec![to_name(2)])
            .union(&NameSet::from_static_names(vec![to_name(1)]))
            .difference(
                &NameSet::from_static_names(vec![to_name(3)])
                    .intersection(&NameSet::from_static_names(vec![to_name(2), to_name(3)])),
            );
        assert_eq!(
            format!("{:?}", r(set.flatten()).unwrap()),
            "<static [0202, 0101]>"
        );
    }

    #[test]
    fn test_ops() {
        let ab: NameSet = "a b".into();
        let bc: NameSet = "b c".into();
        let s = |set: NameSet| -> Vec<String> { shorten_iter(set.iter()) };
        assert_eq!(s(ab.clone() | bc.clone()), ["61", "62", "63"]);
        assert_eq!(s(ab.clone() & bc.clone()), ["62"]);
        assert_eq!(s(ab.clone() - bc.clone()), ["61"]);
    }

    #[test]
    fn test_skip_take_slow_path() {
        let s: NameSet = "a b c d".into();
        let d = |set: NameSet| -> String { format!("{:?}", r(set.flatten_names()).unwrap()) };
        assert_eq!(d(s.take(2)), "<static [a, b]>");
        assert_eq!(d(s.skip(2)), "<static [c, d]>");
        assert_eq!(d(s.skip(1).take(2)), "<static [b, c]>");
    }

    #[test]
    fn test_hints_empty_full_fast_paths() {
        let partial: NameSet = "a".into();
        partial.hints().add_flags(Flags::ID_ASC);
        let empty: NameSet = "".into();
        let full: NameSet = "f".into();
        full.hints().add_flags(Flags::FULL | Flags::ID_DESC);

        assert_eq!(
            hints_ops(&partial, &empty),
            [
                "- Hints(ID_ASC)",
                "  Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "& Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "  Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "| Hints(ID_ASC)",
                "  Hints(ID_ASC)"
            ]
        );
        // Fast paths are not used for "|" because there is no dag associated.
        assert_eq!(
            hints_ops(&partial, &full),
            [
                "- Hints(ID_ASC)",
                "  Hints(ID_DESC)",
                "& Hints(ID_ASC)",
                "  Hints(ID_ASC)",
                "| Hints((empty))",
                "  Hints((empty))"
            ]
        );
        assert_eq!(
            hints_ops(&empty, &full),
            [
                "- Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "  Hints(FULL | ID_DESC | ANCESTORS)",
                "& Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "  Hints(EMPTY | ID_DESC | ID_ASC | TOPO_DESC | ANCESTORS)",
                "| Hints(FULL | ID_DESC | ANCESTORS)",
                "  Hints(FULL | ID_DESC | ANCESTORS)"
            ]
        );
    }

    #[test]
    fn test_hints_full_subset() {
        let mut t = crate::tests::TestDag::new();
        let a = r(t.dag.all()).unwrap(); // [] FULL EMPTY
        t.drawdag("X", &[]);
        let b = r(t.dag.all()).unwrap(); // [X] FULL
        t.drawdag("X--Y--Z", &[]);
        let c = r(t.dag.all()).unwrap(); // [X Y Z] FULL
        let d = r(t.dag.heads(r(t.dag.all()).unwrap())).unwrap(); // [Z]

        let a = move || a.clone();
        let b = move || b.clone();
        let c = move || c.clone();
        let d = move || d.clone();
        let f = |set: NameSet| {
            let s = format!("{:?}", &set);
            let v = set
                .iter()
                .unwrap()
                .map(|i| String::from_utf8(i.unwrap().as_ref().to_vec()).unwrap())
                .collect::<Vec<String>>()
                .join(" ");
            format!("{} = [{}]", s, v)
        };

        assert_eq!(f(a()), "<spans []> = []");
        assert_eq!(f(b()), "<spans [X+N0]> = [X]");
        assert_eq!(f(c()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(d()), "<spans [Z+N2]> = [Z]");

        assert_eq!(f(a() - c()), "<empty> = []");
        assert_eq!(f(a() - d()), "<spans []> = []");
        assert_eq!(f(b() - c()), "<empty> = []");
        assert_eq!(f(b() - d()), "<spans [X+N0]> = [X]");
        assert_eq!(f(c() - b()), "<spans [Y:Z+N1:N2]> = [Z Y]");
        assert_eq!(f(c() - a()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(c() - d()), "<spans [X:Y+N0:N1]> = [Y X]");
        assert_eq!(f(d() - a()), "<spans [Z+N2]> = [Z]");
        assert_eq!(f(d() - b()), "<spans [Z+N2]> = [Z]");
        assert_eq!(f(d() - c()), "<empty> = []");

        assert_eq!(f(a() | c()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(a() | b()), "<spans [X+N0]> = [X]");
        assert_eq!(f(a() | d()), "<spans [Z+N2]> = [Z]");
        assert_eq!(f(b() | a()), "<spans [X+N0]> = [X]");
        assert_eq!(f(b() | c()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(b() | d()), "<spans [Z+N2, X+N0]> = [Z X]");
        assert_eq!(f(c() | a()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(c() | b()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(c() | d()), "<spans [X:Z+N0:N2]> = [Z Y X]");
        assert_eq!(f(d() | a()), "<spans [Z+N2]> = [Z]");
        assert_eq!(f(d() | b()), "<spans [Z+N2, X+N0]> = [Z X]");
        assert_eq!(f(d() | c()), "<spans [X:Z+N0:N2]> = [Z Y X]");

        assert_eq!(f(a() & c()), "<spans []> = []");
        assert_eq!(f(a() & d()), "<empty> = []");
        assert_eq!(f(b() & c()), "<spans [X+N0]> = [X]");
        assert_eq!(f(c() & a()), "<spans []> = []");
        assert_eq!(f(c() & b()), "<spans [X+N0]> = [X]");
        assert_eq!(f(c() & d()), "<spans [Z+N2]> = [Z]");
        assert_eq!(f(d() & a()), "<empty> = []");
        assert_eq!(f(d() & b()), "<spans []> = []");
        assert_eq!(f(d() & c()), "<spans [Z+N2]> = [Z]");
    }

    #[test]
    fn test_hints_min_max_id() {
        let bc: NameSet = "b c".into();
        bc.hints().set_min_id(Id(20));
        bc.hints().add_flags(Flags::ID_DESC);

        let ad: NameSet = "a d".into();
        ad.hints().set_max_id(Id(40));
        ad.hints().add_flags(Flags::ID_ASC);

        assert_eq!(
            hints_ops(&bc, &ad),
            [
                "- Hints(ID_DESC, 20..)",
                "  Hints(ID_ASC, ..=40)",
                "& Hints(ID_DESC, 20..=40)",
                "  Hints(ID_ASC, 20..=40)",
                "| Hints((empty))",
                "  Hints((empty))"
            ]
        );

        ad.hints().set_min_id(Id(10));
        bc.hints().set_max_id(Id(30));
        assert_eq!(
            hints_ops(&bc, &ad),
            [
                "- Hints(ID_DESC, 20..=30)",
                "  Hints(ID_ASC, 10..=40)",
                "& Hints(ID_DESC, 20..=30)",
                "  Hints(ID_ASC, 20..=30)",
                "| Hints((empty))",
                "  Hints((empty))"
            ]
        );
    }

    #[test]
    fn test_hints_ancestors() {
        let a: NameSet = "a".into();
        a.hints().add_flags(Flags::ANCESTORS);

        let b: NameSet = "b".into();
        assert_eq!(
            hints_ops(&a, &b),
            [
                "- Hints((empty))",
                "  Hints((empty))",
                "& Hints((empty))",
                "  Hints((empty))",
                "| Hints((empty))",
                "  Hints((empty))"
            ]
        );

        b.hints().add_flags(Flags::ANCESTORS);
        assert_eq!(
            hints_ops(&a, &b),
            [
                "- Hints((empty))",
                "  Hints((empty))",
                "& Hints(ANCESTORS)",
                "  Hints(ANCESTORS)",
                "| Hints(ANCESTORS)",
                "  Hints(ANCESTORS)"
            ]
        );
    }

    #[test]
    fn test_filter() {
        id_static::tests::with_dag(|dag| {
            let sets: Vec<NameSet> = vec!["C B A".into(), nb(dag.ancestors("C".into())).unwrap()];
            for abc in sets {
                let filter: NameSet = abc.filter(Box::new(|v: &VertexName| {
                    Box::pin(async move { Ok(v.as_ref() != b"A") })
                }));
                check_invariants(filter.0.as_ref()).unwrap();
                assert_eq!(abc.hints().dag_version(), filter.hints().dag_version());
                assert_eq!(
                    abc.hints().id_map_version(),
                    filter.hints().id_map_version()
                );
                assert!(filter.hints().flags().contains(Flags::FILTER));
                assert!(!filter.hints().flags().contains(Flags::ANCESTORS));
                assert_eq!(
                    format!("{:?}", r(filter.flatten_names())),
                    "Ok(<static [C, B]>)"
                );
            }
        })
    }

    // Print hints for &, |, - operations.
    fn hints_ops(lhs: &NameSet, rhs: &NameSet) -> Vec<String> {
        vec![
            (lhs.clone() - rhs.clone(), rhs.clone() - lhs.clone()),
            (lhs.clone() & rhs.clone(), rhs.clone() & lhs.clone()),
            (lhs.clone() | rhs.clone(), rhs.clone() | lhs.clone()),
        ]
        .into_iter()
        .zip("-&|".chars())
        .flat_map(|((set1, set2), ch)| {
            vec![
                format!("{} {:?}", ch, set1.hints()),
                format!("  {:?}", set2.hints()),
            ]
        })
        .collect()
    }

    /// Check consistency of a `AsyncNameSetQuery`, such as `iter().nth(0)` matches
    /// `first()` etc.
    pub(crate) fn check_invariants(query: &dyn AsyncNameSetQuery) -> Result<()> {
        // Collect contains_fast result before calling other functions which might
        // change the internal set state.
        let contains_fast_vec: Vec<Option<bool>> = (0..=0xff)
            .map(|b| {
                let name = VertexName::from(vec![b; 20]);
                nb(query.contains_fast(&name)).unwrap_or(None)
            })
            .collect();
        let is_empty = nb(query.is_empty())?;
        let count = nb(query.count())?;
        let first = nb(query.first())?;
        let last = nb(query.last())?;
        let names: Vec<VertexName> = ni(query.iter())?.collect::<Result<Vec<_>>>()?;
        assert_eq!(
            first,
            names.first().cloned(),
            "first() should match iter().first() (set: {:?})",
            &query
        );
        assert_eq!(
            last,
            names.last().cloned(),
            "last() should match iter().last() (set: {:?})",
            &query
        );
        assert_eq!(
            count,
            names.len(),
            "count() should match iter().count() (set: {:?})",
            &query
        );
        assert_eq!(
            is_empty,
            count == 0,
            "is_empty() should match count() == 0 (set: {:?})",
            &query
        );
        assert!(
            names
                .iter()
                .all(|name| nb(query.contains(name)).ok() == Some(true)),
            "contains() should return true for names returned by iter() (set: {:?})",
            &query
        );
        assert!(
            names
                .iter()
                .all(|name| nb(query.contains_fast(name)).unwrap_or(None) != Some(false)),
            "contains_fast() should not return false for names returned by iter() (set: {:?})",
            &query
        );
        assert!(
            (0..=0xff).all(|b| {
                let name = VertexName::from(vec![b; 20]);
                nb(query.contains(&name)).ok() == Some(names.contains(&name))
            }),
            "contains() should return false for names not returned by iter() (set: {:?})",
            &query
        );
        assert!(
            (0..=0xff)
                .zip(contains_fast_vec.into_iter())
                .all(|(b, old_contains)| {
                    let name = VertexName::from(vec![b; 20]);
                    let contains = nb(query.contains_fast(&name)).unwrap_or(None);
                    old_contains == None || contains == old_contains
                }),
            "contains_fast() should be consistent (set: {:?})",
            &query
        );
        assert!(
            (0..=0xff).all(|b| {
                let name = VertexName::from(vec![b; 20]);
                let contains = nb(query.contains_fast(&name)).unwrap_or(None);
                contains == None || contains == Some(names.contains(&name))
            }),
            "contains_fast() should not return true for names not returned by iter() (set: {:?})",
            &query
        );
        let reversed: Vec<VertexName> = ni(query.iter_rev())?.collect::<Result<Vec<_>>>()?;
        assert_eq!(
            names,
            reversed.into_iter().rev().collect::<Vec<VertexName>>(),
            "iter() should match iter_rev().rev() (set: {:?})",
            &query
        );
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::any::Any;
use std::collections::HashSet;
use std::fmt;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Release;
use std::sync::Arc;

use futures::lock::Mutex;
use futures::StreamExt;
use indexmap::IndexSet;
use tracing::debug;
use tracing::instrument;
use tracing::trace;
use tracing::Level;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use super::NameSet;
use crate::fmt::write_debug;
use crate::Result;
use crate::VertexName;

/// Slice of a set.
#[derive(Clone)]
pub struct SliceSet {
    inner: NameSet,
    hints: Hints,
    skip_count: u64,
    take_count: Option<u64>,

    // Skipped vertexes. Updated during iteration.
    skip_cache: Arc<Mutex<HashSet<VertexName>>>,
    // Taken vertexes. Updated during iteration.
    take_cache: Arc<Mutex<IndexSet<VertexName>>>,
    // If take_cache is complete.
    take_cache_complete: Arc<AtomicBool>,
}

impl SliceSet {
    pub fn new(set: NameSet, skip_count: u64, take_count: Option<u64>) -> Self {
        let hints = set.hints().clone();
        hints.update_flags_with(|mut f| {
            // Only keep compatible flags.
            f &= Flags::ID_DESC
                | Flags::ID_ASC
                | Flags::TOPO_DESC
                | Flags::HAS_MIN_ID
                | Flags::HAS_MAX_ID
                | Flags::EMPTY;
            // Add EMPTY hints if take_count is 0.
            if take_count == Some(0) {
                f |= Flags::EMPTY;
            }
            f
        });
        Self {
            inner: set,
            hints,
            skip_count,
            take_count,
            skip_cache: Default::default(),
            take_cache: Default::default(),
            take_cache_complete: Default::default(),
        }
    }

    fn is_take_cache_complete(&self) -> bool {
        self.take_cache_complete.load(Acquire)
    }

    async fn is_skip_cache_complete(&self) -> bool {
        self.skip_cache.lock().await.len() as u64 == self.skip_count
    }

    #[instrument(level=Level::DEBUG)]
    async fn populate_take_cache(&self) -> Result<()> {
        // See Iter::next. If take_count is not set, the "take" can be unbounded,
        // and take_cache won't be populated.
        assert!(self.take_count.is_some());

        // Use iter() to populate take_cache.
        let mut iter = self.iter().await?;
        while let Some(_) = iter.next().await {}
        assert!(self.is_take_cache_complete());

        Ok(())
    }
}

struct Iter {
    inner_iter: BoxVertexStream,
    set: SliceSet,
    index: u64,
}

const SKIP_CACHE_SIZE_THRESHOLD: u64 = 1000;

impl Iter {
    async fn next(&mut self) -> Option<Result<VertexName>> {
        if self.set.is_take_cache_complete() {
            // Fast path - no need to use inner_iter.
            let index = self.index.max(self.set.skip_count);
            let take_index = index - self.set.skip_count;
            let result = {
                let cache = self.set.take_cache.lock().await;
                cache.get_index(take_index as _).cloned()
            };
            trace!("next(index={}) = {:?} (fast path)", index, &result);
            self.index = index + 1;
            return Ok(result).transpose();
        }

        loop {
            // Slow path - use inner_iter.
            let index = self.index;
            trace!("next(index={})", index);
            let next: Option<VertexName> = match self.inner_iter.next().await {
                Some(Err(e)) => {
                    self.index = u64::MAX;
                    return Some(Err(e));
                }
                Some(Ok(v)) => Some(v),
                None => None,
            };
            self.index += 1;

            // Skip?
            if index < self.set.skip_count {
                if index < SKIP_CACHE_SIZE_THRESHOLD {
                    // Update skip_cache.
                    if let Some(v) = next.as_ref() {
                        let mut cache = self.set.skip_cache.lock().await;
                        cache.insert(v.clone());
                    }
                }
                continue;
            }

            // Take?
            let take_index = index - self.set.skip_count;
            let should_take: bool = match self.set.take_count {
                Some(count) => {
                    if take_index < count {
                        // Update take_cache.
                        let mut cache = self.set.take_cache.lock().await;
                        if take_index == cache.len() as u64 {
                            if let Some(v) = next.as_ref() {
                                cache.insert(v.clone());
                            } else {
                                // No more item in the original set.
                                self.set.take_cache_complete.store(true, Release);
                            }
                        }
                        true
                    } else {
                        self.set.take_cache_complete.store(true, Release);
                        false
                    }
                }
                None => {
                    // Do not update take_cache, since the inner
                    // set can be quite large.
                    true
                }
            };
            if should_take {
                return next.map(Ok);
            } else {
                return None;
            }
        }
    }

    fn into_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |mut state| async move {
            let result = state.next().await;
            result.map(|r| (r, state))
        }))
    }
}

struct TakeCacheRevIter {
    take_cache: Arc<Mutex<IndexSet<VertexName>>>,
    index: usize,
}

impl TakeCacheRevIter {
    async fn next(&mut self) -> Option<Result<VertexName>> {
        let index = self.index;
        self.index += 1;
        let cache = self.take_cache.lock().await;
        if index >= cache.len() {
            None
        } else {
            let index = cache.len() - index - 1;
            cache.get_index(index).cloned().map(Ok)
        }
    }

    fn into_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |mut state| async move {
            let result = state.next().await;
            result.map(|r| (r, state))
        }))
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for SliceSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let inner_iter = self.inner.iter().await?;
        let iter = Iter {
            inner_iter,
            set: self.clone(),
            index: 0,
        };
        Ok(iter.into_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        if let Some(_take) = self.take_count {
            self.populate_take_cache().await?;
            trace!("iter_rev({:0.6?}): use take_cache", self);
            // Use take_cache to answer RevIter. This is probably better
            // than using inner.iter_rev(), if take_count is small:
            //     [<----------------------------]
            //     [skip][take][...(need skip)...]
            let iter = TakeCacheRevIter {
                take_cache: self.take_cache.clone(),
                index: 0,
            };
            Ok(iter.into_stream())
        } else {
            // Unbounded "take_count". Reuse inner.rev_iter().
            //     [<-------------------]
            //     [skip][<---take------]
            trace!("iter_rev({:0.6?}): use inner.iter_rev()", self,);
            let count = self.count().await?;
            let iter = self.inner.iter_rev().await?;
            Ok(Box::pin(iter.take(count)))
        }
    }

    async fn count(&self) -> Result<usize> {
        let count = self.inner.count().await?;
        // consider skip_count
        let count = (count as u64).max(self.skip_count) - self.skip_count;
        // consider take_count
        let count = count.min(self.take_count.unwrap_or(u64::MAX));
        Ok(count as _)
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        if let Some(result) = self.contains_fast(name).await? {
            return Ok(result);
        }

        debug!("SliceSet::contains({:.6?}, {:?}) (slow path)", self, name);
        let mut iter = self.iter().await?;
        while let Some(item) = iter.next().await {
            if &item? == name {
                return Ok(true);
            }
        }
        Ok(false)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        // Check take_cache.
        {
            let take_cache = self.take_cache.lock().await;
            let is_take_cache_complete = self.is_take_cache_complete();
            let contains = take_cache.contains(name);
            match (contains, is_take_cache_complete) {
                (_, true) | (true, _) => return Ok(Some(contains)),
                (false, false) => {}
            }
        }

        // Check skip_cache.
        // Assumes one vertex only occurs once in a set.
        let skip_contains = self.skip_cache.lock().await.contains(name);
        if skip_contains {
            return Ok(Some(false));
        }

        // Check with the original set.
        let result = self.inner.contains_fast(name).await?;
        match (result, self.is_skip_cache_complete().await) {
            // Not in the original set. Slice is a subset. Result: false.
            (Some(false), _) => Ok(Some(false)),
            // In the original set. Skip cache is completed _and_ checked
            // above (name was _not_ skipped). Result: true.
            (Some(true), true) => {
                // skip_cache was checked above
                debug_assert!(!self.skip_cache.lock().await.contains(name));
                Ok(Some(true))
            }
            // Unsure cases.
            (None, false) => Ok(None),
            (Some(true), false) => Ok(None),
            (None, true) => Ok(None),
        }
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

impl fmt::Debug for SliceSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("<slice")?;
        write_debug(f, &self.inner)?;
        f.write_str(" [")?;
        if self.skip_count > 0 {
            write!(f, "{}", self.skip_count)?;
        }
        f.write_str("..")?;
        if let Some(n) = self.take_count {
            write!(f, "{}", self.skip_count + n)?;
        }
        f.write_str("]>")
    }
}

#[cfg(test)]
#[allow(clippy::redundant_clone)]
mod tests {
    use nonblocking::non_blocking_result as r;

    use super::super::tests::*;
    use super::*;

    #[test]
    fn test_basic() -> Result<()> {
        let orig = NameSet::from("a b c d e f g h i");
        let count = r(orig.count())?;

        let set = SliceSet::new(orig.clone(), 0, None);
        assert_eq!(r(set.count())?, count);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 0, Some(0));
        assert_eq!(r(set.count())?, 0);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 4, None);
        assert_eq!(r(set.count())?, count - 4);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 4, Some(0));
        assert_eq!(r(set.count())?, 0);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 0, Some(4));
        assert_eq!(r(set.count())?, 4);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 4, Some(4));
        assert_eq!(r(set.count())?, 4);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 7, Some(4));
        assert_eq!(r(set.count())?, 2);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 20, Some(4));
        assert_eq!(r(set.count())?, 0);
        check_invariants(&set)?;

        let set = SliceSet::new(orig.clone(), 20, Some(0));
        assert_eq!(r(set.count())?, 0);
        check_invariants(&set)?;

        Ok(())
    }

    #[test]
    fn test_debug() {
        let orig = NameSet::from("a b c d e f g h i");
        let set = SliceSet::new(orig.clone(), 0, None);
        assert_eq!(
            format!("{:?}", set),
            "<slice <static [a, b, c] + 6 more> [..]>"
        );
        let set = SliceSet::new(orig.clone(), 4, None);
        assert_eq!(
            format!("{:?}", set),
            "<slice <static [a, b, c] + 6 more> [4..]>"
        );
        let set = SliceSet::new(orig.clone(), 4, Some(4));
        assert_eq!(
            format!("{:?}", set),
            "<slice <static [a, b, c] + 6 more> [4..8]>"
        );
        let set = SliceSet::new(orig.clone(), 0, Some(4));
        assert_eq!(
            format!("{:?}", set),
            "<slice <static [a, b, c] + 6 more> [..4]>"
        );
    }

    quickcheck::quickcheck! {
        fn test_static_quickcheck(skip_and_take: u8) -> bool {
            let skip = (skip_and_take & 0xf) as u64;
            let take = (skip_and_take >> 4) as u64;
            let take = if take > 12 {
                None
            } else {
                Some(take as u64)
            };
            let orig = NameSet::from("a c b d e f g i h j");
            let set = SliceSet::new(orig, skip, take);
            check_invariants(&set).unwrap();
            true
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::any::Any;
use std::fmt;

use indexmap::IndexSet;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use crate::Result;
use crate::VertexName;

/// A set backed by a concrete ordered set.
pub struct StaticSet(pub(crate) IndexSet<VertexName>, Hints);

impl StaticSet {
    pub fn from_names(names: impl IntoIterator<Item = VertexName>) -> Self {
        let names: IndexSet<VertexName> = names.into_iter().collect();
        let hints = Hints::default();
        if names.is_empty() {
            hints.add_flags(Flags::EMPTY);
        }
        Self(names, hints)
    }

    pub fn empty() -> Self {
        let names: IndexSet<VertexName> = Default::default();
        let hints = Hints::default();
        hints.add_flags(Flags::EMPTY);
        Self(names, hints)
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for StaticSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let iter = self.0.clone().into_iter().map(Ok);
        Ok(Box::pin(futures::stream::iter(iter)))
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let iter = self.0.clone().into_iter().rev().map(Ok);
        Ok(Box::pin(futures::stream::iter(iter)))
    }

    async fn count(&self) -> Result<usize> {
        Ok(self.0.len())
    }

    async fn is_empty(&self) -> Result<bool> {
        Ok(self.0.is_empty())
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        Ok(self.0.contains(name))
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        Ok(Some(self.0.contains(name)))
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn hints(&self) -> &Hints {
        &self.1
    }
}

impl fmt::Debug for StaticSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.0.is_empty() {
            return f.write_str("<empty>");
        }
        write!(f, "<static ")?;
        // Only show 3 commits by default.
        let limit = f.width().unwrap_or(3);
        f.debug_list().entries(self.0.iter().take(limit)).finish()?;
        let remaining = self.0.len().max(limit) - limit;
        if remaining > 0 {
            write!(f, " + {} more>", remaining)?;
        } else {
            write!(f, ">")?;
        }
        Ok(())
    }
}

// Test infra is unhappy about 'r#' yet (D20008157).
#[cfg(not(fbcode_build))]
#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::super::tests::*;
    use super::*;

    fn static_set(a: &[u8]) -> StaticSet {
        StaticSet::from_names(a.iter().map(|&b| to_name(b)))
    }

    #[test]
    fn test_static_basic() -> Result<()> {
        let set = static_set(b"\x11\x33\x22\x77\x22\x55\x11");
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["11", "33", "22", "77", "55"]);
        assert_eq!(
            shorten_iter(ni(set.iter_rev())),
            ["55", "77", "22", "33", "11"]
        );
        assert!(!nb(set.is_empty())?);
        assert_eq!(nb(set.count())?, 5);
        assert_eq!(shorten_name(nb(set.first())?.unwrap()), "11");
        assert_eq!(shorten_name(nb(set.last())?.unwrap()), "55");
        Ok(())
    }

    #[test]
    fn test_debug() {
        let set = static_set(b"");
        assert_eq!(format!("{:?}", set), "<empty>");

        let set = static_set(b"\x11\x33\x22");
        assert_eq!(format!("{:?}", set), "<static [1111, 3333, 2222]>");

        let set = static_set(b"\xaa\x00\xaa\xdd\xee\xdd\x11\x22");
        assert_eq!(
            format!("{:?}", &set),
            "<static [aaaa, 0000, dddd] + 3 more>"
        );
        // {:#?} can be used to show commits in multi-line.
        assert_eq!(
            format!("{:#?}", &set),
            "<static [\n    aaaa,\n    0000,\n    dddd,\n] + 3 more>"
        );
        // {:5.2} can be used to control how many commits to show, and their length.
        assert_eq!(
            format!("{:5.2?}", &set),
            "<static [aa, 00, dd, ee, 11] + 1 more>"
        );
    }

    quickcheck::quickcheck! {
        fn test_static_quickcheck(a: Vec<u8>) -> bool {
            let set = static_set(&a);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap();
            assert!(count <= a.len());

            let set2: HashSet<_> = a.iter().cloned().collect();
            assert_eq!(count, set2.len());

            true
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::any::Any;
use std::fmt;

use futures::StreamExt;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use super::NameSet;
use crate::fmt::write_debug;
use crate::Result;
use crate::VertexName;

/// Union of 2 sets.
///
/// The order is preserved. The first set is iterated first, then the second set
/// is iterated, with duplicated names skipped.
pub struct UnionSet {
    sets: [NameSet; 2],
    hints: Hints,
}

impl UnionSet {
    pub fn new(lhs: NameSet, rhs: NameSet) -> Self {
        let hints = Hints::union(&[lhs.hints(), rhs.hints()]);
        if hints.id_map().is_some() {
            if let (Some(id1), Some(id2)) = (lhs.hints().min_id(), rhs.hints().min_id()) {
                hints.set_min_id(id1.min(id2));
            }
            if let (Some(id1), Some(id2)) = (lhs.hints().max_id(), rhs.hints().max_id()) {
                hints.set_max_id(id1.max(id2));
            }
        };
        hints.add_flags(lhs.hints().flags() & rhs.hints().flags() & Flags::ANCESTORS);
        if lhs.hints().contains(Flags::FILTER) || rhs.hints().contains(Flags::FILTER) {
            hints.add_flags(Flags::FILTER);
        }
        Self {
            sets: [lhs, rhs],
            hints,
        }
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for UnionSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        debug_assert_eq!(self.sets.len(), 2);
        let diff = self.sets[1].clone() - self.sets[0].clone();
        let diff_iter = diff.iter().await?;
        let set0_iter = self.sets[0].iter().await?;
        let iter = set0_iter.chain(diff_iter);
        Ok(Box::pin(iter))
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        debug_assert_eq!(self.sets.len(), 2);
        let diff = self.sets[1].clone() - self.sets[0].clone();
        let diff_iter = diff.iter_rev().await?;
        let set0_iter = self.sets[0].iter_rev().await?;
        let iter = diff_iter.chain(set0_iter);
        Ok(Box::pin(iter))
    }

    async fn count(&self) -> Result<usize> {
        debug_assert_eq!(self.sets.len(), 2);
        // This is more efficient if sets[0] is a large set that has a fast path
        // for "count()".
        let mut count = self.sets[0].count().await?;
        let mut iter = self.sets[1].iter().await?;
        while let Some(item) = iter.next().await {
            let name = item?;
            if !self.sets[0].contains(&name).await? {
                count += 1;
            }
        }
        Ok(count)
    }

    async fn is_empty(&self) -> Result<bool> {
        for set in &self.sets {
            if !set.is_empty().await? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        for set in &self.sets {
            if set.contains(name).await? {
                return Ok(true);
            }
        }
        Ok(false)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        for set in &self.sets {
            if let Some(result) = set.contains_fast(name).await? {
                return Ok(Some(result));
            }
        }
        Ok(None)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

impl fmt::Debug for UnionSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<or")?;
        write_debug(f, &self.sets[0])?;
        write_debug(f, &self.sets[1])?;
        write!(f, ">")
    }
}

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

    use super::super::tests::*;
    use super::*;

    fn union(a: &[u8], b: &[u8]) -> UnionSet {
        let a = NameSet::from_query(VecQuery::from_bytes(a));
        let b = NameSet::from_query(VecQuery::from_bytes(b));
        UnionSet::new(a, b)
    }

    #[test]
    fn test_union_basic() -> Result<()> {
        // 'a' overlaps with 'b'. UnionSet should de-duplicate items.
        let set = union(b"\x11\x33\x22", b"\x44\x11\x55\x33");
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["11", "33", "22", "44", "55"]);
        assert_eq!(
            shorten_iter(ni(set.iter_rev())),
            ["55", "44", "22", "33", "11"]
        );
        assert!(!nb(set.is_empty())?);
        assert_eq!(nb(set.count())?, 5);
        assert_eq!(shorten_name(nb(set.first())?.unwrap()), "11");
        assert_eq!(shorten_name(nb(set.last())?.unwrap()), "55");
        for &b in b"\x11\x22\x33\x44\x55".iter() {
            assert!(nb(set.contains(&to_name(b)))?);
        }
        for &b in b"\x66\x77\x88".iter() {
            assert!(!nb(set.contains(&to_name(b)))?);
        }
        Ok(())
    }

    quickcheck::quickcheck! {
        fn test_union_quickcheck(a: Vec<u8>, b: Vec<u8>) -> bool {
            let set = union(&a, &b);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap();
            assert!(count <= a.len() + b.len());

            let set2: HashSet<_> = a.iter().chain(b.iter()).cloned().collect();
            assert_eq!(count, set2.len());

            assert!(a.iter().all(|&b| nb(set.contains(&to_name(b))).ok() == Some(true)));
            assert!(b.iter().all(|&b| nb(set.contains(&to_name(b))).ok() == Some(true)));

            true
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;

use crate::namedag::MemNameDag;
use crate::ops::Parents;
use crate::Result;
use crate::Vertex;

/// Represents a graph from ASCII parsed by `drawdag`.
pub struct DrawDag {
    parents: HashMap<Vertex, Vec<Vertex>>,
}

impl<'a> From<&'a str> for DrawDag {
    fn from(text: &'a str) -> Self {
        let parents = ::drawdag::parse(text);
        let v = |s: String| Vertex::copy_from(s.as_bytes());
        let parents = parents
            .into_iter()
            .map(|(k, vs)| (v(k), vs.into_iter().map(v).collect()))
            .collect();
        Self { parents }
    }
}

impl DrawDag {
    /// Heads in the graph. Vertexes that are not parents of other vertexes.
    pub fn heads(&self) -> Vec<Vertex> {
        let mut heads = self
            .parents
            .keys()
            .collect::<HashSet<_>>()
            .difference(&self.parents.values().flatten().collect())
            .map(|&v| v.clone())
            .collect::<Vec<_>>();
        heads.sort();
        heads
    }
}

#[async_trait::async_trait]
impl Parents for DrawDag {
    async fn parent_names(&self, name: Vertex) -> Result<Vec<Vertex>> {
        Parents::parent_names(&self.parents, name).await
    }

    async fn hint_subdag_for_insertion(&self, heads: &[Vertex]) -> Result<MemNameDag> {
        Parents::hint_subdag_for_insertion(&self.parents, heads).await
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::sync::Arc;

use nonblocking::non_blocking;

use crate::ops::DagAlgorithm;
use crate::NameSet;
use crate::Result;
use crate::VerLink;
use crate::VertexName;

/// The DummyDag implements a DAG that contains all vertexes with no parents.
#[derive(Debug, Clone)]
pub(crate) struct DummyDag {
    version: VerLink,
}

impl DummyDag {
    pub fn new() -> Self {
        Self {
            version: VerLink::new(),
        }
    }
}

#[async_trait::async_trait]
impl DagAlgorithm for DummyDag {
    async fn sort(&self, set: &NameSet) -> Result<NameSet> {
        Ok(set.clone())
    }

    /// Get ordered parent vertexes.
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        let _ = name;
        Ok(Vec::new())
    }

    /// Returns a set that covers all vertexes tracked by this DAG.
    async fn all(&self) -> Result<NameSet> {
        crate::errors::programming("DummyDag does not support all()")
    }

    /// Returns a set that covers all vertexes in the master group.
    async fn master_group(&self) -> Result<NameSet> {
        crate::errors::programming("DummyDag does not support master_group()")
    }

    /// Vertexes buffered in memory, not yet written to disk.
    async fn dirty(&self) -> Result<NameSet> {
        Ok(NameSet::empty())
    }

    /// Calculates all ancestors reachable from any name from the given set.
    async fn ancestors(&self, set: NameSet) -> Result<NameSet> {
        Ok(set)
    }

    /// Calculates parents of the given set.
    async fn parents(&self, set: NameSet) -> Result<NameSet> {
        let _ = set;
        Ok(NameSet::empty())
    }

    /// Calculates the n-th first ancestor.
    async fn first_ancestor_nth(&self, name: VertexName, n: u64) -> Result<Option<VertexName>> {
        if n == 0 {
            Ok(Some(name))
        } else {
            crate::errors::programming("DummyDag does not resolve x~n where n > 1")
        }
    }

    /// Calculates heads of the given set.
    async fn heads(&self, set: NameSet) -> Result<NameSet> {
        Ok(set)
    }

    /// Calculates children of the given set.
    async fn children(&self, set: NameSet) -> Result<NameSet> {
        let _ = set;
        Ok(NameSet::empty())
    }

    /// Calculates roots of the given set.
    async fn roots(&self, set: NameSet) -> Result<NameSet> {
        Ok(set)
    }

    /// Calculates one "greatest common ancestor" of the given set.
    ///
    /// If there are no common ancestors, return None.
    /// If there are multiple greatest common ancestors, pick one arbitrarily.
    /// Use `gca_all` to get all of them.
    async fn gca_one(&self, set: NameSet) -> Result<Option<VertexName>> {
        if non_blocking(set.count())?? == 1 {
            non_blocking(set.first())?
        } else {
            Ok(None)
        }
    }

    /// Calculates all "greatest common ancestor"s of the given set.
    /// `gca_one` is faster if an arbitrary answer is ok.
    async fn gca_all(&self, set: NameSet) -> Result<NameSet> {
        self.common_ancestors(set).await
    }

    /// Calculates all common ancestors of the given set.
    async fn common_ancestors(&self, set: NameSet) -> Result<NameSet> {
        if non_blocking(set.count())?? == 1 {
            Ok(set)
        } else {
            Ok(NameSet::empty())
        }
    }

    /// Tests if `ancestor` is an ancestor of `descendant`.
    async fn is_ancestor(&self, ancestor: VertexName, descendant: VertexName) -> Result<bool> {
        Ok(ancestor == descendant)
    }

    /// Calculates "heads" of the ancestors of the given set. That is,
    /// Find Y, which is the smallest subset of set X, where `ancestors(Y)` is
    /// `ancestors(X)`.
    ///
    /// This is faster than calculating `heads(ancestors(set))`.
    ///
    /// This is different from `heads`. In case set contains X and Y, and Y is
    /// an ancestor of X, but not the immediate ancestor, `heads` will include
    /// Y while this function won't.
    async fn heads_ancestors(&self, set: NameSet) -> Result<NameSet> {
        Ok(set)
    }

    /// Calculates the "dag range" - vertexes reachable from both sides.
    async fn range(&self, roots: NameSet, heads: NameSet) -> Result<NameSet> {
        Ok(roots & heads)
    }

    /// Calculates the descendants of the given set.
    async fn descendants(&self, set: NameSet) -> Result<NameSet> {
        Ok(set)
    }

    fn is_vertex_lazy(&self) -> bool {
        false
    }

    /// Get a snapshot of the current graph.
    fn dag_snapshot(&self) -> Result<Arc<dyn DagAlgorithm + Send + Sync>> {
        Ok(Arc::new(self.clone()))
    }

    fn dag_id(&self) -> &str {
        "dummy_dag"
    }

    fn dag_version(&self) -> &VerLink {
        &self.version
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::sync::Arc;
use std::sync::Mutex;

use futures::StreamExt;
use futures::TryStreamExt;
use nonblocking::non_blocking;
use nonblocking::non_blocking_result;
use tracing::debug;

use crate::ops::CheckIntegrity;
use crate::ops::DagAddHeads;
use crate::ops::DagAlgorithm;
use crate::ops::DagExportCloneData;
use crate::ops::DagExportPullData;
use crate::ops::DagImportCloneData;
use crate::ops::DagImportPullData;
use crate::ops::DagPersistent;
use crate::ops::DagStrip;
use crate::ops::IdConvert;
use crate::protocol;
use crate::protocol::RemoteIdConvertProtocol;
#[cfg(feature = "render")]
use crate::render::render_namedag;
use crate::tests::DrawDag;
use crate::Group;
use crate::Level;
use crate::NameDag;
use crate::Result;
use crate::Set;
use crate::Vertex;
use crate::VertexListWithOptions;

/// Dag structure for testing purpose.
pub struct TestDag {
    pub dag: NameDag,
    pub seg_size: usize,
    pub dir: tempfile::TempDir,
    pub output: Arc<Mutex<Vec<String>>>,
}

impl TestDag {
    /// Creates a `TestDag` for testing.
    /// Side effect of the `TestDag` will be removed on drop.
    pub fn new() -> Self {
        Self::new_with_segment_size(3)
    }

    /// Crates a `TestDag` using the given ASCII.
    ///
    /// This is just `new`, followed by `drawdag`, with an extra rule that
    /// comments like "# master: M" at the end can be used to specify master
    /// heads .
    pub fn draw(text: &str) -> Self {
        let mut dag = Self::new();
        let mut split = text.split("# master:");
        let text = split.next().unwrap_or("");
        let master = match split.next() {
            Some(t) => t.split_whitespace().collect::<Vec<_>>(),
            None => Vec::new(),
        };
        dag.drawdag(text, &master);
        dag
    }

    /// Similar to `draw` but creates a lazy client so all vertexes
    /// in the master group are lazy.
    pub async fn draw_client(text: &str) -> Self {
        let server = Self::draw(text);
        // clone data won't include non-master group.
        let mut client = server.client_cloned_data().await;
        tracing::debug!("CLIENT");
        #[cfg(test)]
        tracing::debug!("CLIENT: {}", client.dump_state().await);
        let non_master_heads = {
            let all = server.dag.all().await.unwrap();
            let non_master = all.difference(&server.dag.master_group().await.unwrap());
            let heads = server.dag.heads(non_master).await.unwrap();
            let iter = heads.iter().await.unwrap();
            iter.try_collect::<Vec<_>>().await.unwrap()
        };
        let heads =
            VertexListWithOptions::from(non_master_heads).with_highest_group(Group::NON_MASTER);
        client
            .dag
            .add_heads_and_flush(&server.dag.dag_snapshot().unwrap(), &heads)
            .await
            .unwrap();
        client
    }

    /// Creates a `TestDag` with a specific segment size.
    pub fn new_with_segment_size(seg_size: usize) -> Self {
        let dir = tempfile::tempdir().unwrap();
        let dag = NameDag::open(dir.path().join("n")).unwrap();
        Self {
            dir,
            dag,
            seg_size,
            output: Default::default(),
        }
    }

    /// Reopen the dag. Drop in-memory state including caches.
    pub fn reopen(&mut self) {
        let mut dag = NameDag::open(self.dir.path().join("n")).unwrap();
        dag.set_remote_protocol(self.dag.get_remote_protocol());
        self.dag = dag;
    }

    /// Add vertexes to the graph. Does not resolve vertexes remotely.
    pub fn drawdag(&mut self, text: &str, master_heads: &[&str]) {
        self.drawdag_with_limited_heads(text, master_heads, None);
    }

    /// Add vertexes to the graph. Async version that might resolve vertexes
    /// remotely on demand.
    pub async fn drawdag_async(&mut self, text: &str, master_heads: &[&str]) {
        // Do not call self.validate to avoid fetching vertexes remotely.
        self.drawdag_with_limited_heads_async(text, master_heads, None, false)
            .await
    }

    /// Add vertexes to the graph.
    ///
    /// If `heads` is set, ignore part of the graph. Only consider specified
    /// heads.
    pub fn drawdag_with_limited_heads(
        &mut self,
        text: &str,
        master_heads: &[&str],
        heads: Option<&[&str]>,
    ) {
        non_blocking(self.drawdag_with_limited_heads_async(text, master_heads, heads, true))
            .unwrap()
    }

    pub async fn drawdag_with_limited_heads_async(
        &mut self,
        text: &str,
        master_heads: &[&str],
        heads: Option<&[&str]>,
        validate: bool,
    ) {
        let (all_heads, parent_func) = get_heads_and_parents_func_from_ascii(text);
        let heads = match heads {
            Some(heads) => heads
                .iter()
                .map(|s| Vertex::copy_from(s.as_bytes()))
                .collect(),
            None => all_heads,
        };
        self.dag.dag.set_new_segment_size(self.seg_size);
        self.dag
            .add_heads(&parent_func, &heads.into())
            .await
            .unwrap();
        if validate {
            self.validate().await;
        }
        let problems = self.dag.check_segments().await.unwrap();
        assert!(
            problems.is_empty(),
            "problems after drawdag: {:?}",
            problems
        );
        let master_heads = master_heads
            .iter()
            .map(|s| Vertex::copy_from(s.as_bytes()))
            .collect::<Vec<_>>();
        let need_flush = !master_heads.is_empty();
        if need_flush {
            let heads = VertexListWithOptions::from(master_heads).with_highest_group(Group::MASTER);
            self.dag.flush(&heads).await.unwrap();
        }
        if validate {
            self.validate().await;
        }
        assert_eq!(self.dag.check_segments().await.unwrap(), [] as [String; 0]);
    }

    /// Flush space-separated master heads.
    pub async fn flush(&mut self, master_heads: &str) {
        let heads: Vec<Vertex> = master_heads
            .split_whitespace()
            .map(|v| Vertex::copy_from(v.as_bytes()))
            .collect();
        let heads = VertexListWithOptions::from(heads).with_highest_group(Group::MASTER);
        self.dag.flush(&heads).await.unwrap();
    }

    /// Replace ASCII with Ids in the graph.
    pub fn annotate_ascii(&self, text: &str) -> String {
        self.dag.map.replace(text)
    }

    /// Render the segments.
    pub fn render_segments(&self) -> String {
        format!("{:?}", &self.dag.dag)
    }

    #[cfg(feature = "render")]
    /// Render the graph.
    pub fn render_graph(&self) -> String {
        render_namedag(&self.dag, |v| {
            Some(
                non_blocking_result(self.dag.vertex_id(v.clone()))
                    .unwrap()
                    .to_string(),
            )
        })
        .unwrap()
    }

    /// Use this DAG as the "server", return the "client" Dag that has lazy Vertexes.
    pub async fn client(&self) -> TestDag {
        let mut client = TestDag::new();
        client.set_remote(&self);
        client
    }

    /// Update remote protocol to use the (updated) server graph.
    pub fn set_remote(&mut self, server_dag: &Self) {
        let remote = server_dag.remote_protocol(self.output.clone());
        self.dag.set_remote_protocol(remote);
    }

    /// Alternative syntax of `set_remote`.
    pub fn with_remote(mut self, server_dag: &Self) -> Self {
        self.set_remote(server_dag);
        self
    }

    /// Similar to `client`, but also clone the Dag from the server.
    pub async fn client_cloned_data(&self) -> TestDag {
        let mut client = self.client().await;
        let data = self.dag.export_clone_data().await.unwrap();
        tracing::debug!("clone data: {:?}", &data);
        client.dag.import_clone_data(data).await.unwrap();
        client
    }

    /// Pull from the server Dag using the master fast forward fast path.
    pub async fn pull_ff_master(
        &mut self,
        server: &Self,
        old_master: impl Into<Vertex>,
        new_master: impl Into<Vertex>,
    ) -> Result<()> {
        self.set_remote(server);
        let new_master = new_master.into();
        let missing = server
            .dag
            .only(new_master.clone().into(), old_master.into().into())
            .await?;
        let data = server.dag.export_pull_data(&missing).await?;
        debug!("pull_ff data: {:?}", &data);
        let heads = VertexListWithOptions::from(vec![new_master]).with_highest_group(Group::MASTER);
        self.dag.import_pull_data(data, &heads).await?;
        Ok(())
    }

    /// Strip space-separated vertexes.
    pub async fn strip(&mut self, names: &'static str) {
        let set = Set::from_static_names(names.split(' ').map(|s| s.into()));
        self.dag.strip(&set).await.unwrap();
        let problems = self.dag.check_segments().await.unwrap();
        assert!(problems.is_empty(), "problems after strip: {:?}", problems);
    }

    /// Remote protocol used to resolve Id <-> Vertex remotely using the test dag
    /// as the "server".
    ///
    /// Logs of the remote access will be written to `output`.
    pub fn remote_protocol(
        &self,
        output: Arc<Mutex<Vec<String>>>,
    ) -> Arc<dyn RemoteIdConvertProtocol> {
        let remote = ProtocolMonitor {
            inner: Box::new(self.dag.try_snapshot().unwrap()),
            output,
        };
        Arc::new(remote)
    }

    /// Describe segments at the given level and group as a string.
    pub fn debug_segments(&self, level: Level, group: Group) -> String {
        let lines = crate::namedag::debug_segments_by_level_group(
            &self.dag.dag,
            &self.dag.map,
            level,
            group,
        );
        lines
            .iter()
            .map(|l| format!("\n        {}", l))
            .collect::<Vec<String>>()
            .concat()
    }

    /// Output of remote protocols since the last call.
    pub fn output(&self) -> Vec<String> {
        let mut result = Vec::new();
        let mut output = self.output.lock().unwrap();
        std::mem::swap(&mut result, &mut *output);
        result
    }

    /// Check that a vertex exists locally.
    pub fn contains_vertex_locally(&self, name: impl Into<Vertex>) -> bool {
        non_blocking_result(self.dag.contains_vertex_name_locally(&[name.into()])).unwrap()[0]
    }

    #[cfg(test)]
    /// Dump Dag state as a string.
    pub async fn dump_state(&self) -> String {
        use crate::iddagstore::tests::dump_store_state;
        use crate::Id;
        let iddag = &self.dag.dag;
        let all = iddag.all().unwrap();
        let iddag_state = dump_store_state(&iddag.store, &all);
        let all_str = format!("{:?}", &self.dag.all().await.unwrap());
        let idmap_state: String = {
            let all: Vec<Id> = all.iter_asc().collect();
            let contains = self.dag.contains_vertex_id_locally(&all).await.unwrap();
            let local_ids: Vec<Id> = all
                .into_iter()
                .zip(contains)
                .filter(|(_, c)| *c)
                .map(|(i, _)| i)
                .collect();
            let local_vertexes = self
                .dag
                .vertex_name_batch(&local_ids)
                .await
                .unwrap()
                .into_iter()
                .collect::<Result<Vec<_>>>()
                .unwrap();
            local_ids
                .into_iter()
                .zip(local_vertexes)
                .map(|(i, v)| format!("{:?}->{:?}", i, v))
                .collect::<Vec<_>>()
                .join(" ")
        };

        format!("{}{}\n{}", all_str, iddag_state, idmap_state)
    }

    #[cfg(test)]
    /// Dump Dag segments as ASCII string.
    pub fn dump_segments_ascii(&self) -> String {
        use std::collections::HashSet;

        use crate::Id;
        use crate::IdSet;
        use crate::IdSpan;

        let span_iter = |span: IdSpan| IdSet::from_spans(vec![span]).into_iter().rev();
        let iddag = &self.dag.dag;
        let all_ids = iddag.all_ids_in_groups(&Group::ALL).unwrap();
        let max_level = iddag.max_level().unwrap();
        let mut output = String::new();
        for level in 0..=max_level {
            output = format!("{}\n        Lv{}:", output.trim_end(), level);
            for span in all_ids.iter_span_asc() {
                output += " |";
                let segments = iddag.segments_in_span_ascending(*span, level).unwrap();
                let segment_ids: HashSet<Id> = segments
                    .iter()
                    .flat_map(|s| span_iter(s.span().unwrap()))
                    .collect();
                let segment_highs: HashSet<Id> =
                    segments.iter().map(|s| s.high().unwrap()).collect();
                for id in span_iter(*span) {
                    let id_str = format!("{:?}", id);
                    if segment_ids.contains(&id) {
                        output += &id_str
                    } else {
                        let space = " ".repeat(id_str.len());
                        output += &space;
                    };
                    output.push(
                        if segment_highs.contains(&id)
                            || (segment_ids.contains(&(id + 1)) && !segment_ids.contains(&id))
                        {
                            '|'
                        } else {
                            ' '
                        },
                    );
                }
            }
        }
        output.trim_end().to_string()
    }

    async fn validate(&self) {
        // All vertexes should be accessible, and round-trip through IdMap.
        let mut iter = self.dag.all().await.unwrap().iter().await.unwrap();
        while let Some(v) = iter.next().await {
            let v = v.unwrap();
            let id = self.dag.vertex_id(v.clone()).await.unwrap();
            let v2 = self.dag.vertex_name(id).await.unwrap();
            assert_eq!(v, v2);
        }
    }
}

pub(crate) struct ProtocolMonitor {
    pub(crate) inner: Box<dyn RemoteIdConvertProtocol>,
    pub(crate) output: Arc<Mutex<Vec<String>>>,
}

#[async_trait::async_trait]
impl RemoteIdConvertProtocol for ProtocolMonitor {
    async fn resolve_names_to_relative_paths(
        &self,
        heads: Vec<Vertex>,
        names: Vec<Vertex>,
    ) -> Result<Vec<(protocol::AncestorPath, Vec<Vertex>)>> {
        let msg = format!("resolve names: {:?}, heads: {:?}", &names, &heads);
        self.output.lock().unwrap().push(msg);
        self.inner
            .resolve_names_to_relative_paths(heads, names)
            .await
    }

    async fn resolve_relative_paths_to_names(
        &self,
        paths: Vec<protocol::AncestorPath>,
    ) -> Result<Vec<(protocol::AncestorPath, Vec<Vertex>)>> {
        let msg = format!("resolve paths: {:?}", &paths);
        self.output.lock().unwrap().push(msg);
        self.inner.resolve_relative_paths_to_names(paths).await
    }
}

fn get_heads_and_parents_func_from_ascii(text: &str) -> (Vec<Vertex>, DrawDag) {
    let dag = DrawDag::from(text);
    let heads = dag.heads();
    (heads, dag)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Tests about discontinuous segments
//!
//! Previously, segments in a group are continuous. In other words, all segments
//! in the master group can be represented using a single span `0..=x`.  With
//! discontinuous segments, a group might be represented as a few spans.
//!
//! The discontinuous spans are designed to better support multiple long-lived
//! long branches. For example:
//!
//! ```plain,ignore
//! 1---2---3--...--999---1000     branch1
//!      \
//!       5000--...--5999---6000   branch2
//! ```
//!
//! Note: discontinuous segments is not designed to support massive amount of
//! branches. It introduces O(branch) factor in complexity in many places.

use crate::ops::DagAddHeads;
use crate::ops::DagAlgorithm;
use crate::ops::DagPersistent;
use crate::tests::DrawDag;
use crate::tests::TestDag;
use crate::Group;
use crate::Vertex;
use crate::VertexListWithOptions;
use crate::VertexOptions;

#[tokio::test]
async fn test_simple_3_branches() {
    let mut dag = TestDag::new();
    let draw = DrawDag::from(
        r#"
            A--B--C--D--E--F
                   \
                    G--H--I
                     \
                      J--K--L"#,
    );

    let heads = VertexListWithOptions::from(vec![
        reserved_head("F", 100),
        reserved_head("I", 50),
        reserved_head("L", 0),
    ]);

    dag.dag.add_heads_and_flush(&draw, &heads).await.unwrap();
    assert_eq!(
        format!("{:?}", &dag.dag),
        r#"Max Level: 0
 Level 0
  Group Master:
   Segments: 3
    J+159 : L+161 [G+106]
    G+106 : I+108 [C+2]
    A+0 : F+5 [] Root OnlyHead
  Group Non-Master:
   Segments: 0
"#
    );

    assert_eq!(
        format!("{:?}", dag.dag.ancestors("I".into()).await.unwrap()),
        "<spans [G:I+106:108, A:C+0:2]>"
    );
    assert_eq!(
        format!("{:?}", dag.dag.descendants("B".into()).await.unwrap()),
        "<spans [J:L+159:161, G:I+106:108, B:F+1:5]>"
    );
    assert_eq!(
        format!("{:?}", dag.dag.range("C".into(), "K".into()).await.unwrap()),
        "<spans [J:K+159:160, G+106, C+2]>"
    );
    assert_eq!(
        format!("{:?}", dag.dag.parents("G".into()).await.unwrap()),
        "<spans [C+2]>"
    );
    assert_eq!(
        format!("{:?}", dag.dag.children("C".into()).await.unwrap()),
        "<spans [G+106, D+3]>"
    );

    let all = dag.dag.all().await.unwrap();
    assert_eq!(
        format!("{:?}", dag.dag.children(all.clone()).await.unwrap()),
        "<spans [J:L+159:161, G:I+106:108, B:F+1:5]>"
    );
    assert_eq!(
        format!("{:?}", dag.dag.parents(all.clone()).await.unwrap()),
        "<spans [J:K+159:160, G:H+106:107, A:E+0:4]>"
    );
    assert_eq!(
        format!(
            "{:?}",
            dag.dag.range(all.clone(), all.clone()).await.unwrap()
        ),
        "<spans [J:L+159:161, G:I+106:108, A:F+0:5]>"
    );
}

#[tokio::test]
async fn test_grow_branches() {
    let mut dag = TestDag::new();
    let draw = DrawDag::from(
        r#"
            A0-A1-A2
            B0-B1-B2
            C0-C1-C2"#,
    );
    let heads = VertexListWithOptions::from(vec![
        reserved_head("A2", 4),
        reserved_head("B2", 4),
        reserved_head("C2", 4),
    ]);
    dag.dag.add_heads_and_flush(&draw, &heads).await.unwrap();
    assert_eq!(
        format!("{:?}", dag.dag.all().await.unwrap()),
        "<spans [C0:C2+14:16, B0:B2+7:9, A0:A2+0:2]>"
    );

    // Grow all branches with larger reservation. Larger reservation
    // is changed to pre-existing reservation to avoid fragmentation.
    let draw = DrawDag::from(
        r#" A2-A3-A4
            B2-B3-B4
            C2-C3-C4"#,
    );
    let heads = VertexListWithOptions::from(vec![
        reserved_head("A4", 20),
        reserved_head("B4", 20),
        reserved_head("C4", 20),
    ]);
    dag.dag.add_heads_and_flush(&draw, &heads).await.unwrap();
    assert_eq!(
        format!("{:?}", dag.dag.all().await.unwrap()),
        "<spans [C0:C4+14:18, B0:B4+7:11, A0:A4+0:4]>"
    );

    // Large reservation is respected when the previous reservation
    // gets used up. Note how B5 and C5 respects reservations from
    // A8 (->C4) and B6.
    let draw = DrawDag::from(
        r#" A4-A5-A6-A7-A8
            B4-B5-B6
            C4-C5-C6-C7-C8"#,
    );
    let heads = VertexListWithOptions::from(vec![
        reserved_head("A8", 20),
        reserved_head("B6", 20),
        reserved_head("C8", 20),
    ]);
    dag.dag.add_heads(&draw, &heads).await.unwrap();
    assert_eq!(
        dag.debug_segments(0, Group::MASTER),
        r#"
        C5+61 : C8+64 [C4+18]
        B5+39 : B6+40 [B4+11]
        C0+14 : C4+18 [] Root
        A7+12 : A8+13 [A6+6]
        B0+7 : B4+11 [] Root
        A0+0 : A6+6 [] Root OnlyHead"#
    );
}

#[tokio::test]
async fn test_reservation_on_existing_vertex() {
    let mut dag = TestDag::new();
    let draw = DrawDag::from(
        r#" A0-A1-A2
            B0-B1-B2"#,
    );
    let heads = VertexListWithOptions::from(vec![reserved_head("A2", 4)]);
    dag.dag.add_heads_and_flush(&draw, &heads).await.unwrap();
    assert_eq!(
        format!("{:?}", dag.dag.all().await.unwrap()),
        "<spans [A0:A2+0:2]>"
    );

    // A2: 4 is respected (3..=6 is reserved).
    let heads = VertexListWithOptions::from(vec![reserved_head("B2", 4), reserved_head("A2", 4)]);
    dag.dag.add_heads(&draw, &heads).await.unwrap();
    assert_eq!(
        format!("{:?}", dag.dag.all().await.unwrap()),
        "<spans [B0:B2+7:9, A0:A2+0:2]>"
    );
}

fn reserved_head(s: &'static str, reserve_size: u32) -> (Vertex, VertexOptions) {
    (
        Vertex::from(s),
        VertexOptions {
            reserve_size,
            highest_group: Group::MASTER,
            ..Default::default()
        },
    )
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::TestDag;
use crate::ops::CheckIntegrity;
use crate::ops::DagAlgorithm;
use crate::Group;

#[tokio::test]
async fn test_isomorphic_graph_with_different_segments() {
    let mut dag1 = TestDag::new();
    let mut dag2 = TestDag::new();

    dag1.drawdag(
        r#"
        A--B--C--D--E    K--L--M
         \     \     \
          F--G--H--I--J"#,
        &["J", "M"],
    );
    dag2.drawdag(
        r#"
        A--B--C--D--E    K--L--M
         \     \     \
          F--G--H--I--J"#,
        &["K", "F", "B", "H", "D", "I", "E", "J", "M"],
    );

    assert_eq!(
        dag1.debug_segments(0, Group::MASTER),
        r#"
        K+10 : M+12 [] Root
        J+9 : J+9 [E+4, I+8] OnlyHead
        H+7 : I+8 [C+2, G+6]
        F+5 : G+6 [A+0]
        A+0 : E+4 [] Root OnlyHead"#
    );
    assert_eq!(
        dag2.debug_segments(0, Group::MASTER),
        r#"
        L+11 : M+12 [K+0]
        J+10 : J+10 [E+9, I+8]
        E+9 : E+9 [D+7]
        I+8 : I+8 [H+6]
        D+7 : D+7 [C+5]
        H+6 : H+6 [C+5, G+4]
        C+5 : C+5 [B+3]
        G+4 : G+4 [F+2]
        B+3 : B+3 [A+1]
        A+1 : F+2 [] Root
        K+0 : K+0 [] Root OnlyHead"#
    );

    let heads = dag1.dag.heads(dag1.dag.all().await.unwrap()).await.unwrap();
    assert_eq!(
        dag1.dag
            .check_isomorphic_graph(&dag2.dag, heads.clone())
            .await
            .unwrap(),
        [] as [&str; 0]
    );
}

#[tokio::test]
async fn test_non_isomorphic_graphs() {
    assert_eq!(
        quick_check_graphs("A--B--C--D", "A--D").await,
        ["range A::D with parents []: length mismatch: 4 != 2"]
    );
    assert_eq!(
        quick_check_graphs("A--B--C", "Z--C A--B--C").await,
        ["range A::C with parents []: merge mismatch: 0 != 1"]
    );
    assert_eq!(
        quick_check_graphs("A--B--C", "Z--B A--B--C").await,
        ["range A::C with parents []: merge mismatch: 0 != 1"]
    );
    assert_eq!(
        quick_check_graphs("A--B--C", "Z--A--B--C").await,
        ["range A::C with parents []: parents mismatch: [] != [Z]"]
    );
    assert_eq!(
        quick_check_graphs("A B C", "A B").await,
        ["range C::C with parents []: cannot resolve range on the other graph: VertexNotFound(C)"]
    );
}

async fn quick_check_graphs(ascii1: &str, ascii2: &str) -> Vec<String> {
    let dag1 = TestDag::draw(ascii1);
    let dag2 = TestDag::draw(ascii2);
    let heads = dag1.dag.heads(dag1.dag.all().await.unwrap()).await.unwrap();
    dag1.dag
        .check_isomorphic_graph(&dag2.dag, heads)
        .await
        .unwrap()
}