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

//! Commit Graph
//!
//! The graph of all commits in the repository.

use std::collections::HashSet;
use std::sync::Arc;

use anyhow::Result;
use borrowed::borrowed;
use commit_graph_types::edges::ChangesetNode;
use commit_graph_types::edges::ChangesetParents;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use futures::stream;
use futures::stream::BoxStream;
use futures::Future;
use futures::StreamExt;
use futures::TryStreamExt;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::Generation;

pub use crate::ancestors_stream::AncestorsStreamBuilder;

mod ancestors_stream;
mod compat;
mod core;
mod frontier;
mod segments;

/// Commit Graph.
///
/// This contains the graph of all commits known to Mononoke for a particular
/// repository.  It provides methods for traversing the commit graph and
/// finding out graph-related information for the changesets contained
/// therein.
#[derive(Clone)]
#[facet::facet]
pub struct CommitGraph {
    /// The storage back-end where the commits are actually stored.
    storage: Arc<dyn CommitGraphStorage>,
}

impl CommitGraph {
    pub fn new(storage: Arc<dyn CommitGraphStorage>) -> CommitGraph {
        CommitGraph { storage }
    }

    /// Add a new changeset to the commit graph.
    ///
    /// Returns true if a new changeset was inserted, or false if the
    /// changeset already existed.
    pub async fn add(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
        parents: ChangesetParents,
    ) -> Result<bool> {
        let parent_edges = self
            .storage
            .fetch_many_edges_required(ctx, &parents, Prefetch::None)
            .await?;

        self.storage
            .add(
                ctx,
                self.build_edges(ctx, cs_id, parents, &parent_edges).await?,
            )
            .await
    }

