/*
 * 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::anyhow;
use anyhow::Error;
use async_trait::async_trait;
use auto_impl::auto_impl;
use context::CoreContext;
use context::PerfCounterType;
use metaconfig_types::CommitSyncConfigVersion;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use sql::mysql;
use sql::mysql_async::prelude::ConvIr;
use sql::mysql_async::prelude::FromValue;
use sql::mysql_async::FromValueError;
use sql::mysql_async::Value;
use sql::Connection;
use sql::Transaction;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;
use stats::prelude::*;
use thiserror::Error;

#[derive(Debug, Eq, Error, PartialEq)]
pub enum ErrorKind {
    #[error(
        "tried to insert inconsistent small bcs id {actual_bcs_id:?} version {actual_config_version:?}, while db has {expected_bcs_id:?} version {expected_config_version:?}"
    )]
    InconsistentWorkingCopyEntry {
        expected_bcs_id: Option<ChangesetId>,
        expected_config_version: Option<CommitSyncConfigVersion>,
        actual_bcs_id: Option<ChangesetId>,
        actual_config_version: Option<CommitSyncConfigVersion>,
    },
    #[error(
        "tried to insert inconsistent version for {large_cs_id} in repo {large_repo_id}: tried to insert {expected_version_name}, found {actual_version_name}"
    )]
    InconsistentLargeRepoCommitVersion {
        large_repo_id: RepositoryId,
        large_cs_id: ChangesetId,
        expected_version_name: CommitSyncConfigVersion,
        actual_version_name: CommitSyncConfigVersion,
    },
}

// TODO(simonfar): Once we've proven the concept, we want to cache these
define_stats! {
    prefix = "mononoke.synced_commit_mapping";
    gets: timeseries(Rate, Sum),
    gets_master: timeseries(Rate, Sum),
    adds: timeseries(Rate, Sum),
    add_many_in_txn: timeseries(Rate, Sum),
    add_bulks: timeseries(Rate, Sum),
    insert_working_copy_eqivalence: timeseries(Rate, Sum),
    get_equivalent_working_copy: timeseries(Rate, Sum),
}

// Repo that originally contained the synced commit
#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, mysql::OptTryFromRowField)]
pub enum SyncedCommitSourceRepo {
    Large,
    Small,
}

impl ConvIr<SyncedCommitSourceRepo> for SyncedCommitSourceRepo {
    fn new(v: Value) -> Result<Self, FromValueError> {
        use SyncedCommitSourceRepo::*;

        match v {
            Value::Bytes(ref b) if b == b"large" => Ok(Large),
            Value::Bytes(ref b) if b == b"small" => Ok(Small),
            v => Err(FromValueError(v)),
        }
    }

    fn commit(self) -> SyncedCommitSourceRepo {
        self
    }

    fn rollback(self) -> Value {
        self.into()
    }
}

impl FromValue for SyncedCommitSourceRepo {
    type Intermediate = SyncedCommitSourceRepo;
}

impl From<SyncedCommitSourceRepo> for Value {
    fn from(source_repo: SyncedCommitSourceRepo) -> Self {
        use SyncedCommitSourceRepo::*;

        match source_repo {
            Small => Value::Bytes(b"small".to_vec()),
            Large => Value::Bytes(b"large".to_vec()),
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct SyncedCommitMappingEntry {
    pub large_repo_id: RepositoryId,
    pub large_bcs_id: ChangesetId,
    pub small_repo_id: RepositoryId,
    pub small_bcs_id: ChangesetId,
    pub version_name: Option<CommitSyncConfigVersion>,
    pub source_repo: Option<SyncedCommitSourceRepo>,
}

impl SyncedCommitMappingEntry {
    pub fn new(
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        small_repo_id: RepositoryId,
        small_bcs_id: ChangesetId,
        version_name: CommitSyncConfigVersion,
        source_repo: SyncedCommitSourceRepo,
    ) -> Self {
        Self {
            large_repo_id,
            large_bcs_id,
            small_repo_id,
            small_bcs_id,
            version_name: Some(version_name),
            source_repo: Some(source_repo),
        }
    }

    fn into_equivalent_working_copy_entry(self) -> EquivalentWorkingCopyEntry {
        let Self {
            large_repo_id,
            large_bcs_id,
            small_repo_id,
            small_bcs_id,
            version_name,
            source_repo: _,
        } = self;

        EquivalentWorkingCopyEntry {
            large_repo_id,
            large_bcs_id,
            small_repo_id,
            small_bcs_id: Some(small_bcs_id),
            version_name,
        }
    }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct EquivalentWorkingCopyEntry {
    pub large_repo_id: RepositoryId,
    pub large_bcs_id: ChangesetId,
    pub small_repo_id: RepositoryId,
    pub small_bcs_id: Option<ChangesetId>,
    pub version_name: Option<CommitSyncConfigVersion>,
}

#[derive(Debug, PartialEq, Eq)]
pub enum WorkingCopyEquivalence {
    /// There's no matching working copy. It can happen if a pre-big-merge commit from one small
    /// repo is mapped into another small repo
    NoWorkingCopy(CommitSyncConfigVersion),
    /// ChangesetId of matching working copy and CommitSyncConfigVersion that was used for mapping
    WorkingCopy(ChangesetId, CommitSyncConfigVersion),
}

#[async_trait]
#[auto_impl(Arc)]
#[facet::facet]
pub trait SyncedCommitMapping: Send + Sync {
    /// Given the full large, small mapping, store it in the DB.
    /// Future resolves to true if the mapping was saved, false otherwise
    async fn add(&self, ctx: &CoreContext, entry: SyncedCommitMappingEntry) -> Result<bool, Error>;

    /// Bulk insert a set of large, small mappings
    /// This is meant for blobimport and similar
    async fn add_bulk(
        &self,
        ctx: &CoreContext,
        entries: Vec<SyncedCommitMappingEntry>,
    ) -> Result<u64, Error>;

    /// Find all the mapping entries for a given source commit and target repo
    async fn get(
        &self,
        ctx: &CoreContext,
        source_repo_id: RepositoryId,
        bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> Result<
        Vec<(
            ChangesetId,
            Option<CommitSyncConfigVersion>,
            Option<SyncedCommitSourceRepo>,
        )>,
        Error,
    >;

    /// Inserts equivalent working copy of a large bcs id. It's similar to mapping entry,
    /// however there are a few differences:
    /// 1) For (large repo, small repo) pair, many large commits can map to the same small commit
    /// 2) Small commit can be null
    ///
    /// If there's a mapping between small and large commits, then equivalent working copy is
    /// the same as the same as the mapping.
    async fn insert_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        entry: EquivalentWorkingCopyEntry,
    ) -> Result<bool, Error>;

    /// Same as previous command, but it overwrites existing value.
    /// This is not intended to be used in production, but just as a debug tool
    async fn overwrite_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        entry: EquivalentWorkingCopyEntry,
    ) -> Result<bool, Error>;

    /// Finds equivalent working copy
    async fn get_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        source_repo_id: RepositoryId,
        source_bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> Result<Option<WorkingCopyEquivalence>, Error>;

    /// Get version for large repo commit
    async fn get_large_repo_commit_version(
        &self,
        ctx: &CoreContext,
        large_repo_id: RepositoryId,
        large_repo_cs_id: ChangesetId,
    ) -> Result<Option<CommitSyncConfigVersion>, Error>;
}

#[derive(Clone)]
pub struct SqlSyncedCommitMapping {
    write_connection: Connection,
    read_connection: Connection,
    read_master_connection: Connection,
}

mononoke_queries! {
    write InsertMapping(values: (
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        small_repo_id: RepositoryId,
        small_bcs_id: ChangesetId,
        sync_map_version_name: Option<CommitSyncConfigVersion>,
        source_repo: Option<SyncedCommitSourceRepo>,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO synced_commit_mapping (large_repo_id, large_bcs_id, small_repo_id, small_bcs_id, sync_map_version_name, source_repo) VALUES {values}"
    }

    read SelectMapping(
        source_repo_id: RepositoryId,
        bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> (RepositoryId, ChangesetId, RepositoryId, ChangesetId, Option<CommitSyncConfigVersion>, Option<SyncedCommitSourceRepo>) {
        "SELECT large_repo_id, large_bcs_id, small_repo_id, small_bcs_id, sync_map_version_name, source_repo
          FROM synced_commit_mapping
          WHERE (large_repo_id = {source_repo_id} AND large_bcs_id = {bcs_id} AND small_repo_id = {target_repo_id}) OR
          (small_repo_id = {source_repo_id} AND small_bcs_id = {bcs_id} AND large_repo_id = {target_repo_id})"
    }

    write InsertWorkingCopyEquivalence(values: (
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        small_repo_id: RepositoryId,
        small_bcs_id: Option<ChangesetId>,
        sync_map_version_name: Option<CommitSyncConfigVersion>,
    )) {
        insert_or_ignore,
        "{insert_or_ignore}
         INTO synced_working_copy_equivalence
         (large_repo_id, large_bcs_id, small_repo_id, small_bcs_id, sync_map_version_name)
         VALUES {values}"
    }

    write ReplaceWorkingCopyEquivalence(values: (
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        small_repo_id: RepositoryId,
        small_bcs_id: Option<ChangesetId>,
        sync_map_version_name: Option<CommitSyncConfigVersion>,
    )) {
        none,
        "REPLACE
         INTO synced_working_copy_equivalence
         (large_repo_id, large_bcs_id, small_repo_id, small_bcs_id, sync_map_version_name)
         VALUES {values}"
    }

    read SelectWorkingCopyEquivalence(
        source_repo_id: RepositoryId,
        bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> (RepositoryId, ChangesetId, RepositoryId, Option<ChangesetId>, Option<CommitSyncConfigVersion>) {
        "SELECT large_repo_id, large_bcs_id, small_repo_id, small_bcs_id, sync_map_version_name
          FROM synced_working_copy_equivalence
          WHERE (large_repo_id = {source_repo_id} AND small_repo_id = {target_repo_id} AND large_bcs_id = {bcs_id})
          OR (large_repo_id = {target_repo_id} AND small_repo_id = {source_repo_id} AND small_bcs_id = {bcs_id})
          ORDER BY mapping_id ASC
          LIMIT 1
          "
    }

    write InsertVersionForLargeRepoCommit(values: (
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        sync_map_version_name: CommitSyncConfigVersion,
    )) {
        insert_or_ignore,
        "{insert_or_ignore}
        INTO version_for_large_repo_commit
        (large_repo_id, large_bcs_id, sync_map_version_name)
        VALUES {values}"
    }

    write ReplaceVersionForLargeRepoCommit(values: (
        large_repo_id: RepositoryId,
        large_bcs_id: ChangesetId,
        sync_map_version_name: CommitSyncConfigVersion,
    )) {
        none,
        "REPLACE
        INTO version_for_large_repo_commit
        (large_repo_id, large_bcs_id, sync_map_version_name)
        VALUES {values}"
    }

    read SelectVersionForLargeRepoCommit(
        large_repo_id: RepositoryId,
        cs_id: ChangesetId,
    ) -> (CommitSyncConfigVersion,) {
        "SELECT sync_map_version_name
          FROM version_for_large_repo_commit
          WHERE large_repo_id = {large_repo_id} AND large_bcs_id = {cs_id}"
    }
}

impl SqlConstruct for SqlSyncedCommitMapping {
    const LABEL: &'static str = "synced_commit_mapping";

    const CREATION_QUERY: &'static str =
        include_str!("../schemas/sqlite-synced-commit-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            write_connection: connections.write_connection,
            read_connection: connections.read_connection,
            read_master_connection: connections.read_master_connection,
        }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlSyncedCommitMapping {}

impl SqlSyncedCommitMapping {
    async fn add_many(
        &self,
        ctx: &CoreContext,
        entries: Vec<SyncedCommitMappingEntry>,
    ) -> Result<u64, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);
        let txn = self.write_connection.start_transaction().await?;
        let (txn, affected_rows) = add_many_in_txn(txn, entries).await?;
        txn.commit().await?;
        Ok(affected_rows)
    }

    async fn insert_or_overwrite_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        entry: EquivalentWorkingCopyEntry,
        should_overwrite: bool,
    ) -> Result<bool, Error> {
        STATS::insert_working_copy_eqivalence.add_value(1);

        let EquivalentWorkingCopyEntry {
            large_repo_id,
            large_bcs_id,
            small_repo_id,
            small_bcs_id,
            version_name,
        } = entry;

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);
        if let Some(ref version_name) = version_name {
            // TODO(stash): make version non-optional
            self.insert_version_for_large_repo_commit(
                ctx,
                &self.write_connection,
                large_repo_id,
                large_bcs_id,
                version_name,
                should_overwrite,
            )
            .await?;
        }
        let result = if should_overwrite {
            ReplaceWorkingCopyEquivalence::query(
                &self.write_connection,
                &[(
                    &large_repo_id,
                    &large_bcs_id,
                    &small_repo_id,
                    &small_bcs_id,
                    &version_name,
                )],
            )
            .await?
        } else {
            InsertWorkingCopyEquivalence::query(
                &self.write_connection,
                &[(
                    &large_repo_id,
                    &large_bcs_id,
                    &small_repo_id,
                    &small_bcs_id,
                    &version_name,
                )],
            )
            .await?
        };

        if result.affected_rows() >= 1 {
            Ok(true)
        } else {
            if !should_overwrite {
                // Check that db stores consistent entry
                let maybe_equivalent_wc = self
                    .get_equivalent_working_copy(ctx, large_repo_id, large_bcs_id, small_repo_id)
                    .await?;

                if let Some(equivalent_wc) = maybe_equivalent_wc {
                    use WorkingCopyEquivalence::*;
                    let (expected_bcs_id, expected_version) = match equivalent_wc {
                        WorkingCopy(wc, mapping) => (Some(wc), mapping),
                        NoWorkingCopy(mapping) => (None, mapping),
                    };
                    let expected_version = Some(expected_version);
                    if (expected_bcs_id != small_bcs_id) || (expected_version != version_name) {
                        let err = ErrorKind::InconsistentWorkingCopyEntry {
                            expected_bcs_id,
                            expected_config_version: expected_version,
                            actual_bcs_id: small_bcs_id,
                            actual_config_version: version_name,
                        };
                        return Err(err.into());
                    }
                }
            }
            Ok(false)
        }
    }

    async fn insert_version_for_large_repo_commit(
        &self,
        ctx: &CoreContext,
        write_connection: &Connection,
        large_repo_id: RepositoryId,
        large_cs_id: ChangesetId,
        version_name: &CommitSyncConfigVersion,
        should_overwrite: bool,
    ) -> Result<bool, Error> {
        let result = if should_overwrite {
            ReplaceVersionForLargeRepoCommit::query(
                write_connection,
                &[(&large_repo_id, &large_cs_id, version_name)],
            )
            .await?
        } else {
            InsertVersionForLargeRepoCommit::query(
                write_connection,
                &[(&large_repo_id, &large_cs_id, version_name)],
            )
            .await?
        };

        if result.affected_rows() >= 1 {
            Ok(true)
        } else {
            if !should_overwrite {
                // Check that db stores consistent entry
                let maybe_large_repo_version = self
                    .get_large_repo_commit_version(ctx, large_repo_id, large_cs_id)
                    .await?;

                if let Some(actual_version_name) = maybe_large_repo_version {
                    if &actual_version_name != version_name {
                        let err = ErrorKind::InconsistentLargeRepoCommitVersion {
                            large_repo_id,
                            large_cs_id,
                            expected_version_name: version_name.clone(),
                            actual_version_name,
                        };
                        return Err(err.into());
                    }
                }
            }
            Ok(false)
        }
    }
}

#[async_trait]
impl SyncedCommitMapping for SqlSyncedCommitMapping {
    async fn add(&self, ctx: &CoreContext, entry: SyncedCommitMappingEntry) -> Result<bool, Error> {
        STATS::adds.add_value(1);

        self.add_many(ctx, vec![entry])
            .await
            .map(|count| count == 1)
    }

    async fn add_bulk(
        &self,
        ctx: &CoreContext,
        entries: Vec<SyncedCommitMappingEntry>,
    ) -> Result<u64, Error> {
        STATS::add_bulks.add_value(1);

        self.add_many(ctx, entries).await
    }

    async fn get(
        &self,
        ctx: &CoreContext,
        source_repo_id: RepositoryId,
        bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> Result<
        Vec<(
            ChangesetId,
            Option<CommitSyncConfigVersion>,
            Option<SyncedCommitSourceRepo>,
        )>,
        Error,
    > {
        STATS::gets.add_value(1);

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let rows = SelectMapping::query(
            &self.read_connection,
            &source_repo_id,
            &bcs_id,
            &target_repo_id,
        )
        .await?;

        let rows = if rows.is_empty() {
            STATS::gets_master.add_value(1);
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            SelectMapping::query(
                &self.read_master_connection,
                &source_repo_id,
                &bcs_id,
                &target_repo_id,
            )
            .await?
        } else {
            rows
        };

        Ok(rows
            .into_iter()
            .map(|row| {
                let (
                    large_repo_id,
                    large_bcs_id,
                    _small_repo_id,
                    small_bcs_id,
                    maybe_version_name,
                    maybe_source_repo,
                ) = row;
                if target_repo_id == large_repo_id {
                    (large_bcs_id, maybe_version_name, maybe_source_repo)
                } else {
                    (small_bcs_id, maybe_version_name, maybe_source_repo)
                }
            })
            .collect())
    }

    async fn insert_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        entry: EquivalentWorkingCopyEntry,
    ) -> Result<bool, Error> {
        self.insert_or_overwrite_equivalent_working_copy(
            ctx, entry, false, /* should overwrite */
        )
        .await
    }

    async fn overwrite_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        entry: EquivalentWorkingCopyEntry,
    ) -> Result<bool, Error> {
        self.insert_or_overwrite_equivalent_working_copy(
            ctx, entry, true, /* should overwrite */
        )
        .await
    }

    async fn get_equivalent_working_copy(
        &self,
        ctx: &CoreContext,
        source_repo_id: RepositoryId,
        source_bcs_id: ChangesetId,
        target_repo_id: RepositoryId,
    ) -> Result<Option<WorkingCopyEquivalence>, Error> {
        STATS::get_equivalent_working_copy.add_value(1);

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);

        let rows = SelectWorkingCopyEquivalence::query(
            &self.read_connection,
            &source_repo_id,
            &source_bcs_id,
            &target_repo_id,
        )
        .await?;
        let maybe_row = if !rows.is_empty() {
            rows.get(0).cloned()
        } else {
            ctx.perf_counters()
                .increment_counter(PerfCounterType::SqlReadsMaster);
            SelectWorkingCopyEquivalence::query(
                &self.read_master_connection,
                &source_repo_id,
                &source_bcs_id,
                &target_repo_id,
            )
            .await
            .map(|rows| rows.get(0).cloned())?
        };

        Ok(match maybe_row {
            Some(row) => {
                let (
                    large_repo_id,
                    large_bcs_id,
                    _small_repo_id,
                    maybe_small_bcs_id,
                    maybe_mapping,
                ) = row;

                let mapping = maybe_mapping.ok_or_else(|| {
                    anyhow!(
                        "unexpected empty mapping for {}, {}->{}",
                        source_bcs_id,
                        source_repo_id,
                        target_repo_id
                    )
                })?;
                if target_repo_id == large_repo_id {
                    Some(WorkingCopyEquivalence::WorkingCopy(large_bcs_id, mapping))
                } else {
                    match maybe_small_bcs_id {
                        Some(small_bcs_id) => {
                            Some(WorkingCopyEquivalence::WorkingCopy(small_bcs_id, mapping))
                        }
                        None => Some(WorkingCopyEquivalence::NoWorkingCopy(mapping)),
                    }
                }
            }
            None => None,
        })
    }

    async fn get_large_repo_commit_version(
        &self,
        ctx: &CoreContext,
        large_repo_id: RepositoryId,
        large_repo_cs_id: ChangesetId,
    ) -> Result<Option<CommitSyncConfigVersion>, Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsReplica);
        let maybe_version = SelectVersionForLargeRepoCommit::query(
            &self.read_connection,
            &large_repo_id,
            &large_repo_cs_id,
        )
        .await?
        .pop()
        .map(|x| x.0);

        if let Some(version) = maybe_version {
            return Ok(Some(version));
        }

        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlReadsMaster);
        Ok(SelectVersionForLargeRepoCommit::query(
            &self.read_master_connection,
            &large_repo_id,
            &large_repo_cs_id,
        )
        .await?
        .pop()
        .map(|x| x.0))
    }
}

