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

//! Dice calculations relating to deferreds

use std::collections::HashMap;
use std::hash::Hash;
use std::pin::Pin;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::deferred::data::DeferredData;
use buck2_artifact::deferred::id::DeferredId;
use buck2_artifact::deferred::key::DeferredKey;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_common::result::ToUnsharedResultExt;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::base_deferred_key::BaseDeferredKeyDyn;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_events::dispatch::create_span;
use buck2_events::dispatch::Span;
use buck2_execute::artifact::artifact_dyn::ArtifactDyn;
use buck2_execute::digest_config::HasDigestConfig;
use buck2_execute::materialize::materializer::HasMaterializer;
use buck2_node::nodes::configured_frontend::ConfiguredTargetNodeCalculation;
use buck2_util::late_binding::LateBinding;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use futures::stream::FuturesUnordered;
use futures::Future;
use futures::FutureExt;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use more_futures::cancellation::CancellationContext;
use once_cell::sync::Lazy;

use crate::actions::artifact::get_artifact_fs::GetArtifactFs;
use crate::analysis::calculation::RuleAnalysisCalculation;
use crate::analysis::AnalysisResult;
use crate::artifact_groups::calculation::ArtifactGroupCalculation;
use crate::artifact_groups::ArtifactGroup;
use crate::bxl::calculation::BXL_CALCULATION_IMPL;
use crate::bxl::result::BxlResult;
use crate::deferred::types::BaseKey;
use crate::deferred::types::DeferredInput;
use crate::deferred::types::DeferredLookup;
use crate::deferred::types::DeferredRegistry;
use crate::deferred::types::DeferredResult;
use crate::deferred::types::DeferredValueAny;
use crate::deferred::types::DeferredValueAnyReady;
use crate::deferred::types::DeferredValueReady;
use crate::deferred::types::ResolveDeferredCtx;

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "ResolveDeferred({})", _0)]
pub struct DeferredResolve(pub DeferredKey);

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "ComputeDeferred({})", _0)]
pub struct DeferredCompute(pub DeferredKey);

#[async_trait]
pub trait DeferredCalculation {
    /// Computes and returns the evaluated value of an 'DeferredData'
    async fn compute_deferred_data<T: Send + Sync + 'static>(
        &self,
        data: &DeferredData<T>,
    ) -> anyhow::Result<DeferredValueReady<T>>;
}

#[async_trait]
impl DeferredCalculation for DiceComputations {
    async fn compute_deferred_data<T: Send + Sync + 'static>(
        &self,
        data: &DeferredData<T>,
    ) -> anyhow::Result<DeferredValueReady<T>> {
        if data.deferred_key().id().is_trivial() {
            let deferred = lookup_deferred(self, data.deferred_key()).await?;
            let deferred = deferred
                .get()?
                .as_trivial()
                .context("Invalid deferred")?
                .dupe();
            return DeferredValueAnyReady::TrivialDeferred(deferred).resolve(data);
        }

        let deferred = resolve_deferred(self, data.deferred_key()).await?;
        deferred.resolve(data)
    }
}

pub static EVAL_ANON_TARGET: LateBinding<
    for<'c> fn(
        &'c DiceComputations,
        Arc<dyn BaseDeferredKeyDyn>,
    ) -> Pin<Box<dyn Future<Output = anyhow::Result<AnalysisResult>> + Send + 'c>>,
> = LateBinding::new("EVAL_ANON_TARGET");

async fn lookup_deferred_inner(
    key: &BaseDeferredKey,
    dice: &DiceComputations,
) -> anyhow::Result<DeferredHolder> {
    match key {
        BaseDeferredKey::TargetLabel(target) => {
            let analysis = dice
                .get_analysis_result(target)
                .await?
                .require_compatible()?;

            Ok(DeferredHolder::Analysis(analysis))
        }
        BaseDeferredKey::BxlLabel(bxl) => {
            let bxl_result = BXL_CALCULATION_IMPL
                .get()?
                .eval_bxl(dice, bxl.dupe())
                .await?
                .bxl_result;

            Ok(DeferredHolder::Bxl(bxl_result))
        }
        BaseDeferredKey::AnonTarget(target) => Ok(DeferredHolder::Analysis(
            (EVAL_ANON_TARGET.get()?)(dice, target.dupe()).await?,
        )),
    }
}

struct PartialLookup {
    holder: DeferredHolder,
    id: DeferredId,
}

impl PartialLookup {
    fn get(&self) -> anyhow::Result<DeferredLookup<'_>> {
        self.holder.lookup_deferred(self.id)
    }
}

/// looks up an deferred
async fn lookup_deferred(
    dice: &DiceComputations,
    key: &DeferredKey,
) -> anyhow::Result<PartialLookup> {
    Ok(match key {
        DeferredKey::Base(target, id) => {
            let holder = lookup_deferred_inner(target, dice).await?;
            PartialLookup { holder, id: *id }
        }
        DeferredKey::Deferred(key, id) => {
            let deferred = compute_deferred(dice, key).await?;
            PartialLookup {
                holder: DeferredHolder::Deferred(deferred),
                id: *id,
            }
        }
    })
}

/// Fully resolve the deferred, including any deferreds it might have return when attempting
/// to calculate it.
async fn resolve_deferred(
    dice: &DiceComputations,
    deferred: &DeferredKey,
) -> anyhow::Result<DeferredValueAnyReady> {
    #[async_trait]
    impl Key for DeferredResolve {
        type Value = SharedResult<DeferredValueAnyReady>;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellation: &CancellationContext,
        ) -> Self::Value {
            let result = compute_deferred(ctx, &self.0).await?;
            match result.value() {
                DeferredValueAny::Ready(value) => Ok(value.dupe()),
                DeferredValueAny::Deferred(key) => resolve_deferred(ctx, key).await.shared_error(),
            }
        }

        fn equality(_: &Self::Value, _: &Self::Value) -> bool {
            // TODO(bobyf) consider if we want deferreds to be eq
            false
        }
    }

    dice.compute(&DeferredResolve(deferred.dupe()))
        .await?
        .unshared_error()
}

/// Computes and returns the untyped deferred at the given key. This does not fully resolve
/// the deferred as another deferred may be returned.
async fn compute_deferred(
    dice: &DiceComputations,
    deferred: &DeferredKey,
) -> anyhow::Result<DeferredResult> {
    #[async_trait]
    impl Key for DeferredCompute {
        type Value = SharedResult<DeferredResult>;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellation: &CancellationContext,
        ) -> Self::Value {
            let deferred = lookup_deferred(ctx, &self.0).await?;
            let deferred = deferred.get()?.as_complex();

            // We'll create the Span lazily when materialization hits it.
            let span = Lazy::new(|| deferred.span().map(create_span));

            let target_node_futs = FuturesUnordered::new();
            let deferreds_futs = FuturesUnordered::new();
            let mut materialized_artifacts = Vec::new();

            deferred.inputs().iter().for_each(|input| match input {
                DeferredInput::ConfiguredTarget(target) => target_node_futs.push(
                    ctx.get_configured_target_node(target)
                        .map(|res| anyhow::Ok((target.dupe(), res?.require_compatible()?))),
                ),
                DeferredInput::Deferred(deferred_key) => deferreds_futs.push(
                    resolve_deferred(ctx, deferred_key)
                        .map(|res| anyhow::Ok((deferred_key.dupe(), res?))),
                ),
                DeferredInput::MaterializedArtifact(artifact) => {
                    materialized_artifacts.push(ArtifactGroup::Artifact(artifact.dupe()));
                }
            });

            let materialized_artifacts_fut =
                self.create_materializer_futs(&materialized_artifacts, ctx, &span);

            // TODO(nga): do we need to compute artifacts?
            let (targets, deferreds, materialized_artifacts) = futures::future::join3(
                futures_pair_to_map(target_node_futs),
                futures_pair_to_map(deferreds_futs),
                materialized_artifacts_fut,
            )
            .await;

            let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(self.0.dupe())));
            let mut ctx = ResolveDeferredCtx::new(
                self.0.dupe(),
                targets?,
                deferreds?,
                materialized_artifacts?,
                &mut registry,
                ctx.global_data().get_io_provider().project_root().dupe(),
                ctx.global_data().get_digest_config(),
            );

            let mut execute = move || deferred.execute(&mut ctx);

            let res = match Lazy::into_value(span).unwrap_or_else(|init| init()) {
                Some(span) => {
                    span.wrap_closure(|| (execute(), buck2_data::DeferredEvaluationEnd {}))
                }
                None => execute(),
            };

            // TODO populate the deferred map
            Ok(DeferredResult::new(res?, registry.take_result()?))
        }

        fn equality(_: &Self::Value, _: &Self::Value) -> bool {
            false
        }

        fn validity(x: &Self::Value) -> bool {
            x.is_ok()
        }
    }

    impl DeferredCompute {
        fn create_materializer_futs<'a>(
            &'a self,
            materialized_artifacts: &'a Vec<ArtifactGroup>,
            ctx: &'a DiceComputations,
            span: &'a Lazy<Option<Span>, impl FnOnce() -> Option<Span>>,
        ) -> impl Future<Output = anyhow::Result<HashMap<Artifact, ProjectRelativePathBuf>>> + 'a
        {
            if materialized_artifacts.is_empty() {
                return async move { Ok(HashMap::new()) }.left_future();
            }
            // This is a bit suboptimal: we wait for all artifacts to be ready in order to
            // materialize any of them. However that is how we execute *all* local actions so in
            // the grand scheme of things that's probably not a huge deal.
            let materialized_artifacts_fut = {
                let artifact_futs = futures::future::try_join_all(
                    materialized_artifacts
                        .iter()
                        .map(|artifact| ctx.ensure_artifact_group(artifact)),
                );

                artifact_futs.and_then(move |_| async move {
                    let materializer = ctx.per_transaction_data().get_materializer();
                    let artifact_fs = ctx.get_artifact_fs().await?;

                    let fut = materialized_artifacts
                        .iter()
                        .map(|artifact| async {
                            let artifact = artifact
                                .unpack_artifact()
                                .expect("we only put Artifacts into this list")
                                .dupe();
                            let path = artifact.resolve_path(&artifact_fs)?;
                            materializer.ensure_materialized(vec![path.clone()]).await?;

                            anyhow::Ok((artifact, path))
                        })
                        .collect::<FuturesUnordered<_>>()
                        .try_collect::<HashMap<_, _>>();

                    match span.as_ref() {
                        Some(span) => {
                            span.create_child(buck2_data::DeferredPreparationStageStart {
                                stage: Some(buck2_data::MaterializedArtifacts {}.into()),
                            })
                            .wrap_future(
                                fut.map(|r| (r, buck2_data::DeferredPreparationStageEnd {})),
                            )
                            .await
                        }
                        None => fut.await,
                    }
                })
            };
            materialized_artifacts_fut.right_future()
        }
    }

    Ok(dice.compute(&DeferredCompute(deferred.dupe())).await??)
}

async fn futures_pair_to_map<K: Eq + Hash, V>(
    mut futs: FuturesUnordered<impl Future<Output = anyhow::Result<(K, V)>>>,
) -> anyhow::Result<HashMap<K, V>> {
    let mut res = HashMap::with_capacity(futs.len());
    while let Some(p) = futs.next().await {
        // terminate immediately once a fail occurs
        // TODO support keep going?

        let (key, value) = p?;
        res.insert(key, value);
    }
    Ok(res)
}