    /// Find all changeset ids with a given prefix.
    pub async fn find_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: ChangesetIdPrefix,
        limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        self.storage.find_by_prefix(ctx, cs_prefix, limit).await
    }

    /// Returns true if the changeset exists.
    pub async fn exists(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<bool> {
        let edges = self.storage.fetch_edges(ctx, cs_id).await?;
        Ok(edges.is_some())
    }

    /// Returns the parents of a single changeset.
    pub async fn changeset_parents(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetParents>> {
        let edges = self.storage.fetch_edges(ctx, cs_id).await?;
        Ok(edges.map(|edges| {
            edges
                .parents
                .into_iter()
                .map(|parent| parent.cs_id)
                .collect()
        }))
    }

    /// Returns the parents of a single changeset that must exist.
    pub async fn changeset_parents_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetParents> {
        let edges = self.storage.fetch_edges_required(ctx, cs_id).await?;
        Ok(edges
            .parents
            .into_iter()
            .map(|parent| parent.cs_id)
            .collect())
    }

    /// Returns the generation number of a single changeset.
    pub async fn changeset_generation(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<Generation>> {
        let edges = self.storage.fetch_edges(ctx, cs_id).await?;
        Ok(edges.map(|edges| edges.node.generation))
    }

    /// Returns the generation number of a single changeset that must exist.
    pub async fn changeset_generation_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Generation> {
        let edges = self.storage.fetch_edges_required(ctx, cs_id).await?;
        Ok(edges.node.generation)
    }

    /// Returns a frontier for the ancestors of heads
    /// that satisfy a given property.
    ///
    /// Note: The property needs to be monotonic i.e. if the
    /// property holds for one changeset then it has to hold
    /// for all its parents.
    pub async fn ancestors_frontier_with<MonotonicProperty, Out>(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        monotonic_property: MonotonicProperty,
    ) -> Result<Vec<ChangesetId>>
    where
        MonotonicProperty: Fn(ChangesetId) -> Out + Send + Sync + 'static,
        Out: Future<Output = Result<bool>>,
    {
        let mut ancestors_frontier = vec![];
        let mut frontier = self.frontier(ctx, heads).await?;

        let monotonic_property = move |node: ChangesetNode| {
            borrowed!(monotonic_property);
            monotonic_property(node.cs_id)
        };

        while let Some(ancestors_frontier_extension) = self
            .lower_frontier_step(ctx, &mut frontier, &monotonic_property, Prefetch::None)
            .await?
        {
            ancestors_frontier.extend(ancestors_frontier_extension);
        }

        Ok(ancestors_frontier.into_iter().collect())
    }

    /// Returns true if the ancestor changeset is an ancestor of the descendant
    /// changeset.
    ///
    /// Ancestry is inclusive: a commit is its own ancestor.
    pub async fn is_ancestor(
        &self,
        ctx: &CoreContext,
        ancestor: ChangesetId,
        descendant: ChangesetId,
    ) -> Result<bool> {
        let (mut frontier, target_gen) = futures::try_join!(
            self.single_frontier(ctx, descendant),
            self.changeset_generation_required(ctx, ancestor)
        )?;
        debug_assert!(!frontier.is_empty(), "frontier should contain descendant");
        self.lower_frontier(ctx, &mut frontier, target_gen).await?;
        Ok(frontier.highest_generation_contains(ancestor, target_gen))
    }

    /// Returns a stream of all ancestors of any changeset in heads,
    /// excluding any ancestor of any changeset in common, in reverse
    /// topological order.
    pub async fn ancestors_difference_stream(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        common: Vec<ChangesetId>,
    ) -> Result<BoxStream<'static, Result<ChangesetId>>> {
        AncestorsStreamBuilder::new(Arc::new(self.clone()), ctx.clone(), heads)
            .exclude_ancestors_of(common)
            .build()
            .await
    }

    /// Returns all ancestors of any changeset in heads, excluding any
    /// ancestor of any changeset in common, in reverse topological order.
    pub async fn ancestors_difference(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        common: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetId>> {
        self.ancestors_difference_stream(ctx, heads, common)
            .await?
            .try_collect()
            .await
    }

    /// Returns a stream of all changesets that are both descendants of
    /// start_id and ancestors of end_id, in topological order.
    pub async fn range_stream(
        &self,
        ctx: &CoreContext,
        start_id: ChangesetId,
        end_id: ChangesetId,
    ) -> Result<BoxStream<'static, ChangesetId>> {
        let range: Vec<_> =
            AncestorsStreamBuilder::new(Arc::new(self.clone()), ctx.clone(), vec![end_id])
                .descendants_of(start_id)
                .build()
                .await?
                .try_collect()
                .await?;

        Ok(stream::iter(range.into_iter().rev()).boxed())
    }

    /// Returns all of the highest generation changesets that
    /// are ancestors of both u and v, sorted by changeset id.
    pub async fn common_base(
        &self,
        ctx: &CoreContext,
        u: ChangesetId,
        v: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        let (mut u_frontier, mut v_frontier) =
            futures::try_join!(self.single_frontier(ctx, u), self.single_frontier(ctx, v))?;

        loop {
            let u_gen = match u_frontier.last_key_value() {
                Some((gen, _)) => *gen,
                // if u_frontier is empty then there are no common ancestors.
                None => return Ok(vec![]),
            };

            // lower v_frontier to the highest generation of u_frontier
            self.lower_frontier(ctx, &mut v_frontier, u_gen).await?;

            // Check if the highest generation of u_frontier intersects with v_frontier
            // and return the intersection if so.
            let mut intersection = u_frontier.highest_generation_intersection(&v_frontier);
            if !intersection.is_empty() {
                intersection.sort();
                return Ok(intersection);
            }

            let u_highest_generation_edges = match u_frontier
                .last_key_value()
                .and_then(|(_, cs_ids)| cs_ids.iter().next())
            {
                Some(cs_id) => self.storage.fetch_edges_required(ctx, *cs_id).await?,
                None => return Ok(vec![]),
            };

            // Try to lower u_frontier to the generation of one of its
            // highest generation changesets' skip tree skew ancestor.
            // This is optimized for the case where u_frontier has only
            // one changeset, but is correct in all cases.
            if let Some(ancestor) = u_highest_generation_edges.skip_tree_skew_ancestor {
                let mut lowered_u_frontier = u_frontier.clone();
                let mut lowered_v_frontier = v_frontier.clone();

                self.lower_frontier(ctx, &mut lowered_u_frontier, ancestor.generation)
                    .await?;
                self.lower_frontier(ctx, &mut lowered_v_frontier, ancestor.generation)
                    .await?;

                // If the two lowered frontier are disjoint then it's safe to lower,
                // otherwise there might be a higher generation common ancestor.
                if lowered_u_frontier.is_disjoint(&lowered_v_frontier) {
                    u_frontier = lowered_u_frontier;
                    v_frontier = lowered_v_frontier;

                    continue;
                }
            }

            // If we could lower u_frontier using the skip tree skew ancestor
            // lower only the highest generation instead.
            self.lower_frontier_highest_generation(ctx, &mut u_frontier)
                .await?;
        }
    }

    /// Slices ancestors of heads into a sequence of slices for processing.
    ///
    /// Each slice contains a frontier of changesets within a generation range, returning
    /// (slice_start, slice_frontier) corresponds to the frontier that has generations numbers
    /// within [slice_start..(slice_start + slice_size)].
    ///
    /// Useful for any type of processing that needs to happen on ancestors of changesets first.
    /// By processing slices one by one we avoid traversing the entire history all at once.
    ///
    /// The returned slices consist only of frontiers which haven't been processed yet
    /// (determined by the provided needs_processing function). Slicing stops once we
    /// reach a frontier with all its changesets processed.
    pub async fn slice_ancestors<NeedsProcessing, Out>(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        needs_processing: NeedsProcessing,
        slice_size: u64,
    ) -> Result<Vec<(u64, Vec<ChangesetId>)>>
    where
        NeedsProcessing: Fn(Vec<ChangesetId>) -> Out,
        Out: Future<Output = Result<HashSet<ChangesetId>>>,
    {
        let mut frontier = self.frontier(ctx, heads).await?;

        let max_gen = match frontier.last_key_value() {
            Some((gen, _)) => gen,
            None => return Ok(vec![]),
        };

        // The start of the slice is largest number in the sequence
        // 1, slice_size + 1, 2 * slice_size + 1 ...
        let mut slice_start = ((max_gen.value() - 1) / slice_size) * slice_size + 1;

        let mut slices = vec![];

        // Loop over slices in decreasing order of start generation.
        loop {
            let needed_cs_ids = needs_processing(frontier.changesets()).await?;
            frontier = frontier
                .into_flat_iter()
                .filter(|(cs_id, _)| needed_cs_ids.contains(cs_id))
                .collect();

            if frontier.is_empty() {
                break;
            }

            // Only push changesets that are in this slice's range.
            // Any remaining changesets will be pushed in the next iterations.
            slices.push((
                slice_start,
                frontier.changesets_in_range(
                    Generation::new(slice_start)..Generation::new(slice_start + slice_size),
                ),
            ));

            if slice_start > 1 {
                // Lower the frontier to the end of the next slice (current slice_start - 1).
                self.lower_frontier(ctx, &mut frontier, Generation::new(slice_start - 1))
                    .await?;
                slice_start -= slice_size;
            } else {
                break;
            }
        }

        Ok(slices.into_iter().rev().collect())
    }

    /// Returns the children of a single changeset.
    pub async fn changeset_children(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        self.storage.fetch_children(ctx, cs_id).await
    }
}


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

use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use commit_graph_types::segments::ChangesetSegment;
use commit_graph_types::segments::ChangesetSegmentFrontier;
use commit_graph_types::segments::ChangesetSegmentLocation;
use commit_graph_types::segments::ChangesetSegmentParent;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use futures_stats::TimedTryFutureExt;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use slog::debug;
use smallvec::SmallVec;

use crate::CommitGraph;

/// A set that stores changeset ids and keeps track of all changesets
/// reachable from them by following skew binary ancestor edges.
#[derive(Default, Debug)]
struct SkewAncestorsSet {
    changesets: BTreeMap<Generation, HashSet<ChangesetId>>,
    skew_ancestors: HashMap<ChangesetId, HashSet<ChangesetId>>,
    skew_ancestors_counts: HashMap<ChangesetId, usize>,
}

impl SkewAncestorsSet {
    /// Adds a changeset to the set.
    pub async fn add(
        &mut self,
        ctx: &CoreContext,
        storage: &Arc<dyn CommitGraphStorage>,
        cs_id: ChangesetId,
        base_generation: Generation,
    ) -> Result<()> {
        let mut edges = storage.fetch_edges_required(ctx, cs_id).await?;

        if self
            .changesets
            .entry(edges.node.generation)
            .or_default()
            .insert(cs_id)
        {
            // if this changeset wasn't already present in the set, increment the count
            // of all changesets reachable by following skew binary ancestors edges.
            loop {
                self.skew_ancestors
                    .entry(cs_id)
                    .or_default()
                    .insert(edges.node.cs_id);
                *self
                    .skew_ancestors_counts
                    .entry(edges.node.cs_id)
                    .or_default() += 1;

                match edges.skip_tree_skew_ancestor {
                    Some(skip_tree_skew_ancestor)
                        if skip_tree_skew_ancestor.generation >= base_generation =>
                    {
                        edges = storage
                            .fetch_edges_required(ctx, skip_tree_skew_ancestor.cs_id)
                            .await?;
                    }
                    _ => break,
                }
            }
        }

        Ok(())
    }

    /// Returns the highest generation of a changeset in the set,
    /// or None if the set is empty.
    pub fn highest_generation(&self) -> Option<Generation> {
        self.changesets
            .last_key_value()
            .map(|(generation, _)| *generation)
    }

    /// Returns whether the given changeset is reachable from any changeset
    /// in the set by following skew binary ancestor edges.
    pub fn contains_ancestor(&self, cs_id: ChangesetId) -> bool {
        self.skew_ancestors_counts
            .get(&cs_id)
            .map_or(false, |count| *count > 0)
    }

    /// Removes and returns the highest generation number from the set and all changesets
    /// having that generation number.
    pub fn pop_last(&mut self) -> Option<(Generation, HashSet<ChangesetId>)> {
        match self.changesets.pop_last() {
            Some((generation, cs_ids)) => {
                for cs_id in cs_ids.iter() {
                    if let Some(skew_ancestors) = self.skew_ancestors.get(cs_id) {
                        for skew_ancestor in skew_ancestors {
                            if let Some(skew_ancestor_count) =
                                self.skew_ancestors_counts.get_mut(skew_ancestor)
                            {
                                *skew_ancestor_count -= 1;
                            }
                        }
                    }
                }
                Some((generation, cs_ids))
            }
            None => None,
        }
    }
}

impl CommitGraph {
    /// Returns a frontier of segments from each of the given changesets to their
    /// corresponding merge ancestor.
    async fn segment_frontier(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<ChangesetSegmentFrontier> {
        let mut frontier: ChangesetSegmentFrontier = Default::default();

        let all_edges = self
            .storage
            .fetch_many_edges_required(ctx, &cs_ids, Prefetch::None)
            .await?;

        for (cs_id, edges) in all_edges {
            let base = edges.merge_ancestor.unwrap_or(edges.node);
            frontier
                .segments
                .entry(base.generation)
                .or_default()
                .entry(base.cs_id)
                .or_default()
                .insert(cs_id);
        }

        Ok(frontier)
    }

    /// Lower a segment frontier to the specified target generation.
    async fn lower_segment_frontier(
        &self,
        ctx: &CoreContext,
        segment_frontier: &mut ChangesetSegmentFrontier,
        target_generation: Generation,
    ) -> Result<()> {
        loop {
            match segment_frontier.segments.last_key_value() {
                None => return Ok(()),
                Some((generation, _)) if *generation <= target_generation => return Ok(()),
                _ => {}
            }

            if let Some((_generation, segments)) = segment_frontier.segments.pop_last() {
                let segment_bases: Vec<_> = segments.into_keys().collect();
                let all_edges = self
                    .storage
                    .fetch_many_edges_required(ctx, &segment_bases, Prefetch::None)
                    .await?;

                let parents: Vec<_> = all_edges
                    .into_iter()
                    .flat_map(|(_cs_id, edges)| edges.parents)
                    .map(|node| node.cs_id)
                    .collect();

                let parent_edges = self
                    .storage
                    .fetch_many_edges_required(ctx, &parents, Prefetch::None)
                    .await?;

                for (cs_id, edges) in parent_edges {
                    let base = edges.merge_ancestor.unwrap_or(edges.node);
                    segment_frontier
                        .segments
                        .entry(base.generation)
                        .or_default()
                        .entry(base.cs_id)
                        .or_default()
                        .insert(cs_id);
                }
            }
        }
    }

    /// Given a list of changesets heads and another list of changesets common, all having
    /// their merge_ancestor pointing to base, returns a list of segments representing all
    /// ancestors of heads, excluding all ancestors of common.
    async fn disjoint_segments(
        &self,
        ctx: &CoreContext,
        base: ChangesetId,
        heads: Vec<ChangesetId>,
        common: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetSegment>> {
        let base_edges = self.storage.fetch_edges_required(ctx, base).await?;

        let mut heads_skew_ancestors_set: SkewAncestorsSet = Default::default();
        let mut common_skew_ancestors_set: SkewAncestorsSet = Default::default();

        for cs_id in heads {
            heads_skew_ancestors_set
                .add(ctx, &self.storage, cs_id, base_edges.node.generation)
                .await?;
        }
        for cs_id in common {
            common_skew_ancestors_set
                .add(ctx, &self.storage, cs_id, base_edges.node.generation)
                .await?;
        }

        #[derive(Copy, Clone, Debug)]
        enum Origin {
            Head {
                cs_id: ChangesetId,
                generation: Generation,
            },
            Common,
        }

        let mut frontier: BTreeMap<Generation, HashMap<ChangesetId, Origin>> = Default::default();
        let mut segments = vec![];

        loop {
            // Add the highest generation head changesets to the frontier, if there's
            // no changeset in the frontier or in common with a higher generation number.
            if let Some(heads_generation) = heads_skew_ancestors_set.highest_generation() {
                if frontier
                    .last_key_value()
                    .map_or(true, |(frontier_generation, _)| {
                        heads_generation >= *frontier_generation
                    })
                    && common_skew_ancestors_set
                        .highest_generation()
                        .map_or(true, |common_generation| {
                            heads_generation >= common_generation
                        })
                {
                    if let Some((generation, heads)) = heads_skew_ancestors_set.pop_last() {
                        for cs_id in heads {
                            frontier
                                .entry(generation)
                                .or_default()
                                .entry(cs_id)
                                .or_insert(Origin::Head { cs_id, generation });
                        }
                    }
                }
            }

            // Add the highest generation common changesets to the frontier, if there's
            // no changeset in the frontier with a higher generation number.
            if let Some(common_generation) = common_skew_ancestors_set.highest_generation() {
                if frontier
                    .last_key_value()
                    .map_or(true, |(frontier_generation, _)| {
                        common_generation >= *frontier_generation
                    })
                {
                    if let Some((generation, common)) = common_skew_ancestors_set.pop_last() {
                        for cs_id in common {
                            frontier
                                .entry(generation)
                                .or_default()
                                .insert(cs_id, Origin::Common);
                        }
                    }
                }
            }

            match frontier.pop_last() {
                Some((_generation, last_changesets)) => {
                    let cs_ids: Vec<_> = last_changesets.keys().copied().collect();
                    let all_edges = self
                        .storage
                        .fetch_many_edges_required(ctx, &cs_ids, Prefetch::None)
                        .await?;

                    // Try to lower the highest generation changesets in the frontier to their
                    // skew binary ancestors, and store any that can't be lowered in either
                    // blocked_heads or blocked_common.

                    let mut immediate_skew_ancestors_count: HashMap<ChangesetId, usize> =
                        Default::default();

                    for cs_id in last_changesets.keys() {
                        let edges = all_edges.get(cs_id).ok_or_else(|| {
                            anyhow!("Missing changeset edges in commit graph for {}", cs_id)
                        })?;

                        if let Some(skew_ancestor) = edges.skip_tree_skew_ancestor {
                            if skew_ancestor.generation >= base_edges.node.generation
                                && !heads_skew_ancestors_set.contains_ancestor(skew_ancestor.cs_id)
                                && !common_skew_ancestors_set.contains_ancestor(skew_ancestor.cs_id)
                            {
                                *immediate_skew_ancestors_count
                                    .entry(skew_ancestor.cs_id)
                                    .or_default() += 1;
                            }
                        }
                    }

                    let mut blocked_heads = vec![];
                    let mut blocked_common = vec![];

                    for (cs_id, origin) in last_changesets.iter() {
                        let edges = all_edges.get(cs_id).ok_or_else(|| {
                            anyhow!("Missing changeset edges in commit graph for {}", cs_id)
                        })?;

                        if let Some(skew_ancestor) = edges.skip_tree_skew_ancestor {
                            if skew_ancestor.generation >= base_edges.node.generation
                                && !heads_skew_ancestors_set.contains_ancestor(skew_ancestor.cs_id)
                                && !common_skew_ancestors_set.contains_ancestor(skew_ancestor.cs_id)
                                && immediate_skew_ancestors_count.get(&skew_ancestor.cs_id)
                                    == Some(&1)
                            {
                                frontier
                                    .entry(skew_ancestor.generation)
                                    .or_default()
                                    .insert(skew_ancestor.cs_id, *origin);
                                continue;
                            }
                        }

                        match origin {
                            Origin::Head {
                                cs_id: origin_cs_id,
                                generation: origin_generation,
                            } => blocked_heads.push((
                                *cs_id,
                                *origin_cs_id,
                                *origin_generation,
                                edges,
                            )),
                            Origin::Common => blocked_common.push((*cs_id, edges)),
                        }
                    }

                    // Lower all blocked common changesets to their immediate parent,
                    // if they are not already at the generation of the base.

                    for (_cs_id, edges) in blocked_common {
                        for parent in edges.parents.iter() {
                            if parent.generation >= base_edges.node.generation {
                                frontier
                                    .entry(parent.generation)
                                    .or_default()
                                    .insert(parent.cs_id, Origin::Common);
                            }
                        }
                    }

                    // Try to lower all blocked head changesets to their immediate parent,
                    // producing a segment for any that can't be lowered due to being blocked
                    // by another changeset in the frontier or a common changeset.

                    for (cs_id, origin_cs_id, origin_generation, edges) in blocked_heads {
                        if edges.node.generation == base_edges.node.generation {
                            segments.push(ChangesetSegment {
                                head: origin_cs_id,
                                base: cs_id,
                                length: origin_generation.value() - edges.node.generation.value()
                                    + 1,
                                parents: edges
                                    .parents
                                    .iter()
                                    .map(|parent| ChangesetSegmentParent {
                                        cs_id: parent.cs_id,
                                        location: Some(ChangesetSegmentLocation {
                                            head: parent.cs_id,
                                            distance: 0,
                                        }),
                                    })
                                    .collect(),
                            });
                            continue;
                        }

                        for parent in edges.parents.iter() {
                            match (
                                frontier
                                    .get(&parent.generation)
                                    .and_then(|segments| segments.get_key_value(&parent.cs_id)),
                                common_skew_ancestors_set.contains_ancestor(parent.cs_id),
                            ) {
                                // Parent is contained in another segment that originates from one of the heads.
                                // Stop extending segment.
                                (
                                    Some((
                                        _,
                                        Origin::Head {
                                            cs_id: parent_segment_origin,
                                            generation: parent_segment_origin_generation,
                                        },
                                    )),
                                    _,
                                ) => segments.push(ChangesetSegment {
                                    head: origin_cs_id,
                                    base: cs_id,
                                    length: origin_generation.value()
                                        - edges.node.generation.value()
                                        + 1,
                                    parents: SmallVec::from(vec![ChangesetSegmentParent {
                                        cs_id: parent.cs_id,
                                        location: Some(ChangesetSegmentLocation {
                                            head: *parent_segment_origin,
                                            distance: parent_segment_origin_generation.value()
                                                - parent.generation.value(),
                                        }),
                                    }]),
                                }),
                                // Parent is an ancestor of common.
                                // Stop extending segment.
                                (Some(_), _) | (_, true) => segments.push(ChangesetSegment {
                                    head: origin_cs_id,
                                    base: cs_id,
                                    length: origin_generation.value()
                                        - edges.node.generation.value()
                                        + 1,
                                    parents: SmallVec::from(vec![ChangesetSegmentParent {
                                        cs_id: parent.cs_id,
                                        location: None,
                                    }]),
                                }),
                                // Parent isn't contained in any other segment, and isn't an ancestor of common.
                                // Continue extending segment.
                                (None, false) => {
                                    frontier.entry(parent.generation).or_default().insert(
                                        parent.cs_id,
                                        Origin::Head {
                                            cs_id: origin_cs_id,
                                            generation: origin_generation,
                                        },
                                    );
                                }
                            }
                        }
                    }
                }
                None => break,
            }
        }

        Ok(segments)
    }

    /// Returns a list of segments representing all ancestors of heads, excluding
    /// all ancestors of common.
    pub async fn ancestors_difference_segments(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        common: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetSegment>> {
        let (mut heads_segment_frontier, mut common_segment_frontier) = futures::try_join!(
            self.segment_frontier(ctx, heads),
            self.segment_frontier(ctx, common)
        )?;

        let mut difference_segments_futures = vec![];

        while let Some((generation, segments)) = heads_segment_frontier.segments.pop_last() {
            self.lower_segment_frontier(ctx, &mut common_segment_frontier, generation)
                .await?;

            let mut bases_not_reachable_from_common = vec![];

            // Go through all the segment bases and calculate the disjoint segments rooted
            // at each base, and for all bases not reachable from common, continue traversing
            // the merge graph.

            for (base, heads) in segments {
                let common = match common_segment_frontier
                    .segments
                    .get(&generation)
                    .and_then(|segments| segments.get(&base))
                {
                    Some(common_segments) => common_segments.iter().copied().collect(),
                    None => {
                        bases_not_reachable_from_common.push(base);
                        vec![]
                    }
                };
                difference_segments_futures.push(self.disjoint_segments(
                    ctx,
                    base,
                    heads.into_iter().collect(),
                    common,
                ));
            }

            let all_edges = self
                .storage
                .fetch_many_edges_required(ctx, &bases_not_reachable_from_common, Prefetch::None)
                .await?;

            let parents: Vec<_> = all_edges
                .into_iter()
                .flat_map(|(_cs_id, edges)| edges.parents)
                .map(|node| node.cs_id)
                .collect();

            let parent_edges = self
                .storage
                .fetch_many_edges_required(ctx, &parents, Prefetch::None)
                .await?;

            for (cs_id, edges) in parent_edges {
                let base = edges.merge_ancestor.unwrap_or(edges.node);
                heads_segment_frontier
                    .segments
                    .entry(base.generation)
                    .or_default()
                    .entry(base.cs_id)
                    .or_default()
                    .insert(cs_id);
            }
        }

        stream::iter(difference_segments_futures)
            .buffered(100)
            .map_ok(|segments| stream::iter(segments).map(Ok))
            .try_flatten()
            .try_collect()
            .await
    }

    /// Returns all changesets in a segment in reverse topological order, verifying
    /// that there are no merge changesets in the segment except potentially base,
    /// and that base is an ancestor of head.
    async fn segment_changesets(
        &self,
        ctx: &CoreContext,
        head: ChangesetId,
        base: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        let mut segment_cs_ids = vec![];
        let mut current_cs_id = head;

        loop {
            segment_cs_ids.push(current_cs_id);

            if current_cs_id == base {
                break;
            }

            let mut parents = self
                .changeset_parents_required(ctx, current_cs_id)
                .await?
                .into_iter();

            match (parents.next(), parents.next()) {
                (_, Some(_)) => {
                    return Err(anyhow!(
                        "Found merge changeset {} before segment base",
                        current_cs_id
                    ));
                }
                (None, _) => {
                    return Err(anyhow!(
                        "Segment base {} is not reachable from head {}",
                        base,
                        head
                    ));
                }
                (Some(parent), None) => current_cs_id = parent,
            }
        }

        Ok(segment_cs_ids)
    }

    /// Same as ancestors_difference_segments, but also verifies that:
    /// - The union of all segments matches exactly the returned changesets from ancestors_difference
    /// - All segments are disjoints, no two segments contain the same changeset.
    /// - The parents of each segment are correct.
    /// - No segment contains a merge changeset except potentially at its base.
    /// - Segments are returned in reverse topological order, each parent of each segment either
    /// belong to a subsequent segment or is an ancestor of common.
    pub async fn verified_ancestors_difference_segments(
        &self,
        ctx: &CoreContext,
        heads: Vec<ChangesetId>,
        common: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetSegment>> {
        let (
            (ancestors_difference_stats, difference_cs_ids),
            (ancestors_difference_segments_stats, difference_segments),
        ) = futures::try_join!(
            self.ancestors_difference(ctx, heads.clone(), common.clone())
                .try_timed(),
            self.ancestors_difference_segments(ctx, heads.clone(), common.clone())
                .try_timed(),
        )?;

        debug!(
            ctx.logger(),
            "ancestors_difference stats {:?}, ancestors_difference_segments stats {:?}",
            ancestors_difference_stats,
            ancestors_difference_segments_stats
        );

        let difference_cs_ids: HashSet<_> = difference_cs_ids.into_iter().collect();

        let mut union_segments_cs_ids: HashMap<_, _> = Default::default();

        for (segment_num, segment) in difference_segments.iter().rev().enumerate() {
            let parents = self.changeset_parents_required(ctx, segment.base).await?;
            let segment_parents: SmallVec<[ChangesetId; 1]> =
                segment.parents.iter().map(|parent| parent.cs_id).collect();

            if segment_parents != parents {
                return Err(anyhow!(
                    "Incorrect segment parents, expected {:?} found {:?} for segment base {}",
                    segment_parents,
                    parents,
                    segment.base
                ));
            }

            for parent in segment.parents.iter() {
                if difference_cs_ids.contains(&parent.cs_id)
                    && !union_segments_cs_ids.contains_key(&parent.cs_id)
                {
                    return Err(anyhow!(
                        "Segments are not in reverse topological order, segment parent {} not found in any subsequent segment and isn't an ancestor of common",
                        parent,
                    ));
                }

                match (
                    parent.location,
                    union_segments_cs_ids.contains_key(&parent.cs_id),
                ) {
                    // If a location is provided, verify that it resolves to the changeset id.
                    (Some(location), _) => {
                        let location_head_depth = self
                            .storage
                            .fetch_edges_required(ctx, location.head)
                            .await?
                            .node
                            .skip_tree_depth;
                        let location_level = match location_head_depth.cmp(&location.distance) {
                            Ordering::Less => {
                                return Err(anyhow!(
                                    "Invalid location {}, location head depth is less than location distance",
                                    location,
                                ));
                            }
                            Ordering::Greater | Ordering::Equal => {
                                location_head_depth - location.distance
                            }
                        };
                        let resolved_location = self
                            .skip_tree_level_ancestor(ctx, location.head, location_level)
                            .await?
                            .ok_or_else(|| anyhow!("While resolving location {}", location))?;

                        if resolved_location.cs_id != parent.cs_id {
                            return Err(anyhow!(
                                "Incorrect location for parent of {}, location {} resolves to {}, expected {}",
                                segment.base,
                                location,
                                resolved_location.cs_id,
                                parent.cs_id,
                            ));
                        }
                    }
                    // If the parent belongs to another segment, a location must be provided.
                    (None, true) => {
                        return Err(anyhow!(
                            "Segment parent {} location is None, but it's contained in a subsequent segment",
                            parent.cs_id,
                        ));
                    }
                    _ => {}
                }
            }

            let segment_cs_ids = self
                .segment_changesets(ctx, segment.head, segment.base)
                .await?;

            for cs_id in segment_cs_ids {
                if !difference_cs_ids.contains(&cs_id) {
                    return Err(anyhow!(
                        "Changeset {} in segment {:?} doesn't belong to ancestors difference",
                        cs_id,
                        segment,
                    ));
                }
                if let Some(other_segment) = union_segments_cs_ids.insert(cs_id, segment) {
                    return Err(anyhow!(
                        "Changeset {} found in two segments: {:?}, {:?}",
                        cs_id,
                        segment,
                        other_segment,
                    ));
                }
            }

            if (segment_num + 1) % 1000 == 0 {
                debug!(
                    ctx.logger(),
                    "finished verifying {} segments",
                    segment_num + 1
                );
            }
        }

        debug!(ctx.logger(), "finished verifying all segments");

        if let Some(cs_id) = difference_cs_ids
            .difference(&union_segments_cs_ids.into_keys().collect::<HashSet<_>>())
            .next()
        {
            return Err(anyhow!(
                "Changeset {} found in ancestors difference but is not contained in any segment",
                cs_id,
            ));
        }

        Ok(difference_segments)
    }
}


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

use std::sync::Arc;

use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use changeset_fetcher::ArcChangesetFetcher;
use changeset_fetcher::SimpleChangesetFetcher;
use changesets::ArcChangesets;
use changesets::ChangesetEntry;
use changesets::ChangesetInsert;
use changesets::Changesets;
use changesets::SortOrder;
use commit_graph::ArcCommitGraph;
use commit_graph_types::edges::ChangesetParents;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::stream::BoxStream;
use futures_stats::TimedTryFutureExt;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::Generation;
use mononoke_types::RepositoryId;
use scuba_ext::MononokeScubaSampleBuilder;
use smallvec::SmallVec;
use vec1::vec1;
use vec1::Vec1;

pub struct ChangesetsCommitGraphCompat {
    changesets: ArcChangesets,
    changeset_fetcher: ArcChangesetFetcher,
    commit_graph: ArcCommitGraph,
    repo_name: String,
    scuba: MononokeScubaSampleBuilder,
}

impl ChangesetsCommitGraphCompat {
    pub fn new(
        fb: FacebookInit,
        changesets: ArcChangesets,
        commit_graph: ArcCommitGraph,
        repo_name: String,
        scuba_table: Option<&str>,
    ) -> Result<Self> {
        let changeset_fetcher = Arc::new(SimpleChangesetFetcher::new(
            changesets.clone(),
            changesets.repo_id(),
        ));

        let scuba = match scuba_table {
            Some(scuba_table) => MononokeScubaSampleBuilder::new(fb, scuba_table).with_context(
                || "Couldn't create scuba sample builder for table mononoke_commit_graph",
            )?,
            None => MononokeScubaSampleBuilder::with_discard(),
        };

        Ok(Self {
            changesets,
            changeset_fetcher,
            commit_graph,
            repo_name,
            scuba,
        })
    }

    async fn maybe_write_to_new_commit_graph(
        &self,
        ctx: &CoreContext,
        css: Vec1<(ChangesetId, ChangesetParents)>,
    ) -> Result<()> {
        let mut scuba = self.scuba.clone();

        scuba.add_common_server_data();
        // Only the last id, which is good enough for logging.
        scuba.add("changeset_id", css.last().0.to_string());
        scuba.add("changeset_count", css.len());
        scuba.add("repo_name", self.repo_name.as_str());

        // We use add_recursive because some parents might be missing
        // from the new commit graph.
        match self
            .commit_graph
            .add_recursive(ctx, self.changeset_fetcher.clone(), css)
            .try_timed()
            .await
        {
            Err(err) => {
                scuba.add("error", err.to_string());
                scuba.log_with_msg("Insertion failed", None);

                Err(err).with_context(
                    || "during commit_graph_compat::Changesets::maybe_write_to_new_commit_graph",
                )?
            }
            Ok((stats, added_to_commit_graph)) => {
                scuba.add("time_s", stats.completion_time.as_secs_f64());
                scuba.add("num_added", added_to_commit_graph);

                if added_to_commit_graph > 0 {
                    scuba.log_with_msg("Insertion succeeded", None);
                } else {
                    scuba.log_with_msg("Changesets already stored", None);
                }

                Ok(())
            }
        }
    }
}

#[async_trait]
impl Changesets for ChangesetsCommitGraphCompat {
    fn repo_id(&self) -> RepositoryId {
        self.changesets.repo_id()
    }

    async fn add(&self, ctx: &CoreContext, cs: ChangesetInsert) -> Result<bool> {
        let (added_to_changesets, _) = futures::try_join!(
            self.changesets.add(ctx, cs.clone()),
            self.maybe_write_to_new_commit_graph(
                ctx,
                vec1![(cs.cs_id, SmallVec::from_vec(cs.parents))],
            )
        )
        .with_context(|| "during commit_graph_compat::Changesets::add")?;
        Ok(added_to_changesets)
    }

    async fn add_many(
        &self,
        ctx: &CoreContext,
        css: Vec1<(ChangesetInsert, Generation)>,
    ) -> Result<()> {
        futures::try_join!(
            self.changesets.add_many(ctx, css.clone()),
            self.maybe_write_to_new_commit_graph(
                ctx,
                css.mapped(|(cs, _gen)| (cs.cs_id, SmallVec::from_vec(cs.parents))),
            )
        )
        .with_context(|| "during commit_graph_compat::Changesets::add")?;
        Ok(())
    }

    async fn get(&self, ctx: &CoreContext, cs_id: ChangesetId) -> Result<Option<ChangesetEntry>> {
        self.changesets.get(ctx, cs_id).await
    }

    async fn get_many(
        &self,
        ctx: &CoreContext,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<Vec<ChangesetEntry>> {
        self.changesets.get_many(ctx, cs_ids).await
    }

    async fn get_many_by_prefix(
        &self,
        ctx: &CoreContext,
        cs_prefix: ChangesetIdPrefix,
        limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        self.changesets
            .get_many_by_prefix(ctx, cs_prefix, limit)
            .await
    }

    fn prime_cache(&self, ctx: &CoreContext, changesets: &[ChangesetEntry]) {
        self.changesets.prime_cache(ctx, changesets)
    }

    async fn enumeration_bounds(
        &self,
        ctx: &CoreContext,
        read_from_master: bool,
        known_heads: Vec<ChangesetId>,
    ) -> Result<Option<(u64, u64)>> {
        self.changesets
            .enumeration_bounds(ctx, read_from_master, known_heads)
            .await
    }

    fn list_enumeration_range(
        &self,
        ctx: &CoreContext,
        min_id: u64,
        max_id: u64,
        sort_and_limit: Option<(SortOrder, u64)>,
        read_from_master: bool,
    ) -> BoxStream<'_, Result<(ChangesetId, u64)>> {
        self.changesets.list_enumeration_range(
            ctx,
            min_id,
            max_id,
            sort_and_limit,
            read_from_master,
        )
    }
}


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

pub mod changesets;

pub use crate::changesets::ChangesetsCommitGraphCompat;


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

use std::collections::HashSet;
use std::str::FromStr;
use std::sync::Arc;

use anyhow::Result;
use cloned::cloned;
use commit_graph::CommitGraph;
use commit_graph_types::storage::CommitGraphStorage;
use context::CoreContext;
use in_memory_commit_graph_storage::InMemoryCommitGraphStorage;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::RepositoryId;
use smallvec::smallvec;
use vec1::vec1;

use crate::utils::*;

pub mod utils;

#[macro_export]
macro_rules! impl_commit_graph_tests {
    ( $test_runner:ident ) => {
        $crate::impl_commit_graph_tests_internal!(
            $test_runner,
            test_storage_store_and_fetch,
            test_skip_tree,
            test_p1_linear_tree,
            test_ancestors_difference,
            test_find_by_prefix,
            test_add_recursive,
            test_add_recursive_many_changesets,
            test_ancestors_frontier_with,
            test_range_stream,
            test_common_base,
            test_slice_ancestors,
            test_children,
            test_ancestors_difference_segments_1,
            test_ancestors_difference_segments_2,
        );
    };
}

#[macro_export]
macro_rules! impl_commit_graph_tests_internal {
    ( $test_runner:ident, $($test_name:ident, )* ) => {
        $(
            #[fbinit::test]
            pub async fn $test_name(fb: FacebookInit) -> Result<()> {
                $test_runner(fb, $crate::$test_name).await
            }
        )*
    }
}

pub async fn test_storage_store_and_fetch(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
             A-B-C-D-G-H-I
              \     /
               E---F
         "##,
        storage.clone(),
    )
    .await?;

    // Check the public API.
    assert!(graph.exists(&ctx, name_cs_id("A")).await?);

    assert!(!graph.exists(&ctx, name_cs_id("nonexistent")).await?);
    assert_eq!(
        graph
            .changeset_generation(&ctx, name_cs_id("G"))
            .await?
            .unwrap()
            .value(),
        5
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("A"))
            .await?
            .unwrap()
            .as_slice(),
        &[]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("E"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("A")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("G"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("D"), name_cs_id("F")]
    );

    assert!(
        graph
            .is_ancestor(&ctx, name_cs_id("C"), name_cs_id("C"))
            .await?
    );
    assert!(
        graph
            .is_ancestor(&ctx, name_cs_id("A"), name_cs_id("H"))
            .await?
    );
    assert!(
        graph
            .is_ancestor(&ctx, name_cs_id("A"), name_cs_id("F"))
            .await?
    );
    assert!(
        graph
            .is_ancestor(&ctx, name_cs_id("F"), name_cs_id("I"))
            .await?
    );
    assert!(
        graph
            .is_ancestor(&ctx, name_cs_id("C"), name_cs_id("I"))
            .await?
    );
    assert!(
        !graph
            .is_ancestor(&ctx, name_cs_id("I"), name_cs_id("A"))
            .await?
    );
    assert!(
        !graph
            .is_ancestor(&ctx, name_cs_id("E"), name_cs_id("D"))
            .await?
    );
    assert!(
        !graph
            .is_ancestor(&ctx, name_cs_id("B"), name_cs_id("E"))
            .await?
    );

    // Check some underlying storage details.
    assert_eq!(
        storage
            .fetch_edges(&ctx, name_cs_id("A"))
            .await?
            .unwrap()
            .merge_ancestor,
        None
    );
    assert_eq!(
        storage
            .fetch_edges(&ctx, name_cs_id("C"))
            .await?
            .unwrap()
            .merge_ancestor,
        Some(name_cs_node("A", 1, 0, 0))
    );
    assert_eq!(
        storage
            .fetch_edges(&ctx, name_cs_id("I"))
            .await?
            .unwrap()
            .merge_ancestor,
        Some(name_cs_node("G", 5, 1, 4))
    );

    Ok(())
}

pub async fn test_skip_tree(ctx: CoreContext, storage: Arc<dyn CommitGraphStorage>) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    assert_eq!(
        storage
            .fetch_edges(&ctx, name_cs_id("K"))
            .await?
            .unwrap()
            .node
            .cs_id,
        name_cs_id("K")
    );

    assert_skip_tree_parent(&storage, &ctx, "G", "B").await?;
    assert_skip_tree_parent(&storage, &ctx, "K", "J").await?;
    assert_skip_tree_parent(&storage, &ctx, "J", "H").await?;
    assert_skip_tree_parent(&storage, &ctx, "H", "G").await?;

    assert_skip_tree_skew_ancestor(&storage, &ctx, "H", "A").await?;
    assert_skip_tree_skew_ancestor(&storage, &ctx, "K", "J").await?;
    assert_skip_tree_skew_ancestor(&storage, &ctx, "U", "T").await?;
    assert_skip_tree_skew_ancestor(&storage, &ctx, "T", "S").await?;
    assert_skip_tree_skew_ancestor(&storage, &ctx, "S", "L").await?;

    assert_skip_tree_level_ancestor(&graph, &ctx, "S", 4, Some("P")).await?;
    assert_skip_tree_level_ancestor(&graph, &ctx, "U", 7, Some("S")).await?;
    assert_skip_tree_level_ancestor(&graph, &ctx, "T", 7, Some("S")).await?;
    assert_skip_tree_level_ancestor(&graph, &ctx, "O", 2, Some("N")).await?;
    assert_skip_tree_level_ancestor(&graph, &ctx, "N", 3, None).await?;
    assert_skip_tree_level_ancestor(&graph, &ctx, "K", 2, Some("G")).await?;

    assert_skip_tree_lowest_common_ancestor(&graph, &ctx, "D", "F", Some("B")).await?;
    assert_skip_tree_lowest_common_ancestor(&graph, &ctx, "K", "I", Some("H")).await?;
    assert_skip_tree_lowest_common_ancestor(&graph, &ctx, "D", "C", Some("C")).await?;
    assert_skip_tree_lowest_common_ancestor(&graph, &ctx, "N", "K", None).await?;
    assert_skip_tree_lowest_common_ancestor(&graph, &ctx, "A", "I", Some("A")).await?;

    Ok(())
}

pub async fn test_p1_linear_tree(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    assert_p1_linear_skew_ancestor(&storage, &ctx, "A", None).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "B", Some("A")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "C", Some("B")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "D", Some("A")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "E", Some("B")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "F", Some("A")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "G", Some("D")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "H", Some("G")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "I", Some("D")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "J", Some("D")).await?;
    assert_p1_linear_skew_ancestor(&storage, &ctx, "K", Some("A")).await?;

    assert_p1_linear_level_ancestor(&graph, &ctx, "S", 4, Some("P")).await?;
    assert_p1_linear_level_ancestor(&graph, &ctx, "U", 7, Some("S")).await?;
    assert_p1_linear_level_ancestor(&graph, &ctx, "T", 7, Some("S")).await?;
    assert_p1_linear_level_ancestor(&graph, &ctx, "O", 2, Some("N")).await?;
    assert_p1_linear_level_ancestor(&graph, &ctx, "N", 3, None).await?;
    assert_p1_linear_level_ancestor(&graph, &ctx, "K", 2, Some("C")).await?;

    assert_p1_linear_lowest_common_ancestor(&graph, &ctx, "D", "F", Some("B")).await?;
    assert_p1_linear_lowest_common_ancestor(&graph, &ctx, "K", "I", Some("H")).await?;
    assert_p1_linear_lowest_common_ancestor(&graph, &ctx, "D", "C", Some("C")).await?;
    assert_p1_linear_lowest_common_ancestor(&graph, &ctx, "N", "K", None).await?;
    assert_p1_linear_lowest_common_ancestor(&graph, &ctx, "A", "I", Some("A")).await?;

    Ok(())
}

pub async fn test_ancestors_difference(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    assert_ancestors_difference(
        &graph,
        &ctx,
        vec!["K"],
        vec![],
        vec!["K", "J", "I", "H", "G", "D", "F", "C", "E", "B", "A"],
    )
    .await?;

    assert_ancestors_difference(
        &graph,
        &ctx,
        vec!["K", "U"],
        vec![],
        vec![
            "U", "T", "S", "R", "Q", "P", "O", "N", "M", "L", "K", "J", "I", "H", "G", "D", "F",
            "C", "E", "B", "A",
        ],
    )
    .await?;

    assert_ancestors_difference(&graph, &ctx, vec!["K"], vec!["G"], vec!["K", "J", "I", "H"])
        .await?;

    assert_ancestors_difference(&graph, &ctx, vec!["K", "I"], vec!["J"], vec!["K"]).await?;

    assert_ancestors_difference(
        &graph,
        &ctx,
        vec!["I"],
        vec!["C"],
        vec!["I", "H", "G", "F", "E", "D"],
    )
    .await?;

    assert_ancestors_difference(
        &graph,
        &ctx,
        vec!["J", "S"],
        vec!["C", "E", "O"],
        vec!["J", "I", "H", "G", "F", "D", "S", "R", "Q", "P"],
    )
    .await?;

    let set1 = ["A", "B", "C", "D", "E", "F", "G", "H", "I"]
        .into_iter()
        .map(name_cs_id)
        .collect::<HashSet<_>>();

    let set1_fn = move |cs_id| {
        cloned!(set1);
        async move { Ok(set1.contains(&cs_id)) }
    };

    assert_ancestors_difference_with(
        &graph,
        &ctx,
        vec!["J", "S"],
        vec!["C", "E", "O"],
        set1_fn.clone(),
        vec!["J", "S", "R", "Q", "P"],
    )
    .await?;

    assert_ancestors_difference_with(
        &graph,
        &ctx,
        vec!["K"],
        vec!["C", "E"],
        set1_fn.clone(),
        vec!["K", "J"],
    )
    .await?;

    let set2 = ["A", "B", "C"]
        .into_iter()
        .map(name_cs_id)
        .collect::<HashSet<_>>();

    let set2_fn = move |cs_id| {
        cloned!(set2);
        async move { Ok(set2.contains(&cs_id)) }
    };

    assert_ancestors_difference_with(
        &graph,
        &ctx,
        vec!["H"],
        vec![],
        set2_fn.clone(),
        vec!["D", "E", "F", "G", "H"],
    )
    .await?;

    assert_ancestors_difference_with(
        &graph,
        &ctx,
        vec!["H"],
        vec!["F"],
        set2_fn.clone(),
        vec!["D", "G", "H"],
    )
    .await?;

    Ok(())
}

pub async fn test_find_by_prefix(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
             J-K-L-LZZ
             M-MA-MAA-MAB-MAC
             M-MB-MBB-MBC
             N-NAA
             O-P-QQ
             a-b-c
         "##,
        storage.clone(),
    )
    .await?;

    assert_eq!(
        graph
            .find_by_prefix(&ctx, ChangesetIdPrefix::from_bytes("Z")?, 10)
            .await?,
        ChangesetIdsResolvedFromPrefix::NoMatch
    );
    assert_eq!(
        graph
            .find_by_prefix(&ctx, ChangesetIdPrefix::from_bytes("Q")?, 10)
            .await?,
        ChangesetIdsResolvedFromPrefix::Single(name_cs_id("QQ"))
    );
    assert_eq!(
        graph
            .find_by_prefix(&ctx, ChangesetIdPrefix::from_bytes("MA")?, 10)
            .await?,
        ChangesetIdsResolvedFromPrefix::Multiple(vec![
            name_cs_id("MA"),
            name_cs_id("MAA"),
            name_cs_id("MAB"),
            name_cs_id("MAC"),
        ])
    );
    assert_eq!(
        graph
            .find_by_prefix(&ctx, ChangesetIdPrefix::from_bytes("M")?, 6)
            .await?,
        ChangesetIdsResolvedFromPrefix::TooMany(vec![
            name_cs_id("M"),
            name_cs_id("MA"),
            name_cs_id("MAA"),
            name_cs_id("MAB"),
            name_cs_id("MAC"),
            name_cs_id("MB"),
        ])
    );
    // Check prefixes that are not a full byte. `P` is `\x50` in ASCII.
    assert_eq!(
        graph
            .find_by_prefix(&ctx, ChangesetIdPrefix::from_str("5")?, 2)
            .await?,
        ChangesetIdsResolvedFromPrefix::Multiple(vec![name_cs_id("P"), name_cs_id("QQ")])
    );

    Ok(())
}

pub async fn test_add_recursive(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let reference_storage = Arc::new(InMemoryCommitGraphStorage::new(RepositoryId::new(1)));

    let reference_graph = Arc::new(
        from_dag(
            &ctx,
            r##"
             A-B-C-D-G-H-I
              \     /
               E---F---J
         "##,
            reference_storage,
        )
        .await?,
    );

    let graph = CommitGraph::new(storage);
    assert_eq!(
        graph
            .add_recursive(
                &ctx,
                reference_graph.clone(),
                vec1![(name_cs_id("I"), smallvec![name_cs_id("H")])],
            )
            .await?,
        9
    );
    assert_eq!(
        graph
            .add_recursive(
                &ctx,
                reference_graph,
                vec1![(name_cs_id("J"), smallvec![name_cs_id("F")])]
            )
            .await?,
        1
    );

    assert!(graph.exists(&ctx, name_cs_id("A")).await?);

    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("E"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("A")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("G"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("D"), name_cs_id("F")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("I"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("H")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("J"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("F")]
    );

    Ok(())
}

pub async fn test_add_recursive_many_changesets(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let reference_storage = Arc::new(InMemoryCommitGraphStorage::new(RepositoryId::new(1)));

    let reference_graph = Arc::new(
        from_dag(
            &ctx,
            r##"
             A-B-C-D-G-H-I
              \     /
               E---F---J
         "##,
            reference_storage,
        )
        .await?,
    );

    let graph = CommitGraph::new(storage);
    assert_eq!(
        graph
            .add_recursive(
                &ctx,
                reference_graph.clone(),
                vec1![
                    (name_cs_id("I"), smallvec![name_cs_id("H")]),
                    (name_cs_id("K"), smallvec![name_cs_id("I")]),
                    (name_cs_id("L"), smallvec![name_cs_id("K")]),
                    (name_cs_id("M"), smallvec![name_cs_id("J")]),
                ],
            )
            .await?,
        13
    );

    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("I"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("H")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("K"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("I")]
    );
    assert_eq!(
        graph
            .changeset_parents(&ctx, name_cs_id("L"))
            .await?
            .unwrap()
            .as_slice(),
        &[name_cs_id("K")]
    );

    assert_eq!(
        graph
            .add_recursive(
                &ctx,
                reference_graph.clone(),
                vec1![
                    (name_cs_id("N"), smallvec![name_cs_id("M")]),
                    (name_cs_id("O"), smallvec![name_cs_id("K"), name_cs_id("N")]),
                ],
            )
            .await?,
        2
    );
    Ok(())
}

pub async fn test_ancestors_frontier_with(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    let set1 = ["A", "B", "C", "D", "E", "F", "G", "H", "I"]
        .into_iter()
        .map(name_cs_id)
        .collect::<HashSet<_>>();

    assert_ancestors_frontier_with(
        &graph,
        &ctx,
        vec!["K", "U"],
        move |cs_id| {
            cloned!(set1);
            async move { Ok(set1.contains(&cs_id)) }
        },
        vec!["H", "I"],
    )
    .await?;

    let set2 = ["A", "B", "C", "E"]
        .into_iter()
        .map(name_cs_id)
        .collect::<HashSet<_>>();

    assert_ancestors_frontier_with(
        &graph,
        &ctx,
        vec!["D", "F"],
        {
            cloned!(set2);
            move |cs_id| {
                cloned!(set2);
                async move { Ok(set2.contains(&cs_id)) }
            }
        },
        vec!["C", "E"],
    )
    .await?;

    assert_ancestors_frontier_with(
        &graph,
        &ctx,
        vec!["G"],
        {
            cloned!(set2);
            move |cs_id| {
                cloned!(set2);
                async move { Ok(set2.contains(&cs_id)) }
            }
        },
        vec!["C", "E"],
    )
    .await?;

    assert_ancestors_frontier_with(
        &graph,
        &ctx,
        vec!["K"],
        {
            cloned!(set2);
            move |cs_id| {
                cloned!(set2);
                async move { Ok(set2.contains(&cs_id)) }
            }
        },
        vec!["C", "E"],
    )
    .await?;

    assert_ancestors_frontier_with(
        &graph,
        &ctx,
        vec!["D"],
        {
            cloned!(set2);
            move |cs_id| {
                cloned!(set2);
                async move { Ok(set2.contains(&cs_id)) }
            }
        },
        vec!["C"],
    )
    .await?;

    Ok(())
}

pub async fn test_range_stream(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    assert_range_stream(
        &graph,
        &ctx,
        "A",
        "K",
        vec!["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"],
    )
    .await?;
    assert_range_stream(&graph, &ctx, "D", "K", vec!["D", "G", "H", "I", "J", "K"]).await?;
    assert_range_stream(&graph, &ctx, "A", "U", vec![]).await?;
    assert_range_stream(&graph, &ctx, "O", "T", vec!["O", "P", "Q", "R", "S", "T"]).await?;

    Ok(())
}

pub async fn test_common_base(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
        A-B-C-D-E-L------N
           \       \    /
            F-G-H   M  /
             \     /  /
              I-J-K--/

        O-P-Q-R-S-T-U-V-W
        "##,
        storage.clone(),
    )
    .await?;

    assert_common_base(&graph, &ctx, "J", "J", vec!["J"]).await?;
    assert_common_base(&graph, &ctx, "K", "J", vec!["J"]).await?;
    assert_common_base(&graph, &ctx, "E", "H", vec!["B"]).await?;
    assert_common_base(&graph, &ctx, "G", "J", vec!["F"]).await?;
    assert_common_base(&graph, &ctx, "M", "N", vec!["K", "L"]).await?;
    assert_common_base(&graph, &ctx, "L", "K", vec!["B"]).await?;
    assert_common_base(&graph, &ctx, "M", "H", vec!["F"]).await?;
    assert_common_base(&graph, &ctx, "A", "B", vec!["A"]).await?;
    assert_common_base(&graph, &ctx, "N", "W", vec![]).await?;
    assert_common_base(&graph, &ctx, "D", "Q", vec![]).await?;

    Ok(())
}

pub async fn test_slice_ancestors(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
         A-B-C-D-G-H---J-K
            \   /   \ /
             E-F     I

         L-M-N-O-P-Q-R-S-T-U
         "##,
        storage.clone(),
    )
    .await?;

    assert_slice_ancestors(
        &graph,
        &ctx,
        vec!["H"],
        |cs_ids| async { Ok(cs_ids.into_iter().collect::<HashSet<_>>()) },
        2,
        vec![(1, vec!["B"]), (3, vec!["D", "F"]), (5, vec!["H"])],
    )
    .await?;

    assert_slice_ancestors(
        &graph,
        &ctx,
        vec!["Q"],
        |cs_ids| async { Ok(cs_ids.into_iter().collect::<HashSet<_>>()) },
        1,
        vec![
            (1, vec!["L"]),
            (2, vec!["M"]),
            (3, vec!["N"]),
            (4, vec!["O"]),
            (5, vec!["P"]),
            (6, vec!["Q"]),
        ],
    )
    .await?;

    assert_slice_ancestors(
        &graph,
        &ctx,
        vec!["Q"],
        |cs_ids| async { Ok(cs_ids.into_iter().collect::<HashSet<_>>()) },
        3,
        vec![(1, vec!["N"]), (4, vec!["Q"])],
    )
    .await?;

    let set1 = ["P", "Q", "R", "S", "T", "U"]
        .into_iter()
        .map(name_cs_id)
        .collect::<HashSet<_>>();

    assert_slice_ancestors(
        &graph,
        &ctx,
        vec!["Q"],
        |_| async { Ok(set1.clone()) },
        1,
        vec![(5, vec!["P"]), (6, vec!["Q"])],
    )
    .await?;

    Ok(())
}

pub async fn test_children(ctx: CoreContext, storage: Arc<dyn CommitGraphStorage>) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
        A-B-C-D-E-L------N
           \       \    /
            F-G-H   M  /
             \     /  /
              I-J-K--/
        "##,
        storage,
    )
    .await?;

    assert_children(&graph, &ctx, "A", vec!["B"]).await?;
    assert_children(&graph, &ctx, "B", vec!["C", "F"]).await?;
    assert_children(&graph, &ctx, "C", vec!["D"]).await?;
    assert_children(&graph, &ctx, "D", vec!["E"]).await?;
    assert_children(&graph, &ctx, "E", vec!["L"]).await?;
    assert_children(&graph, &ctx, "F", vec!["G", "I"]).await?;
    assert_children(&graph, &ctx, "G", vec!["H"]).await?;
    assert_children(&graph, &ctx, "H", vec![]).await?;
    assert_children(&graph, &ctx, "I", vec!["J"]).await?;
    assert_children(&graph, &ctx, "J", vec!["K"]).await?;
    assert_children(&graph, &ctx, "K", vec!["M", "N"]).await?;
    assert_children(&graph, &ctx, "L", vec!["M", "N"]).await?;
    assert_children(&graph, &ctx, "M", vec![]).await?;
    assert_children(&graph, &ctx, "N", vec![]).await?;

    Ok(())
}

