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

use anyhow::format_err;
use anyhow::Result;
use blobstore::Loadable;
use context::CoreContext;
use derived_data_manager::DerivationContext;
use filenodes::FilenodeInfo;
use filenodes::FilenodeResult;
use filenodes::PreparedFilenode;
use futures::future::try_join_all;
use futures::pin_mut;
use futures::stream;
use futures::FutureExt;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use itertools::Either;
use itertools::Itertools;
use manifest::find_intersection_of_diffs_and_parents;
use manifest::Entry;
use mercurial_derivation::MappedHgChangesetId;
use mercurial_types::blobs::File;
use mercurial_types::fetch_manifest_envelope;
use mercurial_types::nodehash::NULL_HASH;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileEnvelope;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestEnvelope;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use mononoke_types::RepoPath;

use crate::mapping::FilenodesOnlyPublic;
use crate::mapping::PreparedRootFilenode;

pub async fn derive_filenodes(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bcs: BonsaiChangeset,
) -> Result<FilenodesOnlyPublic> {
    if tunables::tunables()
        .filenodes_disabled()
        .unwrap_or_default()
    {
        return Ok(FilenodesOnlyPublic::Disabled);
    }
    let (_, public_filenode, non_roots) =
        prepare_filenodes_for_cs(ctx, derivation_ctx, bcs).await?;
    if !non_roots.is_empty() {
        if let FilenodeResult::Disabled = derivation_ctx
            .filenodes()?
            .add_filenodes(ctx, non_roots)
            .await?
        {
            return Ok(FilenodesOnlyPublic::Disabled);
        }
    }
    // In case it got updated while deriving
    if tunables::tunables()
        .filenodes_disabled()
        .unwrap_or_default()
    {
        return Ok(FilenodesOnlyPublic::Disabled);
    }
    Ok(public_filenode)
}

pub async fn derive_filenodes_in_batch(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    batch: Vec<BonsaiChangeset>,
) -> Result<Vec<(ChangesetId, FilenodesOnlyPublic, Vec<PreparedFilenode>)>> {
    stream::iter(
        batch
            .clone()
            .into_iter()
            .map(|bcs| async move { prepare_filenodes_for_cs(ctx, derivation_ctx, bcs).await }),
    )
    .buffered(100)
    .try_collect()
    .await
}

pub async fn prepare_filenodes_for_cs(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bcs: BonsaiChangeset,
) -> Result<(ChangesetId, FilenodesOnlyPublic, Vec<PreparedFilenode>)> {
    let filenodes = generate_all_filenodes(ctx, derivation_ctx, &bcs).await?;
    if filenodes.is_empty() {
        // This commit didn't create any new filenodes, and it's root manifest is the
        // same as one of the parents (that can happen if this commit is empty).
        // In that case we don't need to insert a root filenode - it will be inserted
        // when parent is derived.
        Ok((
            bcs.get_changeset_id(),
            FilenodesOnlyPublic::Present {
                root_filenode: None,
            },
            filenodes,
        ))
    } else {
        let (roots, non_roots): (Vec<_>, Vec<_>) =
            filenodes.into_iter().partition_map(classify_filenode);
        let roots_num = roots.len();
        let mut roots = roots.into_iter();

        match (roots.next(), roots.next()) {
            (Some(root_filenode), None) => Ok((
                bcs.get_changeset_id(),
                FilenodesOnlyPublic::Present {
                    root_filenode: Some(root_filenode),
                },
                non_roots,
            )),
            _ => Err(format_err!(
                "expected exactly one root, found {} for cs_id {}",
                roots_num,
                bcs.get_changeset_id().to_string()
            )),
        }
    }
}

pub async fn generate_all_filenodes(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bcs: &BonsaiChangeset,
) -> Result<Vec<PreparedFilenode>> {
    let blobstore = derivation_ctx.blobstore();
    let hg_id = derivation_ctx
        .derive_dependency::<MappedHgChangesetId>(ctx, bcs.get_changeset_id())
        .await?
        .hg_changeset_id();
    let root_mf = hg_id.load(ctx, &blobstore).await?.manifestid();
    // In case of non-existant manifest (that's created by hg if the first commit in the repo
    // is empty) it's fine to return the empty list of filenodes.
    if root_mf.clone().into_nodehash() == NULL_HASH {
        return Ok(vec![]);
    }
    // Bonsai might have > 2 parents, while mercurial supports at most 2.
    // That's fine for us - we just won't generate filenodes for paths that came from
    // stepparents. That means that linknode for these filenodes will point to a stepparent
    let parents = try_join_all(
        derivation_ctx
            .fetch_parents::<MappedHgChangesetId>(ctx, bcs)
            .await?
            .into_iter()
            .map(|id| async move {
                Result::<_>::Ok(
                    id.hg_changeset_id()
                        .load(ctx, &blobstore)
                        .await?
                        .manifestid(),
                )
            }),
    )
    .await?;
    let linknode = hg_id;

    (async_stream::stream! {
        let s = find_intersection_of_diffs_and_parents(
            ctx.clone(),
            blobstore.clone(),
            root_mf,
            parents.clone(),
        )
        .try_filter_map(|(path, entry, parent_entries)| {
            async move {
                // file entry has file type and file node id. If file type is different but filenode is
                // the same we don't want to create a new filenode, and this filter removes
                // all entries where at least one parent has the same filenode id.
                if let Entry::Leaf((_, hg_filenode_id)) = entry {
                    for parent_entry in parent_entries {
                        if let Entry::Leaf((_, parent_filenode_id)) = parent_entry {
                            if parent_filenode_id == hg_filenode_id {
                                return Ok(None);
                            }
                        }
                    }
                }
                Ok(Some((path, entry)))
            }
        })
        .map_ok(move |(path, entry)| {
            match entry {
                Entry::Tree(hg_mf_id) => fetch_manifest_envelope(ctx, blobstore, hg_mf_id)
                    .map_ok(move |envelope| create_manifest_filenode(path, envelope, linknode))
                    .left_future(),
                Entry::Leaf((_, hg_filenode_id)) => async move {
                    let envelope = hg_filenode_id.load(ctx, blobstore).await?;
                    create_file_filenode(path, envelope, linknode)
                }
                    .right_future(),
            }
        })
        .try_buffer_unordered(100);

        pin_mut!(s);
        while let Some(value) = s.next().await {
            yield value;
        }
    })
    .try_collect()
    .await
}

pub(crate) fn classify_filenode(
    filenode: PreparedFilenode,
) -> Either<PreparedRootFilenode, PreparedFilenode> {
    if filenode.path == RepoPath::RootPath {
        let FilenodeInfo {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        } = filenode.info;

        Either::Left(PreparedRootFilenode {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        })
    } else {
        Either::Right(filenode)
    }
}

fn create_manifest_filenode(
    path: Option<MPath>,
    envelope: HgManifestEnvelope,
    linknode: HgChangesetId,
) -> PreparedFilenode {
    let path = match path {
        Some(path) => RepoPath::DirectoryPath(path),
        None => RepoPath::RootPath,
    };
    let filenode = HgFileNodeId::new(envelope.node_id());
    let (p1, p2) = envelope.parents();
    let p1 = p1.map(HgFileNodeId::new);
    let p2 = p2.map(HgFileNodeId::new);

    PreparedFilenode {
        path,
        info: FilenodeInfo {
            filenode,
            p1,
            p2,
            copyfrom: None,
            linknode,
        },
    }
}