/// Represents an Analysis or Deferred result. Technically, we can treat analysis as a 'Deferred'
/// and get rid of this enum
enum DeferredHolder {
    Analysis(AnalysisResult),
    Bxl(Arc<BxlResult>),
    Deferred(DeferredResult),
}

impl DeferredHolder {
    fn lookup_deferred(&self, id: DeferredId) -> anyhow::Result<DeferredLookup<'_>> {
        match self {
            DeferredHolder::Analysis(result) => result.lookup_deferred(id),
            DeferredHolder::Deferred(result) => result.lookup_deferred(id),
            DeferredHolder::Bxl(result) => result.lookup_deferred(id),
        }
    }
}


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

//! 'DeferredData's are deferred work registered by a rule implementation that can be stored in
//! fields of providers.
//!
//! These are usually slow operations required to generate the information on 'Provider's.
//! For example, information that depend on the execution of some actions, or information that
//! require reading the filesystem.
//!
//! We model these asynchronous data to be stored on fields of 'Provider's as 'DeferredData'. Each
//! 'DeferredData' will hold an 'DeferredKey'. That key that maps to an implementation of
//! 'Deferred', which is the actual work to be ran when execution of the deferred is needed.

pub mod calculation;
pub mod types;


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

use std::any;
use std::any::type_name;
use std::any::Any;
use std::any::Demand;
use std::collections::HashMap;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::marker::PhantomData;
use std::sync::Arc;

use allocative::Allocative;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::deferred::data::DeferredData;
use buck2_artifact::deferred::id::DeferredId;
use buck2_artifact::deferred::key::DeferredKey;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_execute::digest_config::DigestConfig;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use dupe::Dupe;
use gazebo::variants::VariantName;
use indexmap::indexset;
use indexmap::IndexSet;
use itertools::Itertools;
use once_cell::sync::Lazy;
use thiserror::Error;

/// An asynchronous chunk of work that will be executed when requested.
/// The 'Deferred' can have "inputs" which are values that will be guaranteed to be ready to use
/// before the 'Deferred' is actually executed. These can be 'Artifact's, which means that those
/// 'Artifact's will be materialized and its corresponding 'Action's executed, or other
/// 'DeferredData', which means those 'Deferred' will be computed first.
///
/// `any::Provider` can be used to obtain data for introspection. At the moment of writing,
/// only `ProvideOutputs` can be extracted from deferreds with outputs.
pub trait Deferred: Allocative + any::Provider {
    type Output;

    /// the set of 'Deferred's that should be computed first before executing this 'Deferred'
    fn inputs(&self) -> &IndexSet<DeferredInput>;

    /// executes this 'Deferred', assuming all inputs and input artifacts are already computed
    fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValue<Self::Output>>;

    /// An optional stage to wrap execution in.
    fn span(&self) -> Option<buck2_data::span_start_event::Data> {
        None
    }
}

/// The context for executing a 'Deferred'.
pub trait DeferredCtx {
    fn get_configured_target(&self, label: &ConfiguredTargetLabel)
    -> Option<&ConfiguredTargetNode>;

    fn get_action_key(&self) -> String;

    fn get_deferred_data(&self, key: &DeferredKey) -> Option<DeferredValueAnyReady>;

    fn get_materialized_artifact(&self, artifact: &Artifact) -> Option<&ProjectRelativePath>;

    fn registry(&mut self) -> &mut DeferredRegistry;

    fn project_filesystem(&self) -> &ProjectRoot;

    fn digest_config(&self) -> DigestConfig;
}

/// DeferredCtx with already resolved values
pub struct ResolveDeferredCtx<'a> {
    key: DeferredKey,
    configured_targets: HashMap<ConfiguredTargetLabel, ConfiguredTargetNode>,
    deferreds: HashMap<DeferredKey, DeferredValueAnyReady>,
    materialized_artifacts: HashMap<Artifact, ProjectRelativePathBuf>,
    registry: &'a mut DeferredRegistry,
    project_filesystem: ProjectRoot,
    digest_config: DigestConfig,
}

impl<'a> ResolveDeferredCtx<'a> {
    pub fn new(
        key: DeferredKey,
        configured_targets: HashMap<ConfiguredTargetLabel, ConfiguredTargetNode>,
        deferreds: HashMap<DeferredKey, DeferredValueAnyReady>,
        materialized_artifacts: HashMap<Artifact, ProjectRelativePathBuf>,
        registry: &'a mut DeferredRegistry,
        project_filesystem: ProjectRoot,
        digest_config: DigestConfig,
    ) -> Self {
        Self {
            key,
            configured_targets,
            deferreds,
            materialized_artifacts,
            registry,
            project_filesystem,
            digest_config,
        }
    }
}

impl<'a> DeferredCtx for ResolveDeferredCtx<'a> {
    fn get_configured_target(
        &self,
        label: &ConfiguredTargetLabel,
    ) -> Option<&ConfiguredTargetNode> {
        self.configured_targets.get(label)
    }

    fn get_action_key(&self) -> String {
        self.key.action_key()
    }

    fn get_deferred_data(&self, key: &DeferredKey) -> Option<DeferredValueAnyReady> {
        self.deferreds.get(key).map(|b| b.dupe())
    }

    fn get_materialized_artifact(&self, artifact: &Artifact) -> Option<&ProjectRelativePath> {
        self.materialized_artifacts
            .get(artifact)
            .map(|x| x.as_ref())
    }

    fn registry(&mut self) -> &mut DeferredRegistry {
        self.registry
    }

    fn project_filesystem(&self) -> &ProjectRoot {
        &self.project_filesystem
    }

    fn digest_config(&self) -> DigestConfig {
        self.digest_config
    }
}

/// input to a deferred that needs to be computed first before executing
#[derive(Clone, Debug, Eq, PartialEq, Hash, Allocative)]
pub enum DeferredInput {
    ConfiguredTarget(ConfiguredTargetLabel),
    Deferred(DeferredKey),
    MaterializedArtifact(Artifact),
}

/// The base key. We can actually get rid of this and just use 'DeferredKey' if rule analysis is an
/// 'Deferred' itself. This is used to construct the composed 'DeferredKey::Deferred' or
/// 'DeferredKey::Base' type.
#[derive(Allocative)]
pub enum BaseKey {
    Base(BaseDeferredKey),
    // While DeferredKey is Dupe, it has quite a lot of Arc's inside it, so maybe an Arc here makes sense?
    // Maybe not?
    #[cfg_attr(feature = "gazebo_lint", allow(gazebo_lint_arc_on_dupe))]
    Deferred(Arc<DeferredKey>),
}

impl BaseKey {
    fn make_key(&self, id: DeferredId) -> DeferredKey {
        match self {
            BaseKey::Base(base) => DeferredKey::Base(base.dupe(), id),
            BaseKey::Deferred(base) => DeferredKey::Deferred(base.dupe(), id),
        }
    }
}

/// Implemented by all trivial deferreds.
pub trait TrivialDeferred: Allocative + AnyValue + Debug + Send + Sync {
    /// Convert the object to an AnyValue object
    fn as_any_value(&self) -> &dyn AnyValue;

    /// Obtain deferred-specific debug data.
    ///
    /// This function is copied from `any::Provider` trait, which cannot be implemented
    /// for `Arc<RegisteredAction>`.
    fn provide<'a>(&'a self, demand: &mut Demand<'a>);
}

#[derive(Allocative)]
pub struct TrivialDeferredValue(pub Arc<dyn TrivialDeferred>);

impl any::Provider for TrivialDeferredValue {
    fn provide<'a>(&'a self, demand: &mut Demand<'a>) {
        self.0.provide(demand)
    }
}

impl DeferredAny for TrivialDeferredValue {
    fn inputs(&self) -> &IndexSet<DeferredInput> {
        static INPUTS: Lazy<IndexSet<DeferredInput>> = Lazy::new(IndexSet::new);
        &INPUTS
    }

    fn execute(&self, _ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValueAny> {
        Ok(DeferredValueAny::Ready(
            DeferredValueAnyReady::TrivialDeferred(self.0.dupe()),
        ))
    }

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

    fn type_name(&self) -> &str {
        self.0.type_name()
    }

    fn span(&self) -> Option<buck2_data::span_start_event::Data> {
        // No evaluation Span if you never evaluate.
        None
    }
}

#[derive(Allocative)]
pub enum DeferredTableEntry {
    Trivial(TrivialDeferredValue),
    Complex(Box<dyn DeferredAny>),
}

impl any::Provider for DeferredTableEntry {
    fn provide<'a>(&'a self, demand: &mut Demand<'a>) {
        match self {
            Self::Trivial(v) => v.provide(demand),
            Self::Complex(v) => v.provide(demand),
        }
    }
}

impl DeferredAny for DeferredTableEntry {
    fn inputs(&self) -> &IndexSet<DeferredInput> {
        match self {
            Self::Trivial(v) => v.inputs(),
            Self::Complex(v) => v.inputs(),
        }
    }

    fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValueAny> {
        match self {
            Self::Trivial(v) => v.execute(ctx),
            Self::Complex(v) => v.execute(ctx),
        }
    }

    fn as_any(&self) -> &dyn Any {
        match self {
            Self::Trivial(v) => v.as_any(),
            Self::Complex(v) => v.as_any(),
        }
    }

    fn type_name(&self) -> &str {
        match self {
            Self::Trivial(v) => v.type_name(),
            Self::Complex(v) => v.type_name(),
        }
    }

    fn span(&self) -> Option<buck2_data::span_start_event::Data> {
        match self {
            Self::Trivial(..) => None,
            Self::Complex(v) => v.span(),
        }
    }
}

#[derive(Allocative)]
enum DeferredRegistryEntry {
    Set(DeferredTableEntry),
    Pending,
}

/// The registry for creating 'DeferredData's and registering 'Deferred's
#[derive(Allocative)]
pub struct DeferredRegistry {
    base_key: BaseKey,
    registry: Vec<DeferredRegistryEntry>,
}

#[derive(Allocative)]
#[allocative(bound = "")]
pub struct ReservedDeferredData<T>(DeferredData<T>);