pub async fn add_many_in_txn(
    txn: Transaction,
    entries: Vec<SyncedCommitMappingEntry>,
) -> Result<(Transaction, u64), Error> {
    STATS::add_many_in_txn.add_value(1);

    let insert_entries: Vec<_> = entries
        .iter()
        .map(|entry| {
            (
                &entry.large_repo_id,
                &entry.large_bcs_id,
                &entry.small_repo_id,
                &entry.small_bcs_id,
                &entry.version_name,
                &entry.source_repo,
            )
        })
        .collect();

    let (txn, _result) = InsertMapping::query_with_transaction(txn, &insert_entries).await?;
    let owned_entries: Vec<_> = entries
        .into_iter()
        .map(|entry| entry.into_equivalent_working_copy_entry())
        .collect();

    let mut large_repo_commit_versions = vec![];
    for entry in &owned_entries {
        if let Some(version_name) = &entry.version_name {
            large_repo_commit_versions.push((
                &entry.large_repo_id,
                &entry.large_bcs_id,
                version_name,
            ));
        }
    }
    let (txn, _result) =
        InsertVersionForLargeRepoCommit::query_with_transaction(txn, &large_repo_commit_versions)
            .await?;

    let ref_entries: Vec<_> = owned_entries
        .iter()
        .map(|entry| {
            (
                &entry.large_repo_id,
                &entry.large_bcs_id,
                &entry.small_repo_id,
                &entry.small_bcs_id,
                &entry.version_name,
            )
        })
        .collect();

    let (txn, result) =
        InsertWorkingCopyEquivalence::query_with_transaction(txn, &ref_entries).await?;
    Ok((txn, result.affected_rows()))
}


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