fn create_file_filenode(
    path: Option<MPath>,
    envelope: HgFileEnvelope,
    linknode: HgChangesetId,
) -> Result<PreparedFilenode> {
    let path = match path {
        Some(path) => RepoPath::FilePath(path),
        None => {
            return Err(format_err!("unexpected empty file path"));
        }
    };
    let filenode = envelope.node_id();
    let (p1, p2) = envelope.parents();
    let copyfrom = File::extract_copied_from(envelope.metadata())?
        .map(|(path, node)| (RepoPath::FilePath(path), node));

    Ok(PreparedFilenode {
        path,
        info: FilenodeInfo {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        },
    })
}

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

    use anyhow::anyhow;
    use anyhow::Context;
    use anyhow::Result;
    use async_trait::async_trait;
    use bonsai_hg_mapping::BonsaiHgMapping;
    use bookmarks::Bookmarks;
    use changeset_fetcher::ChangesetFetcher;
    use changesets::Changesets;
    use cloned::cloned;
    use derived_data_manager::BatchDeriveOptions;
    use fbinit::FacebookInit;
    use filenodes::FilenodeRange;
    use filenodes::Filenodes;
    use filestore::FilestoreConfig;
    use fixtures::Linear;
    use fixtures::TestRepoFixture;
    use futures::compat::Stream01CompatExt;
    use manifest::ManifestOps;
    use maplit::hashmap;
    use mercurial_derivation::DeriveHgChangeset;
    use mononoke_types::FileType;
    use repo_blobstore::RepoBlobstore;
    use repo_derived_data::RepoDerivedData;
    use repo_derived_data::RepoDerivedDataRef;
    use revset::AncestorsNodeStream;
    use slog::info;
    use test_repo_factory::TestRepoFactory;
    use tests_utils::resolve_cs_id;
    use tests_utils::CreateCommitContext;
    use tunables::with_tunables;
    use tunables::MononokeTunables;

    use super::*;

    #[facet::container]
    struct TestRepo {
        #[facet]
        bonsai_hg_mapping: dyn BonsaiHgMapping,
        #[facet]
        bookmarks: dyn Bookmarks,
        #[facet]
        repo_blobstore: RepoBlobstore,
        #[facet]
        repo_derived_data: RepoDerivedData,
        #[facet]
        filestore_config: FilestoreConfig,
        #[facet]
        changeset_fetcher: dyn ChangesetFetcher,
        #[facet]
        changesets: dyn Changesets,
        #[facet]
        filenodes: dyn Filenodes,
    }

    async fn verify_filenodes(
        ctx: &CoreContext,
        repo: &TestRepo,
        cs_id: ChangesetId,
        expected_paths: Vec<RepoPath>,
    ) -> Result<()> {
        let bonsai = cs_id.load(ctx, &repo.repo_blobstore).await?;
        let filenodes = generate_all_filenodes(
            ctx,
            &repo.repo_derived_data().manager().derivation_context(None),
            &bonsai,
        )
        .await?;

        assert_eq!(filenodes.len(), expected_paths.len());
        for path in expected_paths {
            assert!(filenodes.iter().any(|filenode| filenode.path == path));
        }

        let linknode = repo.derive_hg_changeset(ctx, cs_id).await?;

        for filenode in filenodes {
            assert_eq!(filenode.info.linknode, linknode);
        }
        Ok(())
    }

    async fn test_generate_filenodes_simple(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let filename = "path";
        let commit = CreateCommitContext::new_root(&ctx, &repo)
            .add_file(filename, "content")
            .commit()
            .await?;

        // Two filenodes - one for root manifest, another for a file
        verify_filenodes(
            &ctx,
            &repo,
            commit,
            vec![RepoPath::RootPath, RepoPath::file(filename)?],
        )
        .await?;

        Ok(())
    }

    #[fbinit::test]
    fn generate_filenodes_simple(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_generate_filenodes_simple(fb))
    }

    async fn test_generate_filenodes_merge(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let first_p1 = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path1", "content")
            .commit()
            .await?;

        let first_p2 = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path2", "content")
            .commit()
            .await?;

        let merge = CreateCommitContext::new(&ctx, &repo, vec![first_p1, first_p2])
            .commit()
            .await?;

        // Only root filenode was added - other filenodes were reused from parents
        verify_filenodes(&ctx, &repo, merge, vec![RepoPath::RootPath]).await?;

        Ok(())
    }

    #[fbinit::test]
    fn generate_filenodes_merge(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_generate_filenodes_merge(fb))
    }

    async fn test_generate_type_change(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let parent = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path", "content")
            .commit()
            .await?;

        let child = CreateCommitContext::new(&ctx, &repo, vec![parent])
            .add_file_with_type("path", "content", FileType::Executable)
            .commit()
            .await?;

        // Only root filenode should be changed - change of file type doesn't change filenode
        verify_filenodes(&ctx, &repo, child, vec![RepoPath::RootPath]).await?;

        Ok(())
    }

    #[fbinit::test]
    fn generate_filenodes_type_change(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_generate_type_change(fb))
    }

    async fn test_many_parents(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let p1 = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path1", "content")
            .commit()
            .await?;
        let p2 = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path2", "content")
            .commit()
            .await?;
        let p3 = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path3", "content")
            .commit()
            .await?;

        let merge = CreateCommitContext::new(&ctx, &repo, vec![p1, p2, p3])
            .commit()
            .await?;

        info!(ctx.logger(), "checking filenodes for {}", p3);
        verify_filenodes(
            &ctx,
            &repo,
            p3,
            vec![RepoPath::RootPath, RepoPath::file("path3")?],
        )
        .await?;

        info!(ctx.logger(), "checking filenodes for {}", merge);
        verify_filenodes(&ctx, &repo, merge, vec![RepoPath::RootPath]).await?;

        Ok(())
    }

    #[fbinit::test]
    fn many_parents(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_many_parents(fb))
    }

    async fn test_derive_empty_commits(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let parent_empty = CreateCommitContext::new_root(&ctx, &repo).commit().await?;

        let child_empty = CreateCommitContext::new(&ctx, &repo, vec![parent_empty])
            .add_file("file", "content")
            .commit()
            .await?;

        let manager = repo.repo_derived_data().manager();

        manager
            .derive::<FilenodesOnlyPublic>(&ctx, child_empty, None)
            .await?;

        // Make sure they are in the mapping
        let maps = manager
            .fetch_derived_batch::<FilenodesOnlyPublic>(&ctx, vec![parent_empty, child_empty], None)
            .await?;

        assert_eq!(maps.len(), 2);
        Ok(())
    }

    #[fbinit::test]
    fn derive_empty_commits(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_derive_empty_commits(fb))
    }

    async fn test_derive_only_empty_commits(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let parent_empty = CreateCommitContext::new_root(&ctx, &repo).commit().await?;
        let child_empty = CreateCommitContext::new(&ctx, &repo, vec![parent_empty])
            .commit()
            .await?;

        let manager = repo.repo_derived_data().manager();
        manager
            .derive::<FilenodesOnlyPublic>(&ctx, child_empty, None)
            .await?;

        // Make sure they are in the mapping
        let maps = manager
            .fetch_derived_batch::<FilenodesOnlyPublic>(&ctx, vec![child_empty, parent_empty], None)
            .await?;
        assert_eq!(maps.len(), 2);
        Ok(())
    }

    #[fbinit::test]
    fn derive_only_empty_commits(fb: FacebookInit) -> Result<()> {
        let runtime = tokio::runtime::Runtime::new()?;
        runtime.block_on(test_derive_only_empty_commits(fb))
    }

    #[fbinit::test]
    fn derive_disabled_filenodes(fb: FacebookInit) -> Result<()> {
        let tunables = MononokeTunables::default();
        tunables.update_bools(&hashmap! {"filenodes_disabled".to_string() => true});

        with_tunables(tunables, || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_time()
                .build()?;
            runtime.block_on(test_derive_disabled_filenodes(fb))
        })
    }

    async fn test_derive_disabled_filenodes(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        let cs = CreateCommitContext::new_root(&ctx, &repo).commit().await?;
        let derived = repo
            .repo_derived_data()
            .derive::<FilenodesOnlyPublic>(&ctx, cs)
            .await?;
        assert_eq!(derived, FilenodesOnlyPublic::Disabled);

        assert_eq!(
            repo.repo_derived_data()
                .fetch_derived::<FilenodesOnlyPublic>(&ctx, cs)
                .await?
                .unwrap(),
            FilenodesOnlyPublic::Disabled
        );

        Ok(())
    }

    #[fbinit::test]
    async fn verify_batch_and_sequential_derive(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let repo1: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        Linear::initrepo(fb, &repo1).await;
        let repo2: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;
        Linear::initrepo(fb, &repo2).await;
        let master_cs_id = resolve_cs_id(&ctx, &repo1, "master").await?;
        let mut cs_ids =
            AncestorsNodeStream::new(ctx.clone(), &repo1.changeset_fetcher, master_cs_id)
                .compat()
                .try_collect::<Vec<_>>()
                .await?;
        cs_ids.reverse();

        let manager1 = repo1.repo_derived_data().manager();
        manager1
            .derive_exactly_batch::<FilenodesOnlyPublic>(
                &ctx,
                cs_ids.clone(),
                BatchDeriveOptions::Parallel { gap_size: None },
                None,
            )
            .await?;
        let batch = manager1
            .fetch_derived_batch::<FilenodesOnlyPublic>(&ctx, cs_ids.clone(), None)
            .await?;

        let sequential = {
            let mut res = HashMap::new();
            for cs in cs_ids.clone() {
                let root_filenode = repo2
                    .repo_derived_data()
                    .derive::<FilenodesOnlyPublic>(&ctx, cs)
                    .await?;
                res.insert(cs, root_filenode);
            }
            res
        };

        assert_eq!(batch, sequential);
        for cs in cs_ids {
            compare_filenodes(&ctx, &repo1, &repo2, cs).await?;
        }
        Ok(())
    }

    #[fbinit::test]
    async fn derive_parents_before_children(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let filenodes_cs_id = Arc::new(Mutex::new(None));
        let mut factory = TestRepoFactory::new(fb)?;
        let repo: TestRepo = factory
            .with_filenodes_override({
                cloned!(filenodes_cs_id);
                move |inner| {
                    Arc::new(FilenodesWrapper {
                        inner,
                        cs_id: filenodes_cs_id.clone(),
                    })
                }
            })
            .build()
            .await?;
        Linear::initrepo(fb, &repo).await;
        let commit8 =
            resolve_cs_id(&ctx, &repo, "a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157").await?;
        let commit8 = repo.derive_hg_changeset(&ctx, commit8).await?;
        *filenodes_cs_id.lock().unwrap() = Some(commit8);
        let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;
        let mut cs_ids =
            AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher, master_cs_id)
                .compat()
                .try_collect::<Vec<_>>()
                .await?;
        cs_ids.reverse();

        let manager = repo.repo_derived_data().manager();

        match manager
            .derive_exactly_batch::<FilenodesOnlyPublic>(
                &ctx,
                cs_ids.clone(),
                BatchDeriveOptions::Parallel { gap_size: None },
                None,
            )
            .await
        {
            Ok(_) => {}
            Err(_) => {}
        };

        // FilenodesWrapper prevents writing of root filenode for a9473beb2eb03ddb1cccc3fbaeb8a4820f9cd157 (8th commit in repo)
        // so all children (9, 10, 11) should not have root_filenodes written
        for cs_id in cs_ids.into_iter().skip(8) {
            let filenode = manager
                .fetch_derived::<FilenodesOnlyPublic>(&ctx, cs_id, None)
                .await?;
            assert_eq!(filenode, None);
        }
        Ok(())
    }

    struct FilenodesWrapper {
        inner: Arc<dyn Filenodes + Send + Sync>,
        cs_id: Arc<Mutex<Option<HgChangesetId>>>,
    }

    #[async_trait]
    impl Filenodes for FilenodesWrapper {
        async fn add_filenodes(
            &self,
            ctx: &CoreContext,
            info: Vec<PreparedFilenode>,
        ) -> Result<FilenodeResult<()>> {
            // compare PreparedFilenode::Info::Linknode
            if let Some(cs_id) = *self.cs_id.lock().unwrap() {
                if info.iter().any(|filenode| filenode.info.linknode == cs_id) {
                    return Err(anyhow!("filenodes for {} are prohibited", cs_id));
                }
            }
            self.inner.add_filenodes(ctx, info).await
        }

        async fn add_or_replace_filenodes(
            &self,
            ctx: &CoreContext,
            info: Vec<PreparedFilenode>,
        ) -> Result<FilenodeResult<()>> {
            self.inner.add_or_replace_filenodes(ctx, info).await
        }

        async fn get_filenode(
            &self,
            ctx: &CoreContext,
            path: &RepoPath,
            filenode: HgFileNodeId,
        ) -> Result<FilenodeResult<Option<FilenodeInfo>>> {
            self.inner.get_filenode(ctx, path, filenode).await
        }

        async fn get_all_filenodes_maybe_stale(
            &self,
            ctx: &CoreContext,
            path: &RepoPath,
            limit: Option<u64>,
        ) -> Result<FilenodeResult<FilenodeRange>> {
            self.inner
                .get_all_filenodes_maybe_stale(ctx, path, limit)
                .await
        }

        fn prime_cache(&self, ctx: &CoreContext, filenodes: &[PreparedFilenode]) {
            self.inner.prime_cache(ctx, filenodes)
        }
    }

    async fn compare_filenodes(
        ctx: &CoreContext,
        repo: &TestRepo,
        backup_repo: &TestRepo,
        cs: ChangesetId,
    ) -> Result<()> {
        let manifest = repo
            .derive_hg_changeset(ctx, cs)
            .await?
            .load(ctx, &repo.repo_blobstore)
            .await
            .with_context(|| format!("while fetching manifest from prod for cs {:?}", cs))?
            .manifestid();
        manifest
            .list_all_entries(ctx.clone(), repo.repo_blobstore.clone())
            .map_ok(|(path, entry)| {
                async move {
                    let (path, node) = match (path, entry) {
                        (Some(path), Entry::Leaf((_, id))) => (RepoPath::FilePath(path), id),
                        (Some(path), Entry::Tree(id)) => (
                            RepoPath::DirectoryPath(path),
                            HgFileNodeId::new(id.into_nodehash()),
                        ),
                        (None, Entry::Leaf((_, id))) => (RepoPath::RootPath, id),
                        (None, Entry::Tree(id)) => {
                            (RepoPath::RootPath, HgFileNodeId::new(id.into_nodehash()))
                        }
                    };
                    let prod = repo.filenodes
                        .get_filenode(ctx, &path, node)
                        .await
                        .with_context(|| format!("while get prod filenode for cs {:?}", cs))?;
                    let backup = backup_repo.filenodes
                        .get_filenode(ctx, &path, node)
                        .await
                        .with_context(|| format!("while get backup filenode for cs {:?}", cs))?;
                    match (prod, backup) {
                        (FilenodeResult::Present(prod), FilenodeResult::Present(backup)) => {
                            assert!(prod == backup, "Differernt filenode for cs {} with path {:?}\nfilenode in prod repo {:?}\nfilenode in backup repo {:?}", cs, path, prod, backup);
                            Ok(())
                        }
                        (FilenodeResult::Disabled, FilenodeResult::Disabled) => Ok(()),
                        (_, _) => Err(anyhow!("filenodes results different for cs: {:?}", cs)),
                    }
                }
            })
            .try_buffer_unordered(100)
            .try_for_each(|_| async { Ok(()) })
            .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.
 */

mod derive;
mod mapping;

pub use derive::generate_all_filenodes;
pub use mapping::FilenodesOnlyPublic;
pub use mapping::PreparedRootFilenode;


/*
 * 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 anyhow::anyhow;
use anyhow::bail;
use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Loadable;
use context::CoreContext;
use derived_data::impl_bonsai_derived_via_manager;
use derived_data_manager::dependencies;
use derived_data_manager::BonsaiDerivable;
use derived_data_manager::DerivableType;
use derived_data_manager::DerivationContext;
use derived_data_service_if::types as thrift;
use filenodes::FilenodeInfo;
use filenodes::FilenodeResult;
use filenodes::PreparedFilenode;
use mercurial_derivation::MappedHgChangesetId;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::NULL_HASH;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::RepoPath;

use crate::derive::derive_filenodes;
use crate::derive::derive_filenodes_in_batch;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PreparedRootFilenode {
    pub filenode: HgFileNodeId,
    pub p1: Option<HgFileNodeId>,
    pub p2: Option<HgFileNodeId>,
    pub copyfrom: Option<(RepoPath, HgFileNodeId)>,
    pub linknode: HgChangesetId,
}

impl From<PreparedRootFilenode> for PreparedFilenode {
    fn from(root_filenode: PreparedRootFilenode) -> Self {
        let PreparedRootFilenode {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        } = root_filenode;

        PreparedFilenode {
            path: RepoPath::RootPath,
            info: FilenodeInfo {
                filenode,
                p1,
                p2,
                copyfrom,
                linknode,
            },
        }
    }
}

impl TryFrom<PreparedFilenode> for PreparedRootFilenode {
    type Error = Error;

    fn try_from(filenode: PreparedFilenode) -> Result<Self> {
        let PreparedFilenode { path, info } = filenode;

        let FilenodeInfo {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        } = info;

        if path != RepoPath::RootPath {
            return Err(format_err!("unexpected path for root filenode: {:?}", path));
        }
        Ok(Self {
            filenode,
            p1,
            p2,
            copyfrom,
            linknode,
        })
    }
}

/// Derives filenodes that are stores in Filenodes object (usually in a database).
/// Note: that should be derived only for public commits!
///
/// Filenodes might be disabled, in that case FilenodesOnlyPublic will always return
/// FilenodesOnlyPublic::Disabled enum variant.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FilenodesOnlyPublic {
    Present {
        root_filenode: Option<PreparedRootFilenode>,
    },
    Disabled,
}

#[async_trait]
impl BonsaiDerivable for FilenodesOnlyPublic {
    const VARIANT: DerivableType = DerivableType::FileNodes;

    type Dependencies = dependencies![MappedHgChangesetId];

    async fn derive_single(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsai: BonsaiChangeset,
        _parents: Vec<Self>,
    ) -> Result<Self> {
        derive_filenodes(ctx, derivation_ctx, bonsai).await
    }

    async fn derive_batch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsais: Vec<BonsaiChangeset>,
        _gap_size: Option<usize>,
    ) -> Result<HashMap<ChangesetId, Self>> {
        let filenodes = derivation_ctx.filenodes()?;
        let prepared = derive_filenodes_in_batch(ctx, derivation_ctx, bonsais).await?;
        let mut res = HashMap::with_capacity(prepared.len());
        for (cs_id, public_filenode, non_roots) in prepared.into_iter() {
            let filenode = match public_filenode {
                FilenodesOnlyPublic::Present { root_filenode } => match root_filenode {
                    Some(filenode) if !non_roots.is_empty() => {
                        match filenodes.add_filenodes(ctx, non_roots).await? {
                            FilenodeResult::Disabled => FilenodesOnlyPublic::Disabled,
                            FilenodeResult::Present(()) => FilenodesOnlyPublic::Present {
                                root_filenode: Some(filenode),
                            },
                        }
                    }
                    _ => FilenodesOnlyPublic::Present { root_filenode },
                },
                FilenodesOnlyPublic::Disabled => FilenodesOnlyPublic::Disabled,
            };
            res.insert(cs_id, filenode.clone());
        }
        Ok(res)
    }

    async fn store_mapping(
        self,
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        _changeset_id: ChangesetId,
    ) -> Result<()> {
        let root_filenode = match self {
            FilenodesOnlyPublic::Present { root_filenode } => match root_filenode {
                Some(root_filenode) => root_filenode,
                None => return Ok(()),
            },
            FilenodesOnlyPublic::Disabled => return Ok(()),
        };

        match derivation_ctx
            .filenodes()?
            .add_filenodes(ctx, vec![root_filenode.into()])
            .await?
        {
            FilenodeResult::Present(()) => Ok(()),
            FilenodeResult::Disabled => {
                // Filenodes got disabled just after we finished deriving them
                // but before we stored the mapping.  Ideally we would return
                // FilenodesMaybePublic::Disabled to the caller, but in this
                // very small window there is no way to do that. Instead we
                // must fail the request.
                bail!("filenodes were disabled after being successfully derived")
            }
        }
    }

    async fn fetch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<Option<Self>> {
        if tunables::tunables()
            .filenodes_disabled()
            .unwrap_or_default()
        {
            return Ok(Some(FilenodesOnlyPublic::Disabled));
        }
        let filenode_res = fetch_root_filenode(ctx, derivation_ctx, changeset_id).await?;
        let maybe_root_filenode = match filenode_res {
            FilenodeResult::Present(maybe_root_filenode) => maybe_root_filenode,
            FilenodeResult::Disabled => {
                return Ok(Some(FilenodesOnlyPublic::Disabled));
            }
        };

        Ok(
            maybe_root_filenode.map(move |filenode| FilenodesOnlyPublic::Present {
                root_filenode: Some(filenode),
            }),
        )
    }

    fn from_thrift(data: thrift::DerivedData) -> Result<Self> {
        match data {
            thrift::DerivedData::filenode(thrift::DerivedDataFilenode::filenode_present(
                filenode,
            )) => match filenode.root_filenode {
                None => Ok(FilenodesOnlyPublic::Present {
                    root_filenode: None,
                }),
                Some(data) => Ok(FilenodesOnlyPublic::Present {
                    root_filenode: Some(
                        FilenodeInfo::from_thrift(data)
                            .map(|info| PreparedFilenode {
                                path: RepoPath::RootPath,
                                info,
                            })?
                            .try_into()?,
                    ),
                }),
            },
            thrift::DerivedData::filenode(thrift::DerivedDataFilenode::filenode_disabled(_)) => {
                Ok(FilenodesOnlyPublic::Disabled)
            }
            _ => Err(anyhow!(
                "Can't convert {} from provided thrift::DerivedData",
                Self::NAME.to_string(),
            )),
        }
    }

    fn into_thrift(data: Self) -> Result<thrift::DerivedData> {
        match data {
            FilenodesOnlyPublic::Present {
                root_filenode: Some(root),
            } => Ok(thrift::DerivedData::filenode(
                thrift::DerivedDataFilenode::filenode_present(thrift::DerivedDataFilenodePresent {
                    root_filenode: Some(PreparedFilenode::from(root).info.into_thrift()),
                }),
            )),
            FilenodesOnlyPublic::Present {
                root_filenode: None,
            } => Ok(thrift::DerivedData::filenode(
                thrift::DerivedDataFilenode::filenode_present(thrift::DerivedDataFilenodePresent {
                    root_filenode: None,
                }),
            )),
            FilenodesOnlyPublic::Disabled => Ok(thrift::DerivedData::filenode(
                thrift::DerivedDataFilenode::filenode_disabled(thrift::DisabledFilenodes {}),
            )),
        }
    }
}

async fn fetch_root_filenode(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    cs_id: ChangesetId,
) -> Result<FilenodeResult<Option<PreparedRootFilenode>>> {
    // If hg changeset is not generated, then root filenode can't possible be generated
    // Check it and return None if hg changeset is not generated
    let maybe_hg_cs_id = derivation_ctx
        .bonsai_hg_mapping()?
        .get_hg_from_bonsai(ctx, cs_id)
        .await?;
    let hg_cs_id = if let Some(hg_cs_id) = maybe_hg_cs_id {
        hg_cs_id
    } else {
        return Ok(FilenodeResult::Present(None));
    };

    let mf_id = hg_cs_id
        .load(ctx, &derivation_ctx.blobstore())
        .await?
        .manifestid();

    // Special case null manifest id if we run into it
    let mf_id = mf_id.into_nodehash();
    if mf_id == NULL_HASH {
        Ok(FilenodeResult::Present(Some(PreparedRootFilenode {
            filenode: HgFileNodeId::new(NULL_HASH),
            p1: None,
            p2: None,
            copyfrom: None,
            linknode: HgChangesetId::new(NULL_HASH),
        })))
    } else {
        let filenode_res = derivation_ctx
            .filenodes()?
            .get_filenode(ctx, &RepoPath::RootPath, HgFileNodeId::new(mf_id))
            .await?;

        match filenode_res {
            FilenodeResult::Present(maybe_info) => {
                let info = maybe_info
                    .map(|info| {
                        PreparedRootFilenode::try_from(PreparedFilenode {
                            path: RepoPath::RootPath,
                            info,
                        })
                    })
                    .transpose()?;
                Ok(FilenodeResult::Present(info))
            }
            FilenodeResult::Disabled => Ok(FilenodeResult::Disabled),
        }
    }
}

impl_bonsai_derived_via_manager!(FilenodesOnlyPublic);


/*
 * 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 anyhow::Context;
use anyhow::Error;
use blobstore::Loadable;
use borrowed::borrowed;
use context::CoreContext;
use derived_data::batch::split_batch_in_linear_stacks;
use derived_data::batch::FileConflicts;
use derived_data::batch::StackItem;
use derived_data_manager::BonsaiDerivable;
use derived_data_manager::DerivationContext;
use futures::stream::FuturesOrdered;
use futures::stream::TryStreamExt;
use itertools::Itertools;
use mononoke_types::ChangesetId;
use mononoke_types::FsnodeId;
use stats::prelude::*;

use crate::derive::derive_fsnodes_stack;
use crate::RootFsnodeId;

define_stats! {
    prefix = "mononoke.derived_data.fsnodes";
    new_parallel: timeseries(Rate, Sum),
}

/// Derive a batch of fsnodes, potentially doing it faster than deriving fsnodes sequentially.
/// The primary purpose of this is to be used while backfilling fsnodes for a large repository.
///
/// The best results are achieved if a batch is a linear stack (i.e. no merges) of commits where batch[i-1] is a parent
/// of batch[i]. However if it's not the case then using derive_fsnode_in_batch shouldn't be much slower than a sequential
/// derivation of the same commits.
///
/// `derive_fsnode_in_batch` proceed in a few stages:
/// 1) Split `batch` in a a few linear stacks (there are certain rules about how it can be done, see `split_batch_in_linear_stacks` for more details)
/// 2) Stacks are processed one after another (i.e. we get benefits from parallel execution only if two commits are in the same stack)
/// 3) For each commit stack derive fsnode commits in parallel. This is done by calling `derive_fsnode()`
///    with parents of the first commit in the stack, and all bonsai file changes since first commit in the stack. See example below:
///
///   Stack:
///     Commit 1 - Added "file1" with content "A", has parent commit 0
///     Commit 2 - Added "file2" with content "B", has parent commit 1
///     Commit 3 - Modified "file1" with content "C", has parent commit 2
///
///   We make three derive_fsnode() calls in parallel with these parameters:
///      derive_fsnode([commit0], {"file1" => "A"})
///      derive_fsnode([commit0], {"file1" => "A", "file2" => "B"})
///      derive_fsnode([commit0], {"file1" => "C", "file2" => "B"})
///
/// So effectively we combine the changes from all commits in the stack. Note that it's not possible to do it for unodes
/// because unodes depend on the order of the changes.
///
/// Fsnode derivation can be cpu-bounded, and the speed up is achieved by spawning derivation on different
/// tokio tasks - this allows us to use more cpu.
pub async fn derive_fsnode_in_batch(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    batch: Vec<ChangesetId>,
    gap_size: Option<usize>,
) -> Result<HashMap<ChangesetId, RootFsnodeId>, Error> {
    let linear_stacks = split_batch_in_linear_stacks(
        ctx,
        derivation_ctx.blobstore(),
        batch,
        FileConflicts::ChangeDelete.into(),
    )
    .await?;
    let mut res: HashMap<ChangesetId, RootFsnodeId> = HashMap::new();
    for linear_stack in linear_stacks {
        // Fetch the parent fsnodes, either from a previous iteration of this
        // loop (which will have stored the mapping in `res`), or from the
        // main mapping, where they should already be derived.
        let parent_fsnodes = linear_stack
            .parents
            .into_iter()
            .map(|p| {
                borrowed!(res);
                async move {
                    anyhow::Result::<_>::Ok(
                        match res.get(&p) {
                            Some(fsnode_id) => fsnode_id.clone(),
                            None => derivation_ctx.fetch_dependency(ctx, p).await?,
                        }
                        .into_fsnode_id(),
                    )
                }
            })
            .collect::<FuturesOrdered<_>>()
            .try_collect::<Vec<_>>()
            .await?;

        STATS::new_parallel.add_value(1);
        new_batch_derivation(
            ctx,
            derivation_ctx,
            parent_fsnodes,
            gap_size,
            linear_stack.stack_items,
            &mut res,
        )
        .await?;
    }

    Ok(res)
}

pub async fn new_batch_derivation(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    parent_fsnode_manifests: Vec<FsnodeId>,
    gap_size: Option<usize>,
    file_changes: Vec<StackItem>,
    already_derived: &mut HashMap<ChangesetId, RootFsnodeId>,
) -> Result<(), Error> {
    if parent_fsnode_manifests.len() > 1 {
        // we can't derive stack for a merge commit,
        // so let's derive it without batching
        for item in file_changes {
            let bonsai = item.cs_id.load(ctx, derivation_ctx.blobstore()).await?;
            let parents = derivation_ctx
                .fetch_unknown_parents(ctx, Some(already_derived), &bonsai)
                .await?;
            let derived = RootFsnodeId::derive_single(ctx, derivation_ctx, bonsai, parents).await?;
            already_derived.insert(item.cs_id, derived);
        }
    } else {
        let first = file_changes.first().map(|item| item.cs_id);
        let last = file_changes.last().map(|item| item.cs_id);

        let file_changes: Vec<_> = match gap_size {
            Some(gap_size) => file_changes
                .into_iter()
                .chunks(gap_size)
                .into_iter()
                .filter_map(|chunk| {
                    // We are deriving with gaps - that means we are deriving
                    // just for the top commit out of `gap_size` linear stack
                    // of commits. To do that we need to combine all file changes
                    // for a given linear stack together
                    let mut combined_file_changes = BTreeMap::new();
                    let mut last_cs_id = None;
                    for item in chunk {
                        combined_file_changes.extend(item.per_commit_file_changes);
                        last_cs_id = Some(item.cs_id);
                    }
                    Some((last_cs_id?, combined_file_changes))
                })
                .collect(),
            None => file_changes
                .into_iter()
                .map(|item| (item.cs_id, item.per_commit_file_changes))
                .collect(),
        };

        let derived = derive_fsnodes_stack(
            ctx,
            derivation_ctx,
            file_changes,
            parent_fsnode_manifests.get(0).copied(),
        )
        .await
        .with_context(|| format!("failed deriving stack of {:?} to {:?}", first, last,))?;

        already_derived.extend(
            derived
                .into_iter()
                .map(|(csid, mf_id)| (csid, RootFsnodeId(mf_id))),
        );
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use bonsai_hg_mapping::BonsaiHgMapping;
    use bookmarks::Bookmarks;
    use changeset_fetcher::ChangesetFetcher;
    use changeset_fetcher::ChangesetFetcherArc;
    use changesets::Changesets;
    use derived_data_manager::BatchDeriveOptions;
    use fbinit::FacebookInit;
    use filestore::FilestoreConfig;
    use fixtures::Linear;
    use fixtures::TestRepoFixture;
    use futures::compat::Stream01CompatExt;
    use repo_blobstore::RepoBlobstore;
    use repo_derived_data::RepoDerivedData;
    use repo_derived_data::RepoDerivedDataRef;
    use revset::AncestorsNodeStream;
    use test_repo_factory::TestRepoFactory;
    use tests_utils::bookmark;
    use tests_utils::drawdag::create_from_dag;
    use tests_utils::resolve_cs_id;

    use super::*;

    #[facet::container]
    #[derive(Clone)]
    struct TestRepo {
        #[facet]
        bonsai_hg_mapping: dyn BonsaiHgMapping,
        #[facet]
        bookmarks: dyn Bookmarks,
        #[facet]
        repo_derived_data: RepoDerivedData,
        #[facet]
        changeset_fetcher: dyn ChangesetFetcher,
        #[facet]
        repo_blobstore: RepoBlobstore,
        #[facet]
        filestore_config: FilestoreConfig,
        #[facet]
        changesets: dyn Changesets,
    }

    #[fbinit::test]
    async fn batch_derive(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let new_batch = {
            let repo = Linear::getrepo(fb).await;
            let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;

            let mut cs_ids =
                AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), master_cs_id)
                    .compat()
                    .try_collect::<Vec<_>>()
                    .await?;
            cs_ids.reverse();

            let manager = repo.repo_derived_data().manager();

            manager
                .derive_exactly_batch::<RootFsnodeId>(
                    &ctx,
                    cs_ids,
                    BatchDeriveOptions::Parallel { gap_size: None },
                    None,
                )
                .await?;

            manager
                .fetch_derived::<RootFsnodeId>(&ctx, master_cs_id, None)
                .await?
                .unwrap()
                .into_fsnode_id()
        };

        let sequential = {
            let repo = Linear::getrepo(fb).await;
            let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;
            repo.repo_derived_data()
                .manager()
                .derive::<RootFsnodeId>(&ctx, master_cs_id, None)
                .await?
                .into_fsnode_id()
        };

        assert_eq!(new_batch, sequential);
        Ok(())
    }

    #[fbinit::test]
    async fn batch_derive_with_merge(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let new_batch = {
            let repo = repo_with_merge(&ctx).await?;
            let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;

            let mut cs_ids =
                AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), master_cs_id)
                    .compat()
                    .try_collect::<Vec<_>>()
                    .await?;
            cs_ids.reverse();

            let manager = repo.repo_derived_data().manager();
            manager
                .derive_exactly_batch::<RootFsnodeId>(
                    &ctx,
                    cs_ids,
                    BatchDeriveOptions::Parallel { gap_size: None },
                    None,
                )
                .await?;

            manager
                .fetch_derived::<RootFsnodeId>(&ctx, master_cs_id, None)
                .await?
                .unwrap()
                .into_fsnode_id()
        };

        let sequential = {
            let repo = repo_with_merge(&ctx).await?;
            let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;
            repo.repo_derived_data()
                .manager()
                .derive::<RootFsnodeId>(&ctx, master_cs_id, None)
                .await?
                .into_fsnode_id()
        };

        assert_eq!(new_batch, sequential);
        Ok(())
    }

    async fn repo_with_merge(ctx: &CoreContext) -> Result<TestRepo, Error> {
        let repo: TestRepo = TestRepoFactory::new(ctx.fb)?.build().await?;

        let commit_map = create_from_dag(
            ctx,
            &repo,
            r##"
            A-M
             /
            B
            "##,
        )
        .await?;

        let m = commit_map.get(&"M".to_string()).unwrap();
        bookmark(ctx, &repo, "master").set_to(*m).await?;

        Ok(repo)
    }

    #[fbinit::test]
    async fn batch_derive_with_gaps(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);

        let repo = Linear::getrepo(fb).await;
        let master_cs_id = resolve_cs_id(&ctx, &repo, "master").await?;
        let ddm = repo.repo_derived_data().manager();
        ddm.derive::<RootFsnodeId>(&ctx, master_cs_id, None)
            .await?
            .into_fsnode_id();

        for gap_size in 1..12 {
            let new_batch_with_gaps = derive_new_batch(fb, &ctx, gap_size).await?;
            for (cs_id, derived_with_gaps) in new_batch_with_gaps {
                let derived_sequential =
                    ddm.fetch_derived::<RootFsnodeId>(&ctx, cs_id, None).await?;
                assert_eq!(derived_with_gaps, derived_sequential.unwrap());
            }
        }

        Ok(())
    }

    async fn derive_new_batch(
        fb: FacebookInit,
        ctx: &CoreContext,
        gap_size: usize,
    ) -> Result<HashMap<ChangesetId, RootFsnodeId>, Error> {
        let repo = Linear::getrepo(fb).await;
        let master_cs_id = resolve_cs_id(ctx, &repo, "master").await?;

        let mut cs_ids =
            AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), master_cs_id)
                .compat()
                .try_collect::<Vec<_>>()
                .await?;
        cs_ids.reverse();
        let manager = repo.repo_derived_data().manager();
        manager
            .derive_exactly_batch::<RootFsnodeId>(
                ctx,
                cs_ids.clone(),
                BatchDeriveOptions::Parallel {
                    gap_size: Some(gap_size),
                },
                None,
            )
            .await?;

        let derived = cs_ids
            .chunks(gap_size)
            .filter_map(|chunk| chunk.last().cloned())
            .collect();

        let derived = manager
            .fetch_derived_batch::<RootFsnodeId>(ctx, derived, None)
            .await?;
        for cs_id in cs_ids {
            if !derived.contains_key(&cs_id) {
                assert_eq!(
                    manager
                        .fetch_derived::<RootFsnodeId>(ctx, cs_id, None)
                        .await?,
                    None
                );
            }
        }

        Ok(derived)
    }
}


/*
 * 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::sync::Arc;

use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use ascii::AsciiString;
use blobstore::Blobstore;
use blobstore::Loadable;
use borrowed::borrowed;
use cloned::cloned;
use context::CoreContext;
use derived_data_manager::DerivationContext;
use digest::Digest;
use filestore::get_metadata;
use filestore::FetchKey;
use futures::channel::mpsc;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::stream::FuturesOrdered;
use futures::stream::FuturesUnordered;
use futures::stream::TryStreamExt;
use manifest::derive_manifest_with_io_sender;
use manifest::derive_manifests_for_simple_stack_of_commits;
use manifest::Entry;
use manifest::LeafInfo;
use manifest::ManifestChanges;
use manifest::TreeInfo;
use mononoke_types::fsnode::Fsnode;
use mononoke_types::fsnode::FsnodeDirectory;
use mononoke_types::fsnode::FsnodeEntry;
use mononoke_types::fsnode::FsnodeFile;
use mononoke_types::fsnode::FsnodeSummary;
use mononoke_types::hash::Sha1;
use mononoke_types::hash::Sha256;
use mononoke_types::BlobstoreKey;
use mononoke_types::BlobstoreValue;
use mononoke_types::ChangesetId;
use mononoke_types::ContentId;
use mononoke_types::ContentMetadataV2;
use mononoke_types::FileType;
use mononoke_types::FsnodeId;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use sorted_vector_map::SortedVectorMap;

use crate::FsnodeDerivationError;

pub(crate) async fn derive_fsnodes_stack(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    file_changes: Vec<(ChangesetId, BTreeMap<MPath, Option<(ContentId, FileType)>>)>,
    parent: Option<FsnodeId>,
) -> Result<HashMap<ChangesetId, FsnodeId>, Error> {
    let blobstore = derivation_ctx.blobstore();

    let content_ids = file_changes
        .iter()
        .flat_map(|(_cs_id, per_commit_file_changes)| {
            per_commit_file_changes
                .iter()
                .filter_map(|(_mpath, content_id_and_file_type)| {
                    content_id_and_file_type.map(|(content_id, _file_type)| content_id)
                })
        })
        .collect();

    let prefetched_content_metadata =
        Arc::new(prefetch_content_metadata(ctx, &blobstore, content_ids).await?);

    let manifest_changes = file_changes
        .into_iter()
        .map(|(cs_id, file_changes)| ManifestChanges {
            cs_id,
            changes: file_changes.into_iter().collect(),
        })
        .collect::<Vec<_>>();

    let res = derive_manifests_for_simple_stack_of_commits(
        ctx.clone(),
        blobstore.clone(),
        parent,
        manifest_changes,
        {
            cloned!(blobstore, ctx, prefetched_content_metadata);
            move |tree_info, _cs_id| {
                cloned!(blobstore, ctx, prefetched_content_metadata);
                async move {
                    create_fsnode(
                        &ctx,
                        &blobstore,
                        None,
                        prefetched_content_metadata,
                        tree_info,
                    )
                    .await
                }
            }
        },
        |leaf_info, _cs_id| check_fsnode_leaf(leaf_info),
    )
    .await?;

    Ok(res.into_iter().collect())
}

/// Derives fsnodes for bonsai_changeset `cs_id` given parent fsnodes. Note
/// that `derive_manifest()` does a lot of the heavy lifting for us, and this
/// crate has to provide only functions to create a single fsnode, and check
/// that the leaf entries (which are `(ContentId, FileType)` pairs) are valid
/// during merges.
pub(crate) async fn derive_fsnode(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    parents: Vec<FsnodeId>,
    changes: Vec<(MPath, Option<(ContentId, FileType)>)>,
) -> Result<FsnodeId> {
    let blobstore = derivation_ctx.blobstore();
    let content_ids = changes
        .iter()
        .filter_map(|(_mpath, content_id_and_file_type)| {
            content_id_and_file_type.map(|(content_id, _file_type)| content_id)
        })
        .collect();

    let prefetched_content_metadata =
        Arc::new(prefetch_content_metadata(ctx, &blobstore, content_ids).await?);

    // We must box and store the derivation future, otherwise lifetime
    // analysis is unable to see that the blobstore lasts long enough.
    let derive_fut = derive_manifest_with_io_sender(
        ctx.clone(),
        blobstore.clone(),
        parents.clone(),
        changes,
        {
            cloned!(blobstore, ctx, prefetched_content_metadata);
            move |tree_info, sender| {
                cloned!(blobstore, ctx, prefetched_content_metadata);
                async move {
                    create_fsnode(
                        &ctx,
                        &blobstore,
                        Some(sender),
                        prefetched_content_metadata,
                        tree_info,
                    )
                    .await
                }
            }
        },
        |leaf_info, _sender| check_fsnode_leaf(leaf_info),
    )
    .boxed();
    let maybe_tree_id = derive_fut.await?;

    match maybe_tree_id {
        Some(tree_id) => Ok(tree_id),
        None => {
            // All files have been deleted, generate empty fsnode
            let tree_info = TreeInfo {
                path: None,
                parents,
                subentries: Default::default(),
            };
            let (_, tree_id) =
                create_fsnode(ctx, blobstore, None, prefetched_content_metadata, tree_info).await?;
            Ok(tree_id)
        }
    }
}

// Prefetch metadata for all content IDs introduced by a changeset.
pub async fn prefetch_content_metadata(
    ctx: &CoreContext,
    blobstore: &impl Blobstore,
    content_ids: HashSet<ContentId>,
) -> Result<HashMap<ContentId, ContentMetadataV2>> {
    content_ids
        .into_iter()
        .map({
            move |content_id| async move {
                match get_metadata(blobstore, ctx, &FetchKey::Canonical(content_id)).await? {
                    Some(metadata) => Ok(Some((content_id, metadata))),
                    None => Ok(None),
                }
            }
        })
        .collect::<FuturesUnordered<_>>()
        .try_filter_map(|maybe_metadata| async move { Ok(maybe_metadata) })
        .try_collect()
        .await
}

/// Collect all the subentries for a new fsnode, re-using entries the parent
/// fsnodes to avoid fetching too much.
async fn collect_fsnode_subentries(
    ctx: &CoreContext,
    blobstore: &impl Blobstore,
    prefetched_content_metadata: &HashMap<ContentId, ContentMetadataV2>,
    parents: Vec<FsnodeId>,
    subentries: BTreeMap<
        MPathElement,
        (
            Option<Option<FsnodeSummary>>,
            Entry<FsnodeId, (ContentId, FileType)>,
        ),
    >,
) -> Result<Vec<(MPathElement, FsnodeEntry)>> {
    // Load the parent fsnodes and collect their entries into a cache
    let mut file_cache = HashMap::new();
    let mut dir_cache = HashMap::new();
    let mut parent_fsnodes = parents
        .into_iter()
        .map({
            move |fsnode_id| async move {
                fsnode_id
                    .load(ctx, blobstore)
                    .await
                    .context(FsnodeDerivationError::MissingParent(fsnode_id))
            }
        })
        .collect::<FuturesUnordered<_>>();
    while let Some(parent_fsnode) = parent_fsnodes.try_next().await? {
        for (_elem, entry) in parent_fsnode.list() {
            match entry {
                FsnodeEntry::File(file) => {
                    file_cache
                        .entry((*file.content_id(), *file.file_type()))
                        .or_insert_with(|| file.clone());
                }
                FsnodeEntry::Directory(dir) => {
                    dir_cache.entry(*dir.id()).or_insert_with(|| dir.clone());
                }
            }
        }
    }

    // Find (from the traversal or the cache) or fetch (from the blobstore)
    // the `FsnodeEntry` for each of the subentries.
    borrowed!(file_cache, dir_cache);
    subentries
        .into_iter()
        .map(move |(elem, (summary, entry))| {
            async move {
                match entry {
                    Entry::Tree(fsnode_id) => {
                        if let Some(Some(summary)) = summary {
                            // The subdirectory was just created. Use the
                            // summary we just calculated.
                            Ok((
                                elem.clone(),
                                FsnodeEntry::Directory(FsnodeDirectory::new(fsnode_id, summary)),
                            ))
                        } else if let Some(entry) = dir_cache.get(&fsnode_id) {
                            // The subdirectory was already in this
                            // directory. Use the cached entry.
                            Ok((elem.clone(), FsnodeEntry::Directory(entry.clone())))
                        } else {
                            // Some other directory is being used. Fetch its
                            // summary from the blobstore.
                            let fsnode = fsnode_id.load(ctx, blobstore).await.with_context({
                                || {
                                    FsnodeDerivationError::MissingSubentry(
                                        String::from_utf8_lossy(elem.as_ref()).to_string(),
                                        fsnode_id,
                                    )
                                }
                            })?;

                            let entry = FsnodeEntry::Directory(FsnodeDirectory::new(
                                fsnode_id,
                                fsnode.summary().clone(),
                            ));
                            Ok((elem.clone(), entry))
                        }
                    }
                    Entry::Leaf(content_id_and_file_type) => {
                        if let Some(entry) = file_cache.get(&content_id_and_file_type) {
                            // The file was already in this directory. Use
                            // the cached entry.
                            Ok((elem.clone(), FsnodeEntry::File(entry.clone())))
                        } else {
                            // Some other file is being used. Use the
                            // metadata we prefetched to create a new entry.
                            let (content_id, file_type) = content_id_and_file_type.clone();
                            if let Some(metadata) = prefetched_content_metadata.get(&content_id) {
                                let entry = FsnodeEntry::File(FsnodeFile::new(
                                    content_id,
                                    file_type,
                                    metadata.total_size,
                                    metadata.sha1,
                                    metadata.sha256,
                                ));
                                Ok((elem.clone(), entry))
                            } else {
                                Err(FsnodeDerivationError::MissingContent(content_id).into())
                            }
                        }
                    }
                }
            }
        })
        .collect::<FuturesOrdered<_>>()
        .try_collect()
        .await
}

/// Create a new fsnode for the tree described by `tree_info`.
async fn create_fsnode(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    sender: Option<mpsc::UnboundedSender<BoxFuture<'static, Result<(), Error>>>>,
    prefetched_content_metadata: Arc<HashMap<ContentId, ContentMetadataV2>>,
    tree_info: TreeInfo<FsnodeId, (ContentId, FileType), Option<FsnodeSummary>>,
) -> Result<(Option<FsnodeSummary>, FsnodeId)> {
    let entries = collect_fsnode_subentries(
        ctx,
        &blobstore,
        prefetched_content_metadata.as_ref(),
        tree_info.parents,
        tree_info.subentries,
    )
    .await?;

    // Build a summary of the entries and store it as the new fsnode.
    let entries: SortedVectorMap<_, _> = entries.into_iter().collect();
    let simple_format_sha1 = {
        let digest = generate_simple_format_digest(
            sha1::Sha1::new(),
            &entries,
            |fsnode_file| fsnode_file.content_sha1().to_hex(),
            |fsnode_dir| fsnode_dir.summary().simple_format_sha1.to_hex(),
        );
        let bytes = digest.finalize().into();
        Sha1::from_byte_array(bytes)
    };
    let simple_format_sha256 = {
        let digest = generate_simple_format_digest(
            sha2::Sha256::new(),
            &entries,
            |fsnode_file| fsnode_file.content_sha256().to_hex(),
            |fsnode_dir| fsnode_dir.summary().simple_format_sha256.to_hex(),
        );
        let bytes = digest.finalize().into();
        Sha256::from_byte_array(bytes)
    };
    let mut summary = FsnodeSummary {
        simple_format_sha1,
        simple_format_sha256,
        child_files_count: 0,
        child_files_total_size: 0,
        child_dirs_count: 0,
        descendant_files_count: 0,
        descendant_files_total_size: 0,
    };
    for (_elem, entry) in entries.iter() {
        match entry {
            FsnodeEntry::File(fsnode_file) => {
                let size = fsnode_file.size();
                summary.child_files_count += 1;
                summary.child_files_total_size += size;
                summary.descendant_files_count += 1;
                summary.descendant_files_total_size += size;
            }
            FsnodeEntry::Directory(fsnode_dir) => {
                let subdir_summary = fsnode_dir.summary();
                summary.child_dirs_count += 1;
                summary.descendant_files_count += subdir_summary.descendant_files_count;
                summary.descendant_files_total_size += subdir_summary.descendant_files_total_size;
            }
        }
    }
    let fsnode = Fsnode::new(entries, summary.clone());
    let blob = fsnode.into_blob();
    let fsnode_id = *blob.id();
    let key = fsnode_id.blobstore_key();
    cloned!(blobstore, ctx);
    let f = async move { blobstore.put(&ctx, key, blob.into()).await };

    match sender {
        Some(sender) => sender
            .unbounded_send(f.boxed())
            .map_err(|err| format_err!("failed to send fsnode future {}", err))?,
        None => f.await?,
    };
    Ok((Some(summary), fsnode_id))
}

/// Generate the simple format hash for a directory. See
/// `mononoke_types/src/fsnodes.rs` for a description of what the simple
/// format is.
fn generate_simple_format_digest<H, F, D>(
    mut digest: H,
    dir: &SortedVectorMap<MPathElement, FsnodeEntry>,
    get_file_hash: F,
    get_dir_hash: D,
) -> H
where
    H: Digest,
    F: Fn(&FsnodeFile) -> AsciiString,
    D: Fn(&FsnodeDirectory) -> AsciiString,
{
    for (elem, entry) in dir.iter() {
        match entry {
            FsnodeEntry::File(file) => {
                digest.update(get_file_hash(file).as_bytes());
                digest.update(match file.file_type() {
                    FileType::Regular => b" file ",
                    FileType::Executable => b" exec ",
                    FileType::Symlink => b" link ",
                });
            }
            FsnodeEntry::Directory(dir) => {
                digest.update(get_dir_hash(dir).as_bytes());
                digest.update(b" tree ");
            }
        }
        digest.update(elem.as_ref());
        digest.update(b"\0");
    }
    digest
}

/// There are no leaves stored for fsnodes, however we still need to check
/// that any merge operations have resulted in valid fsnodes, where, for each
/// file, either all the parents have the same file contents, or the
/// changeset includes a change for that file.
async fn check_fsnode_leaf(
    leaf_info: LeafInfo<(ContentId, FileType), (ContentId, FileType)>,
) -> Result<(Option<FsnodeSummary>, (ContentId, FileType))> {
    if let Some(content_id_and_file_type) = leaf_info.leaf {
        Ok((None, content_id_and_file_type))
    } else {
        // This bonsai changeset is a merge. If all content IDs and file
        // types match for this file, then the content ID is valid. Check
        // this is so.
        if leaf_info.parents.len() < 2 {
            return Err(FsnodeDerivationError::InvalidBonsai(
                "no change is provided, but file has only one parent".to_string(),
            )
            .into());
        }
        let mut iter = leaf_info.parents.into_iter();
        let fsnode_file = iter.next().and_then(|first_elem| {
            if iter.all(|next_elem| next_elem == first_elem) {
                Some(first_elem)
            } else {
                None
            }
        });
        if let Some(fsnode_file) = fsnode_file {
            Ok((None, fsnode_file))
        } else {
            Err(FsnodeDerivationError::InvalidBonsai(
                "no change is provided, but file content or type is different".to_string(),
            )
            .into())
        }
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use derived_data_test_utils::bonsai_changeset_from_hg;
    use derived_data_test_utils::iterate_all_manifest_entries;
    use fbinit::FacebookInit;
    use fixtures::Linear;
    use fixtures::ManyFilesDirs;
    use fixtures::TestRepoFixture;
    use repo_blobstore::RepoBlobstoreRef;
    use repo_derived_data::RepoDerivedDataRef;
    use tokio::runtime::Runtime;

    use super::*;
    use crate::mapping::get_file_changes;

    #[fbinit::test]
    fn flat_linear_test(fb: FacebookInit) {
        let runtime = Runtime::new().unwrap();
        let repo = runtime.block_on(Linear::getrepo(fb));
        let derivation_ctx = repo.repo_derived_data().manager().derivation_context(None);

        let ctx = CoreContext::test_mock(fb);
        let parent_fsnode_id = {
            let parent_hg_cs = "2d7d4ba9ce0a6ffd222de7785b249ead9c51c536";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, parent_hg_cs))
                .unwrap();

            let f = derive_fsnode(&ctx, &derivation_ctx, vec![], get_file_changes(&bcs));

            let root_fsnode_id = runtime.block_on(f).unwrap();

            // Make sure it's saved in the blobstore.
            let root_fsnode = runtime
                .block_on(root_fsnode_id.load(&ctx, repo.repo_blobstore()))
                .unwrap();

            // Make sure the fsnodes describe the full manifest.
            let all_fsnodes: BTreeMap<_, _> = runtime
                .block_on(
                    iterate_all_manifest_entries(&ctx, &repo, Entry::Tree(root_fsnode_id))
                        .try_collect(),
                )
                .unwrap();
            assert_eq!(
                all_fsnodes.keys().collect::<Vec<_>>(),
                vec![
                    &None,
                    &Some(MPath::new("1").unwrap()),
                    &Some(MPath::new("files").unwrap())
                ]
            );

            // Make sure the root fsnode is correct.
            assert_eq!(
                root_fsnode.summary(),
                &FsnodeSummary {
                    simple_format_sha1: Sha1::from_str("b23c1ffbd30ffa92942ec67748ea28084bbbb100")
                        .unwrap(),
                    simple_format_sha256: Sha256::from_str(
                        "6639d6877355e0b26b135cc0bf90b8cb901c85f43e71fbf4e8a5202ff1a7f868"
                    )
                    .unwrap(),
                    child_files_count: 2,
                    child_files_total_size: 4,
                    child_dirs_count: 0,
                    descendant_files_count: 2,
                    descendant_files_total_size: 4,
                }
            );
            root_fsnode_id
        };

        {
            let child_hg_cs = "3e0e761030db6e479a7fb58b12881883f9f8c63f";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, child_hg_cs))
                .unwrap();

            let f = derive_fsnode(
                &ctx,
                &derivation_ctx,
                vec![parent_fsnode_id.clone()],
                get_file_changes(&bcs),
            );

            let root_fsnode_id = runtime.block_on(f).unwrap();

            // Make sure it's saved in the blobstore
            let root_fsnode = runtime
                .block_on(root_fsnode_id.load(&ctx, repo.repo_blobstore()))
                .unwrap();

            // Make sure the fsnodes describe the full manifest.
            let all_fsnodes: BTreeMap<_, _> = runtime
                .block_on(
                    iterate_all_manifest_entries(&ctx, &repo, Entry::Tree(root_fsnode_id))
                        .try_collect(),
                )
                .unwrap();
            assert_eq!(
                all_fsnodes.keys().collect::<Vec<_>>(),
                vec![
                    &None,
                    &Some(MPath::new("1").unwrap()),
                    &Some(MPath::new("2").unwrap()),
                    &Some(MPath::new("files").unwrap())
                ]
            );

            // Make sure the root fsnode is correct.
            assert_eq!(
                root_fsnode.summary(),
                &FsnodeSummary {
                    simple_format_sha1: Sha1::from_str("e1c394ce3584f6c69680807b93997e46099e3f84")
                        .unwrap(),
                    simple_format_sha256: Sha256::from_str(
                        "bae96aa31c752e678f22bd1d051ba0a2def7992cf6857c1a1736fd6517106c08"
                    )
                    .unwrap(),
                    child_files_count: 3,
                    child_files_total_size: 8,
                    child_dirs_count: 0,
                    descendant_files_count: 3,
                    descendant_files_total_size: 8,
                }
            );
        }
    }

    #[fbinit::test]
    fn nested_directories_test(fb: FacebookInit) {
        let runtime = Runtime::new().unwrap();
        let repo = runtime.block_on(ManyFilesDirs::getrepo(fb));
        let derivation_ctx = repo.repo_derived_data().manager().derivation_context(None);

        let ctx = CoreContext::test_mock(fb);

        // Derive fsnodes for the first two commits.  We will look at commit 3 and 4.
        let parent_fsnode_id = {
            let parent_hg_cs = "5a28e25f924a5d209b82ce0713d8d83e68982bc8";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, parent_hg_cs))
                .unwrap();
            let f = derive_fsnode(&ctx, &derivation_ctx, vec![], get_file_changes(&bcs));
            runtime.block_on(f).unwrap()
        };

        let parent_fsnode_id = {
            let parent_hg_cs = "2f866e7e549760934e31bf0420a873f65100ad63";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, parent_hg_cs))
                .unwrap();
            let f = derive_fsnode(
                &ctx,
                &derivation_ctx,
                vec![parent_fsnode_id.clone()],
                get_file_changes(&bcs),
            );
            runtime.block_on(f).unwrap()
        };

        let parent_fsnode_id = {
            let parent_hg_cs = "d261bc7900818dea7c86935b3fb17a33b2e3a6b4";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, parent_hg_cs))
                .unwrap();

            let f = derive_fsnode(
                &ctx,
                &derivation_ctx,
                vec![parent_fsnode_id.clone()],
                get_file_changes(&bcs),
            );

            let root_fsnode_id = runtime.block_on(f).unwrap();

            // Make sure it's saved in the blobstore.
            let root_fsnode = runtime
                .block_on(root_fsnode_id.load(&ctx, repo.repo_blobstore()))
                .unwrap();

            // Make sure the fsnodes describe the full manifest.
            let all_fsnodes: BTreeMap<_, _> = runtime
                .block_on(
                    iterate_all_manifest_entries(&ctx, &repo, Entry::Tree(root_fsnode_id))
                        .try_collect(),
                )
                .unwrap();
            assert_eq!(
                all_fsnodes.keys().collect::<Vec<_>>(),
                vec![
                    &None,
                    &Some(MPath::new("1").unwrap()),
                    &Some(MPath::new("2").unwrap()),
                    &Some(MPath::new("dir1").unwrap()),
                    &Some(MPath::new("dir1/file_1_in_dir1").unwrap()),
                    &Some(MPath::new("dir1/file_2_in_dir1").unwrap()),
                    &Some(MPath::new("dir1/subdir1").unwrap()),
                    &Some(MPath::new("dir1/subdir1/file_1").unwrap()),
                    &Some(MPath::new("dir1/subdir1/subsubdir1").unwrap()),
                    &Some(MPath::new("dir1/subdir1/subsubdir1/file_1").unwrap()),
                    &Some(MPath::new("dir1/subdir1/subsubdir2").unwrap()),
                    &Some(MPath::new("dir1/subdir1/subsubdir2/file_1").unwrap()),
                    &Some(MPath::new("dir1/subdir1/subsubdir2/file_2").unwrap()),
                    &Some(MPath::new("dir2").unwrap()),
                    &Some(MPath::new("dir2/file_1_in_dir2").unwrap()),
                ]
            );

            // Make sure the root fsnode is correct.  Test the simple format hashes are as expected.
            let simple_format_sha1 = {
                let text = concat!(
                    "e5fa44f2b31c1fb553b6021e7360d07d5d91ff5e file 1\0",
                    "7448d8798a4380162d4b56f9b452e2f6f9e24e7a file 2\0",
                    "be0dface7b74d8a69b39fb4691ef9eee36077ede tree dir1\0",
                    "ad02b5a5f778d9ad6afd42fcc8e0b889254b5215 tree dir2\0",
                );
                let mut digest = sha1::Sha1::new();
                digest.update(text);
                let bytes = digest.finalize().into();
                Sha1::from_byte_array(bytes)
            };
            let simple_format_sha256 = {
                let text = concat!(
                    "4355a46b19d348dc2f57c046f8ef63d4538ebb936000f3c9ee954a27460dd865 file 1\0",
                    "53c234e5e8472b6ac51c1ae1cab3fe06fad053beb8ebfd8977b010655bfdd3c3 file 2\0",
                    "eebf7e41f348db6b31c11fe7adf577bd0951300436a1fd37e53d628127b3517e tree dir1\0",
                    "583c3d388efb78eb9dec46626662d6657bb53706c1ee10770c0fb3e859bd36e1 tree dir2\0",
                );
                let mut digest = sha2::Sha256::new();
                digest.update(text);
                let bytes = digest.finalize().into();
                Sha256::from_byte_array(bytes)
            };

            assert_eq!(
                root_fsnode.summary(),
                &FsnodeSummary {
                    simple_format_sha1,
                    simple_format_sha256,
                    child_files_count: 2,
                    child_files_total_size: 4,
                    child_dirs_count: 2,
                    descendant_files_count: 9,
                    descendant_files_total_size: 67,
                }
            );

            // Check one of the deeper fsnodes.
            let deep_fsnode_id = match all_fsnodes.get(&Some(MPath::new("dir1/subdir1").unwrap())) {
                Some(Entry::Tree(fsnode_id)) => fsnode_id,
                _ => panic!("dir1/subdir1 fsnode should be a tree"),
            };
            let deep_fsnode = runtime
                .block_on(deep_fsnode_id.load(&ctx, repo.repo_blobstore()))
                .unwrap();
            let deep_fsnode_entries: Vec<_> = deep_fsnode.list().collect();
            assert_eq!(
                deep_fsnode_entries,
                vec![
                    (
                        &MPathElement::new(b"file_1".to_vec()).unwrap(),
                        &FsnodeEntry::File(FsnodeFile::new(
                            ContentId::from_str(
                                "271119c630193464c55cf7feb72c16dd18b953e65ee85047d4d956381cdb96d9",
                            )
                            .unwrap(),
                            FileType::Regular,
                            9,
                            Sha1::from_str("eba7b799b1e587506677c108eaa51ca273ddcfdc").unwrap(),
                            Sha256::from_str(
                                "8bfab644c313b0227b1862786426697a9b5283eea8eb6066cc8d8e134ce0daa6",
                            )
                            .unwrap(),
                        )),
                    ),
                    (
                        &MPathElement::new(b"subsubdir1".to_vec()).unwrap(),
                        &FsnodeEntry::Directory(FsnodeDirectory::new(
                            FsnodeId::from_str(
                                "a8deafcaa9fdafad84447d754b1e24140f1113e68b3f7ec17c55a4fd9c90fdc7",
                            )
                            .unwrap(),
                            FsnodeSummary {
                                simple_format_sha1: Sha1::from_str(
                                    "b58f168a508c4388f496b21acee468b13b89ac50",
                                )
                                .unwrap(),
                                simple_format_sha256: Sha256::from_str(
                                    "dc19c9a55b59d97dbb2912e5acb1bad17c1aa2835c4cf288f2d47bb0b0b539bc",
                                )
                                .unwrap(),
                                child_files_count: 1,
                                child_files_total_size: 9,
                                child_dirs_count: 0,
                                descendant_files_count: 1,
                                descendant_files_total_size: 9,
                            },
                        )),
                    ),
                    (
                        &MPathElement::new(b"subsubdir2".to_vec()).unwrap(),
                        &FsnodeEntry::Directory(FsnodeDirectory::new(
                            FsnodeId::from_str(
                                "3daadf48916e4b25e28e21a15de4352188154fcaa749207eb00fb72dbd812343",
                            )
                            .unwrap(),
                            FsnodeSummary {
                                simple_format_sha1: Sha1::from_str(
                                    "89babd6eab901598b3df0a0c3318b83bd2edfa4b",
                                )
                                .unwrap(),
                                simple_format_sha256: Sha256::from_str(
                                    "5fa3074b0a96f8ac8f5ab31266d1ce95b1c84a35e8b3f4b8108385881e9e9f15",
                                )
                                .unwrap(),
                                child_files_count: 2,
                                child_files_total_size: 18,
                                child_dirs_count: 0,
                                descendant_files_count: 2,
                                descendant_files_total_size: 18,
                            },
                        )),
                    ),
                ]
            );
            assert_eq!(
                deep_fsnode.summary(),
                &FsnodeSummary {
                    simple_format_sha1: Sha1::from_str("56daa486e90686b5ead6f3ce360a6cbb7b73f23f")
                        .unwrap(),
                    simple_format_sha256: Sha256::from_str(
                        "8c24cf135f74933775b0ac1453c127c987874beb6fde5454ccf9c6c8c726e032"
                    )
                    .unwrap(),
                    child_files_count: 1,
                    child_files_total_size: 9,
                    child_dirs_count: 2,
                    descendant_files_count: 4,
                    descendant_files_total_size: 36,
                }
            );

            root_fsnode_id
        };

        {
            let child_hg_cs = "051946ed218061e925fb120dac02634f9ad40ae2";
            let (_bcs_id, bcs) = runtime
                .block_on(bonsai_changeset_from_hg(&ctx, &repo, child_hg_cs))
                .unwrap();

            let f = derive_fsnode(
                &ctx,
                &derivation_ctx,
                vec![parent_fsnode_id.clone()],
                get_file_changes(&bcs),
            );

            let root_fsnode_id = runtime.block_on(f).unwrap();

            // Make sure it's saved in the blobstore
            let root_fsnode = runtime
                .block_on(root_fsnode_id.load(&ctx, repo.repo_blobstore()))
                .unwrap();

            // Make sure the fsnodes describe the full manifest.
            let all_fsnodes: BTreeMap<_, _> = runtime
                .block_on(
                    iterate_all_manifest_entries(&ctx, &repo, Entry::Tree(root_fsnode_id))
                        .try_collect(),
                )
                .unwrap();
            assert_eq!(
                all_fsnodes.keys().collect::<Vec<_>>(),
                vec![
                    &None,
                    &Some(MPath::new("1").unwrap()),
                    &Some(MPath::new("2").unwrap()),
                    &Some(MPath::new("dir1").unwrap()),
                    &Some(MPath::new("dir2").unwrap()),
                    &Some(MPath::new("dir2/file_1_in_dir2").unwrap()),
                ]
            );

            // Make sure the root fsnode is correct.
            assert_eq!(
                root_fsnode.summary(),
                &FsnodeSummary {
                    simple_format_sha1: Sha1::from_str("753b41c7bf23bda7eabb71095867c2ddf3e485df")
                        .unwrap(),
                    simple_format_sha256: Sha256::from_str(
                        "d86a29b0e92bcbff179e20a27fde18aacedc225a6c4884a7760d2880794c9e51"
                    )
                    .unwrap(),
                    child_files_count: 3,
                    child_files_total_size: 16,
                    child_dirs_count: 1,
                    descendant_files_count: 4,
                    descendant_files_total_size: 25,
                }
            );
        }
    }
}


/*
 * 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 mononoke_types::ContentId;
use mononoke_types::FsnodeId;
use thiserror::Error;

mod batch;
mod derive;
mod mapping;

pub use derive::prefetch_content_metadata;
pub use mapping::format_key;
pub use mapping::RootFsnodeId;

#[derive(Debug, Error)]
pub enum FsnodeDerivationError {
    #[error("Invalid bonsai changeset: {0}")]
    InvalidBonsai(String),
    #[error("Missing content: {0}")]
    MissingContent(ContentId),
    #[error("Missing fsnode parent: {0}")]
    MissingParent(FsnodeId),
    #[error("Missing fsnode subentry for '{0}': {1}")]
    MissingSubentry(String, FsnodeId),
}


/*
 * 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 anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use bytes::Bytes;
use context::CoreContext;
use derived_data::impl_bonsai_derived_via_manager;
use derived_data_manager::dependencies;
use derived_data_manager::BonsaiDerivable;
use derived_data_manager::DerivableType;
use derived_data_manager::DerivationContext;
use derived_data_service_if::types as thrift;
use mononoke_types::BlobstoreBytes;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::ContentId;
use mononoke_types::FileType;
use mononoke_types::FsnodeId;
use mononoke_types::MPath;

use crate::batch::derive_fsnode_in_batch;
use crate::derive::derive_fsnode;

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct RootFsnodeId(pub(crate) FsnodeId);

impl RootFsnodeId {
    pub fn fsnode_id(&self) -> &FsnodeId {
        &self.0
    }
    pub fn into_fsnode_id(self) -> FsnodeId {
        self.0
    }
}

impl TryFrom<BlobstoreBytes> for RootFsnodeId {
    type Error = Error;

    fn try_from(blob_bytes: BlobstoreBytes) -> Result<Self> {
        FsnodeId::from_bytes(&blob_bytes.into_bytes()).map(RootFsnodeId)
    }
}

impl TryFrom<BlobstoreGetData> for RootFsnodeId {
    type Error = Error;

    fn try_from(blob_get_data: BlobstoreGetData) -> Result<Self> {
        blob_get_data.into_bytes().try_into()
    }
}

impl From<RootFsnodeId> for BlobstoreBytes {
    fn from(root_fsnode_id: RootFsnodeId) -> Self {
        BlobstoreBytes::from_bytes(Bytes::copy_from_slice(root_fsnode_id.0.blake2().as_ref()))
    }
}

pub fn format_key(derivation_ctx: &DerivationContext, changeset_id: ChangesetId) -> String {
    let root_prefix = "derived_root_fsnode.";
    let key_prefix = derivation_ctx.mapping_key_prefix::<RootFsnodeId>();
    format!("{}{}{}", root_prefix, key_prefix, changeset_id)
}

#[async_trait]
impl BonsaiDerivable for RootFsnodeId {
    const VARIANT: DerivableType = DerivableType::Fsnodes;

    type Dependencies = dependencies![];

    async fn derive_single(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsai: BonsaiChangeset,
        parents: Vec<Self>,
    ) -> Result<Self, Error> {
        let fsnode_id = derive_fsnode(
            ctx,
            derivation_ctx,
            parents
                .into_iter()
                .map(|root_fsnode_id| root_fsnode_id.into_fsnode_id())
                .collect(),
            get_file_changes(&bonsai),
        )
        .await?;
        Ok(RootFsnodeId(fsnode_id))
    }

    async fn derive_batch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsais: Vec<BonsaiChangeset>,
        gap_size: Option<usize>,
    ) -> Result<HashMap<ChangesetId, Self>> {
        derive_fsnode_in_batch(
            ctx,
            derivation_ctx,
            bonsais.into_iter().map(|b| b.get_changeset_id()).collect(),
            gap_size,
        )
        .await
    }

    async fn store_mapping(
        self,
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<()> {
        let key = format_key(derivation_ctx, changeset_id);
        derivation_ctx.blobstore().put(ctx, key, self.into()).await
    }

    async fn fetch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<Option<Self>> {
        let key = format_key(derivation_ctx, changeset_id);
        Ok(derivation_ctx
            .blobstore()
            .get(ctx, &key)
            .await?
            .map(TryInto::try_into)
            .transpose()?)
    }

    fn from_thrift(data: thrift::DerivedData) -> Result<Self> {
        if let thrift::DerivedData::fsnode(thrift::DerivedDataFsnode::root_fsnode_id(id)) = data {
            FsnodeId::from_thrift(id).map(Self)
        } else {
            Err(anyhow!(
                "Can't convert {} from provided thrift::DerivedData",
                Self::NAME.to_string(),
            ))
        }
    }

    fn into_thrift(data: Self) -> Result<thrift::DerivedData> {
        Ok(thrift::DerivedData::fsnode(
            thrift::DerivedDataFsnode::root_fsnode_id(data.into_fsnode_id().into_thrift()),
        ))
    }
}

impl_bonsai_derived_via_manager!(RootFsnodeId);

pub(crate) fn get_file_changes(
    bcs: &BonsaiChangeset,
) -> Vec<(MPath, Option<(ContentId, FileType)>)> {
    bcs.file_changes()
        .map(|(mpath, file_change)| {
            (
                mpath.clone(),
                file_change
                    .simplify()
                    .map(|bc| (bc.content_id(), bc.file_type())),
            )
        })
        .collect()
}

#[cfg(test)]
mod test {
    use blobstore::Loadable;
    use bookmarks::BookmarkKey;
    use bookmarks::BookmarksRef;
    use borrowed::borrowed;
    use changeset_fetcher::ChangesetFetcherArc;
    use derived_data::BonsaiDerived;
    use derived_data_test_utils::iterate_all_manifest_entries;
    use fbinit::FacebookInit;
    use fixtures::BranchEven;
    use fixtures::BranchUneven;
    use fixtures::BranchWide;
    use fixtures::Linear;
    use fixtures::ManyDiamonds;
    use fixtures::ManyFilesDirs;
    use fixtures::MergeEven;
    use fixtures::MergeUneven;
    use fixtures::TestRepoFixture;
    use fixtures::UnsharedMergeEven;
    use fixtures::UnsharedMergeUneven;
    use futures::compat::Stream01CompatExt;
    use futures::future::Future;
    use futures::stream::Stream;
    use futures::stream::TryStreamExt;
    use futures::try_join;
    use manifest::Entry;
    use mercurial_derivation::DeriveHgChangeset;
    use mercurial_types::HgChangesetId;
    use mercurial_types::HgManifestId;
    use repo_blobstore::RepoBlobstoreRef;
    use repo_derived_data::RepoDerivedDataRef;
    use revset::AncestorsNodeStream;
    use tokio::runtime::Runtime;

    use super::*;

    async fn fetch_manifest_by_cs_id(
        ctx: &CoreContext,
        repo: &impl RepoBlobstoreRef,
        hg_cs_id: HgChangesetId,
    ) -> Result<HgManifestId> {
        Ok(hg_cs_id
            .load(ctx, repo.repo_blobstore())
            .await?
            .manifestid())
    }

    async fn verify_fsnode(
        ctx: &CoreContext,
        repo: &(impl RepoBlobstoreRef + RepoDerivedDataRef + Send + Sync),
        bcs_id: ChangesetId,
        hg_cs_id: HgChangesetId,
    ) -> Result<()> {
        let root_fsnode_id = RootFsnodeId::derive(ctx, repo, bcs_id)
            .await?
            .into_fsnode_id();

        let fsnode_entries = iterate_all_manifest_entries(ctx, repo, Entry::Tree(root_fsnode_id))
            .map_ok(|(path, _)| path)
            .try_collect::<Vec<_>>();

        let root_mf_id = fetch_manifest_by_cs_id(ctx, repo, hg_cs_id).await?;

        let filenode_entries = iterate_all_manifest_entries(ctx, repo, Entry::Tree(root_mf_id))
            .map_ok(|(path, _)| path)
            .try_collect::<Vec<_>>();

        let (mut fsnode_entries, mut filenode_entries) =
            try_join!(fsnode_entries, filenode_entries)?;
        fsnode_entries.sort();
        filenode_entries.sort();
        assert_eq!(fsnode_entries, filenode_entries);
        Ok(())
    }

    async fn all_commits<'a>(
        ctx: &'a CoreContext,
        repo: &'a (impl BookmarksRef + ChangesetFetcherArc + RepoDerivedDataRef + Send + Sync),
    ) -> Result<impl Stream<Item = Result<(ChangesetId, HgChangesetId)>> + 'a> {
        let master_book = BookmarkKey::new("master").unwrap();
        let bcs_id = repo
            .bookmarks()
            .get(ctx.clone(), &master_book)
            .await?
            .unwrap();

        Ok(
            AncestorsNodeStream::new(ctx.clone(), &repo.changeset_fetcher_arc(), bcs_id.clone())
                .compat()
                .and_then(move |new_bcs_id| async move {
                    let hg_cs_id = repo.derive_hg_changeset(ctx, new_bcs_id).await?;
                    Ok((new_bcs_id, hg_cs_id))
                }),
        )
    }

    fn verify_repo<F, R>(fb: FacebookInit, repo: F, runtime: &Runtime)
    where
        F: Future<Output = R>,
        R: BookmarksRef + RepoBlobstoreRef + RepoDerivedDataRef + ChangesetFetcherArc + Send + Sync,
    {
        let ctx = CoreContext::test_mock(fb);
        let repo = runtime.block_on(repo);
        borrowed!(ctx, repo);

        runtime
            .block_on(async move {
                all_commits(ctx, repo)
                    .await
                    .unwrap()
                    .try_for_each(move |(bcs_id, hg_cs_id)| async move {
                        verify_fsnode(ctx, repo, bcs_id, hg_cs_id).await
                    })
                    .await
            })
            .unwrap();
    }

    #[fbinit::test]
    fn test_derive_data(fb: FacebookInit) {
        let runtime = Runtime::new().unwrap();
        verify_repo(fb, Linear::getrepo(fb), &runtime);
        verify_repo(fb, BranchEven::getrepo(fb), &runtime);
        verify_repo(fb, BranchUneven::getrepo(fb), &runtime);
        verify_repo(fb, BranchWide::getrepo(fb), &runtime);
        verify_repo(fb, ManyDiamonds::getrepo(fb), &runtime);
        verify_repo(fb, ManyFilesDirs::getrepo(fb), &runtime);
        verify_repo(fb, MergeEven::getrepo(fb), &runtime);
        verify_repo(fb, MergeUneven::getrepo(fb), &runtime);
        verify_repo(fb, UnsharedMergeEven::getrepo(fb), &runtime);
        verify_repo(fb, UnsharedMergeUneven::getrepo(fb), &runtime);
    }
}