impl<'a, T> ReservedDeferredData<T>
where
    T: Clone + Send + Sync + 'static,
{
    fn new(id: DeferredData<T>) -> Self {
        Self(id)
    }

    pub fn data(&'a self) -> &'a DeferredData<T> {
        &self.0
    }
}

#[derive(Allocative)]
#[allocative(bound = "")]
pub struct ReservedTrivialDeferredData<T>(DeferredData<T>);

impl<'a, T> ReservedTrivialDeferredData<T>
where
    T: Clone + Send + Sync + 'static,
{
    fn new(id: DeferredData<T>) -> Self {
        Self(id)
    }

    pub fn data(&'a self) -> &'a DeferredData<T> {
        &self.0
    }
}

impl DeferredRegistry {
    pub fn new(base_key: BaseKey) -> Self {
        Self {
            base_key,
            registry: Vec::new(),
        }
    }

    /// Reserves a 'DeferredData', with it's underlying key, on the promise that it should be bound
    /// to a 'Deferred' before 'take_result' is called.
    pub fn reserve<D>(&mut self) -> ReservedDeferredData<D>
    where
        D: Clone + Send + Sync + 'static,
    {
        let id = DeferredId {
            id: self.registry.len().try_into().unwrap(),
            trivial: false,
        };
        self.registry.push(DeferredRegistryEntry::Pending);
        ReservedDeferredData::new(DeferredData::unchecked_new(self.base_key.make_key(id)))
    }

    /// Reserves a 'DeferredData', with it's underlying key, on the promise that it should be bound
    /// to a 'Deferred' before 'take_result' is called.
    pub fn reserve_trivial<D>(&mut self) -> ReservedTrivialDeferredData<D>
    where
        D: Clone + Send + Sync + 'static,
    {
        let id = DeferredId {
            id: self.registry.len().try_into().unwrap(),
            trivial: true,
        };
        self.registry.push(DeferredRegistryEntry::Pending);
        ReservedTrivialDeferredData::new(DeferredData::unchecked_new(self.base_key.make_key(id)))
    }

    /// binds a reserved 'ReservedDeferredData' to a 'Deferred'
    pub fn bind<D, T>(&mut self, reserved: ReservedDeferredData<T>, d: D) -> DeferredData<T>
    where
        D: Deferred<Output = T> + Send + Sync + 'static,
        T: Allocative + Clone + Debug + Send + Sync + 'static,
    {
        let id = reserved.0.deferred_key().id().as_usize();

        match self.registry.get_mut(id) {
            Some(entry @ DeferredRegistryEntry::Pending) => {
                *entry = DeferredRegistryEntry::Set(DeferredTableEntry::Complex(Box::new(d)));
            }
            _ => {
                panic!("the reserved should always be in pending");
            }
        }

        reserved.0
    }

    /// binds a reserved 'ReservedDeferredData' to a 'Deferred'
    pub fn bind_trivial<D>(
        &mut self,
        reserved: ReservedTrivialDeferredData<D>,
        d: D,
    ) -> DeferredData<D>
    where
        D: TrivialDeferred + Clone + 'static,
    {
        let id = reserved.0.deferred_key().id().as_usize();

        match self.registry.get_mut(id) {
            Some(entry @ DeferredRegistryEntry::Pending) => {
                *entry = DeferredRegistryEntry::Set(DeferredTableEntry::Trivial(
                    TrivialDeferredValue(Arc::new(d) as _),
                ));
            }
            _ => {
                panic!("the reserved should always be in pending");
            }
        }

        reserved.0
    }

    /// creates a new 'DeferredData'
    pub fn defer<
        D: Deferred<Output = T> + Send + Sync + 'static,
        T: Allocative + Clone + Debug + Send + Sync + 'static,
    >(
        &mut self,
        d: D,
    ) -> DeferredData<T> {
        let id = DeferredId {
            id: self.registry.len().try_into().unwrap(),
            trivial: false,
        };
        self.registry
            .push(DeferredRegistryEntry::Set(DeferredTableEntry::Complex(
                Box::new(d),
            )));
        DeferredData::unchecked_new(self.base_key.make_key(id))
    }

    /// creates a new 'DeferredData'
    pub fn defer_trivial<D>(&mut self, d: D) -> DeferredData<D>
    where
        D: TrivialDeferred + Clone + 'static,
    {
        let id = DeferredId {
            id: self.registry.len().try_into().unwrap(),
            trivial: true,
        };
        self.registry
            .push(DeferredRegistryEntry::Set(DeferredTableEntry::Trivial(
                TrivialDeferredValue(Arc::new(d) as _),
            )));
        DeferredData::unchecked_new(self.base_key.make_key(id))
    }

    /// performs a mapping function over an 'DeferredData'
    pub fn map<T, U, F>(&mut self, orig: &DeferredData<T>, f: F) -> DeferredData<U>
    where
        T: Allocative + Clone + Debug + Send + Sync + 'static,
        U: Allocative + Clone + Debug + Send + Sync + 'static,
        F: Fn(&T, &mut dyn DeferredCtx) -> DeferredValue<U> + Send + Sync + 'static,
    {
        #[derive(Allocative)]
        #[allocative(bound = "")]
        struct Map<T, U, F> {
            orig: IndexSet<DeferredInput>,
            #[allocative(skip)]
            f: F,
            p: PhantomData<(T, U)>,
        }

        impl<T, U, F> any::Provider for Map<T, U, F>
        where
            T: Allocative + Send + Sync + 'static,
            F: Fn(&T, &mut dyn DeferredCtx) -> DeferredValue<U> + 'static,
            U: Allocative + 'static,
        {
            fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
        }

        impl<T, U, F> Deferred for Map<T, U, F>
        where
            T: Allocative + Send + Sync + 'static,
            F: Fn(&T, &mut dyn DeferredCtx) -> DeferredValue<U> + 'static,
            U: Allocative + 'static,
        {
            type Output = U;

            fn inputs(&self) -> &IndexSet<DeferredInput> {
                &self.orig
            }

            fn execute(
                &self,
                ctx: &mut dyn DeferredCtx,
            ) -> anyhow::Result<DeferredValue<Self::Output>> {
                let orig = match self.orig.iter().exactly_one() {
                    Ok(DeferredInput::Deferred(orig)) => orig,
                    _ => unreachable!(),
                };
                Ok((self.f)(
                    ctx.get_deferred_data(orig)
                        .unwrap()
                        .downcast::<T>()
                        .unwrap(),
                    ctx,
                ))
            }
        }

        self.defer(Map {
            orig: indexset![DeferredInput::Deferred(orig.deferred_key().dupe())],
            f,
            p: PhantomData,
        })
    }

    pub fn take_result(self) -> anyhow::Result<Vec<DeferredTableEntry>> {
        self.registry
            .into_iter()
            .enumerate()
            .map(|(i, e)| match e {
                DeferredRegistryEntry::Set(e) => anyhow::Ok(e),
                DeferredRegistryEntry::Pending => {
                    Err(DeferredErrors::UnboundReservedDeferred(i).into())
                }
            })
            .collect()
    }
}

#[derive(Debug, Error)]
pub enum DeferredErrors {
    #[error("no deferred found for deferred id `{0}`")]
    DeferredNotFound(u32),
    #[error("reserved deferred id of `{0:?}` was never bound")]
    UnboundReservedDeferred(usize),
}

pub enum DeferredLookup<'a> {
    Trivial(&'a TrivialDeferredValue),
    Complex(&'a (dyn DeferredAny + 'static)),
}

impl<'a> DeferredLookup<'a> {
    pub fn as_complex(&self) -> &'a (dyn DeferredAny + 'static) {
        match self {
            Self::Trivial(v) => *v as _,
            Self::Complex(v) => *v,
        }
    }

    pub fn as_trivial(&self) -> Option<&'a Arc<dyn TrivialDeferred>> {
        match self {
            Self::Trivial(v) => Some(&v.0),
            Self::Complex(..) => None,
        }
    }
}

/// Contains all the deferreds generated by analyzing a particular rule implementation
#[derive(Clone, Debug, Dupe, Allocative)]
pub struct DeferredResult(Arc<DeferredResultData>);

#[derive(Allocative)]
struct DeferredResultData {
    deferreds: Vec<DeferredTableEntry>,
    value: DeferredValueAny,
}

#[derive(Clone, Dupe, Allocative)]
pub struct DeferredTable(Arc<Vec<DeferredTableEntry>>);

impl Debug for DeferredResultData {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "DeferredResult(value=`{:?}`, {} deferreds)",
            self.value,
            self.deferreds.len(),
        )
    }
}

impl Debug for DeferredTable {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "DeferredTable({} deferreds)", self.0.len())
    }
}

impl DeferredTable {
    pub fn new(deferreds: Vec<DeferredTableEntry>) -> Self {
        Self(Arc::new(deferreds))
    }

    /// looks up an 'Deferred' given the id
    pub fn lookup_deferred(&self, id: DeferredId) -> anyhow::Result<DeferredLookup<'_>> {
        match self.0.get(id.as_usize()) {
            Some(DeferredTableEntry::Complex(deferred)) => Ok(DeferredLookup::Complex(&**deferred)),
            Some(DeferredTableEntry::Trivial(value)) => Ok(DeferredLookup::Trivial(value)),
            None => Err(anyhow::anyhow!(DeferredErrors::DeferredNotFound(id.id))),
        }
    }

    /// Iterator on the DeferredTable which converts a `DeferredTableEntry` to a `DeferredLookup`
    pub fn iter(&self) -> impl Iterator<Item = DeferredLookup<'_>> {
        self.0.iter().map(|deferred| match deferred {
            DeferredTableEntry::Complex(deferred) => DeferredLookup::Complex(&**deferred),
            DeferredTableEntry::Trivial(value) => DeferredLookup::Trivial(value),
        })
    }
}

impl DeferredResult {
    pub fn new(value: DeferredValueAny, deferreds: Vec<DeferredTableEntry>) -> Self {
        Self(Arc::new(DeferredResultData { deferreds, value }))
    }

    /// looks up an 'Deferred' given the id
    pub fn lookup_deferred(&self, id: DeferredId) -> anyhow::Result<DeferredLookup<'_>> {
        match self.0.deferreds.get(id.as_usize()) {
            Some(DeferredTableEntry::Complex(deferred)) => Ok(DeferredLookup::Complex(&**deferred)),
            Some(DeferredTableEntry::Trivial(value)) => Ok(DeferredLookup::Trivial(value)),
            None => Err(anyhow::anyhow!(DeferredErrors::DeferredNotFound(id.id))),
        }
    }

    pub fn value(&self) -> &DeferredValueAny {
        &self.0.value
    }
}

/// typed value computed by a deferred. This can either be a completed calculation, in which case
/// a value is returned, or defer to another 'Deferred' computation.
#[derive(Allocative)]
pub enum DeferredValue<T> {
    Ready(T),
    Deferred(DeferredData<T>),
}

/// Enum of AnyValue or TrivialDeferreds.
#[derive(Allocative, Debug, Dupe, Clone)]
pub enum DeferredValueAnyReady {
    AnyValue(Arc<dyn AnyValue>),
    TrivialDeferred(Arc<dyn TrivialDeferred>),
}

impl DeferredValueAnyReady {
    pub fn downcast<T: Send + 'static>(&self) -> anyhow::Result<&T> {
        match self {
            Self::AnyValue(v) => v.downcast(),
            Self::TrivialDeferred(v) => v.as_any_value().downcast(),
        }
    }

    pub fn downcast_into<T: Send + 'static>(self) -> anyhow::Result<DeferredValueReady<T>> {
        // Check if it can downcast to T
        self.downcast::<T>()?;

        Ok(DeferredValueReady {
            inner: self,
            _type: PhantomData,
        })
    }

    pub fn resolve<T: Send + 'static>(
        self,
        _data: &DeferredData<T>,
    ) -> anyhow::Result<DeferredValueReady<T>> {
        self.downcast_into()
    }
}

/// This is a `Any` that has been checked to contain a T and can therefore provide &T infallibly
#[derive(Allocative, Debug, Dupe, Clone)]
pub struct DeferredValueReady<T> {
    inner: DeferredValueAnyReady,
    _type: PhantomData<T>,
}