pub async fn test_ancestors_difference_segments_1(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
        A-B-C-D-E---L------N----O
           \         \    /
            F-G-H     M  /
             \       /  /
              I-J---K--/
                 \
                  \---------P
        "##,
        storage,
    )
    .await?;

    assert_ancestors_difference_segments(&ctx, &graph, vec!["N"], vec![], 3).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["N"], vec!["D"], 3).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["H"], vec!["G"], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["M"], vec![], 3).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["M"], vec!["H"], 3).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["N"], vec!["E", "J"], 3).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["O", "P"], vec![], 4).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["O", "P"], vec!["H"], 4).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["O", "P"], vec!["D", "I"], 4).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["F"], vec!["H"], 0).await?;

    Ok(())
}

pub async fn test_ancestors_difference_segments_2(
    ctx: CoreContext,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<()> {
    let graph = from_dag(
        &ctx,
        r##"
        A--B------C----E---J---K
         \  \      \
          \  \--D   \-----F----L
           \  \            \
            \  \--G---H     \--M
             \     \
              \-P   \--I--N----O
        "##,
        storage,
    )
    .await?;

    assert_ancestors_difference_segments(&ctx, &graph, vec!["K"], vec![], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["L"], vec![], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["M"], vec![], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["O"], vec![], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["K", "L"], vec![], 2).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["K", "L", "M", "O"], vec![], 4).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["K", "L"], vec!["M"], 2).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["K", "L", "H"], vec!["M", "O"], 3)
        .await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["C"], vec!["M"], 0).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["A", "B", "E"], vec![], 1).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["B", "H", "O"], vec!["D"], 2).await?;
    assert_ancestors_difference_segments(&ctx, &graph, vec!["E", "L", "K"], vec!["J"], 2).await?;
    assert_ancestors_difference_segments(
        &ctx,
        &graph,
        vec![
            "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P",
        ],
        vec![],
        7,
    )
    .await?;

    Ok(())
}


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

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