//! Tests for the synced commits mapping.

use anyhow::Error;
use context::CoreContext;
use fbinit::FacebookInit;
use metaconfig_types::CommitSyncConfigVersion;
use mononoke_types_mocks::changesetid as bonsai;
use mononoke_types_mocks::repo::REPO_ONE;
use mononoke_types_mocks::repo::REPO_ZERO;
use sql_construct::SqlConstruct;
use synced_commit_mapping::EquivalentWorkingCopyEntry;
use synced_commit_mapping::SqlSyncedCommitMapping;
use synced_commit_mapping::SyncedCommitMapping;
use synced_commit_mapping::SyncedCommitMappingEntry;
use synced_commit_mapping::SyncedCommitSourceRepo;
use synced_commit_mapping::WorkingCopyEquivalence;

async fn add_and_get<M: SyncedCommitMapping>(fb: FacebookInit, mapping: M) {
    let version_name = CommitSyncConfigVersion("TEST_VERSION_NAME".to_string());
    let ctx = CoreContext::test_mock(fb);
    let entry = SyncedCommitMappingEntry::new(
        REPO_ZERO,
        bonsai::ONES_CSID,
        REPO_ONE,
        bonsai::TWOS_CSID,
        version_name.clone(),
        SyncedCommitSourceRepo::Large,
    );
    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );
    assert!(
        !mapping
            .add(&ctx, entry)
            .await
            .expect("Adding same entry failed")
    );

    let res = mapping
        .get_equivalent_working_copy(&ctx, REPO_ZERO, bonsai::ONES_CSID, REPO_ONE)
        .await
        .expect("get equivalent wc failed, should succeed");

    assert_eq!(
        res,
        Some(WorkingCopyEquivalence::WorkingCopy(
            bonsai::TWOS_CSID,
            version_name.clone(),
        ))
    );

    // insert again
    let entry = SyncedCommitMappingEntry::new(
        REPO_ZERO,
        bonsai::THREES_CSID,
        REPO_ONE,
        bonsai::FOURS_CSID,
        version_name.clone(),
        SyncedCommitSourceRepo::Large,
    );
    assert!(
        mapping
            .add(&ctx, entry.clone())
            .await
            .expect("Adding new entry failed")
    );

    let res = mapping
        .get_equivalent_working_copy(&ctx, REPO_ZERO, bonsai::THREES_CSID, REPO_ONE)
        .await
        .expect("get equivalent wc failed, should succeed");

    assert_eq!(
        res,
        Some(WorkingCopyEquivalence::WorkingCopy(
            bonsai::FOURS_CSID,
            version_name.clone(),
        ))
    );

    let result = mapping
        .get(&ctx, REPO_ZERO, bonsai::ONES_CSID, REPO_ONE)
        .await
        .expect("Get failed")
        .into_iter()
        .next()
        .expect("Unexpectedly, mapping is absent");
    assert_eq!(
        result,
        (
            bonsai::TWOS_CSID,
            Some(version_name.clone()),
            Some(SyncedCommitSourceRepo::Large)
        )
    );
    let result = mapping
        .get(&ctx, REPO_ONE, bonsai::TWOS_CSID, REPO_ZERO)
        .await
        .expect("Get failed")
        .into_iter()
        .next()
        .expect("Unexpectedly, mapping is absent");
    assert_eq!(
        result,
        (
            bonsai::ONES_CSID,
            Some(version_name),
            Some(SyncedCommitSourceRepo::Large)
        )
    );
}

async fn missing<M: SyncedCommitMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);
    let result = mapping
        .get(&ctx, REPO_ONE, bonsai::TWOS_CSID, REPO_ZERO)
        .await
        .expect("Failed to fetch missing changeset (should succeed with None instead)");
    assert!(result.is_empty());
}