impl<T> std::ops::Deref for DeferredValueReady<T>
where
    T: Send + 'static,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        // This was checked earlier
        self.inner.downcast::<T>().unwrap()
    }
}

/// untyped value computed by the deferred. This is same as 'DeferredValue', but with 'T' as an
/// 'ValueAny'
#[derive(Debug, VariantName, Clone, Dupe, Allocative)]
pub enum DeferredValueAny {
    Ready(DeferredValueAnyReady),
    Deferred(DeferredKey),
}

impl DeferredValueAny {
    fn ready<T: Allocative + Clone + Debug + Send + Sync + 'static>(t: T) -> Self {
        Self::Ready(DeferredValueAnyReady::AnyValue(Arc::new(t)))
    }

    fn defer<T>(k: DeferredData<T>) -> Self {
        Self::Deferred(k.into_deferred_key())
    }
}

/// An 'Any' that is the return type of a 'Deferred'. This is box cloneable, and castable.
pub trait AnyValue: Allocative + Any + Debug + Send + Sync {
    fn into_any(&self) -> &(dyn Any);

    fn type_name(&self) -> &str;
}

impl<T> AnyValue for T
where
    T: Allocative + Any + Debug + Send + Sync,
{
    fn into_any(&self) -> &dyn Any {
        self
    }

    fn type_name(&self) -> &str {
        type_name::<T>()
    }
}

impl dyn AnyValue {
    pub(crate) fn downcast<T: Send + 'static>(&self) -> anyhow::Result<&T> {
        match self.into_any().downcast_ref::<T>() {
            Some(t) => Ok(t),
            None => Err(anyhow::anyhow!(
                "Cannot cast Deferred of value type `{}` into type `{}`",
                self.type_name(),
                type_name::<T>()
            )),
        }
    }
}

/// untyped deferred
pub trait DeferredAny: Allocative + any::Provider + Send + Sync {
    /// the set of 'Deferred's that should be computed first before executing this 'Deferred'
    fn inputs(&self) -> &IndexSet<DeferredInput>;

    /// executes this 'Deferred', assuming all inputs and input artifacts are already computed
    fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValueAny>;

    fn as_any(&self) -> &dyn Any;

    fn type_name(&self) -> &str;

    /// An optional stage to wrap execution in.
    fn span(&self) -> Option<buck2_data::span_start_event::Data>;
}

impl dyn DeferredAny {
    pub fn downcast<T: Deferred + Send + 'static>(&self) -> anyhow::Result<&T> {
        match self.as_any().downcast_ref::<T>() {
            Some(t) => Ok(t),
            None => Err(anyhow::anyhow!(
                "Cannot cast Deferred of value type `{}` into type `{}`",
                self.type_name(),
                type_name::<T>()
            )),
        }
    }
}

impl<D, T> DeferredAny for D
where
    D: Deferred<Output = T> + Send + Sync + Any + 'static,
    T: Allocative + Clone + Debug + Send + Sync + 'static,
{
    fn inputs(&self) -> &IndexSet<DeferredInput> {
        self.inputs()
    }

    fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValueAny> {
        match self.execute(ctx)? {
            DeferredValue::Ready(t) => Ok(DeferredValueAny::ready(t)),
            DeferredValue::Deferred(d) => Ok(DeferredValueAny::defer(d)),
        }
    }

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

    fn type_name(&self) -> &str {
        type_name::<D>()
    }

    fn span(&self) -> Option<buck2_data::span_start_event::Data> {
        D::span(self)
    }
}

pub mod testing {

    use buck2_artifact::deferred::key::DeferredKey;
    use gazebo::variants::VariantName;

    use crate::deferred::types::AnyValue;
    use crate::deferred::types::DeferredResult;
    use crate::deferred::types::DeferredTable;
    use crate::deferred::types::DeferredTableEntry;
    use crate::deferred::types::DeferredValueAny;
    use crate::deferred::types::DeferredValueAnyReady;

    pub trait DeferredValueAnyExt {
        fn assert_ready(self) -> DeferredValueAnyReady;
        fn assert_deferred(self) -> DeferredKey;
    }

    impl DeferredValueAnyExt for DeferredValueAny {
        fn assert_ready(self) -> DeferredValueAnyReady {
            match self {
                DeferredValueAny::Ready(v) => v,
                x => panic!("Expected deferred to be Ready but was {}", x.variant_name()),
            }
        }

        fn assert_deferred(self) -> DeferredKey {
            match self {
                DeferredValueAny::Deferred(k) => k,
                x => panic!(
                    "Expected deferred to be Deferred but was {}",
                    x.variant_name()
                ),
            }
        }
    }

    pub trait AnyValueExt {
        /// tests if the any is equal to 't'
        fn eq<T: PartialEq + 'static>(&self, t: &T) -> bool;
    }

    impl AnyValueExt for dyn AnyValue {
        fn eq<T: PartialEq + 'static>(&self, t: &T) -> bool {
            self.into_any()
                .downcast_ref::<T>()
                .map_or(false, |x| x == t)
        }
    }

    pub trait DeferredAnalysisResultExt {
        fn get_registered(&self) -> &Vec<DeferredTableEntry>;
    }

    impl DeferredAnalysisResultExt for DeferredResult {
        fn get_registered(&self) -> &Vec<DeferredTableEntry> {
            &self.0.deferreds
        }
    }

    impl DeferredAnalysisResultExt for DeferredTable {
        fn get_registered(&self) -> &Vec<DeferredTableEntry> {
            &self.0
        }
    }
}

#[cfg(test)]
mod tests {
    use std::any;
    use std::any::Demand;
    use std::fmt;
    use std::fmt::Debug;
    use std::fmt::Formatter;
    use std::sync::Arc;

    use allocative::Allocative;
    use buck2_artifact::deferred::data::DeferredData;
    use buck2_artifact::deferred::id::DeferredId;
    use buck2_artifact::deferred::key::DeferredKey;
    use buck2_core::base_deferred_key::BaseDeferredKey;
    use buck2_core::configuration::data::ConfigurationData;
    use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
    use buck2_core::fs::project::ProjectRoot;
    use buck2_core::target::label::ConfiguredTargetLabel;
    use buck2_execute::digest_config::DigestConfig;
    use dupe::Dupe;
    use indexmap::indexset;
    use indexmap::IndexSet;

    use super::AnyValue;
    use super::TrivialDeferred;
    use crate::deferred::types::testing::DeferredValueAnyExt;
    use crate::deferred::types::BaseKey;
    use crate::deferred::types::Deferred;
    use crate::deferred::types::DeferredAny;
    use crate::deferred::types::DeferredCtx;
    use crate::deferred::types::DeferredInput;
    use crate::deferred::types::DeferredRegistry;
    use crate::deferred::types::DeferredTable;
    use crate::deferred::types::DeferredValue;
    use crate::deferred::types::DeferredValueAny;
    use crate::deferred::types::DeferredValueAnyReady;
    use crate::deferred::types::ResolveDeferredCtx;

    #[derive(Clone, PartialEq, Eq, Allocative)]
    #[allocative(bound = "")]
    struct FakeDeferred<T> {
        inputs: IndexSet<DeferredInput>,
        #[allocative(skip)]
        val: T,
    }