use anyhow::Result;
use commit_graph::AncestorsStreamBuilder;
use commit_graph::CommitGraph;
use commit_graph_types::edges::ChangesetNode;
use commit_graph_types::storage::CommitGraphStorage;
use context::CoreContext;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::Future;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

/// Generate a fake changeset id for graph testing purposes by using the raw
/// bytes of the changeset name, padded with zeroes.
pub fn name_cs_id(name: &str) -> ChangesetId {
    let mut bytes = [0; 32];
    bytes[..name.len()].copy_from_slice(name.as_bytes());
    ChangesetId::from_bytes(bytes).expect("Changeset ID should be valid")
}

/// Generate a fake changeset node for graph testing purposes by using the raw
/// bytes of the changeset name, padded with zeroes.
pub fn name_cs_node(
    name: &str,
    gen: u64,
    skip_tree_depth: u64,
    p1_linear_depth: u64,
) -> ChangesetNode {
    let cs_id = name_cs_id(name);
    let generation = Generation::new(gen);
    ChangesetNode {
        cs_id,
        generation,
        skip_tree_depth,
        p1_linear_depth,
    }
}

/// Build a commit graph from an ASCII-art dag.
pub async fn from_dag(
    ctx: &CoreContext,
    dag: &str,
    storage: Arc<dyn CommitGraphStorage>,
) -> Result<CommitGraph> {
    let mut added: BTreeMap<String, ChangesetId> = BTreeMap::new();
    let dag = drawdag::parse(dag);
    let graph = CommitGraph::new(storage.clone());

    while added.len() < dag.len() {
        let mut made_progress = false;
        for (name, parents) in dag.iter() {
            if added.contains_key(name) {
                // This node was already added.
                continue;
            }

            if parents.iter().any(|parent| !added.contains_key(parent)) {
                // This node still has unadded parents.
                continue;
            }

            let parent_ids = parents.iter().map(|parent| added[parent].clone()).collect();

            let cs_id = name_cs_id(name);
            graph.add(ctx, cs_id, parent_ids).await?;
            added.insert(name.clone(), cs_id);
            made_progress = true;
        }
        if !made_progress {
            anyhow::bail!("Graph contains cycles");
        }
    }
    Ok(graph)
}