async fn equivalent_working_copy<M: SyncedCommitMapping>(fb: FacebookInit, mapping: M) {
    let ctx = CoreContext::test_mock(fb);
    let version_name = CommitSyncConfigVersion("TEST_VERSION_NAME".to_string());
    let result = mapping
        .get_equivalent_working_copy(&ctx, REPO_ONE, bonsai::TWOS_CSID, REPO_ZERO)
        .await
        .expect("Failed to fetch equivalent working copy (should succeed with None instead)");
    assert_eq!(result, None);

    let entry = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::ONES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: Some(bonsai::TWOS_CSID),
        version_name: Some(version_name.clone()),
    };
    let result = mapping
        .insert_equivalent_working_copy(&ctx, entry.clone())
        .await
        .expect("Failed to insert working copy");
    assert!(result);

    let result = mapping
        .insert_equivalent_working_copy(&ctx, entry)
        .await
        .expect("Failed to insert working copy");
    assert!(!result);

    let res = mapping
        .get_equivalent_working_copy(&ctx, REPO_ZERO, bonsai::ONES_CSID, REPO_ONE)
        .await
        .expect("get equivalent wc failed, should succeed");

    assert_eq!(
        res,
        Some(WorkingCopyEquivalence::WorkingCopy(
            bonsai::TWOS_CSID,
            version_name
        ))
    );

    let null_entry = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::THREES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: None,
        version_name: None,
    };

    let result = mapping
        .insert_equivalent_working_copy(&ctx, null_entry)
        .await
        .expect("Failed to insert working copy");
    assert!(result);

    let should_fail = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::THREES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: Some(bonsai::TWOS_CSID),
        version_name: None,
    };
    assert!(
        mapping
            .insert_equivalent_working_copy(&ctx, should_fail)
            .await
            .is_err()
    );
}

#[fbinit::test]
async fn test_add_and_get(fb: FacebookInit) {
    add_and_get(fb, SqlSyncedCommitMapping::with_sqlite_in_memory().unwrap()).await;
}

#[fbinit::test]
async fn test_missing(fb: FacebookInit) {
    missing(fb, SqlSyncedCommitMapping::with_sqlite_in_memory().unwrap()).await
}

#[fbinit::test]
async fn test_equivalent_working_copy(fb: FacebookInit) {
    equivalent_working_copy(fb, SqlSyncedCommitMapping::with_sqlite_in_memory().unwrap()).await
}

#[fbinit::test]
async fn test_version_for_large_repo_commit(fb: FacebookInit) -> Result<(), Error> {
    let mapping = SqlSyncedCommitMapping::with_sqlite_in_memory()?;
    let ctx = CoreContext::test_mock(fb);
    // Check that at first we don't have a version for a given large repo commit
    assert!(
        mapping
            .get_large_repo_commit_version(&ctx, REPO_ZERO, bonsai::ONES_CSID)
            .await?
            .is_none()
    );

    // Insert working copy equivalence, version for the large commit
    let version_name = CommitSyncConfigVersion("TEST_VERSION_NAME".to_string());
    let entry = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::ONES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: Some(bonsai::TWOS_CSID),
        version_name: Some(version_name.clone()),
    };
    let result = mapping
        .insert_equivalent_working_copy(&ctx, entry.clone())
        .await?;
    assert!(result);

    assert_eq!(
        mapping
            .get_large_repo_commit_version(&ctx, REPO_ZERO, bonsai::ONES_CSID)
            .await?,
        Some(version_name.clone())
    );

    // Now try bulk insertion

    let new_version_name = CommitSyncConfigVersion("NEW_TEST_VERSION_NAME".to_string());
    let ctx = CoreContext::test_mock(fb);
    let entry_1 = SyncedCommitMappingEntry::new(
        REPO_ZERO,
        bonsai::THREES_CSID,
        REPO_ONE,
        bonsai::THREES_CSID,
        new_version_name.clone(),
        SyncedCommitSourceRepo::Large,
    );
    let entry_2 = SyncedCommitMappingEntry::new(
        REPO_ZERO,
        bonsai::FOURS_CSID,
        REPO_ONE,
        bonsai::FOURS_CSID,
        new_version_name.clone(),
        SyncedCommitSourceRepo::Large,
    );

    mapping.add_bulk(&ctx, vec![entry_1, entry_2]).await?;

    assert_eq!(
        mapping
            .get_large_repo_commit_version(&ctx, REPO_ZERO, bonsai::THREES_CSID)
            .await?,
        Some(new_version_name.clone())
    );
    assert_eq!(
        mapping
            .get_large_repo_commit_version(&ctx, REPO_ZERO, bonsai::FOURS_CSID)
            .await?,
        Some(new_version_name)
    );

    Ok(())
}

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

    // Insert working copy equivalence, version for the large commit
    let version_name = CommitSyncConfigVersion("TEST_VERSION_NAME".to_string());
    let entry = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::ONES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: Some(bonsai::TWOS_CSID),
        version_name: Some(version_name.clone()),
    };
    let result = mapping
        .insert_equivalent_working_copy(&ctx, entry.clone())
        .await?;
    assert!(result);

    let new_version_name = CommitSyncConfigVersion("NEW_TEST_VERSION_NAME".to_string());
    let entry = EquivalentWorkingCopyEntry {
        large_repo_id: REPO_ZERO,
        large_bcs_id: bonsai::ONES_CSID,
        small_repo_id: REPO_ONE,
        small_bcs_id: Some(bonsai::TWOS_CSID),
        version_name: Some(new_version_name.clone()),
    };
    let result = mapping
        .overwrite_equivalent_working_copy(&ctx, entry.clone())
        .await?;
    assert!(result);

    let res = mapping
        .get_equivalent_working_copy(&ctx, REPO_ZERO, bonsai::ONES_CSID, REPO_ONE)
        .await
        .expect("get equivalent wc failed, should succeed");

    assert_eq!(
        res,
        Some(WorkingCopyEquivalence::WorkingCopy(
            bonsai::TWOS_CSID,
            new_version_name.clone(),
        ))
    );

    Ok(())
}


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

use std::alloc::GlobalAlloc;
use std::alloc::Layout;
use std::alloc::System;
use std::cell::UnsafeCell;
use std::mem;
use std::thread_local;

use anyhow::Error;

thread_local! {
    // NOTE: We use an UnsafeCell here: ALLOCATION_STATS is a thread local that is never borrowed
    // by more than one callsite at a time.
    static ALLOCATION_STATS: UnsafeCell<AllocationStats> = UnsafeCell::new(AllocationStats::default());
}

struct TracingAllocator;

// NOTE: This code assumes that we don't try to allocate more than a u64 at a time: that's probably
// a reasonable assumption to make.
// NOTE: thread_local initialization doesn't allocate, which is why we can use a thread local here.
// if it did allocate, any allocation with TracingAllocator would hang forever.
unsafe impl GlobalAlloc for TracingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ret = System.alloc(layout);

        if !ret.is_null() {
            ALLOCATION_STATS.with(|cell| {
                (*cell.get()).allocated += layout.size() as u64;
            });
        }

        ret
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ret = System.alloc_zeroed(layout);

        if !ret.is_null() {
            ALLOCATION_STATS.with(|cell| {
                (*cell.get()).allocated += layout.size() as u64;
            });
        }

        ret
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let ret = System.realloc(ptr, layout, new_size);

        if !ret.is_null() {
            ALLOCATION_STATS.with(|cell| {
                (*cell.get()).freed += layout.size() as u64;
                (*cell.get()).allocated += new_size as u64;
            });
        }

        ret
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let ret = System.dealloc(ptr, layout);

        ALLOCATION_STATS.with(|cell| {
            (*cell.get()).freed += layout.size() as u64;
        });

        ret
    }
}

#[global_allocator]
static GLOBAL: TracingAllocator = TracingAllocator;

#[derive(Debug, Copy, Clone, Default, Hash, PartialEq, Eq)]
pub struct AllocationStats {
    allocated: u64,
    freed: u64,
}

impl AllocationStats {
    pub fn delta(&self) -> Result<i64, Error> {
        let allocated: i64 = self.allocated.try_into()?;
        let freed: i64 = self.freed.try_into()?;
        Ok(allocated - freed)
    }
}

/// trace_allocations returns AllocationStats representing memory that was allocated and freed
/// during the execution of f. Note that the two numbers are somewhat independent. For example:
/// - Objects allocated in f that are moved out of f won't be freed by f.
/// - Objects moved into f might be freed by f even though they were not allocated by f.
pub fn trace_allocations<T, F>(f: F) -> (T, AllocationStats)
where
    F: FnOnce() -> T,
{
    let new = AllocationStats::default();
    let curr = ALLOCATION_STATS.with(|cell| mem::replace(unsafe { &mut *cell.get() }, new));
    let ret = f();
    let stats = ALLOCATION_STATS.with(|cell| mem::replace(unsafe { &mut *cell.get() }, curr));
    (ret, stats)
}

#[cfg(test)]
#[allow(dead_code)]
mod test {
    use super::*;