    impl<T> Debug for FakeDeferred<T> {
        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
            write!(f, "FakeDeferred")
        }
    }

    impl<T: Clone> any::Provider for FakeDeferred<T> {
        fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
    }

    impl<T: Clone> Deferred for FakeDeferred<T> {
        type Output = T;

        fn inputs(&self) -> &IndexSet<DeferredInput> {
            &self.inputs
        }

        fn execute(&self, _ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValue<T>> {
            Ok(DeferredValue::Ready(self.val.clone()))
        }
    }

    impl TrivialDeferred for FakeDeferred<i32> {
        fn as_any_value(&self) -> &dyn AnyValue {
            self
        }

        fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
    }

    #[derive(Clone, Debug, PartialEq, Eq, Allocative)]
    #[allocative(bound = "")]
    struct DeferringDeferred<T> {
        inputs: IndexSet<DeferredInput>,
        #[allocative(skip)]
        defer: FakeDeferred<T>,
    }

    impl<T: Clone + Debug + Allocative + Send + Sync + 'static> any::Provider for DeferringDeferred<T> {
        fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
    }

    impl<T: Clone + Debug + Allocative + Send + Sync + 'static> Deferred for DeferringDeferred<T> {
        type Output = T;

        fn inputs(&self) -> &IndexSet<DeferredInput> {
            &self.inputs
        }

        fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValue<T>> {
            Ok(DeferredValue::Deferred(
                ctx.registry().defer(self.defer.clone()),
            ))
        }
    }

    fn make_resolved<T: Clone + Debug + Allocative + Send + Sync + 'static>(
        data: &DeferredData<T>,
        deferred: &FakeDeferred<T>,
    ) -> (DeferredKey, DeferredValueAnyReady) {
        (
            data.deferred_key().dupe(),
            DeferredValueAnyReady::AnyValue(Arc::new(deferred.val.clone())),
        )
    }

    fn dummy_base() -> BaseDeferredKey {
        BaseDeferredKey::TargetLabel(ConfiguredTargetLabel::testing_parse(
            "cell//pkg:foo",
            ConfigurationData::testing_new(),
        ))
    }

    fn dummy_project_filesystem() -> ProjectRoot {
        let cwd = if cfg!(windows) {
            AbsNormPath::new("c:/tmp").unwrap().to_owned()
        } else {
            AbsNormPath::new("/dev/null").unwrap().to_owned()
        };
        ProjectRoot::new_unchecked(cwd)
    }

    #[test]
    fn register_deferred() -> anyhow::Result<()> {
        let target = dummy_base();
        let mut registry = DeferredRegistry::new(BaseKey::Base(target.dupe()));

        let deferred = FakeDeferred {
            inputs: IndexSet::new(),
            val: 2,
        };

        let deferred_data = registry.defer(deferred);

        let result = registry.take_result()?;

        let mut ctx = DeferredRegistry::new(BaseKey::Deferred(Arc::new(
            deferred_data.deferred_key().dupe(),
        )));

        assert_eq!(
            *result
                .get(deferred_data.deferred_key().id().as_usize())
                .unwrap()
                .execute(&mut ResolveDeferredCtx::new(
                    deferred_data.deferred_key().dupe(),
                    Default::default(),
                    Default::default(),
                    Default::default(),
                    &mut ctx,
                    dummy_project_filesystem(),
                    DigestConfig::testing_default(),
                ))
                .unwrap()
                .assert_ready()
                .resolve(&deferred_data)?,
            2
        );

        Ok(())
    }

    #[test]
    fn mapping_async_data() -> anyhow::Result<()> {
        let base = BaseDeferredKey::TargetLabel(ConfiguredTargetLabel::testing_parse(
            "cell//pkg:foo",
            ConfigurationData::testing_new(),
        ));
        let mut registry = DeferredRegistry::new(BaseKey::Base(base.dupe()));

        let deferred = FakeDeferred {
            inputs: IndexSet::new(),
            val: 0,
        };

        let deferred_data = registry.defer(deferred.clone());
        let mapped = registry.map(&deferred_data, |x, _ctx| DeferredValue::Ready(x + 1));

        let result = registry.take_result()?;
        let mapped_deferred = result.get(mapped.deferred_key().id().as_usize()).unwrap();

        assert_eq!(
            mapped_deferred.inputs(),
            &indexset![DeferredInput::Deferred(deferred_data.deferred_key().dupe())]
        );

        let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(
            deferred_data.deferred_key().dupe(),
        )));
        let mut resolved = ResolveDeferredCtx::new(
            deferred_data.deferred_key().dupe(),
            Default::default(),
            vec![make_resolved(&deferred_data, &deferred)]
                .into_iter()
                .collect(),
            Default::default(),
            &mut registry,
            dummy_project_filesystem(),
            DigestConfig::testing_default(),
        );

        assert_eq!(
            *mapped_deferred
                .execute(&mut resolved)
                .unwrap()
                .assert_ready()
                .downcast::<i32>()?,
            1
        );

        Ok(())
    }

    #[test]
    fn register_nested_deferred() -> anyhow::Result<()> {
        let target =
            ConfiguredTargetLabel::testing_parse("cell//pkg:foo", ConfigurationData::testing_new());
        let id = DeferredId {
            id: 1,
            trivial: false,
        };

        let base = DeferredKey::Base(BaseDeferredKey::TargetLabel(target.dupe()), id);
        let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(base)));

        let deferred = FakeDeferred {
            inputs: IndexSet::new(),
            val: 2,
        };

        let deferring_deferred = DeferringDeferred {
            inputs: Default::default(),
            defer: deferred,
        };

        let deferring_deferred_data = registry.defer(deferring_deferred);
        let result = registry.take_result()?;

        let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(
            deferring_deferred_data.deferred_key().dupe(),
        )));
        let exec_result = result
            .get(deferring_deferred_data.deferred_key().id().as_usize())
            .unwrap()
            .execute(&mut ResolveDeferredCtx::new(
                deferring_deferred_data.deferred_key().dupe(),
                Default::default(),
                Default::default(),
                Default::default(),
                &mut registry,
                dummy_project_filesystem(),
                DigestConfig::testing_default(),
            ))
            .unwrap();

        let deferred_key = match exec_result {
            DeferredValueAny::Ready(_) => panic!("expected a deferred"),
            DeferredValueAny::Deferred(deferred) => deferred,
        };

        assert_eq!(
            deferred_key,
            DeferredKey::Deferred(
                Arc::new(deferring_deferred_data.deferred_key().dupe()),
                DeferredId {
                    id: 0,
                    trivial: false
                }
            )
        );

        let result = registry.take_result()?;
        let deferred = result.get(deferred_key.id().as_usize()).unwrap();

        let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(deferred_key.dupe())));
        assert_eq!(
            *deferred
                .execute(&mut ResolveDeferredCtx::new(
                    deferred_key,
                    Default::default(),
                    Default::default(),
                    Default::default(),
                    &mut registry,
                    dummy_project_filesystem(),
                    DigestConfig::testing_default(),
                ))
                .unwrap()
                .assert_ready()
                .downcast::<i32>()?,
            2
        );

        Ok(())
    }

    #[test]
    fn reserving_deferred() -> anyhow::Result<()> {
        let base = BaseDeferredKey::TargetLabel(ConfiguredTargetLabel::testing_parse(
            "cell//pkg:foo",
            ConfigurationData::testing_new(),
        ));
        let mut registry = DeferredRegistry::new(BaseKey::Base(base));

        let deferred = FakeDeferred {
            inputs: IndexSet::new(),
            val: 2,
        };

        let reserved = registry.reserve();
        let reserved_deferred_data = reserved.data().dupe();

        let deferred_data = registry.bind(reserved, deferred);

        assert_eq!(deferred_data, reserved_deferred_data);

        let result = registry.take_result()?;

        let mut registry = DeferredRegistry::new(BaseKey::Deferred(Arc::new(
            deferred_data.deferred_key().dupe(),
        )));

        let key = deferred_data.deferred_key().dupe();
        assert_eq!(
            *result
                .get(deferred_data.deferred_key().id().as_usize())
                .unwrap()
                .execute(&mut ResolveDeferredCtx::new(
                    key,
                    Default::default(),
                    Default::default(),
                    Default::default(),
                    &mut registry,
                    dummy_project_filesystem(),
                    DigestConfig::testing_default(),
                ))
                .unwrap()
                .assert_ready()
                .resolve(&deferred_data)?,
            2
        );

        Ok(())
    }

    #[test]
    fn reserving_deferred_unbound() {
        let base = BaseDeferredKey::TargetLabel(ConfiguredTargetLabel::testing_parse(
            "cell//pkg:foo",
            ConfigurationData::testing_new(),
        ));
        let mut registry = DeferredRegistry::new(BaseKey::Base(base));

        let _reserved = registry.reserve::<FakeDeferred<()>>();
        let _reserved1 = registry.reserve::<FakeDeferred<()>>();

        assert_eq!(registry.take_result().is_err(), true);
    }

    #[test]
    fn trivial_deferred() -> anyhow::Result<()> {
        let base = BaseDeferredKey::TargetLabel(ConfiguredTargetLabel::testing_parse(
            "cell//pkg:foo",
            ConfigurationData::testing_new(),
        ));
        let mut registry = DeferredRegistry::new(BaseKey::Base(base));

        let deferred0 = FakeDeferred {
            inputs: IndexSet::new(),
            val: 123,
        };

        let deferred1 = FakeDeferred {
            inputs: IndexSet::new(),
            val: 456,
        };
        let deferred_data0 = registry.defer_trivial(deferred0.clone());
        let deferred_data1 = registry.reserve_trivial();
        let deferred_data1 = registry.bind_trivial(deferred_data1, deferred1.clone());

        let result = DeferredTable::new(registry.take_result()?);

        assert_eq!(
            *DeferredValueAnyReady::TrivialDeferred(
                result
                    .lookup_deferred(deferred_data0.deferred_key().id())?
                    .as_trivial()
                    .unwrap()
                    .dupe()
            )
            .resolve(&deferred_data0)?,
            deferred0
        );

        assert_eq!(
            *DeferredValueAnyReady::TrivialDeferred(
                result
                    .lookup_deferred(deferred_data1.deferred_key().id())?
                    .as_trivial()
                    .unwrap()
                    .dupe()
            )
            .resolve(&deferred_data1)?,
            deferred1
        );

        Ok(())
    }
}


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

use std::any;
use std::any::Demand;
use std::mem;
use std::sync::Arc;

use allocative::Allocative;
use buck2_artifact::actions::key::ActionKey;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_artifact::artifact::provide_outputs::ProvideOutputs;
use buck2_artifact::deferred::data::DeferredData;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::ProvidersName;
use buck2_events::dispatch::get_dispatcher;
use buck2_interpreter::print_handler::EventDispatcherPrintHandler;
use buck2_interpreter::types::label::Label;
use dupe::Dupe;
use gazebo::prelude::*;
use indexmap::indexset;
use indexmap::IndexSet;
use starlark::collections::SmallMap;
use starlark::environment::Module;
use starlark::eval::Evaluator;
use starlark::values::dict::Dict;
use starlark::values::tuple::TupleRef;
use starlark::values::OwnedFrozenValue;
use thiserror::Error;

use crate::actions::key::ActionKeyExt;
use crate::actions::RegisteredAction;
use crate::analysis::registry::AnalysisRegistry;
use crate::bxl;
use crate::deferred::types::BaseKey;
use crate::deferred::types::Deferred;
use crate::deferred::types::DeferredCtx;
use crate::deferred::types::DeferredInput;
use crate::deferred::types::DeferredRegistry;
use crate::deferred::types::DeferredValue;
use crate::interpreter::rule_defs::artifact::associated::AssociatedArtifacts;
use crate::interpreter::rule_defs::artifact::StarlarkArtifact;
use crate::interpreter::rule_defs::artifact::StarlarkArtifactValue;
use crate::interpreter::rule_defs::artifact::StarlarkDeclaredArtifact;
use crate::interpreter::rule_defs::context::AnalysisContext;

/// The artifacts that are returned are dynamic actions, which depend on the `DynamicLambda`
/// to get their real `RegisteredAction`.
#[derive(Clone, Debug, Allocative)]
pub(crate) struct DynamicAction {
    // A singleton pointing at a DynamicLambda
    inputs: IndexSet<DeferredInput>,
    index: usize,
}

impl DynamicAction {
    pub fn new(deferred: &DeferredData<DynamicLambdaOutput>, index: usize) -> Self {
        Self {
            inputs: indexset![DeferredInput::Deferred(deferred.deferred_key().dupe())],
            index,
        }
    }
}

/// The lambda captured by `dynamic_output`, alongside the other required data.
#[derive(Clone, Debug, Allocative)]
pub(crate) struct DynamicLambda {
    /// the owner that defined this lambda
    owner: BaseDeferredKey,
    /// Things required by the lambda (wrapped in DeferredInput)
    dynamic: IndexSet<DeferredInput>,
    /// Things I am allowed to use as inputs, but don't wait for
    inputs: IndexSet<Artifact>,
    /// Things I produce
    outputs: Vec<BuildArtifact>,
    /// A Starlark pair of the attributes and a lambda function that binds the outputs given a context
    attributes_lambda: OwnedFrozenValue,
}

impl DynamicLambda {
    pub(crate) fn new(
        owner: BaseDeferredKey,
        dynamic: IndexSet<Artifact>,
        inputs: IndexSet<Artifact>,
        outputs: Vec<BuildArtifact>,
    ) -> Self {
        let mut depends = IndexSet::with_capacity(dynamic.len() + 1);
        match &owner {
            BaseDeferredKey::TargetLabel(target) => {
                depends.insert(DeferredInput::ConfiguredTarget(target.dupe()));
            }
            BaseDeferredKey::BxlLabel(_) => {
                // do nothing. This is for grabbing the execution platform, which for bxl, we
                // hard code to a local execution.
            }
            BaseDeferredKey::AnonTarget(_) => {
                // This will return an error later, so doesn't need to have the dependency
            }
        }
        depends.extend(dynamic.into_iter().map(DeferredInput::MaterializedArtifact));
        Self {
            owner,
            dynamic: depends,
            inputs,
            outputs,
            attributes_lambda: Default::default(),
        }
    }

    pub(crate) fn bind(&mut self, attributes_lambda: OwnedFrozenValue) {
        self.attributes_lambda = attributes_lambda;
    }
}

/// The `Output` from `DynamicLambda`.
#[derive(Clone, Debug, Allocative)]
pub(crate) struct DynamicLambdaOutput {
    /// The actions the DynamicLambda produces, in the right order.
    /// `DynamicAction.index` is an index into this Vec.
    output: Vec<ActionKey>,
}

impl any::Provider for DynamicAction {
    fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
}

impl Deferred for DynamicAction {
    type Output = Arc<RegisteredAction>;

    fn inputs(&self) -> &IndexSet<DeferredInput> {
        &self.inputs
    }