pub async fn assert_skip_tree_parent(
    storage: &Arc<dyn CommitGraphStorage>,
    ctx: &CoreContext,
    u: &str,
    u_skip_tree_parent: &str,
) -> Result<()> {
    assert_eq!(
        storage
            .fetch_edges(ctx, name_cs_id(u))
            .await?
            .unwrap()
            .skip_tree_parent
            .map(|node| node.cs_id),
        Some(name_cs_id(u_skip_tree_parent))
    );
    Ok(())
}

pub async fn assert_skip_tree_skew_ancestor(
    storage: &Arc<dyn CommitGraphStorage>,
    ctx: &CoreContext,
    u: &str,
    u_skip_tree_skew_ancestor: &str,
) -> Result<()> {
    assert_eq!(
        storage
            .fetch_edges(ctx, name_cs_id(u))
            .await?
            .unwrap()
            .skip_tree_skew_ancestor
            .map(|node| node.cs_id),
        Some(name_cs_id(u_skip_tree_skew_ancestor))
    );
    Ok(())
}

pub async fn assert_skip_tree_level_ancestor(
    graph: &CommitGraph,
    ctx: &CoreContext,
    u: &str,
    target_depth: u64,
    u_level_ancestor: Option<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .skip_tree_level_ancestor(ctx, name_cs_id(u), target_depth,)
            .await?
            .map(|node| node.cs_id),
        u_level_ancestor.map(name_cs_id)
    );
    Ok(())
}

