/*
 * 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 Edges

use std::num::NonZeroU32;

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Result;
use commit_graph_thrift as thrift;
use futures::Future;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use smallvec::SmallVec;

#[derive(Abomonation, Copy, Clone, Debug, Eq, PartialEq)]
pub struct ChangesetNode {
    /// The id of the changeset.
    pub cs_id: ChangesetId,

    /// The changeset's generation: the inclusive number of commits between
    /// this commit and the farthest root commit.  Root commits have a
    /// generation of 1.
    pub generation: Generation,

    /// The changeset's depth in the skip tree.
    pub skip_tree_depth: u64,

    /// The changeset's depth in the p1 tree
    pub p1_linear_depth: u64,
}

impl ChangesetNode {
    pub fn to_thrift(&self) -> thrift::ChangesetNode {
        thrift::ChangesetNode {
            cs_id: self.cs_id.into_thrift(),
            generation: thrift::Generation(self.generation.value() as i64),
            skip_tree_depth: self.skip_tree_depth as i64,
            p1_linear_depth: self.p1_linear_depth as i64,
        }
    }

    pub fn from_thrift(node: thrift::ChangesetNode) -> Result<Self> {
        Ok(Self {
            cs_id: ChangesetId::from_thrift(node.cs_id)?,
            generation: Generation::new(node.generation.0 as u64),
            skip_tree_depth: node.skip_tree_depth as u64,
            p1_linear_depth: node.p1_linear_depth as u64,
        })
    }
}
/// The parents of a changeset.
///
/// This uses a smallvec, as there is usually exactly one.
pub type ChangesetParents = SmallVec<[ChangesetId; 1]>;

/// The parents of a changeset node.
///
/// This uses a smallvec, as there is usually exactly one.  Unlike
/// `ChangesetParents`, this includes the generation number of the parents.
pub type ChangesetNodeParents = SmallVec<[ChangesetNode; 1]>;

/// Outgoing edges from a changeset node.
#[derive(Abomonation, Clone, Debug, Eq, PartialEq)]
pub struct ChangesetEdges {
    /// The starting changeset for this set of edges.
    pub node: ChangesetNode,

    /// The changeset's immediate parents.
    pub parents: ChangesetNodeParents,

    /// For root and merge commits, this is `None`.
    ///
    /// For single-parent commits, this is the merge ancestor: the most recent
    /// ancestor that is a merge or root.
    pub merge_ancestor: Option<ChangesetNode>,

    /// The skip tree parent: this is the most recent single common ancestor
    /// of this commit's parents
    pub skip_tree_parent: Option<ChangesetNode>,

    /// The skip tree skew ancestor: this is some ancestor of the common
    /// ancestors of this commit's parents, which provides a skew-binary
    /// search tree over the commit graph, if such an ancestor exists.
    pub skip_tree_skew_ancestor: Option<ChangesetNode>,

    /// The p1-linear skew ancestor: this is some ancestor of the first
    /// parent of this commit, which provides a skew-binary search tree over
    /// the linear first-parent history of this commit, if such an ancestor
    /// exists.
    ///
    /// Note: this excludes any merged in branches, so should only be used for
    /// history-lossy operations.
    pub p1_linear_skew_ancestor: Option<ChangesetNode>,
}

impl ChangesetEdges {
    pub fn to_thrift(&self) -> thrift::ChangesetEdges {
        thrift::ChangesetEdges {
            node: self.node.to_thrift(),
            parents: self.parents.iter().map(ChangesetNode::to_thrift).collect(),
            merge_ancestor: self.merge_ancestor.as_ref().map(ChangesetNode::to_thrift),
            skip_tree_parent: self.skip_tree_parent.as_ref().map(ChangesetNode::to_thrift),
            skip_tree_skew_ancestor: self
                .skip_tree_skew_ancestor
                .as_ref()
                .map(ChangesetNode::to_thrift),
            p1_linear_skew_ancestor: self
                .p1_linear_skew_ancestor
                .as_ref()
                .map(ChangesetNode::to_thrift),
        }
    }

    pub fn from_thrift(edges: thrift::ChangesetEdges) -> Result<Self> {
        Ok(Self {
            node: ChangesetNode::from_thrift(edges.node)?,
            parents: edges
                .parents
                .into_iter()
                .map(ChangesetNode::from_thrift)
                .collect::<Result<ChangesetNodeParents>>()?,
            merge_ancestor: edges
                .merge_ancestor
                .map(ChangesetNode::from_thrift)
                .transpose()?,
            skip_tree_parent: edges
                .skip_tree_parent
                .map(ChangesetNode::from_thrift)
                .transpose()?,
            skip_tree_skew_ancestor: edges
                .skip_tree_skew_ancestor
                .map(ChangesetNode::from_thrift)
                .transpose()?,
            p1_linear_skew_ancestor: edges
                .p1_linear_skew_ancestor
                .map(ChangesetNode::from_thrift)
                .transpose()?,
        })
    }

    /// Returns the lowest skip tree edge (skip_tree_parent or skip_tree_skew_ancestor)
    /// that satisfies the given property, or None if neither does.
    pub async fn lowest_skip_tree_edge_with<Property, Out>(
        &self,
        property: Property,
    ) -> Result<Option<ChangesetNode>>
    where
        Property: Fn(ChangesetNode) -> Out,
        Out: Future<Output = Result<bool>>,
    {
        if let Some(skip_tree_skew_ancestor) = self.skip_tree_skew_ancestor {
            if property(skip_tree_skew_ancestor).await? {
                return Ok(Some(skip_tree_skew_ancestor));
            }
        }

        if let Some(skip_tree_parent) = self.skip_tree_parent {
            if property(skip_tree_parent).await? {
                return Ok(Some(skip_tree_parent));
            }
        }

        Ok(None)
    }
}

/// A smaller version of ChangesetEdges for use in cases where
/// space efficiency matters (e.g. preloading the commit graph).
///
/// Outgoing edges are replaced by u32 ids identifying a changeset.
#[derive(Debug)]
pub struct CompactChangesetEdges {
    pub generation: u32,
    pub skip_tree_depth: u32,
    pub p1_linear_depth: u32,
    pub parents: SmallVec<[NonZeroU32; 2]>,
    pub merge_ancestor: Option<NonZeroU32>,
    pub skip_tree_parent: Option<NonZeroU32>,
    pub skip_tree_skew_ancestor: Option<NonZeroU32>,
    pub p1_linear_skew_ancestor: Option<NonZeroU32>,
}

impl CompactChangesetEdges {
    pub fn to_thrift(
        &self,
        cs_id: ChangesetId,
        unique_id: NonZeroU32,
    ) -> thrift::CompactChangesetEdges {
        thrift::CompactChangesetEdges {
            cs_id: cs_id.into_thrift(),
            unique_id: unique_id.get() as i32,
            generation: self.generation as i32,
            skip_tree_depth: self.skip_tree_depth as i32,
            p1_linear_depth: self.p1_linear_depth as i32,
            parents: self
                .parents
                .iter()
                .copied()
                .map(|id| id.get() as i32)
                .collect(),
            merge_ancestor: self.merge_ancestor.map(|id| id.get() as i32),
            skip_tree_parent: self.skip_tree_parent.map(|id| id.get() as i32),
            skip_tree_skew_ancestor: self.skip_tree_skew_ancestor.map(|id| id.get() as i32),
            p1_linear_skew_ancestor: self.p1_linear_skew_ancestor.map(|id| id.get() as i32),
        }
    }

    pub fn from_thrift(edges: thrift::CompactChangesetEdges) -> Result<Self> {
        Ok(Self {
            generation: edges.generation as u32,
            skip_tree_depth: edges.skip_tree_depth as u32,
            p1_linear_depth: edges.p1_linear_depth as u32,
            parents: edges
                .parents
                .into_iter()
                .map(|id| {
                    NonZeroU32::new(id as u32)
                        .ok_or_else(|| anyhow!("Couldn't convert parent id to NonZeroU32"))
                })
                .collect::<Result<_>>()?,
            merge_ancestor: edges
                .merge_ancestor
                .map(|id| {
                    NonZeroU32::new(id as u32)
                        .ok_or_else(|| anyhow!("Couldn't convert merge ancestor id to NonZeroU32"))
                })
                .transpose()?,
            skip_tree_parent: edges
                .skip_tree_parent
                .map(|id| {
                    NonZeroU32::new(id as u32).ok_or_else(|| {
                        anyhow!("Couldn't convert skip tree parent id to NonZeroU32")
                    })
                })
                .transpose()?,
            skip_tree_skew_ancestor: edges
                .skip_tree_skew_ancestor
                .map(|id| {
                    NonZeroU32::new(id as u32).ok_or_else(|| {
                        anyhow!("Couldn't convert skip tree skew ancestor id to NonZeroU32")
                    })
                })
                .transpose()?,
            p1_linear_skew_ancestor: edges
                .p1_linear_skew_ancestor
                .map(|id| {
                    NonZeroU32::new(id as u32).ok_or_else(|| {
                        anyhow!("Couldn't convert p1 linear skew ancestor id to NonZeroU32")
                    })
                })
                .transpose()?,
        })
    }
}


/*
 * 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::iter::Extend;
use std::iter::IntoIterator;
use std::iter::Iterator;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ops::RangeBounds;

use maplit::btreemap;
use maplit::hashset;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;

/// A frontier of changesets ordered by generation number.
#[derive(Clone, Debug)]
pub struct ChangesetFrontier(BTreeMap<Generation, HashSet<ChangesetId>>);

impl ChangesetFrontier {
    pub fn new() -> Self {
        Self(Default::default())
    }

    pub fn new_single(cs_id: ChangesetId, generation: Generation) -> Self {
        Self(btreemap! { generation => hashset! { cs_id }})
    }

    pub fn highest_generation_contains(&self, cs_id: ChangesetId, generation: Generation) -> bool {
        match self.last_key_value() {
            None => false,
            Some((highest_frontier_generation, cs_ids)) => {
                *highest_frontier_generation == generation && cs_ids.contains(&cs_id)
            }
        }
    }

    pub fn highest_generation_intersection(
        &self,
        other_frontier: &ChangesetFrontier,
    ) -> Vec<ChangesetId> {
        match self.last_key_value() {
            None => vec![],
            Some((highest_frontier_generation, cs_ids)) => {
                match other_frontier.get(highest_frontier_generation) {
                    None => vec![],
                    Some(other_cs_ids) => cs_ids.intersection(other_cs_ids).copied().collect(),
                }
            }
        }
    }

    pub fn is_disjoint(&self, other_frontier: &ChangesetFrontier) -> bool {
        for (gen, cs_ids) in self.iter().rev() {
            if let Some(other_cs_ids) = other_frontier.get(gen) {
                if !cs_ids.is_disjoint(other_cs_ids) {
                    return false;
                }
            }
        }
        true
    }

    /// Return an iterator over tuples of each changeset in the frontier
    /// together with its generation number.
    pub fn into_flat_iter(self) -> impl Iterator<Item = (ChangesetId, Generation)> {
        self.0
            .into_iter()
            .flat_map(|(gen, cs_ids)| cs_ids.into_iter().map(move |cs_id| (cs_id, gen)))
    }

    /// Returns a vec of all changesets in the frontier.
    pub fn changesets(&self) -> Vec<ChangesetId> {
        self.iter()
            .flat_map(|(_, cs_ids)| cs_ids.iter())
            .copied()
            .collect()
    }

    /// Returns a vec of all changesets in the frontier inside
    /// of the given range.
    pub fn changesets_in_range(&self, range: impl RangeBounds<Generation>) -> Vec<ChangesetId> {
        self.range(range)
            .flat_map(|(_, cs_ids)| cs_ids.iter())
            .copied()
            .collect()
    }
}

impl Deref for ChangesetFrontier {
    type Target = BTreeMap<Generation, HashSet<ChangesetId>>;

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

impl DerefMut for ChangesetFrontier {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl FromIterator<(ChangesetId, Generation)> for ChangesetFrontier {
    fn from_iter<T: IntoIterator<Item = (ChangesetId, Generation)>>(iter: T) -> Self {
        let mut frontier = Self::new();

        for (cs_id, gen) in iter {
            frontier.entry(gen).or_default().insert(cs_id);
        }

        frontier
    }
}

impl Extend<(ChangesetId, Generation)> for ChangesetFrontier {
    fn extend<T: IntoIterator<Item = (ChangesetId, Generation)>>(&mut self, iter: T) {
        for (cs_id, gen) in iter {
            self.entry(gen).or_default().insert(cs_id);
        }
    }
}


/*
 * 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 edges;
pub mod frontier;
pub mod segments;
pub mod storage;


/*
 * 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::HashMap;
use std::collections::HashSet;
use std::fmt::Display;

use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use smallvec::SmallVec;

/// A representation of a segment of changesets.
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct ChangesetSegment {
    /// The highest generation changeset in the segment.
    pub head: ChangesetId,
    /// The lowest generation changeset in the segment.
    pub base: ChangesetId,
    /// The number of changesets in the segment.
    pub length: u64,
    /// Parents of the lowest generation changeset in the segment.
    pub parents: SmallVec<[ChangesetSegmentParent; 1]>,
}

/// A representation of a parent of a segment.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ChangesetSegmentParent {
    pub cs_id: ChangesetId,
    pub location: Option<ChangesetSegmentLocation>,
}

/// A location of a changeset in a segment, represented as the ancestor
/// of the head of the segment that's at the specified distance (i.e. head~distance).
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ChangesetSegmentLocation {
    pub head: ChangesetId,
    pub distance: u64,
}

impl Display for ChangesetSegment {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "{} -> {}, length: {}, parents: {}",
            self.head,
            self.base,
            self.length,
            self.parents
                .iter()
                .map(|cs_id| cs_id.to_string())
                .collect::<Vec<_>>()
                .join(", "),
        )
    }
}

impl Display for ChangesetSegmentParent {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.cs_id)?;

        if let Some(location) = &self.location {
            write!(f, " ({})", location)?;
        }

        Ok(())
    }
}

impl Display for ChangesetSegmentLocation {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}~{}", self.head, self.distance)
    }
}

/// A frontier of changeset segments ordered by the generation
/// number of their base. For each generation the segments are
/// represented as a map from their base to the set of their heads.
#[derive(Default, Debug)]
pub struct ChangesetSegmentFrontier {
    pub segments: BTreeMap<Generation, HashMap<ChangesetId, HashSet<ChangesetId>>>,
}


/*
 * 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 Storage
//!
//! Trait for the storage back-end for the commit graph.

use std::collections::HashMap;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::Generation;
use mononoke_types::RepositoryId;
use mononoke_types::FIRST_GENERATION;
use vec1::Vec1;

use crate::edges::ChangesetEdges;

/// Indication of the kind of edge to traverse for prefetch.
#[derive(Copy, Clone, Debug)]
pub enum PrefetchEdge {
    /// Prefetch a linear range of commits by following the first parent
    FirstParent,

    /// Prefetch along the maximum skip tree distance by following the skip
    /// tree skew ancestor, or first parent if the commit does not have
    /// a skip tree skew ancestor
    SkipTreeSkewAncestor,
}

/// Where to prefetch to.
#[derive(Copy, Clone, Debug)]
pub struct PrefetchTarget {
    /// Prefetch along this edge.
    pub edge: PrefetchEdge,

    /// Prefetch as far back as this generation.
    pub generation: Generation,

    /// Prefetch up to this many steps.
    pub steps: u64,
}

/// Indication for additional changesets to be fetched for subsequent
/// traversals.
///
/// If efficient to do so, implementors should use this hint to fetch
/// additional edges that will be useful for skew-binary traversal
/// to the target generation.
#[derive(Copy, Clone, Debug)]
pub enum Prefetch {
    /// No prefetch is required.
    None,

    /// Prefetch is permitted with the given hint, but additional items are
    /// not to be returned.
    Hint(PrefetchTarget),

    /// Prefetch if possible, and included prefetched items in the result.
    Include(PrefetchTarget),
}

impl Prefetch {
    /// Prepare prefetching for skew-binary traversal over the skip tree.
    pub fn for_skip_tree_traversal(generation: Generation) -> Self {
        // We are prefetching mostly along the skew ancestor edge, which
        // should typically be O(log(N)) in length, except that for merge
        // commits without a common ancestor we follow the p1 parent, so limit
        // to 32 steps so that we don't follow the p1 ancestry too far.
        Prefetch::Hint(PrefetchTarget {
            edge: PrefetchEdge::SkipTreeSkewAncestor,
            generation,
            steps: 32,
        })
    }

    /// Prepare prefetching for linear traversal of the p1 history.
    pub fn for_p1_linear_traversal() -> Self {
        // Prefetch linear ranges of 128 commits.  This is arbitrary, but is a
        // balance between not overfetching for the cache and reducing the
        // number of sequential steps.
        Prefetch::Hint(PrefetchTarget {
            edge: PrefetchEdge::FirstParent,
            generation: FIRST_GENERATION,
            steps: 128,
        })
    }

    pub fn is_hint(&self) -> bool {
        matches!(self, Prefetch::Hint(..))
    }

    pub fn is_include(&self) -> bool {
        matches!(self, Prefetch::Include(..))
    }

    /// Indicate that prefetching should be included if it has been hinted.
    ///
    /// This is called when the caching layer determines that it is able to
    /// store any prefetched values, and so values for any prefetch hint
    /// should be included.
    pub fn include_hint(self) -> Prefetch {
        match self {
            Prefetch::None => Prefetch::None,
            Prefetch::Hint(target) | Prefetch::Include(target) => Prefetch::Include(target),
        }
    }

    /// Target to prefetch to, if prefetching should be included.
    ///
    /// If prefetching is merely hinted, this won't return the target, as
    /// prefetching should not be performed.
    pub fn target(self) -> Option<PrefetchTarget> {
        match self {
            Prefetch::None | Prefetch::Hint(..) => None,
            Prefetch::Include(target) => Some(target),
        }
    }
}

/// Commit Graph Storage.
#[async_trait]
pub trait CommitGraphStorage: Send + Sync {
    /// The repository this commit graph storage is for.
    fn repo_id(&self) -> RepositoryId;

    /// Add a new changeset to the commit graph.
    ///
    /// Returns true if a new changeset was inserted, or false if the
    /// changeset already existed.
    async fn add(&self, ctx: &CoreContext, edges: ChangesetEdges) -> Result<bool>;

    /// Add many changesets at once. Used for low level stuff like backfilling.
    async fn add_many(&self, ctx: &CoreContext, many_edges: Vec1<ChangesetEdges>) -> Result<usize>;

    /// Returns the changeset graph edges for this changeset.
    async fn fetch_edges(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetEdges>>;

    /// Returns the changeset graph edges for this changeset, or an error of
    /// this changeset is missing in the commit graph.
    async fn fetch_edges_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetEdges>;

    /// Returns the changeset graph edges for multiple changesets.
    ///
    /// Prefetch indicates that this request is part of a larger request
    /// involving commits down to a particular generation number, and so
    /// prefetching more nodes into any internal caches would be beneficial.
    async fn fetch_many_edges(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>>;

    /// Same as fetch_many_edges but returns an error if any of
    /// the changesets are missing in the commit graph.
    async fn fetch_many_edges_required(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>>;

    /// Find all changeset ids with a given prefix.
    async fn find_by_prefix(
        &self,
        _ctx: &CoreContext,
        _cs_prefix: ChangesetIdPrefix,
        _limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix>;

    // Fetch all children of a changeset.
    async fn fetch_children(
        &self,
        _ctx: &CoreContext,
        _cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>>;
}


/*
 * 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::HashMap;
use std::collections::HashSet;

use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::RepositoryId;
use parking_lot::RwLock;
use vec1::Vec1;

/// In-memory commit graph storage.
pub struct InMemoryCommitGraphStorage {
    repo_id: RepositoryId,
    changesets: RwLock<BTreeMap<ChangesetId, ChangesetEdges>>,
    children: RwLock<BTreeMap<ChangesetId, HashSet<ChangesetId>>>,
}

impl InMemoryCommitGraphStorage {
    pub fn new(repo_id: RepositoryId) -> Self {
        InMemoryCommitGraphStorage {
            repo_id,
            changesets: Default::default(),
            children: Default::default(),
        }
    }

    pub fn drain(&self) -> Vec<ChangesetEdges> {
        let mut changesets = self.changesets.write();
        let many_edges = changesets.iter().map(|(_, edges)| edges).cloned().collect();
        changesets.clear();
        many_edges
    }

    pub fn len(&self) -> usize {
        self.changesets.read().len()
    }

    pub fn is_empty(&self) -> bool {
        self.changesets.read().is_empty()
    }
}

#[async_trait]
impl CommitGraphStorage for InMemoryCommitGraphStorage {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn add(&self, _ctx: &CoreContext, edges: ChangesetEdges) -> Result<bool> {
        {
            let mut children = self.children.write();
            for parent in edges.parents.iter() {
                children
                    .entry(parent.cs_id)
                    .or_default()
                    .insert(edges.node.cs_id);
            }
        }

        Ok(self
            .changesets
            .write()
            .insert(edges.node.cs_id, edges)
            .is_none())
    }

    async fn add_many(
        &self,
        _ctx: &CoreContext,
        many_edges: Vec1<ChangesetEdges>,
    ) -> Result<usize> {
        let mut added = 0;

        {
            let mut changesets = self.changesets.write();
            let mut children = self.children.write();
            for edges in many_edges {
                if changesets.insert(edges.node.cs_id, edges.clone()).is_none() {
                    added += 1;

                    for parent in edges.parents.iter() {
                        children
                            .entry(parent.cs_id)
                            .or_default()
                            .insert(edges.node.cs_id);
                    }
                }
            }
        }

        Ok(added)
    }

    async fn fetch_edges(
        &self,
        _ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetEdges>> {
        Ok(self.changesets.read().get(&cs_id).cloned())
    }

    async fn fetch_edges_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetEdges> {
        self.fetch_edges(ctx, cs_id).await?.ok_or_else(|| {
            anyhow!(
                "Missing changeset from in-memory commit graph storage: {}",
                cs_id
            )
        })
    }

    async fn fetch_many_edges(
        &self,
        _ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        _prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let mut result = HashMap::with_capacity(cs_ids.len());
        let changesets = self.changesets.read();
        for cs_id in cs_ids {
            if let Some(edges) = changesets.get(cs_id) {
                result.insert(*cs_id, edges.clone());
            }
        }
        Ok(result)
    }

    async fn fetch_many_edges_required(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let edges = self.fetch_many_edges(ctx, cs_ids, prefetch).await?;
        let missing_changesets: Vec<_> = cs_ids
            .iter()
            .filter(|cs_id| !edges.contains_key(cs_id))
            .collect();

        if !missing_changesets.is_empty() {
            Err(anyhow!(
                "Missing changesets from in-memory commit graph storage: {}",
                missing_changesets
                    .into_iter()
                    .map(|cs_id| format!("{}, ", cs_id))
                    .collect::<String>()
            ))
        } else {
            Ok(edges)
        }
    }

    async fn find_by_prefix(
        &self,
        _ctx: &CoreContext,
        cs_prefix: ChangesetIdPrefix,
        limit: usize,
    ) -> Result<ChangesetIdsResolvedFromPrefix> {
        let changesets = self.changesets.read();
        let (min, max) = (cs_prefix.min_bound(), cs_prefix.max_bound());
        let matches: Vec<_> = changesets
            .range(min..=max)
            .take(limit.saturating_add(1))
            .map(|(cs_id, _)| *cs_id)
            .collect();
        Ok(ChangesetIdsResolvedFromPrefix::from_vec_and_limit(
            matches, limit,
        ))
    }

    async fn fetch_children(
        &self,
        _ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Vec<ChangesetId>> {
        Ok(self
            .children
            .read()
            .get(&cs_id)
            .cloned()
            .unwrap_or_default()
            .into_iter()
            .collect())
    }
}

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

    use commit_graph_testlib::*;
    use commit_graph_types::storage::CommitGraphStorage;
    use context::CoreContext;
    use fbinit::FacebookInit;

    use super::*;

    async fn run_test<Fut>(
        fb: FacebookInit,
        test_function: impl FnOnce(CoreContext, Arc<dyn CommitGraphStorage>) -> Fut,
    ) -> Result<()>
    where
        Fut: Future<Output = Result<()>>,
    {
        let ctx = CoreContext::test_mock(fb);
        let storage = Arc::new(InMemoryCommitGraphStorage::new(RepositoryId::new(1)));
        test_function(ctx, storage).await
    }

    impl_commit_graph_tests!(run_test);
}


/*
 * 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::HashMap;
use std::num::NonZeroU32;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use bytes::Bytes;
use cloned::cloned;
use commit_graph_thrift as thrift;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::edges::ChangesetNode;
use commit_graph_types::edges::CompactChangesetEdges;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use fbthrift::compact_protocol;
use mononoke_types::ChangesetId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::Generation;
use mononoke_types::RepositoryId;
use reloader::Loader;
use reloader::Reloader;
use slog::info;
use vec1::Vec1;

#[cfg(test)]
mod tests;

const DEFAULT_RELOADING_INTERVAL_SECS: u64 = 60 * 60;

/// A commit graph storage that wraps another storage and periodically preloads
/// the changeset edges of the commit graph from the blobstore. Writes are passed
/// to the underlying storage, while reads first search the preloaded changeset
/// edges and fall over to the underlying storage for the missing changesets.
///
/// Useful for commit graphs with complicated structure that are small enough
/// to fit into memory.
pub struct PreloadedCommitGraphStorage {
    repo_id: RepositoryId,
    preloaded_edges: Reloader<PreloadedEdges>,
    persistent_storage: Arc<dyn CommitGraphStorage>,
}

pub struct PreloadedEdgesLoader {
    ctx: CoreContext,
    blobstore_without_cache: Arc<dyn Blobstore>,
    blobstore_key: String,
}

#[derive(Debug, Default)]
pub struct PreloadedEdges {
    pub cs_id_to_edges: HashMap<ChangesetId, CompactChangesetEdges>,
    pub unique_id_to_cs_id: HashMap<NonZeroU32, ChangesetId>,
    pub max_sql_id: Option<u64>,
}

impl PreloadedEdges {
    pub fn to_thrift(&self) -> Result<thrift::PreloadedEdges> {
        Ok(thrift::PreloadedEdges {
            edges: self
                .unique_id_to_cs_id
                .iter()
                .map(|(unique_id, cs_id)| {
                    Ok(self
                        .cs_id_to_edges
                        .get(cs_id)
                        .ok_or_else(|| anyhow!("Missing changeset edges for {}", cs_id))?
                        .to_thrift(*cs_id, *unique_id))
                })
                .collect::<Result<_>>()?,
            max_sql_id: self.max_sql_id.map(|id| id as i64),
        })
    }

    pub fn from_thrift(preloaded_edges: thrift::PreloadedEdges) -> Result<Self> {
        let unique_id_to_cs_id = preloaded_edges
            .edges
            .iter()
            .map(|edges| {
                Ok((
                    NonZeroU32::new(edges.unique_id as u32)
                        .ok_or_else(|| anyhow!("Couldn't convert unique_id to NonZeroU32"))?,
                    ChangesetId::from_thrift(edges.cs_id.clone())?,
                ))
            })
            .collect::<Result<_>>()?;
        Ok(Self {
            unique_id_to_cs_id,
            cs_id_to_edges: preloaded_edges
                .edges
                .into_iter()
                .map(|edges| {
                    Ok((
                        ChangesetId::from_thrift(edges.cs_id.clone())?,
                        CompactChangesetEdges::from_thrift(edges)?,
                    ))
                })
                .collect::<Result<_>>()?,
            max_sql_id: preloaded_edges.max_sql_id.map(|id| id as u64),
        })
    }

    fn get_node(&self, unique_id: NonZeroU32) -> Result<ChangesetNode> {
        let cs_id = *self
            .unique_id_to_cs_id
            .get(&unique_id)
            .ok_or_else(|| anyhow!("Missing changeset id for unique id: {}", unique_id))?;
        let edges = self
            .cs_id_to_edges
            .get(&cs_id)
            .ok_or_else(|| anyhow!("Missing changeset edges for {}", cs_id))?;

        Ok(ChangesetNode {
            cs_id,
            generation: Generation::new(edges.generation as u64),
            skip_tree_depth: edges.skip_tree_depth as u64,
            p1_linear_depth: edges.p1_linear_depth as u64,
        })
    }

    pub fn get(&self, cs_id: &ChangesetId) -> Result<Option<ChangesetEdges>> {
        let compact_edges = match self.cs_id_to_edges.get(cs_id) {
            Some(edges) => edges,
            None => return Ok(None),
        };

        Ok(Some(ChangesetEdges {
            node: ChangesetNode {
                cs_id: *cs_id,
                generation: Generation::new(compact_edges.generation as u64),
                skip_tree_depth: compact_edges.skip_tree_depth as u64,
                p1_linear_depth: compact_edges.p1_linear_depth as u64,
            },
            parents: compact_edges
                .parents
                .iter()
                .map(|parent_id| self.get_node(*parent_id))
                .collect::<Result<_>>()?,
            merge_ancestor: compact_edges
                .merge_ancestor
                .map(|merge_ancestor| self.get_node(merge_ancestor))
                .transpose()?,
            skip_tree_parent: compact_edges
                .skip_tree_parent
                .map(|skip_tree_parent| self.get_node(skip_tree_parent))
                .transpose()?,
            skip_tree_skew_ancestor: compact_edges
                .skip_tree_skew_ancestor
                .map(|skip_tree_skew_ancestor| self.get_node(skip_tree_skew_ancestor))
                .transpose()?,
            p1_linear_skew_ancestor: compact_edges
                .p1_linear_skew_ancestor
                .map(|p1_linear_skew_ancestor| self.get_node(p1_linear_skew_ancestor))
                .transpose()?,
        }))
    }
}

#[derive(Debug, Default)]
pub struct ExtendablePreloadedEdges {
    preloaded_edges: PreloadedEdges,
    cs_id_to_unique_id: HashMap<ChangesetId, NonZeroU32>,
}

impl ExtendablePreloadedEdges {
    pub fn from_preloaded_edges(preloaded_edges: PreloadedEdges) -> Self {
        let cs_id_to_unique_id = preloaded_edges
            .unique_id_to_cs_id
            .iter()
            .map(|(unique_id, cs_id)| (*cs_id, *unique_id))
            .collect();
        Self {
            preloaded_edges,
            cs_id_to_unique_id,
        }
    }

    pub fn into_preloaded_edges(self) -> PreloadedEdges {
        self.preloaded_edges
    }

    pub fn unique_id(&mut self, cs_id: ChangesetId) -> NonZeroU32 {
        match self.cs_id_to_unique_id.get(&cs_id) {
            Some(unique_id) => *unique_id,
            None => {
                let unique_id = NonZeroU32::new(self.cs_id_to_unique_id.len() as u32 + 1).unwrap();
                self.cs_id_to_unique_id.insert(cs_id, unique_id);
                self.preloaded_edges
                    .unique_id_to_cs_id
                    .insert(unique_id, cs_id);
                unique_id
            }
        }
    }

    pub fn add(&mut self, edges: ChangesetEdges) -> Result<()> {
        let _unique_id = self.unique_id(edges.node.cs_id);
        let parents = edges
            .parents
            .iter()
            .map(|parent| self.unique_id(parent.cs_id))
            .collect();
        let merge_ancestor = edges
            .merge_ancestor
            .map(|merge_ancestor| self.unique_id(merge_ancestor.cs_id));
        let skip_tree_parent = edges
            .skip_tree_parent
            .map(|skip_tree_parent| self.unique_id(skip_tree_parent.cs_id));
        let skip_tree_skew_ancestor = edges
            .skip_tree_skew_ancestor
            .map(|skip_tree_skew_ancestor| self.unique_id(skip_tree_skew_ancestor.cs_id));
        let p1_linear_skew_ancestor = edges
            .p1_linear_skew_ancestor
            .map(|p1_linear_skew_ancestor| self.unique_id(p1_linear_skew_ancestor.cs_id));

        match self.preloaded_edges.cs_id_to_edges.insert(
            edges.node.cs_id,
            CompactChangesetEdges {
                generation: edges.node.generation.value() as u32,
                skip_tree_depth: edges.node.skip_tree_depth as u32,
                p1_linear_depth: edges.node.p1_linear_depth as u32,
                parents,
                merge_ancestor,
                skip_tree_parent,
                skip_tree_skew_ancestor,
                p1_linear_skew_ancestor,
            },
        ) {
            Some(old_edges) => Err(anyhow!("Duplicate changeset edges found: {:?}", old_edges)),
            None => Ok(()),
        }
    }

    pub fn update_max_sql_id(&mut self, max_sql_id: u64) {
        self.preloaded_edges.max_sql_id = Some(max_sql_id);
    }
}

pub fn deserialize_preloaded_edges(bytes: Bytes) -> Result<PreloadedEdges> {
    let preloaded_edges: thrift::PreloadedEdges = compact_protocol::deserialize(bytes)?;

    PreloadedEdges::from_thrift(preloaded_edges)
}

#[async_trait]
impl Loader<PreloadedEdges> for PreloadedEdgesLoader {
    async fn load(&mut self) -> Result<Option<PreloadedEdges>> {
        tokio::task::spawn({
            cloned!(self.ctx, self.blobstore_without_cache, self.blobstore_key);
            async move {
                info!(ctx.logger(), "Started preloading commit graph");
                let maybe_bytes = blobstore_without_cache.get(&ctx, &blobstore_key).await?;
                match maybe_bytes {
                    Some(bytes) => {
                        let bytes = bytes.into_raw_bytes();
                        let preloaded_edges =
                            tokio::task::spawn_blocking(move || deserialize_preloaded_edges(bytes))
                                .await??;
                        info!(
                            ctx.logger(),
                            "Finished preloading commit graph ({} changesets)",
                            preloaded_edges.cs_id_to_edges.len()
                        );
                        Ok(Some(preloaded_edges))
                    }
                    None => Ok(Some(Default::default())),
                }
            }
        })
        .await?
    }
}

impl PreloadedCommitGraphStorage {
    pub async fn from_blobstore(
        ctx: &CoreContext,
        repo_id: RepositoryId,
        blobstore_without_cache: Arc<dyn Blobstore>,
        preloaded_edges_blobstore_key: String,
        persistent_storage: Arc<dyn CommitGraphStorage>,
    ) -> Result<Arc<Self>> {
        let loader = PreloadedEdgesLoader {
            ctx: ctx.clone(),
            blobstore_key: preloaded_edges_blobstore_key,
            blobstore_without_cache,
        };

        let reloader = Reloader::reload_periodically(
            ctx.clone(),
            move || {
                std::time::Duration::from_secs(
                    tunables::tunables()
                        .preloaded_commit_graph_reloading_interval_secs()
                        .map_or(DEFAULT_RELOADING_INTERVAL_SECS, |secs| secs as u64),
                )
            },
            loader,
        )
        .await?;
        Ok(Arc::new(Self {
            repo_id,
            preloaded_edges: reloader,
            persistent_storage,
        }))
    }
}

#[async_trait]
impl CommitGraphStorage for PreloadedCommitGraphStorage {
    fn repo_id(&self) -> RepositoryId {
        self.repo_id
    }

    async fn add(&self, ctx: &CoreContext, edges: ChangesetEdges) -> Result<bool> {
        self.persistent_storage.add(ctx, edges).await
    }

    async fn add_many(&self, ctx: &CoreContext, many_edges: Vec1<ChangesetEdges>) -> Result<usize> {
        self.persistent_storage.add_many(ctx, many_edges).await
    }

    async fn fetch_edges(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<Option<ChangesetEdges>> {
        match self.preloaded_edges.load().get(&cs_id)? {
            Some(edges) => Ok(Some(edges)),
            None => self.persistent_storage.fetch_edges(ctx, cs_id).await,
        }
    }

    async fn fetch_edges_required(
        &self,
        ctx: &CoreContext,
        cs_id: ChangesetId,
    ) -> Result<ChangesetEdges> {
        match self.preloaded_edges.load().get(&cs_id)? {
            Some(edges) => Ok(edges),
            None => {
                self.persistent_storage
                    .fetch_edges_required(ctx, cs_id)
                    .await
            }
        }
    }

    async fn fetch_many_edges(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let preloaded_edges = self.preloaded_edges.load();
        let mut fetched_edges: HashMap<_, _> = cs_ids
            .iter()
            .filter_map(|cs_id| {
                preloaded_edges
                    .get(cs_id)
                    .map(|edges| edges.map(|edges| (*cs_id, edges)))
                    .transpose()
            })
            .collect::<Result<_>>()?;

        let unfetched_ids: Vec<_> = cs_ids
            .iter()
            .filter(|cs_id| !fetched_edges.contains_key(cs_id))
            .copied()
            .collect();

        if !unfetched_ids.is_empty() {
            fetched_edges.extend(
                self.persistent_storage
                    .fetch_many_edges(ctx, unfetched_ids.as_slice(), prefetch)
                    .await?
                    .into_iter(),
            )
        }

        Ok(fetched_edges)
    }

    async fn fetch_many_edges_required(
        &self,
        ctx: &CoreContext,
        cs_ids: &[ChangesetId],
        prefetch: Prefetch,
    ) -> Result<HashMap<ChangesetId, ChangesetEdges>> {
        let edges = self.fetch_many_edges(ctx, cs_ids, prefetch).await?;
        if let Some(missing_changeset) = cs_ids.iter().find(|cs_id| !edges.contains_key(cs_id)) {
            Err(anyhow!(
                "Missing changeset from preloaded commit graph storage: {}",
                missing_changeset,
            ))
        } else {
            Ok(edges)
        }
    }

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

    async fn fetch_children(&self, ctx: &CoreContext, cs: ChangesetId) -> Result<Vec<ChangesetId>> {
        self.persistent_storage.fetch_children(ctx, cs).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::sync::Arc;

use anyhow::Result;
use commit_graph::CommitGraph;
use commit_graph_testlib::utils::from_dag;
use commit_graph_testlib::utils::name_cs_id;
use commit_graph_types::edges::ChangesetEdges;
use commit_graph_types::storage::CommitGraphStorage;
use commit_graph_types::storage::Prefetch;
use context::CoreContext;
use fbinit::FacebookInit;
use fbthrift::compact_protocol;
use in_memory_commit_graph_storage::InMemoryCommitGraphStorage;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use reloader::Reloader;

use crate::deserialize_preloaded_edges;
use crate::ExtendablePreloadedEdges;
use crate::PreloadedCommitGraphStorage;

impl PreloadedCommitGraphStorage {
    /// Constructs PreloadedCommitGraphStorage from a fixed list of edges
    /// for unit tests.
    pub fn from_edges(
        repo_id: RepositoryId,
        edges_vec: Vec<ChangesetEdges>,
        persistent_storage: Arc<dyn CommitGraphStorage>,
    ) -> Result<Arc<Self>> {
        let mut extendable_preloaded_edges: ExtendablePreloadedEdges = Default::default();

        for edges in edges_vec {
            extendable_preloaded_edges.add(edges)?;
        }

        // Serialize the preloaded edges and then deserialize,
        // to make sure no information is lost.
        let bytes = compact_protocol::serialize(
            &extendable_preloaded_edges
                .into_preloaded_edges()
                .to_thrift()?,
        );
        let preloaded_edges = deserialize_preloaded_edges(bytes)?;

        Ok(Arc::new(Self {
            repo_id,
            preloaded_edges: Reloader::fixed(preloaded_edges),
            persistent_storage,
        }))
    }
}

async fn test_equivalent_storages(
    ctx: &CoreContext,
    first_storage: Arc<dyn CommitGraphStorage>,
    second_storage: Arc<dyn CommitGraphStorage>,
    cs_ids: Vec<ChangesetId>,
) -> Result<()> {
    for cs_id in cs_ids.iter() {
        assert_eq!(
            first_storage.fetch_edges(ctx, *cs_id).await?,
            second_storage.fetch_edges(ctx, *cs_id).await?,
        );
        assert_eq!(
            first_storage.fetch_edges_required(ctx, *cs_id).await?,
            second_storage.fetch_edges_required(ctx, *cs_id).await?,
        );
    }

    assert_eq!(
        first_storage
            .fetch_many_edges(ctx, &cs_ids, Prefetch::None)
            .await?,
        second_storage
            .fetch_many_edges(ctx, &cs_ids, Prefetch::None)
            .await?
    );
    assert_eq!(
        first_storage
            .fetch_many_edges_required(ctx, &cs_ids, Prefetch::None)
            .await?,
        second_storage
            .fetch_many_edges_required(ctx, &cs_ids, Prefetch::None)
            .await?
    );

    Ok(())
}

#[fbinit::test]
async fn test_preloaded_commit_graph_storage(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let underlying_storage = Arc::new(InMemoryCommitGraphStorage::new(RepositoryId::new(1)));
    let _ = from_dag(
        &ctx,
        r##"
            A-B-C-D-G-H-I
             \     /
              E---F
        "##,
        underlying_storage.clone(),
    )
    .await?;

    let edges_vec = vec![
        underlying_storage
            .fetch_edges_required(&ctx, name_cs_id("A"))
            .await?,
        underlying_storage
            .fetch_edges_required(&ctx, name_cs_id("B"))
            .await?,
        underlying_storage
            .fetch_edges_required(&ctx, name_cs_id("C"))
            .await?,
        underlying_storage
            .fetch_edges_required(&ctx, name_cs_id("D"))
            .await?,
        underlying_storage
            .fetch_edges_required(&ctx, name_cs_id("E"))
            .await?,
    ];

    let preloaded_storage = PreloadedCommitGraphStorage::from_edges(
        RepositoryId::new(1),
        edges_vec.clone(),
        underlying_storage.clone(),
    )?;

    let graph = CommitGraph::new(preloaded_storage.clone());
    graph
        .add(&ctx, name_cs_id("J"), [name_cs_id("I")].into())
        .await?;

    // Test that fetching from the preloaded storage is equivalent
    // to fetching directly from the underlying storage.
    test_equivalent_storages(
        &ctx,
        underlying_storage.clone(),
        preloaded_storage,
        ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
            .into_iter()
            .map(name_cs_id)
            .collect(),
    )
    .await?;

    let preloaded_storage_with_empty_underlying_storage = PreloadedCommitGraphStorage::from_edges(
        RepositoryId::new(1),
        edges_vec.clone(),
        Arc::new(InMemoryCommitGraphStorage::new(RepositoryId::new(1))),
    )?;

    // Test that fetching any of the preloaded edges doesn't
    // need to be fetched from the underlying storage by using
    // an empty underlying storage.
    test_equivalent_storages(
        &ctx,
        underlying_storage,
        preloaded_storage_with_empty_underlying_storage.clone(),
        ["A", "B", "C", "D", "E"]
            .into_iter()
            .map(name_cs_id)
            .collect(),
    )
    .await?;

    assert_eq!(
        preloaded_storage_with_empty_underlying_storage
            .fetch_edges(&ctx, name_cs_id("F"))
            .await?,
        None
    );

    Ok(())
}