    fn execute(&self, ctx: &mut dyn DeferredCtx) -> anyhow::Result<DeferredValue<Self::Output>> {
        let id = match self.inputs.iter().into_singleton() {
            Some(DeferredInput::Deferred(x)) => x,
            _ => unreachable!("DynamicAction must have a single Deferred as as inputs"),
        };
        let val = ctx.get_deferred_data(id).unwrap();
        let key = val
            .downcast::<DynamicLambdaOutput>()?
            .output
            .get(self.index)
            .map_or_else(
                || {
                    Err(anyhow::anyhow!(
                        "Unexpected index in DynamicAction (internal error)"
                    ))
                },
                Ok,
            )?;
        Ok(DeferredValue::Deferred(key.deferred_data().dupe()))
    }
}

#[derive(Debug, Error)]
enum DynamicLambdaError {
    #[error("dynamic_output and anon_target cannot be used together (yet)")]
    AnonTargetIncompatible,
}

impl any::Provider for DynamicLambda {
    fn provide<'a>(&'a self, demand: &mut Demand<'a>) {
        demand.provide_value_with(|| ProvideOutputs(Ok(self.outputs.clone())));
    }
}

impl Deferred for DynamicLambda {
    type Output = DynamicLambdaOutput;

    fn inputs(&self) -> &IndexSet<DeferredInput> {
        &self.dynamic
    }

    fn execute(
        &self,
        deferred_ctx: &mut dyn DeferredCtx,
    ) -> anyhow::Result<DeferredValue<Self::Output>> {
        let env = Module::new();
        let heap = env.heap();
        let print = EventDispatcherPrintHandler(get_dispatcher());
        let mut outputs = SmallMap::with_capacity(self.outputs.len());
        let mut declared_outputs = IndexSet::with_capacity(self.outputs.len());

        let analysis_registry = {
            let mut eval = Evaluator::new(&env);
            eval.set_print_handler(&print);

            let data = TupleRef::from_value(self.attributes_lambda.owned_value(env.frozen_heap()))
                .unwrap();
            assert_eq!(data.len(), 2);
            let attributes = data.content()[0];
            let lambda = data.content()[1];

            let execution_platform = {
                match &self.owner {
                    BaseDeferredKey::TargetLabel(target) => {
                        let configured_target = deferred_ctx.get_configured_target(target).unwrap();

                        configured_target.execution_platform_resolution().dupe()
                    }
                    BaseDeferredKey::BxlLabel(_) => {
                        (*bxl::execution_platform::EXECUTION_PLATFORM).dupe()
                    }
                    BaseDeferredKey::AnonTarget(_) => {
                        return Err(DynamicLambdaError::AnonTargetIncompatible.into());
                    }
                }
            };

            // The DeferredCtx has a registry it wants us to use as &mut.
            // The AnalysisRegistry wants ownership of a registry.
            // To overcome the difference, we create a fake registry, swap it with the one in deferred,
            // and swap back after AnalysisRegistry completes.
            let fake_registry = DeferredRegistry::new(BaseKey::Base(self.owner.dupe()));

            let deferred = mem::replace(deferred_ctx.registry(), fake_registry);
            let mut registry = AnalysisRegistry::new_from_owner_and_deferred(
                self.owner.dupe(),
                execution_platform,
                deferred,
            )?;
            registry.set_action_key(Arc::from(deferred_ctx.get_action_key()));

            let mut artifacts = SmallMap::with_capacity(self.inputs.len());
            let fs = deferred_ctx.project_filesystem();
            for x in &self.dynamic {
                let x = match x {
                    DeferredInput::MaterializedArtifact(x) => x,
                    DeferredInput::ConfiguredTarget(_) => continue,
                    _ => unreachable!("DynamicLambda only depends on artifact and target"),
                };
                let k = heap.alloc(StarlarkArtifact::new(x.dupe()));
                let path = deferred_ctx.get_materialized_artifact(x).unwrap();
                let v = heap.alloc(StarlarkArtifactValue::new(
                    x.dupe(),
                    path.to_owned(),
                    fs.dupe(),
                ));
                artifacts.insert_hashed(k.get_hashed()?, v);
            }

            for x in &self.outputs {
                let k = heap.alloc(StarlarkArtifact::new(Artifact::from(x.dupe())));
                let declared =
                    registry.declare_dynamic_output(x.get_path().dupe(), x.output_type());
                declared_outputs.insert(declared.dupe());
                let v = heap.alloc(StarlarkDeclaredArtifact::new(
                    None,
                    declared,
                    AssociatedArtifacts::new(),
                ));
                outputs.insert_hashed(k.get_hashed()?, v);
            }

            let artifacts = Dict::new(artifacts);
            let outputs = Dict::new(outputs);

            let ctx = heap.alloc_typed(AnalysisContext::new(
                heap,
                attributes,
                match &self.owner {
                    BaseDeferredKey::TargetLabel(target) => Some(heap.alloc_typed(Label::new(
                        ConfiguredProvidersLabel::new(target.dupe(), ProvidersName::Default),
                    ))),
                    BaseDeferredKey::BxlLabel(target) | BaseDeferredKey::AnonTarget(target) => {
                        target.configured_label().map(|configured_target_label| {
                            heap.alloc_typed(Label::new(ConfiguredProvidersLabel::new(
                                configured_target_label,
                                ProvidersName::Default,
                            )))
                        })
                    }
                },
                registry,
                deferred_ctx.digest_config(),
            ));

            eval.eval_function(
                lambda,
                &[ctx.to_value(), heap.alloc(artifacts), heap.alloc(outputs)],
                &[],
            )?;
            ctx.assert_no_promises()?;

            ctx.take_state()
        };

        let (_frozen_env, deferred) = analysis_registry.finalize(&env)?(env)?;
        let _fake_registry = mem::replace(deferred_ctx.registry(), deferred);

        // TODO(ndmitchell): Check we don't use anything not in `inputs`

        let output: anyhow::Result<Vec<_>> = declared_outputs
            .into_iter()
            .map(|x| anyhow::Ok(x.ensure_bound()?.action_key().dupe()))
            .collect();
        Ok(DeferredValue::Ready(DynamicLambdaOutput {
            output: output?,
        }))
    }

    fn span(&self) -> Option<buck2_data::span_start_event::Data> {
        let owner = self.owner.to_proto().into();
        Some(buck2_data::DynamicLambdaStart { owner: Some(owner) }.into())
    }
}


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

pub(crate) mod deferred;
pub(crate) mod registry;


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

use allocative::Allocative;
use anyhow::Context;
use buck2_artifact::actions::key::ActionKey;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::artifact_type::OutputArtifact;
use buck2_artifact::deferred::id::DeferredId;
use buck2_core::base_deferred_key::BaseDeferredKey;
use dupe::Dupe;
use indexmap::IndexSet;

use crate::actions::key::ActionKeyExt;
use crate::analysis::registry::AnalysisValueFetcher;
use crate::deferred::types::DeferredRegistry;
use crate::deferred::types::ReservedDeferredData;
use crate::dynamic::deferred::DynamicAction;
use crate::dynamic::deferred::DynamicLambda;
use crate::dynamic::deferred::DynamicLambdaOutput;

#[derive(Allocative)]
pub(crate) struct DynamicRegistry {
    owner: BaseDeferredKey,
    pending: Vec<(ReservedDeferredData<DynamicLambdaOutput>, DynamicLambda)>,
}

impl DynamicRegistry {
    pub fn new(owner: BaseDeferredKey) -> Self {
        Self {
            owner,
            pending: Vec::new(),
        }
    }

    pub fn register(
        &mut self,
        dynamic: IndexSet<Artifact>,
        inputs: IndexSet<Artifact>,
        outputs: IndexSet<OutputArtifact>,
        registry: &mut DeferredRegistry,
    ) -> anyhow::Result<DeferredId> {
        let reserved = registry.reserve::<DynamicLambdaOutput>();
        let outputs = outputs
            .iter()
            .enumerate()
            .map(|(i, output)| {
                let output_id = registry.defer(DynamicAction::new(reserved.data(), i));
                let bound = output
                    .bind(ActionKey::new(output_id))?
                    .as_base_artifact()
                    .dupe();
                Ok(bound)
            })
            .collect::<anyhow::Result<_>>()?;
        let lambda = DynamicLambda::new(self.owner.dupe(), dynamic, inputs, outputs);
        let lambda_id = reserved.data().deferred_key().id();
        self.pending.push((reserved, lambda));
        Ok(lambda_id)
    }

    pub fn ensure_bound(
        self,
        registry: &mut DeferredRegistry,
        analysis_value_fetcher: &AnalysisValueFetcher,
    ) -> anyhow::Result<()> {
        for (key, mut data) in self.pending {
            let id = key.data().deferred_key().id();

            let fv = analysis_value_fetcher
                .get(id)?
                .with_context(|| format!("Key is missing in AnalysisValueFetcher: {:?}", id))?;

            data.bind(fv);
            registry.bind(key, data);
        }
        Ok(())
    }
}


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

use buck2_interpreter::build_context::STARLARK_PATH_FROM_BUILD_CONTEXT;
use buck2_interpreter::path::StarlarkPath;
use either::Either;
use itertools::Itertools;
use starlark::collections::SmallMap;
use starlark::docs::DocString;
use starlark::docs::DocStringKind;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::values::Value;
use starlark_map::small_set::SmallSet;

use crate::interpreter::rule_defs::provider::callable::UserProviderCallable;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetDefinition;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetError;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetOperations;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetProjectionKind;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetProjectionSpec;

#[derive(Debug, thiserror::Error)]
enum NativesError {
    #[error("non-unique field names: [{}]", .0.iter().map(|s| format!("`{}`", s)).join(", "))]
    NonUniqueFields(Vec<String>),
    #[error("`transitive_set()` can only be used in `bzl` files")]
    TransitiveSetOnlyInBzl,
}

#[starlark_module]
pub fn register_provider(builder: &mut GlobalsBuilder) {
    /// Create a `"provider"` type that can be returned from `rule` implementations.
    /// Used to pass information from a rule to the things that depend on it.
    /// Typically named with an `Info` suffix.
    ///
    /// ```python
    /// GroovyLibraryInfo(fields = [
    ///     "objects",  # a list of artifacts
    ///     "options",  # a string containing compiler options
    /// ])
    /// ```
    ///
    /// Given a dependency you can obtain the provider with `my_dep[GroovyLibraryInfo]`
    /// which returns either `None` or a value of type `GroovyLibraryInfo`.
    ///
    /// For providers that accumulate upwards a transitive set is often a good choice.
    fn provider(
        #[starlark(require=named, default = "")] doc: &str,
        #[starlark(require=named)] fields: Either<Vec<String>, SmallMap<&str, &str>>,
        eval: &mut Evaluator,
    ) -> anyhow::Result<UserProviderCallable> {
        let docstring = DocString::from_docstring(DocStringKind::Starlark, doc);
        let path = (STARLARK_PATH_FROM_BUILD_CONTEXT.get()?)(eval)?.path();

        let (field_names, field_docs) = match fields {
            Either::Left(f) => {
                let docs = vec![None; f.len()];
                let field_names: SmallSet<String> = f.iter().cloned().collect();
                if field_names.len() != f.len() {
                    return Err(NativesError::NonUniqueFields(f).into());
                }
                (field_names, docs)
            }
            Either::Right(fields_with_docs) => {
                let mut field_names = SmallSet::with_capacity(fields_with_docs.len());
                let mut field_docs = Vec::with_capacity(fields_with_docs.len());
                for (name, docs) in fields_with_docs {
                    let inserted = field_names.insert(name.to_owned());
                    assert!(inserted);
                    field_docs.push(DocString::from_docstring(DocStringKind::Starlark, docs));
                }
                (field_names, field_docs)
            }
        };
        Ok(UserProviderCallable::new(
            path.into_owned(),
            docstring,
            field_docs,
            field_names,
        ))
    }
}