    #[derive(Copy, Clone, Default, Debug)]
    struct TestStruct1 {
        a: u64,
    }

    #[derive(Copy, Clone, Default, Debug)]
    struct TestStruct2 {
        a: u64,
        b: u64,
    }

    #[derive(Copy, Clone, Default, Debug)]
    struct TestStruct3 {
        a: u64,
        b: u64,
        c: u64,
    }

    #[derive(Copy, Clone, Default, Debug)]
    struct TestStruct4 {
        a: u64,
        b: u64,
        c: u64,
        d: u64,
    }

    #[derive(Copy, Clone, Default, Debug)]
    struct TestStruct5 {
        a: u64,
        b: u64,
        c: u64,
        d: u64,
        e: u64,
    }

    fn alloc_some<T: Default>() -> Box<T> {
        Box::<T>::default()
    }

    #[test]
    fn test_baseline_allocation() -> Result<(), Error> {
        let (_, baseline) = trace_allocations(|| ());
        assert_eq!(baseline.delta()?, 0);
        Ok(())
    }

    #[test]
    fn test_trace_allocations() -> Result<(), Error> {
        fn run_test<T: Default>() -> Result<(), Error> {
            let (b, stats) = trace_allocations(alloc_some::<T>);
            mem::drop(b);
            assert_eq!(stats.delta()?, mem::size_of::<T>() as i64);
            Ok(())
        }

        run_test::<TestStruct1>()?;
        run_test::<TestStruct2>()?;
        run_test::<TestStruct3>()?;
        run_test::<TestStruct4>()?;
        run_test::<TestStruct5>()?;
        Ok(())
    }

    #[test]
    fn test_trace_frees() -> Result<(), Error> {
        fn run_test<T: Default>() -> Result<(), Error> {
            let b = alloc_some::<T>();
            let (_, stats) = trace_allocations(move || mem::drop(b));
            assert_eq!(stats.delta()?, -(mem::size_of::<T>() as i64));
            Ok(())
        }

        run_test::<TestStruct1>()?;
        run_test::<TestStruct2>()?;
        run_test::<TestStruct3>()?;
        run_test::<TestStruct4>()?;
        run_test::<TestStruct5>()?;
        Ok(())
    }

    #[test]
    fn test_nested() -> Result<(), Error> {
        fn run_test<T: Default>() -> Result<(), Error> {
            let ((b, stats1), stats2) = trace_allocations(|| {
                let (b1, stats1) = trace_allocations(alloc_some::<T>);
                (b1, stats1)
            });
            mem::drop(b);
            let delta = stats1.delta()? + stats2.delta()?;
            assert_eq!(delta, mem::size_of::<T>() as i64);
            Ok(())
        }

        run_test::<TestStruct1>()?;
        run_test::<TestStruct2>()?;
        run_test::<TestStruct3>()?;
        run_test::<TestStruct4>()?;
        run_test::<TestStruct5>()?;
        Ok(())
    }

    #[test]
    fn test_realloc() -> Result<(), Error> {
        fn run_test<T: Default>() -> Result<(), Error> {
            let (_, stats) = trace_allocations(|| {
                let mut v = Vec::<T>::with_capacity(2);
                v.push(T::default());
                v.push(T::default());
                v.push(T::default());
                v.push(T::default());
                v.push(T::default());
                v.pop();
                v.shrink_to_fit();
                v
            });

            assert_eq!(stats.delta()?, (mem::size_of::<T>() * 4) as i64);
            Ok(())
        }

        run_test::<TestStruct1>()?;
        run_test::<TestStruct2>()?;
        run_test::<TestStruct3>()?;
        run_test::<TestStruct4>()?;
        run_test::<TestStruct5>()?;
        Ok(())
    }
}


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

use std::cmp;

use anyhow::Error;
use futures_old::Async;
use futures_old::Future;
use futures_old::Stream;

use crate::trace_allocations;

#[must_use = "futures and streams do nothing unless polled"]
pub struct AllocationTraced<T> {
    inner: T,
    delta: i64,
    high: i64,
}

impl<T> AllocationTraced<T> {
    pub fn new(inner: T) -> Self {
        Self {
            inner,
            delta: 0,
            high: 0,
        }
    }
}

impl<T> AllocationTraced<T> {
    fn do_poll<I, E, P>(&mut self, poll: P) -> Result<Async<I>, E>
    where
        P: FnOnce(&mut T) -> Result<Async<I>, E>,
        E: From<Error>,
    {
        let (ret, stats) = trace_allocations(|| poll(&mut self.inner));

        let delta = match stats.delta() {
            Ok(delta) => delta,
            Err(e) => {
                return Err(e.into());
            }
        };

        let new_delta = self.delta + delta;

        self.delta = new_delta;
        self.high = cmp::max(self.high, new_delta);

        ret
    }
}

/// AllocationTraced<F> returns a Future that yields the item or error from F, along with the high
/// watermark for bytes allocated throughout the execution of F.
impl<T> Future for AllocationTraced<T>
where
    T: Future,
    T::Error: From<Error>,
{
    type Item = (T::Item, i64);
    type Error = (T::Error, i64);

    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        match self.do_poll(T::poll) {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready(r)) => Ok(Async::Ready((r, self.high))),
            Err(e) => Err((e, self.high)),
        }
    }
}

/// AllocationTraced<S> returns a Strem that yields items or the error from S, along with the high
/// watermark for bytes allocated throughout the execution of S.
impl<T> Stream for AllocationTraced<T>
where
    T: Stream,
    T::Error: From<Error>,
{
    type Item = (T::Item, i64);
    type Error = (T::Error, i64);

    fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
        match self.do_poll(T::poll) {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready(Some(r))) => Ok(Async::Ready(Some((r, self.high)))),
            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
            Err(e) => Err((e, self.high)),
        }
    }
}

pub trait AllocationTracingFutureExt: Future + Sized {
    fn allocation_traced(self) -> AllocationTraced<Self> {
        AllocationTraced::new(self)
    }
}

impl<T> AllocationTracingFutureExt for T where T: Future {}

pub trait AllocationTracingStreamExt: Stream + Sized {
    fn allocation_traced(self) -> AllocationTraced<Self> {
        AllocationTraced::new(self)
    }
}

impl<T> AllocationTracingStreamExt for T where T: Stream {}