pub async fn assert_skip_tree_lowest_common_ancestor(
    graph: &CommitGraph,
    ctx: &CoreContext,
    u: &str,
    v: &str,
    lca: Option<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .skip_tree_lowest_common_ancestor(ctx, name_cs_id(u), name_cs_id(v),)
            .await?
            .map(|node| node.cs_id),
        lca.map(name_cs_id)
    );
    Ok(())
}

pub async fn assert_ancestors_difference_with<Property, Out>(
    graph: &CommitGraph,
    ctx: &CoreContext,
    heads: Vec<&str>,
    common: Vec<&str>,
    property_fn: Property,
    ancestors_difference: Vec<&str>,
) -> Result<()>
where
    Property: Fn(ChangesetId) -> Out + Send + Sync + 'static,
    Out: Future<Output = Result<bool>> + Send + Sync + 'static,
{
    let heads = heads.into_iter().map(name_cs_id).collect();
    let common = common.into_iter().map(name_cs_id).collect();

    assert_eq!(
        AncestorsStreamBuilder::new(Arc::new(graph.clone()), ctx.clone(), heads)
            .exclude_ancestors_of(common)
            .without(property_fn)
            .build()
            .await?
            .try_collect::<HashSet<_>>()
            .await?,
        ancestors_difference
            .into_iter()
            .map(name_cs_id)
            .collect::<HashSet<_>>()
    );
    Ok(())
}