#[starlark_module]
pub fn register_transitive_set(builder: &mut GlobalsBuilder) {
    fn transitive_set<'v>(
        args_projections: Option<SmallMap<String, Value<'v>>>,
        json_projections: Option<SmallMap<String, Value<'v>>>,
        reductions: Option<SmallMap<String, Value<'v>>>,
        eval: &mut Evaluator,
    ) -> anyhow::Result<TransitiveSetDefinition<'v>> {
        // TODO(cjhopman): Reductions could do similar signature checking.
        let projections: SmallMap<_, _> = args_projections
            .into_iter()
            .flat_map(|v| v.into_iter())
            .map(|(k, v)| {
                (
                    k,
                    TransitiveSetProjectionSpec {
                        kind: TransitiveSetProjectionKind::Args,
                        projection: v,
                    },
                )
            })
            .chain(
                json_projections
                    .into_iter()
                    .flat_map(|v| v.into_iter())
                    .map(|(k, v)| {
                        (
                            k,
                            TransitiveSetProjectionSpec {
                                kind: TransitiveSetProjectionKind::Json,
                                projection: v,
                            },
                        )
                    }),
            )
            .collect();

        // Both kinds of projections take functions with the same signature.
        for (name, spec) in projections.iter() {
            // We should probably be able to require that the projection returns a parameters_spec, but
            // we don't depend on this type-checking and we'd just error out later when calling it if it
            // were wrong.
            if let Some(v) = spec.projection.parameters_spec() {
                if v.len() != 1 {
                    return Err(TransitiveSetError::ProjectionSignatureError {
                        name: name.clone(),
                    }
                    .into());
                }
            };
        }

        let starlark_path: StarlarkPath = (STARLARK_PATH_FROM_BUILD_CONTEXT.get()?)(eval)?;
        Ok(TransitiveSetDefinition::new(
            match starlark_path {
                StarlarkPath::LoadFile(import_path) => import_path.clone(),
                _ => return Err(NativesError::TransitiveSetOnlyInBzl.into()),
            },
            TransitiveSetOperations {
                projections,
                reductions: reductions.unwrap_or_default(),
            },
        ))
    }
}


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

pub mod build_defs;
pub mod more;
pub mod rule_defs;


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

use buck2_interpreter::functions::more::MoreFunctions;
use buck2_interpreter::functions::more::MORE_FUNCTIONS;

use crate::interpreter::build_defs::register_provider;
use crate::interpreter::build_defs::register_transitive_set;
use crate::interpreter::rule_defs::cmd_args::register_cmd_args;
use crate::interpreter::rule_defs::command_executor_config::register_command_executor_config;
use crate::interpreter::rule_defs::provider::registration::register_builtin_providers;
use crate::interpreter::rule_defs::register_rule_defs;

pub(crate) fn init_more_functions() {
    MORE_FUNCTIONS.init(MoreFunctions {
        register_builtin_providers,
        register_cmd_args,
        register_command_executor_config,
        register_provider,
        register_rule_defs,
        register_transitive_set,
    });
}


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

use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use buck2_common::executor_config::CacheUploadBehavior;
use buck2_common::executor_config::CommandExecutorConfig;
use buck2_common::executor_config::CommandGenerationOptions;
use buck2_common::executor_config::Executor;
use buck2_common::executor_config::HybridExecutionLevel;
use buck2_common::executor_config::LocalExecutorOptions;
use buck2_common::executor_config::PathSeparatorKind;
use buck2_common::executor_config::RemoteEnabledExecutor;
use buck2_common::executor_config::RemoteExecutorOptions;
use buck2_common::executor_config::RemoteExecutorUseCase;
use derive_more::Display;
use starlark::any::ProvidesStaticType;
use starlark::environment::GlobalsBuilder;
use starlark::values::dict::DictRef;
use starlark::values::none::NoneOr;
use starlark::values::none::NoneType;
use starlark::values::starlark_value;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::Value;
use starlark::values::ValueLike;
use thiserror::Error;

#[derive(Debug, Error)]
enum CommandExecutorConfigErrors {
    #[error("expected a dict, got `{0}` (type `{1}`)")]
    RePropertiesNotADict(String, String),
    #[error("expected `{0}` to be set")]
    MissingField(&'static str),
    #[error("invalid value in `{0}`")]
    InvalidField(&'static str),
    #[error(
        "executor config must specify at least `local_enabled = True` or `remote_enabled = True`"
    )]
    NoExecutor,
}

#[derive(Debug, Display, NoSerialize, ProvidesStaticType, Allocative)]
#[display(fmt = "{:?}", _0)]
pub struct StarlarkCommandExecutorConfig(pub Arc<CommandExecutorConfig>);

starlark_simple_value!(StarlarkCommandExecutorConfig);

#[starlark_value(type = "command_executor_config")]
impl<'v> StarlarkValue<'v> for StarlarkCommandExecutorConfig {}

#[starlark_module]
pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
    /// Contains configurations for how actions should be executed
    ///
    /// * `local_enabled` : Whether to use local execution for this execution platform.
    /// If both remote_enabled and local_enabled are `True`, we will use the hybrid executor
    /// * `remote_enabled`: Whether to use remote execution for this execution platform
    /// * `remote_cache_enabled`: Whether to query RE caches
    /// * `remote_execution_properties`: Properties for remote execution for this platform
    /// * `remote_execution_action_key`: A component to inject into the action key
    /// This should typically used to inject variability into the action key so that
    /// it's different across e.g. build modes (RE uses the action key for things like expected memory utilization)
    /// * `remote_execution_max_input_files_mebibytes`: The maximum input file size (in bytes) that remote execution can support
    /// * `remote_execution_queue_time_threshold_s`: The maximum time in seconds we are willing to wait
    /// in the RE queue for remote execution to start running our action
    /// * `remote_execution_use_case`: The use case to use when communicating with RE
    /// * `use_limited_hybrid`: Whether to use the limited hybrid executor
    /// * `allow_limited_hybrid_fallbacks`: Whether to allow fallbacks
    /// * `allow_hybrid_fallbacks_on_failure`: Whether to allow fallbacks when the result is failure (i.e. the command failed on the primary, but the infra worked)
    /// * `use_windows_path_separators`: Whether to use Windows path separators in command line arguments
    /// * `use_persistent workers`: Whether to use persistent workers for local execution if they are available
    /// * `allow_cache_uploads`: Whether to upload local actions to the RE cache
    /// * `max_cache_upload_mebibytes`: Maximum size to upload in cache uploads
    /// * `experimental_low_pass_filter`: Whether to use the experimental low pass filter
    /// * `remote_output_paths`: How to express output paths to RE
    #[starlark(as_type = StarlarkCommandExecutorConfig)]
    fn CommandExecutorConfig<'v>(
        #[starlark(require = named)] local_enabled: bool,
        #[starlark(require = named)] remote_enabled: bool,
        #[starlark(default = NoneOr::None, require = named)] remote_cache_enabled: NoneOr<bool>,
        #[starlark(default = NoneType, require = named)] remote_execution_properties: Value<'v>,
        #[starlark(default = NoneType, require = named)] remote_execution_action_key: Value<'v>,
        #[starlark(default = NoneOr::None, require = named)]
        remote_execution_max_input_files_mebibytes: NoneOr<i32>,
        #[starlark(default = NoneOr::None, require = named)]
        remote_execution_queue_time_threshold_s: NoneOr<i32>,
        #[starlark(default = NoneType, require = named)] remote_execution_use_case: Value<'v>,
        #[starlark(default = false, require = named)] use_limited_hybrid: bool,
        #[starlark(default = false, require = named)] allow_limited_hybrid_fallbacks: bool,
        #[starlark(default = false, require = named)] allow_hybrid_fallbacks_on_failure: bool,
        #[starlark(default = false, require = named)] use_windows_path_separators: bool,
        #[starlark(default = false, require = named)] use_persistent_workers: bool,
        #[starlark(default = false, require = named)] allow_cache_uploads: bool,
        #[starlark(default = NoneOr::None, require = named)] max_cache_upload_mebibytes: NoneOr<
            i32,
        >,
        #[starlark(default = false, require = named)] experimental_low_pass_filter: bool,
        #[starlark(default = NoneOr::None, require = named)] remote_output_paths: NoneOr<&str>,
    ) -> anyhow::Result<StarlarkCommandExecutorConfig> {
        let command_executor_config = {
            let remote_execution_max_input_files_mebibytes =
                remote_execution_max_input_files_mebibytes.into_option();

            let remote_execution_queue_time_threshold_s =
                remote_execution_queue_time_threshold_s.into_option();

            let max_cache_upload_mebibytes = max_cache_upload_mebibytes.into_option();

            let re_properties = if remote_execution_properties.is_none() {
                None
            } else {
                let re_properties = DictRef::from_value(remote_execution_properties.to_value())
                    .ok_or_else(|| {
                        CommandExecutorConfigErrors::RePropertiesNotADict(
                            remote_execution_properties.to_value().to_repr(),
                            remote_execution_properties.to_value().get_type().to_owned(),
                        )
                    })?;

                Some(
                    re_properties
                        .iter()
                        .map(|(k, v)| (k.to_str(), v.to_str()))
                        .collect(),
                )
            };

            let re_use_case = if remote_execution_use_case.is_none() {
                None
            } else {
                let re_use_case = remote_execution_use_case
                    .unpack_str()
                    .context("remote_execution_use_case is not a string")?;
                Some(RemoteExecutorUseCase::new(re_use_case.to_owned()))
            };

            let re_action_key = if remote_execution_action_key.is_none() {
                None
            } else {
                let re_action_key = remote_execution_action_key
                    .unpack_str()
                    .context("remote_execution_action_key is not a string")?;
                Some(re_action_key.to_owned())
            };

            let local_options = if local_enabled {
                Some(LocalExecutorOptions {
                    use_persistent_workers,
                })
            } else {
                None
            };
            let remote_options = if remote_enabled {
                let re_max_input_files_bytes = remote_execution_max_input_files_mebibytes
                    .map(u64::try_from)
                    .transpose()
                    .context("remote_execution_max_input_files_mebibytes is negative")?
                    .map(|b| b * 1024 * 1024);

                let re_max_queue_time_ms = remote_execution_queue_time_threshold_s
                    .map(u64::try_from)
                    .transpose()
                    .context("remote_execution_queue_time_threshold_s is negative")?
                    .map(|t| t * 1000);

                Some(RemoteExecutorOptions {
                    re_max_input_files_bytes,
                    re_max_queue_time_ms,
                })
            } else {
                None
            };

            let fallback_on_failure = allow_hybrid_fallbacks_on_failure;

            let hybrid_level = match (use_limited_hybrid, allow_limited_hybrid_fallbacks) {
                (true, true) => HybridExecutionLevel::Fallback {
                    fallback_on_failure,
                },
                (true, false) => HybridExecutionLevel::Limited,
                (false, _) => HybridExecutionLevel::Full {
                    fallback_on_failure,
                    low_pass_filter: experimental_low_pass_filter,
                },
            };

            let max_cache_upload_bytes = max_cache_upload_mebibytes
                .map(u64::try_from)
                .transpose()
                .context("max_cache_upload_mebibytes is negative")?
                .map(|b| b * 1024 * 1024);

            let cache_upload_behavior = if allow_cache_uploads {
                CacheUploadBehavior::Enabled {
                    max_bytes: max_cache_upload_bytes,
                }
            } else {
                CacheUploadBehavior::Disabled
            };

            // FIXME: This should probably default to `remote_enabled` and not `true`, but
            // historically this has been `true`, so we should probably migrate our defs to set
            // `remote_cache_enabled = True` explicitly first.
            let remote_cache_default = if buck2_core::is_open_source() {
                remote_enabled
            } else {
                true
            };
            let remote_cache_enabled = remote_cache_enabled
                .into_option()
                .unwrap_or(remote_cache_default);

            let executor = match (local_options, remote_options, remote_cache_enabled) {
                (local, Some(remote), remote_cache_enabled) => {
                    let executor = match local {
                        Some(local) => RemoteEnabledExecutor::Hybrid {
                            local,
                            remote,
                            level: hybrid_level,
                        },
                        None => RemoteEnabledExecutor::Remote(remote),
                    };

                    Executor::RemoteEnabled {
                        executor,
                        re_properties: re_properties.context(
                            CommandExecutorConfigErrors::MissingField(
                                "remote_execution_properties",
                            ),
                        )?,
                        re_use_case: re_use_case
                            .context(CommandExecutorConfigErrors::MissingField("re_use_case"))?,
                        re_action_key,
                        cache_upload_behavior,
                        remote_cache_enabled,
                    }
                }
                (Some(local), None, true) => Executor::RemoteEnabled {
                    executor: RemoteEnabledExecutor::Local(local),
                    // FIXME: We need a migration flip the default for remote_cache_enabled to
                    // remote_enabled first.
                    re_properties: re_properties.unwrap_or_default(),
                    re_use_case: re_use_case.unwrap_or_else(RemoteExecutorUseCase::buck2_default),
                    re_action_key,
                    cache_upload_behavior,
                    remote_cache_enabled: true,
                },
                (Some(local), None, false) => Executor::Local(local),
                (None, None, _) => {
                    return Err(CommandExecutorConfigErrors::NoExecutor.into());
                }
            };

            let output_paths_behavior = remote_output_paths
                .into_option()
                .map(|s| s.parse())
                .transpose()
                .context(CommandExecutorConfigErrors::InvalidField(
                    "remote_output_paths",
                ))?
                .unwrap_or_default();

            CommandExecutorConfig {
                executor,
                options: CommandGenerationOptions {
                    path_separator: if use_windows_path_separators {
                        PathSeparatorKind::Windows
                    } else {
                        PathSeparatorKind::Unix
                    },
                    output_paths_behavior,
                },
            }
        };

        Ok(StarlarkCommandExecutorConfig(Arc::new(
            command_executor_config,
        )))
    }
}


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