/*
 * 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 core;
mod futures_v1;

pub use crate::core::trace_allocations;
pub use crate::core::AllocationStats;
pub use crate::futures_v1::AllocationTraced;
pub use crate::futures_v1::AllocationTracingFutureExt;
pub use crate::futures_v1::AllocationTracingStreamExt;


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

#![feature(unboxed_closures)]
#![feature(fn_traits)]
#![feature(type_alias_impl_trait)]

use std::marker::PhantomData;

use futures::future;
use futures::stream::FusedStream;
use futures::Future;
use futures::FutureExt;
use futures::Stream;
use futures::TryStream;

mod next_step;

pub use next_step::NextStep;

pub trait AssemblyLine: Stream + Sized {
    fn next_step<F>(self, step_fn: F) -> NextStep<Self, F>
    where
        F: FnMut<(Self::Item,)>,
        F::Output: Future,
        Self: FusedStream,
    {
        NextStep::new(self, step_fn)
    }
}

impl<S: Stream> AssemblyLine for S {}

/// This overly complicated struct and FnMut implementations are used for
/// TryAssemblyLine because we need to name some types as we can't use
/// impl Trait in trait outputs in Rust yet.
pub struct ResultWrapper<F, I, T, Err>(F, PhantomData<(I, T, Err)>)
where
    F: FnMut<(I,)>,
    F::Output: Future<Output = Result<T, Err>>;

impl<F, I, T, Err> FnOnce<(Result<I, Err>,)> for ResultWrapper<F, I, T, Err>
where
    F: FnMut<(I,)>,
    F::Output: Future<Output = Result<T, Err>>,
{
    type Output = future::Either<F::Output, future::Ready<Result<T, Err>>>;
    extern "rust-call" fn call_once(mut self, args: (Result<I, Err>,)) -> Self::Output {
        self.call_mut(args)
    }
}

impl<F, I, T, Err> FnMut<(Result<I, Err>,)> for ResultWrapper<F, I, T, Err>
where
    F: FnMut<(I,)>,
    F::Output: Future<Output = Result<T, Err>>,
{
    extern "rust-call" fn call_mut(&mut self, (res,): (Result<I, Err>,)) -> Self::Output {
        match res {
            Ok(ok) => self.0.call_mut((ok,)).left_future(),
            Err(err) => future::ready(Err(err)).right_future(),
        }
    }
}

pub trait TryAssemblyLine: TryStream + Sized {
    fn try_next_step<F, O>(
        self,
        step_fn: F,
    ) -> NextStep<Self, ResultWrapper<F, Self::Ok, O, Self::Error>>
    where
        F: FnMut<(Self::Ok,)>,
        F::Output: Future<Output = Result<O, Self::Error>>,
        Self: FusedStream,
        // This is always true, not sure why I need this bound
        Self: Stream<Item = Result<Self::Ok, Self::Error>>,
    {
        NextStep::new(self, ResultWrapper(step_fn, PhantomData))
    }
}

impl<S: TryStream> TryAssemblyLine for S {}


/*
 * 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::pin::Pin;

use futures::stream::FusedStream;
use futures::stream::Stream;
use futures::task::Context;
use futures::task::Poll;
use futures::Future;
use pin_project::pin_project;

#[pin_project]
pub struct NextStep<S, F>
where
    S: FusedStream,
    F: FnMut<(S::Item,)>,
    F::Output: Future,
{
    f: F,
    next_in_line: Option<S::Item>,
    running: Option<Pin<Box<F::Output>>>,
    #[pin]
    inner: S,
}

/// Like stream::Then, it applies a future to the output. The difference is
/// that it continues waiting for the next item in the stream WHILE the current
/// one is being processed.
impl<S, F> NextStep<S, F>
where
    S: FusedStream,
    F: FnMut<(S::Item,)>,
    F::Output: Future,
{
    pub fn new(inner: S, f: F) -> Self {
        Self {
            f,
            inner,
            next_in_line: None,
            running: None,
        }
    }
}

impl<S, F> Stream for NextStep<S, F>
where
    S: FusedStream,
    F: FnMut<(S::Item,)>,
    F::Output: Future,
{
    type Item = <F::Output as Future>::Output;

    fn poll_next(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();
        if !this.inner.is_terminated() && this.next_in_line.is_none() {
            if let Poll::Ready(Some(next_in_line)) = this.inner.as_mut().poll_next(ctx) {
                *this.next_in_line = Some(next_in_line);
            }
        }
        if this.running.is_none() {
            if let Some(next_in_line) = this.next_in_line.take() {
                let running = (this.f)(next_in_line);
                *this.running = Some(Box::pin(running));
            }
        }
        if let Some(running) = this.running.as_mut() {
            if let Poll::Ready(processed) = running.as_mut().poll(ctx) {
                *this.running = None;
                return Poll::Ready(Some(processed));
            }
        } else if this.inner.is_terminated() {
            // Nothing running, inner stream is empty, we're done
            return Poll::Ready(None);
        }

        Poll::Pending
    }
}

impl<S, F> FusedStream for NextStep<S, F>
where
    S: FusedStream,
    F: FnMut<(S::Item,)>,
    F::Output: Future,
{
    fn is_terminated(&self) -> bool {
        self.running.is_none() && self.next_in_line.is_none() && self.inner.is_terminated()
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering::Relaxed;

    use futures::stream;
    use futures::StreamExt;

    use super::*;

    struct TestFuture {
        poll_count: AtomicUsize,
        value: AtomicUsize,
        set: AtomicBool,
    }
    impl Default for TestFuture {
        fn default() -> Self {
            Self {
                poll_count: 0.into(),
                value: 0.into(),
                set: false.into(),
            }
        }
    }

    impl TestFuture {
        fn poll_count(&self) -> usize {
            self.poll_count.load(Relaxed)
        }
        fn set_value(&self, x: usize) {
            self.value.store(x, Relaxed);
            self.set.store(true, Relaxed);
        }
    }

    impl Future for &TestFuture {
        type Output = usize;

        fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
            self.poll_count.fetch_add(1, Relaxed);
            if self.set.load(Relaxed) {
                Poll::Ready(self.value.load(Relaxed))
            } else {
                ctx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }

    #[tokio::test]
    async fn test_test_future() {
        let futs: [TestFuture; 3] = Default::default();
        let mut stream = stream::iter(&futs).buffered(1);
        assert_eq!(futs[0].poll_count(), 0);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(futs[0].poll_count(), 1);
        assert_eq!(futs[1].poll_count(), 0);
        futs[0].set_value(12);
        assert_eq!(stream.next().await, Some(12));
        assert_eq!(futs[0].poll_count(), 2);
        assert_eq!(futs[1].poll_count(), 0);
    }

    #[tokio::test]
    async fn test_stream_ext() {
        let futs: [TestFuture; 3] = Default::default();
        let then_futs: [TestFuture; 3] = Default::default();
        let mut stream = stream::iter(&futs).buffered(1).then(|x| &then_futs[x]);
        futs[0].set_value(0);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(then_futs[0].poll_count(), 2);
        assert_eq!(futs[1].poll_count(), 0);
        futs[1].set_value(1);
        futs[2].set_value(2);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        // Still waiting for then_fut[0], doesn't try fut[1]
        assert_eq!(futs[1].poll_count(), 0);
        assert_eq!(futs[2].poll_count(), 0);
    }

    #[tokio::test]
    async fn test_next_step() {
        let futs: [TestFuture; 3] = Default::default();
        let then_futs: [TestFuture; 3] = Default::default();
        let stream = stream::iter(&futs).buffered(1);
        let mut stream = NextStep::new(stream.fuse(), |x| &then_futs[x]);
        futs[0].set_value(0);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(then_futs[0].poll_count(), 2);
        assert_eq!(futs[1].poll_count(), 1);
        assert!(stream.running.is_some());
        assert!(stream.next_in_line.is_none());
        futs[1].set_value(1);
        futs[2].set_value(2);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(futures::poll!(stream.next()), Poll::Pending);
        assert_eq!(stream.next_in_line, Some(1));
        // Awaits for fut[1] while then_fut[0] is being processed
        assert_eq!(futs[1].poll_count(), 2);
        // Only keeps a single thing in queue
        assert_eq!(futs[2].poll_count(), 0);
    }
}


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

use anyhow::Error;
use async_limiter::AsyncLimiter;
use chrono::Local;
use futures::future::join_all;
use governor::Quota;
use governor::RateLimiter;
use nonzero_ext::nonzero;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let limiter = RateLimiter::direct(Quota::per_second(nonzero!(5u32)));
    let limiter = AsyncLimiter::new(limiter).await;

    let futs = (0..10).map(|i| {
        let limiter = limiter.clone();
        async move {
            loop {
                limiter.access().await?;
                println!("[{}] {}", i, Local::now().format("%H:%M:%S%.3f"));
            }
        }
    });

    join_all(futs)
        .await
        .into_iter()
        .collect::<Result<Vec<()>, Error>>()?;

    Ok(())
}


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

use anyhow::Error;
use futures::channel::mpsc;
use futures::channel::oneshot;
use futures::future;
use futures::future::Future;
use futures::stream::StreamExt;
use governor::clock::ReasonablyRealtime;
use governor::state::direct::DirectStateStore;
use governor::state::direct::NotKeyed;
use governor::state::direct::StreamRateLimitExt;
use governor::RateLimiter;

use crate::ErrorKind;

/// A shared asynchronous rate limiter.
#[derive(Clone)]
pub struct AsyncLimiter {
    dispatch: mpsc::UnboundedSender<oneshot::Sender<()>>,
    cancel: mpsc::Sender<()>,
}

impl AsyncLimiter {
    // NOTE: This function is async because it requires a Tokio runtme to spawn things. The best
    // way to require a Tokio runtime to be present is to just make the function async.
    pub async fn new<S, C>(limiter: RateLimiter<NotKeyed, S, C>) -> Self
    where
        S: DirectStateStore + Send + Sync + 'static,
        C: ReasonablyRealtime + Send + Sync + 'static,
    {
        let (dispatch, dispatch_recv) = mpsc::unbounded();
        let (cancel, cancel_recv) = mpsc::channel(1);

        tokio_shim::task::spawn(async move {
            let worker = dispatch_recv
                .zip(futures::stream::select(
                    cancel_recv,
                    futures::stream::repeat(()).ratelimit_stream(&limiter),
                ))
                .for_each(|(reply, _): (oneshot::Sender<()>, _)| {
                    let _ = reply.send(());
                    future::ready(())
                });
            worker.await
        });

        Self { dispatch, cancel }
    }

    /// access() returns a result of a future that returns once the rate limiter reports that it is
    /// OK to let one client proceed. It may return an error if the runtime is shutting down.
    /// Access is granted on a first-come first-serve basis (based on the order in which access()
    /// was called). If a caller doesnot await the future returned by access to completion, then
    /// the rate limiter's internal state will be updated nonetheless. Memory usage is proportional
    /// to the number of pending accesses. Note that this isn't an async fn so as to not capture a
    /// refernce to &self in the future returned by this method, which makes it more suitable for
    /// use in e.g. a futures 0.1 context.
    pub fn access(&self) -> impl Future<Output = Result<(), Error>> + 'static + Send + Sync {
        let (send, recv) = oneshot::channel();
        let dispatch = self.dispatch.clone();

        async move {
            // NOTE: We do the dispatch in this future here, which effectively makes this lazy.
            // This ensures that if you create a future, but don't poll it immediately, it only
            // tries to enter the queue once it's polled.
            dispatch
                .unbounded_send(send)
                .map_err(|_| ErrorKind::RuntimeShuttingDown)?;

            recv.await.map_err(|_| ErrorKind::RuntimeShuttingDown)?;

            Ok(())
        }
    }

    /// cancel() allows a caller to return a token they didn't use. The token will be passed
    /// on to any active waiter. If there are no active waiters, the token is discarded.
    pub fn cancel(&self) {
        let _ = self.cancel.clone().try_send(());
    }
}

#[cfg(test)]
mod test {
    use std::time::Duration;
    use std::time::Instant;

    use governor::Quota;
    use governor::RateLimiter;
    use nonzero_ext::nonzero;

    use super::*;

    #[tokio::test]
    async fn test_access_enters_queue_lazily() -> Result<(), Error> {
        let limiter = RateLimiter::direct(Quota::per_second(nonzero!(5u32)));
        let limiter = AsyncLimiter::new(limiter).await;

        for _ in 0..10 {
            std::mem::drop(limiter.access());
        }

        let now = Instant::now();
        limiter.access().await?;
        limiter.access().await?;

        assert!(now.elapsed() < Duration::from_millis(100));
        Ok(())
    }

    #[tokio::test]
    async fn test_cancel() -> Result<(), Error> {
        let limiter = RateLimiter::direct(Quota::per_second(nonzero!(1u32)));
        let limiter = AsyncLimiter::new(limiter).await;

        let now = Instant::now();

        for _ in 0..100 {
            limiter.access().await?;
            limiter.cancel();
        }

        assert!(now.elapsed() < Duration::from_millis(100));
        Ok(())
    }
}


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

use thiserror::Error;

#[derive(Error, Debug)]
pub enum ErrorKind {
    #[error("Runtime is shutting down")]
    RuntimeShuttingDown,
}


/*
 * 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 async_limiter_;
mod errors;

pub use async_limiter_::AsyncLimiter;
pub use errors::ErrorKind;


/*
 * 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::future::Future;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use either::Either;
use futures::ready;

/// Return value from `unfold` callbacks for ordered traversals.  Each element
/// in the unfolded vector can be either an item to output from the traversal,
/// or a recursive step with an associated weight.
///
/// The associated weight should be an estimate of the number of eventual
/// output items this recursive step will expand into.
pub enum OrderedTraversal<Out, In> {
    Output(Out),
    Recurse(usize, In),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct NodeLocation<Index> {
    pub node_index: Index,  // node index inside execution tree
    pub child_index: usize, // index inside parents children list
}

impl<Index> NodeLocation<Index> {
    pub fn new(node_index: Index, child_index: usize) -> Self {
        NodeLocation {
            node_index,
            child_index,
        }
    }
}

/// Equivalent of `futures::future::Either` but with heterogeneous output
/// types using `either::Either`.
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub(crate) enum Either2<A, B> {
    Left(A),
    Right(B),
}

impl<A, B> Future for Either2<A, B>
where
    A: Future,
    B: Future,
{
    type Output = Either<A::Output, B::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        // see `impl<Left, Right> Future for Either<Left, Right>`
        unsafe {
            let result = match self.get_unchecked_mut() {
                Either2::Left(future) => Either::Left(ready!(Pin::new_unchecked(future).poll(cx))),
                Either2::Right(future) => {
                    Either::Right(ready!(Pin::new_unchecked(future).poll(cx)))
                }
            };
            Poll::Ready(result)
        }
    }
}


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

use std::collections::HashMap;
use std::collections::VecDeque;
use std::future::Future;
use std::hash::Hash;
use std::mem;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use either::Either;
use futures::future::join;
use futures::future::ready;
use futures::future::BoxFuture;
use futures::future::Join;
use futures::future::Ready;
use futures::ready;
use futures::stream::FuturesUnordered;
use futures::stream::StreamExt;

use super::common::Either2;
use super::common::NodeLocation;
use super::Iter;

/// `bounded_traversal_dag` traverses implicit asynchronous DAG specified by `init`
/// and `unfold` arguments, and it also does backward pass with `fold` operation.
/// All `unfold` and `fold` operations are executed in parallel if they do not
/// depend on each other (not related by ancestor-descendant relation in implicit DAG)
/// with amount of concurrency constrained by `scheduled_max`.
///
/// ## Difference between `bounded_traversal_dag` and `bounded_traversal`
/// Obvious difference is that `bounded_traversal_dag` correctly handles DAGs
/// (`bounded_traversal` treats all children references as distinct and its execution time
/// is proportional to number of paths from the root, since DAG can be constructed to contain
/// `O(exp(N))` path it might cause problems) but it comes with a price:
///  - `bounded_traversal_dag` keeps `Out` result of computation for all the nodes
///     but `bounded_traversal` only keeps results for nodes that have not been completely
///     evaluatated
///  - `In` has additional constraints to be `Eq + Hash + Clone`
///  - `Out` has additional constraint to be `Clone`
///
/// ## `init: In`
/// Is the root of the implicit tree to be traversed
///
/// ## `unfold: FnMut(In) -> impl Future<Output = Result<(OutCtx, impl IntoIterator<Item = In>), Err>>`
/// Asynchronous function which given input value produces list of its children. And context
/// associated with current node. If this list is empty, it is a leaf of the tree, and `fold`
/// will be run on this node.
///
/// ## `fold: FnMut(OutCtx, impl Iterator<Out>) -> impl Future<Item = Result<Out, Err>>`
/// Aynchronous function which given node context and output of `fold` for its chidlren
/// should produce new output value.
///
/// ## return value `impl Future<Output = Result<Option<Out>, Err>>`
/// Result of running fold operation on the root of the tree. `None` indiciate that cycle
/// has been found.
///
pub fn bounded_traversal_dag<'caller, Err, In, Ins, Out, OutCtx, Unfold, Fold>(
    scheduled_max: usize,
    init: In,
    unfold: Unfold,
    fold: Fold,
) -> impl Future<Output = Result<Option<Out>, Err>> + 'caller
where
    Err: 'caller,
    In: Eq + Hash + Clone + 'caller,
    Out: Clone + 'caller,
    OutCtx: 'caller,
    // We use BoxFuture here because the `Unfold` future can be very large.
    // As a result, it's more efficient to keep it in one place (the heap)
    // than to move it around on the stack all the time.
    // https://fburl.com/m3cdcdko
    Unfold: FnMut(In) -> BoxFuture<'caller, Result<(OutCtx, Ins), Err>> + 'caller,
    Ins: IntoIterator<Item = In> + 'caller,
    Fold: FnMut(OutCtx, Iter<Out>) -> BoxFuture<'caller, Result<Out, Err>> + 'caller,
{
    BoundedTraversalDAG::new(scheduled_max, init, unfold, fold, None)
}

/// As `bounded_traversal_dag`, but will stop unfolding once enough nodes
/// have been visited.
pub fn bounded_traversal_dag_limited<'caller, Err, In, Ins, Out, OutCtx, Unfold, Fold>(
    scheduled_max: usize,
    init: In,
    unfold: Unfold,
    fold: Fold,
    limit: impl Into<Option<u64>>,
) -> impl Future<Output = Result<Option<Out>, Err>> + 'caller
where
    Err: 'caller,
    In: Eq + Hash + Clone + 'caller,
    Out: Clone + 'caller,
    OutCtx: 'caller,
    // We use BoxFuture here because the `Unfold` future can be very large.
    // As a result, it's more efficient to keep it in one place (the heap)
    // than to move it around on the stack all the time.
    // https://fburl.com/m3cdcdko
    Unfold: FnMut(In) -> BoxFuture<'caller, Result<(OutCtx, Ins), Err>> + 'caller,
    Ins: IntoIterator<Item = In> + 'caller,
    Fold: FnMut(OutCtx, Iter<Out>) -> BoxFuture<'caller, Result<Out, Err>> + 'caller,
{
    BoundedTraversalDAG::new(scheduled_max, init, unfold, fold, limit.into())
}

struct Children<Out, OutCtx> {
    context: OutCtx,
    children: Vec<Option<Out>>,
    children_left: usize,
}

enum Node<In, Out, OutCtx> {
    Pending {
        parents: Vec<NodeLocation<In>>, // nodes blocked by current node
        children: Option<Children<Out, OutCtx>>, // present if node waits for children to be computed
    },
    Done(Out),
}

#[must_use = "futures do nothing unless polled"]
struct BoundedTraversalDAG<In, Out, OutCtx, Unfold, UFut, Fold, FFut>
where
    UFut: Future,
    FFut: Future,
{
    init: In,
    unfold: Unfold,
    fold: Fold,
    scheduled_max: usize,
    /// Jobs being executed to traverse DAG nodes
    scheduled: FuturesUnordered<Join<Ready<In>, Either2<UFut, FFut>>>,
    /// Unscheduled traversal jobs - these are ready to be executed, but are blocked due to scheduled_max
    unscheduled: VecDeque<Join<Ready<In>, Either2<UFut, FFut>>>,
    /// Tree tracking execution progress
    execution_tree: HashMap<In, Node<In, Out, OutCtx>>,
    /// Maximum number of nodes to visit. Once we reach the limit, we stop scheduling `unfold`s.
    unfold_limit: Option<u64>,
}

impl<Err, In, Ins, Out, OutCtx, Unfold, UFut, Fold, FFut>
    BoundedTraversalDAG<In, Out, OutCtx, Unfold, UFut, Fold, FFut>
where
    In: Clone + Eq + Hash,
    Out: Clone,
    Unfold: FnMut(In) -> UFut,
    UFut: Future<Output = Result<(OutCtx, Ins), Err>>,
    Ins: IntoIterator<Item = In>,
    Fold: FnMut(OutCtx, Iter<Out>) -> FFut,
    FFut: Future<Output = Result<Out, Err>>,
{
    fn new(
        scheduled_max: usize,
        init: In,
        unfold: Unfold,
        fold: Fold,
        unfold_limit: Option<u64>,
    ) -> Self {
        let mut this = Self {
            init: init.clone(),
            unfold,
            fold,
            scheduled_max,
            scheduled: FuturesUnordered::new(),
            unscheduled: VecDeque::new(),
            execution_tree: HashMap::new(),
            unfold_limit,
        };
        let init_out = this.enqueue_unfold(
            NodeLocation {
                node_index: init.clone(),
                child_index: 0,
            },
            init,
        );
        // can not be resolved since execution tree is empty
        debug_assert!(init_out.is_none());
        this
    }

    fn enqueue_unfold(&mut self, parent: NodeLocation<In>, value: In) -> Option<Out> {
        match self.execution_tree.get_mut(&value) {
            None => {
                // schedule unfold for previously unseen `value`
                self.execution_tree.insert(
                    value.clone(),
                    Node::Pending {
                        parents: vec![parent],
                        children: None,
                    },
                );
                self.unscheduled.push_front(join(
                    ready(value.clone()),
                    Either2::Left((self.unfold)(value)),
                ));
                None
            }
            Some(Node::Pending { parents, .. }) => {
                // we already have a node associated with the same input value,
                // register as a dependency for this node.
                parents.push(parent);
                None
            }
            Some(Node::Done(result)) => Some(result.clone()),
        }
    }

    fn enqueue_fold(&mut self, value: In, context: OutCtx, children: Iter<Out>) {
        self.unscheduled.push_front(join(
            ready(value),
            Either2::Right((self.fold)(context, children)),
        ));
    }

    fn process_unfold(&mut self, value: In, (context, children): (OutCtx, Ins)) {
        // limit children based on the number of nodes we're able to unfold
        let children = if let Some(limit) = self.unfold_limit {
            let unfold_budget = limit
                .try_into()
                .unwrap_or(usize::MAX)
                .saturating_sub(self.execution_tree.len());
            children.into_iter().take(unfold_budget)
        } else {
            children.into_iter().take(usize::MAX)
        };
        // schedule unfold for node's children
        let mut children_left = 0;
        let children: Vec<_> = children
            .enumerate()
            .map(|(child_index, child)| {
                let out = self.enqueue_unfold(
                    NodeLocation {
                        node_index: value.clone(),
                        child_index,
                    },
                    child,
                );
                if out.is_none() {
                    children_left += 1;
                }
                out
            })
            .collect();

        if children_left != 0 {
            // update pending node with `wait` state
            let node = self
                .execution_tree
                .get_mut(&value)
                .expect("unfold referenced invalid node");
            match node {
                Node::Pending { children: wait, .. } => {
                    *wait = Some(Children {
                        context,
                        children,
                        children_left,
                    });
                }
                _ => unreachable!("running unfold for Node::Done"),
            }
        } else {
            // do not have any dependencies (leaf node), schedule fold immediately
            self.enqueue_fold(value, context, children.into_iter().flatten());
        }
    }

    fn process_fold(&mut self, value: In, result: Out) {
        // mark node as done
        let node = self
            .execution_tree
            .get_mut(&value)
            .expect("fold referenced invalid node");
        let parents = match mem::replace(node, Node::Done(result.clone())) {
            Node::Pending { parents, .. } => parents,
            _ => unreachable!("running fold for Node::Done"),
        };

        // update all the parents wait for this result
        for parent in parents {
            self.update_location(parent, result.clone());
        }
    }

    fn update_location(&mut self, loc: NodeLocation<In>, result: Out) {
        let node = self
            .execution_tree
            .get_mut(&loc.node_index)
            .expect("`update_location` referenced invalid node");
        let children = match node {
            Node::Pending { children, .. } => children,
            _ => unreachable!("updating already resolved parent node"),
        };
        let no_children_left = {
            // update parent
            let children = children
                .as_mut()
                .expect("`update_location` referenced not blocked node");
            debug_assert!(children.children[loc.child_index].is_none());
            children.children[loc.child_index] = Some(result);
            children.children_left -= 1;
            children.children_left == 0
        };
        if no_children_left {
            // all parents children have been completed, so we need
            // to schedule fold operation for it
            let Children {
                context, children, ..
            } = children
                .take()
                .expect("`update_location` reference node without children");
            self.enqueue_fold(loc.node_index, context, children.into_iter().flatten());
        }
    }
}

impl<Err, In, Ins, Out, OutCtx, Unfold, UFut, Fold, FFut> Future
    for BoundedTraversalDAG<In, Out, OutCtx, Unfold, UFut, Fold, FFut>
where
    In: Eq + Hash + Clone,
    Out: Clone,
    Unfold: FnMut(In) -> UFut,
    UFut: Future<Output = Result<(OutCtx, Ins), Err>>,
    Ins: IntoIterator<Item = In>,
    Fold: FnMut(OutCtx, Iter<Out>) -> FFut,
    FFut: Future<Output = Result<Out, Err>>,
{
    type Output = Result<Option<Out>, Err>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = unsafe { self.get_unchecked_mut() };
        loop {
            if this.unscheduled.is_empty() && this.scheduled.is_empty() {
                // we have not received result of with `value == init` and
                // nothing is scheduled or unscheduled, it means that we have
                // cycle dependency somewhere inside input graph
                return Poll::Ready(Ok(None));
            }

            // schedule as many jobs as possible
            for job in this.unscheduled.drain(
                ..std::cmp::min(
                    this.unscheduled.len(),
                    this.scheduled_max - this.scheduled.len(),
                ),
            ) {
                this.scheduled.push(job);
            }

            // execute scheduled until it is blocked or done
            if let Some(job_result) = ready!(this.scheduled.poll_next_unpin(cx)) {
                match job_result {
                    (value, Either::Left(result)) => this.process_unfold(value, result?),
                    (value, Either::Right(result)) => {
                        // we have computed value associated with `init` node
                        if value == this.init {
                            // all jobs have to be completed and execution_tree empty
                            assert!(this.unscheduled.is_empty());
                            assert!(this.scheduled.is_empty());
                            return Poll::Ready(Ok(Some(result?)));
                        }
                        this.process_fold(value, result?);
                    }
                }
            }
        }
    }
}


/*
 * 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 thiserror::Error;

#[derive(Debug, Error)]
pub enum BoundedTraversalError {
    #[error("Programming error at {file}:{line}: {desc}")]
    ProgrammingError {
        desc: String,
        file: &'static str,
        line: u32,
    },
}

macro_rules! programming_error {
    ( $( $args:tt )* ) => {
        $crate::error::BoundedTraversalError::ProgrammingError {
            desc: format!( $( $args )* ),
            file: file!(),
            line: line!(),
        }
    };
}


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

//! Read the documentation of [bounded_traversal](crate::bounded_traversal),
//! [bounded_traversal_dag](crate::bounded_traversal_dag) and
//! [bounded_traversal_stream](crate::bounded_traversal_stream)

#[macro_use]
mod error;
pub use error::BoundedTraversalError;

mod tree;
pub use tree::bounded_traversal;

mod dag;
pub use dag::bounded_traversal_dag;
pub use dag::bounded_traversal_dag_limited;

mod stream;
pub use stream::bounded_traversal_stream;
pub use stream::limited_by_key_shardable;

mod ordered_stream;
pub use ordered_stream::bounded_traversal_limited_ordered_stream;
pub use ordered_stream::bounded_traversal_ordered_stream;

mod common;
pub use common::OrderedTraversal;

#[cfg(test)]
mod tests;

/// A type used frequently in fold-like invocations inside this module
pub type Iter<Out> = std::iter::Flatten<std::vec::IntoIter<Option<Out>>>;