pub async fn assert_ancestors_difference(
    graph: &CommitGraph,
    ctx: &CoreContext,
    heads: Vec<&str>,
    common: Vec<&str>,
    ancestors_difference: Vec<&str>,
) -> Result<()> {
    let heads = heads.into_iter().map(name_cs_id).collect();
    let common = common.into_iter().map(name_cs_id).collect();

    assert_eq!(
        graph
            .ancestors_difference(ctx, heads, common)
            .await?
            .into_iter()
            .collect::<HashSet<_>>(),
        ancestors_difference
            .into_iter()
            .map(name_cs_id)
            .collect::<HashSet<_>>()
    );
    Ok(())
}

pub async fn assert_topological_order(
    graph: &CommitGraph,
    ctx: &CoreContext,
    cs_ids: &Vec<ChangesetId>,
) -> Result<()> {
    let all_cs_ids: HashSet<ChangesetId> = cs_ids.iter().copied().collect();
    let mut previous_cs_ids: HashSet<ChangesetId> = Default::default();

    for cs_id in cs_ids {
        let parents = graph.changeset_parents_required(ctx, *cs_id).await?;
        // Check that each parent of cs_id either isn't contained in cs_ids
        // or is found before cs_id.
        assert!(
            parents
                .into_iter()
                .all(|parent| !all_cs_ids.contains(&parent) || previous_cs_ids.contains(&parent))
        );
        previous_cs_ids.insert(*cs_id);
    }

    Ok(())
}