use std::cell::RefCell;
use std::cell::RefMut;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;

use allocative::Allocative;
use buck2_execute::digest_config::DigestConfig;
use buck2_interpreter::types::label::Label;
use buck2_util::late_binding::LateBinding;
use derive_more::Display;
use dice::DiceComputations;
use starlark::any::ProvidesStaticType;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::eval::Evaluator;
use starlark::typing::Ty;
use starlark::values::starlark_value;
use starlark::values::structs::StructRef;
use starlark::values::type_repr::StarlarkTypeRepr;
use starlark::values::AllocValue;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOfUnchecked;
use starlark::values::ValueTyped;

use crate::analysis::registry::AnalysisRegistry;

/// Functions to allow users to interact with the Actions registry.
///
/// Accessed via `ctx.actions.<function>`
#[derive(
    ProvidesStaticType,
    Debug,
    Display,
    Trace,
    NoSerialize,
    Allocative,
    StarlarkDocs
)]
#[display(fmt = "<ctx.actions>")]
pub struct AnalysisActions<'v> {
    /// Use a RefCell/Option so when we are done with it, without obtaining exclusive access,
    /// we can take the internal state without having to clone it.
    pub state: RefCell<Option<AnalysisRegistry<'v>>>,
    /// Copies from the ctx, so we can capture them for `dynamic`.
    pub attributes: ValueOfUnchecked<'v, StructRef<'v>>,
    /// Digest configuration to use when interpreting digests passed in analysis.
    pub digest_config: DigestConfig,
}

impl<'v> AnalysisActions<'v> {
    pub fn state(&self) -> RefMut<AnalysisRegistry<'v>> {
        RefMut::map(self.state.borrow_mut(), |x| {
            x.as_mut().expect("state to be present during execution")
        })
    }

    pub async fn run_promises(
        &self,
        dice: &DiceComputations,
        eval: &mut Evaluator<'v, '_>,
        description: String,
    ) -> anyhow::Result<()> {
        // We need to loop here because running the promises evaluates promise.map, which might produce more promises.
        // We keep going until there are no promises left.
        loop {
            let promises = self.state().take_promises();
            if let Some(promises) = promises {
                promises
                    .run_promises(dice, eval, description.clone())
                    .await?;
            } else {
                break;
            }
        }
        Ok(())
    }
}

#[starlark_value(type = "actions", StarlarkTypeRepr, UnpackValue)]
impl<'v> StarlarkValue<'v> for AnalysisActions<'v> {
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(|builder| {
            (ANALYSIS_ACTIONS_METHODS_ACTIONS.get().unwrap())(builder);
            (ANALYSIS_ACTIONS_METHODS_ANON_TARGET.get().unwrap())(builder);
        })
    }
}

impl<'v> AllocValue<'v> for AnalysisActions<'v> {
    fn alloc_value(self, heap: &'v Heap) -> Value<'v> {
        heap.alloc_complex_no_freeze(self)
    }
}

struct RefAnalysisAction<'v>(&'v AnalysisActions<'v>);

impl<'v> StarlarkTypeRepr for RefAnalysisAction<'v> {
    fn starlark_type_repr() -> Ty {
        AnalysisActions::starlark_type_repr()
    }
}

impl<'v> UnpackValue<'v> for RefAnalysisAction<'v> {
    fn unpack_value(value: Value<'v>) -> Option<Self> {
        Some(RefAnalysisAction(
            value.downcast_ref::<AnalysisActions>().unwrap(),
        ))
    }
}

#[derive(
    ProvidesStaticType,
    Debug,
    Trace,
    NoSerialize,
    Allocative,
    StarlarkDocs
)]
pub struct AnalysisContext<'v> {
    attrs: ValueOfUnchecked<'v, StructRef<'v>>,
    pub actions: ValueTyped<'v, AnalysisActions<'v>>,
    /// Only `None` when running a `dynamic_output` action from Bxl.
    label: Option<ValueTyped<'v, Label>>,
}

impl<'v> Display for AnalysisContext<'v> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "<ctx")?;
        if let Some(label) = &self.label {
            write!(f, " label=\"{}\"", label)?;
        }
        write!(f, " attrs=...")?;
        write!(f, " actions=...")?;
        write!(f, ">")?;
        Ok(())
    }
}

impl<'v> AnalysisContext<'v> {
    /// The context that is provided to users' UDR implementation functions. Comprised of things like attribute values, actions, etc
    pub fn new(
        heap: &'v Heap,
        attrs: Value<'v>,
        label: Option<ValueTyped<'v, Label>>,
        registry: AnalysisRegistry<'v>,
        digest_config: DigestConfig,
    ) -> Self {
        let attrs = ValueOfUnchecked::new_checked(attrs).unwrap();

        Self {
            attrs,
            actions: heap.alloc_typed(AnalysisActions {
                state: RefCell::new(Some(registry)),
                attributes: attrs,
                digest_config,
            }),
            label,
        }
    }

    pub(crate) fn assert_no_promises(&self) -> anyhow::Result<()> {
        self.actions.state().assert_no_promises()
    }

    /// Must take an `AnalysisContext` which has never had `take_state` called on it before.
    pub fn take_state(&self) -> AnalysisRegistry<'v> {
        self.actions
            .state
            .borrow_mut()
            .take()
            .expect("nothing to have stolen state yet")
    }
}

#[starlark_value(type = "context")]
impl<'v> StarlarkValue<'v> for AnalysisContext<'v> {
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(register_context)
    }
}

impl<'v> AllocValue<'v> for AnalysisContext<'v> {
    fn alloc_value(self, heap: &'v Heap) -> Value<'v> {
        heap.alloc_complex_no_freeze(self)
    }
}

struct RefAnalysisContext<'v>(&'v AnalysisContext<'v>);

impl<'v> StarlarkTypeRepr for RefAnalysisContext<'v> {
    fn starlark_type_repr() -> Ty {
        AnalysisContext::starlark_type_repr()
    }
}

impl<'v> UnpackValue<'v> for RefAnalysisContext<'v> {
    fn unpack_value(value: Value<'v>) -> Option<Self> {
        Some(RefAnalysisContext(
            value.downcast_ref::<AnalysisContext>().unwrap(),
        ))
    }
}

/// The type used for defining rules, usually bound as `ctx`.
/// Usually the sole argument to the `impl` argument of the `rule` function.
///
/// ```python
/// def _impl_my_rule(ctx: "context") -> ["provider"]:
///     return [DefaultInfo()]
/// my_rule = rule(impl = _impl_my_rule, attrs = {})
/// ```
#[starlark_module]
fn register_context(builder: &mut MethodsBuilder) {
    /// Returns the attributes of the target as a Starlark struct with a field for each attribute, which varies per rule.
    /// As an example, given a rule with the `attrs` argument of `{"foo": attrs.string()}`, this field will be
    /// a `struct` containing a field `foo` of type string.
    #[starlark(attribute)]
    fn attrs<'v>(this: RefAnalysisContext) -> anyhow::Result<ValueOfUnchecked<'v, StructRef<'v>>> {
        Ok(this.0.attrs)
    }

    /// Returns an `actions` value containing functions to define actual actions that are run.
    /// See the `actions` type for the operations that are available.
    #[starlark(attribute)]
    fn actions<'v>(
        this: RefAnalysisContext,
    ) -> anyhow::Result<ValueTyped<'v, AnalysisActions<'v>>> {
        Ok(this.0.actions)
    }

    /// Returns a `label` representing the target, or `None` if being invoked from a
    /// `dynamic_output` in Bxl.
    #[starlark(attribute)]
    fn label<'v>(this: RefAnalysisContext) -> anyhow::Result<Option<ValueTyped<'v, Label>>> {
        Ok(this.0.label)
    }
}

pub static ANALYSIS_ACTIONS_METHODS_ACTIONS: LateBinding<fn(&mut MethodsBuilder)> =
    LateBinding::new("ANALYSIS_ACTIONS_METHODS_ACTIONS");
pub static ANALYSIS_ACTIONS_METHODS_ANON_TARGET: LateBinding<fn(&mut MethodsBuilder)> =
    LateBinding::new("ANALYSIS_ACTIONS_METHODS_ANON_TARGET");