pub async fn assert_range_stream(
    graph: &CommitGraph,
    ctx: &CoreContext,
    start: &str,
    end: &str,
    range: Vec<&str>,
) -> Result<()> {
    let start_id = name_cs_id(start);
    let end_id = name_cs_id(end);

    let range_stream_cs_ids = graph
        .range_stream(ctx, start_id, end_id)
        .await?
        .collect::<Vec<_>>()
        .await;

    assert_topological_order(graph, ctx, &range_stream_cs_ids).await?;

    assert_eq!(
        range_stream_cs_ids.into_iter().collect::<HashSet<_>>(),
        range.into_iter().map(name_cs_id).collect::<HashSet<_>>()
    );
    Ok(())
}

pub async fn assert_ancestors_frontier_with<Property, Out>(
    graph: &CommitGraph,
    ctx: &CoreContext,
    heads: Vec<&str>,
    property_fn: Property,
    ancestors_frontier: Vec<&str>,
) -> Result<()>
where
    Property: Fn(ChangesetId) -> Out + Send + Sync + 'static,
    Out: Future<Output = Result<bool>>,
{
    let heads = heads.into_iter().map(name_cs_id).collect();

    assert_eq!(
        graph
            .ancestors_frontier_with(ctx, heads, property_fn)
            .await?
            .into_iter()
            .collect::<HashSet<_>>(),
        ancestors_frontier
            .into_iter()
            .map(name_cs_id)
            .collect::<HashSet<_>>()
    );
    Ok(())
}

pub async fn assert_p1_linear_skew_ancestor(
    storage: &Arc<dyn CommitGraphStorage>,
    ctx: &CoreContext,
    u: &str,
    u_p1_linear_skew_ancestor: Option<&str>,
) -> Result<()> {
    assert_eq!(
        storage
            .fetch_edges(ctx, name_cs_id(u))
            .await?
            .unwrap()
            .p1_linear_skew_ancestor
            .map(|node| node.cs_id),
        u_p1_linear_skew_ancestor.map(name_cs_id)
    );
    Ok(())
}

pub async fn assert_p1_linear_level_ancestor(
    graph: &CommitGraph,
    ctx: &CoreContext,
    u: &str,
    target_depth: u64,
    u_level_ancestor: Option<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .p1_linear_level_ancestor(ctx, name_cs_id(u), target_depth)
            .await?
            .map(|node| node.cs_id),
        u_level_ancestor.map(name_cs_id)
    );
    Ok(())
}

pub async fn assert_p1_linear_lowest_common_ancestor(
    graph: &CommitGraph,
    ctx: &CoreContext,
    u: &str,
    v: &str,
    lca: Option<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .p1_linear_lowest_common_ancestor(ctx, name_cs_id(u), name_cs_id(v))
            .await?
            .map(|node| node.cs_id),
        lca.map(name_cs_id)
    );
    Ok(())
}

pub async fn assert_common_base(
    graph: &CommitGraph,
    ctx: &CoreContext,
    u: &str,
    v: &str,
    common_base: Vec<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .common_base(ctx, name_cs_id(u), name_cs_id(v))
            .await?
            .into_iter()
            .collect::<HashSet<_>>(),
        common_base
            .into_iter()
            .map(name_cs_id)
            .collect::<HashSet<_>>()
    );
    Ok(())
}

pub async fn assert_slice_ancestors<NeedsProcessing, Out>(
    graph: &CommitGraph,
    ctx: &CoreContext,
    heads: Vec<&str>,
    needs_processing: NeedsProcessing,
    slice_size: u64,
    slices: Vec<(u64, Vec<&str>)>,
) -> Result<()>
where
    NeedsProcessing: Fn(Vec<ChangesetId>) -> Out,
    Out: Future<Output = Result<HashSet<ChangesetId>>>,
{
    let heads = heads.into_iter().map(name_cs_id).collect();
    assert_eq!(
        graph
            .slice_ancestors(ctx, heads, needs_processing, slice_size)
            .await?
            .into_iter()
            .map(|(gen_group, cs_ids)| (gen_group, cs_ids.into_iter().collect::<HashSet<_>>()))
            .collect::<Vec<_>>(),
        slices
            .into_iter()
            .map(|(gen_group, cs_ids)| (
                gen_group,
                cs_ids.into_iter().map(name_cs_id).collect::<HashSet<_>>()
            ))
            .collect::<Vec<_>>(),
    );
    Ok(())
}

pub async fn assert_children(
    graph: &CommitGraph,
    ctx: &CoreContext,
    cs_id: &str,
    children: Vec<&str>,
) -> Result<()> {
    assert_eq!(
        graph
            .changeset_children(ctx, name_cs_id(cs_id))
            .await?
            .into_iter()
            .collect::<HashSet<_>>(),
        children.into_iter().map(name_cs_id).collect::<HashSet<_>>(),
    );
    Ok(())
}

pub async fn assert_ancestors_difference_segments(
    ctx: &CoreContext,
    graph: &CommitGraph,
    heads: Vec<&str>,
    common: Vec<&str>,
    segments_count: usize,
) -> Result<()> {
    let heads: Vec<_> = heads.into_iter().map(name_cs_id).collect();
    let common: Vec<_> = common.into_iter().map(name_cs_id).collect();

    assert!(
        graph
            .verified_ancestors_difference_segments(ctx, heads, common)
            .await?
            .len()
            == segments_count
    );

    Ok(())
}


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "commit_graph crate
mononoke_types_thrift mononoke_types_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=commit_graph_thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "commit_graph.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();