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

//!
//! The incrementality module of BUCK
//!
//! This is responsible for performing incremental caching and invalidations
//! with multiple versions in-flight at the same time.
//!

use std::fmt::Debug;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::task::Poll;

use allocative::Allocative;
use assert_matches::assert_matches;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use futures::pin_mut;
use futures::poll;
use more_futures::cancellation::CancellationContext;
use sorted_vector_map::sorted_vector_set;

use crate::api::computations::DiceComputations;
use crate::api::data::DiceData;
use crate::api::key::Key;
use crate::api::storage_type::StorageType;
use crate::api::user_data::NoOpTracker;
use crate::api::user_data::UserComputationData;
use crate::arc::Arc;
use crate::impls::core::graph::history::testing::CellHistoryExt;
use crate::impls::core::graph::history::CellHistory;
use crate::impls::core::graph::types::VersionedGraphKey;
use crate::impls::core::state::StateRequest;
use crate::impls::core::versions::VersionEpoch;
use crate::impls::dice::DiceModern;
use crate::impls::evaluator::AsyncEvaluator;
use crate::impls::events::DiceEventDispatcher;
use crate::impls::incremental::testing::DidDepsChangeExt;
use crate::impls::incremental::IncrementalEngine;
use crate::impls::key::DiceKey;
use crate::impls::key::ParentKey;
use crate::impls::task::sync_dice_task;
use crate::impls::task::PreviouslyCancelledTask;
use crate::impls::transaction::ActiveTransactionGuard;
use crate::impls::transaction::ChangeType;
use crate::impls::user_cycle::UserCycleDetectorData;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceKeyValue;
use crate::impls::value::DiceValidValue;
use crate::impls::value::MaybeValidDiceValue;
use crate::impls::worker::state::DiceWorkerStateCheckingDeps;
use crate::versions::testing::VersionRangesExt;
use crate::versions::VersionNumber;
use crate::versions::VersionRange;
use crate::versions::VersionRanges;

#[derive(Allocative, Clone, Dupe, Debug, Display, PartialEq, Eq, Hash)]
struct K;

#[async_trait]
impl Key for K {
    type Value = usize;

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        unimplemented!("test")
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        x == y
    }
}

#[derive(Allocative, Clone, Debug, Display)]
#[display(fmt = "{:?}", self)]
struct IsRan(Arc<AtomicBool>);

#[async_trait]
impl Key for IsRan {
    type Value = ();

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        self.0.store(true, Ordering::SeqCst);
    }

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

impl PartialEq for IsRan {
    fn eq(&self, _other: &Self) -> bool {
        true
    }
}
impl Eq for IsRan {}
impl Hash for IsRan {
    fn hash<H: Hasher>(&self, _state: &mut H) {}
}

#[derive(Allocative, Clone, Dupe, Debug, Display, PartialEq, Eq, Hash)]
struct Finish;

#[async_trait]
impl Key for Finish {
    type Value = ();

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
    }

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

#[tokio::test]
async fn test_detecting_changed_dependencies() -> anyhow::Result<()> {
    let dice = DiceModern::new(DiceData::new());
    let engine = IncrementalEngine::new(dice.state_handle.dupe(), VersionEpoch::testing_new(0));

    let user_data = std::sync::Arc::new(UserComputationData::new());

    let (ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(1)).await;
    ctx.inject(
        DiceKey { index: 100 },
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::testing_new(&[VersionNumber::new(1)], &[])),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    assert!(
        engine
            .compute_whether_dependencies_changed(
                ParentKey::None,
                eval.dupe(),
                &VersionRanges::testing_new(sorted_vector_set![VersionRange::bounded(
                    VersionNumber::new(0),
                    VersionNumber::new(1)
                )]),
                Arc::new(vec![DiceKey { index: 100 }]),
                &DiceWorkerStateCheckingDeps::testing()
            )
            .await?
            .is_changed()
    );

    let (ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(2)).await;
    ctx.inject(
        DiceKey { index: 100 },
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::testing_new(&[VersionNumber::new(1)], &[])),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    assert!(
        !engine
            .compute_whether_dependencies_changed(
                ParentKey::None,
                eval.dupe(),
                &VersionRanges::testing_new(sorted_vector_set![VersionRange::bounded(
                    VersionNumber::new(1),
                    VersionNumber::new(2)
                )]),
                Arc::new(vec![DiceKey { index: 100 }]),
                &DiceWorkerStateCheckingDeps::testing()
            )
            .await?
            .is_changed()
    );

    // Now we also check that when deps have transients and such.
    // for legacy, this would deal with cycles, but modern dice will detect cycles through post
    // processing and rely on the user cycle detector for now (which returns errors via the result.
    let (ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(2)).await;
    ctx.inject(
        DiceKey { index: 200 },
        DiceComputedValue::new(
            MaybeValidDiceValue::transient(std::sync::Arc::new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::testing_new(&[VersionNumber::new(2)], &[])),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    assert!(
        engine
            .compute_whether_dependencies_changed(
                ParentKey::None,
                eval.dupe(),
                &VersionRanges::testing_new(sorted_vector_set![VersionRange::bounded(
                    VersionNumber::new(1),
                    VersionNumber::new(2)
                )]),
                Arc::new(vec![DiceKey { index: 200 }]),
                &DiceWorkerStateCheckingDeps::testing()
            )
            .await?
            .is_changed()
    );

    Ok(())
}

#[tokio::test]
async fn test_values_gets_reevaluated_when_deps_change() -> anyhow::Result<()> {
    let dice = DiceModern::new(DiceData::new());

    let user_data = std::sync::Arc::new(UserComputationData::new());
    let events = DiceEventDispatcher::new(user_data.tracker.dupe(), dice.dupe());

    let is_ran = Arc::new(AtomicBool::new(false));
    let key = dice.key_index.index_key(IsRan(is_ran.dupe()));

    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::CtxAtVersion {
        version: VersionNumber::new(0),
        guard: ActiveTransactionGuard::new(VersionNumber::new(0), dice.state_handle.dupe()),
        resp: tx,
    });
    let (ctx, guard) = rx.await.unwrap();

    // set the initial state
    let (tx, _rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateComputed {
        key: VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 100 }),
        epoch: ctx.testing_get_epoch(),
        storage: StorageType::LastN(1),
        value: DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)),
        deps: Arc::new(vec![]),
        resp: tx,
    });
    let (tx, _rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateComputed {
        key: VersionedGraphKey::new(VersionNumber::new(0), key.dupe()),
        epoch: ctx.testing_get_epoch(),
        storage: StorageType::LastN(1),
        value: DiceValidValue::testing_new(DiceKeyValue::<IsRan>::new(())),
        deps: Arc::new(vec![DiceKey { index: 100 }]),
        resp: tx,
    });

    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateState {
        changes: vec![(key.dupe(), ChangeType::TestingSoftDirty)],
        resp: tx,
    });
    let v = rx.await.unwrap();
    drop(guard);
    drop(ctx);

    let (ctx, _guard) = dice.testing_shared_ctx(v).await;
    ctx.inject(
        DiceKey { index: 100 },
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::verified(VersionNumber::new(0))),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    let task = IncrementalEngine::spawn_for_key(
        key.dupe(),
        ctx.testing_get_epoch(),
        eval.dupe(),
        UserCycleDetectorData::testing_new(),
        events.dupe(),
        None,
    );
    let res = task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await?;
    assert_eq!(
        res.history().get_verified_ranges(),
        VersionRanges::testing_new(sorted_vector_set![VersionRange::begins_with(v)])
    );
    assert!(!is_ran.load(Ordering::SeqCst));

    // next version
    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateState {
        changes: vec![(key.dupe(), ChangeType::TestingSoftDirty)],
        resp: tx,
    });
    let v = rx.await.unwrap();

    let (ctx, _guard) = dice.testing_shared_ctx(v).await;
    ctx.inject(
        DiceKey { index: 100 },
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::verified(v)),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    let task = IncrementalEngine::spawn_for_key(
        key.dupe(),
        ctx.testing_get_epoch(),
        eval.dupe(),
        UserCycleDetectorData::testing_new(),
        events.dupe(),
        None,
    );
    let res = task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await?;
    assert_eq!(
        res.history().get_verified_ranges(),
        VersionRanges::testing_new(sorted_vector_set![VersionRange::begins_with(v)])
    );
    assert_eq!(is_ran.load(Ordering::SeqCst), true);
    is_ran.store(false, Ordering::SeqCst);

    // now force the dependency to have version numbers [1, 2]
    // also force dirty the root node so we actually check its deps since the above would
    // short circuit dirtying due to the dep value actually being equal.

    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateState {
        changes: vec![(key.dupe(), ChangeType::TestingSoftDirty)],
        resp: tx,
    });
    let new_v = rx.await.unwrap();

    let (ctx, _guard) = dice.testing_shared_ctx(v).await;

    // TODO(nga): `inject` violates `SharedCache` invariant:
    //   value computed should not be downgraded to not computed.
    if true {
        return Ok(());
    }

    ctx.inject(
        DiceKey { index: 100 },
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
            Arc::new(CellHistory::testing_new(&[v, new_v], &[])),
        ),
    );
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    let task = IncrementalEngine::spawn_for_key(
        key.dupe(),
        ctx.testing_get_epoch(),
        eval.dupe(),
        UserCycleDetectorData::testing_new(),
        events.dupe(),
        None,
    );
    let res = task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await?;
    assert_eq!(
        res.history().get_verified_ranges(),
        VersionRanges::testing_new(sorted_vector_set![VersionRange::bounded(v, new_v)])
    );
    assert_eq!(is_ran.load(Ordering::SeqCst), false);

    Ok(())
}

#[tokio::test]
async fn when_equal_return_same_instance() -> anyhow::Result<()> {
    let dice = DiceModern::new(DiceData::new());

    let user_data = std::sync::Arc::new(UserComputationData::new());
    let events = DiceEventDispatcher::new(user_data.tracker.dupe(), dice.dupe());

    let instance = Arc::new(AtomicUsize::new(0));

    #[derive(Clone, Dupe, Allocative)]
    struct InstanceEqual {
        instance_count: usize,
    }

    impl PartialEq for InstanceEqual {
        fn eq(&self, _other: &Self) -> bool {
            true
        }
    }

    #[derive(Allocative, Clone, Debug, Display)]
    #[display(fmt = "{:?}", self)]
    struct InstanceEqualKey(Arc<AtomicUsize>);

    #[async_trait]
    impl Key for InstanceEqualKey {
        type Value = InstanceEqual;

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            InstanceEqual {
                instance_count: self.0.fetch_add(1, Ordering::SeqCst),
            }
        }

        fn equality(x: &Self::Value, y: &Self::Value) -> bool {
            x == y
        }
    }
    impl PartialEq for InstanceEqualKey {
        fn eq(&self, _other: &Self) -> bool {
            true
        }
    }
    impl Eq for InstanceEqualKey {}
    impl Hash for InstanceEqualKey {
        fn hash<H: Hasher>(&self, _state: &mut H) {}
    }

    let key = dice.key_index.index_key(InstanceEqualKey(instance.dupe()));

    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateState {
        changes: vec![],
        resp: tx,
    });
    let v = rx.await.unwrap();

    let (ctx, _guard) = dice.testing_shared_ctx(v).await;
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    let task = IncrementalEngine::spawn_for_key(
        key.dupe(),
        ctx.testing_get_epoch(),
        eval.dupe(),
        UserCycleDetectorData::testing_new(),
        events.dupe(),
        None,
    );
    let res = task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await?;

    let (tx, rx) = tokio::sync::oneshot::channel();
    dice.state_handle.request(StateRequest::UpdateState {
        changes: vec![(key.dupe(), ChangeType::Invalidate)],
        resp: tx,
    });
    let v = rx.await.unwrap();

    let (ctx, _guard) = dice.testing_shared_ctx(v).await;
    let eval = AsyncEvaluator {
        per_live_version_ctx: ctx.dupe(),
        user_data: user_data.dupe(),
        dice: dice.dupe(),
    };

    let task = IncrementalEngine::spawn_for_key(
        key.dupe(),
        ctx.testing_get_epoch(),
        eval.dupe(),
        UserCycleDetectorData::testing_new(),
        events.dupe(),
        None,
    );
    let res2 = task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await?;

    // verify that we incremented the total instance counter
    assert_eq!(instance.load(Ordering::SeqCst), 2);

    assert_eq!(
        res.history().get_verified_ranges(),
        VersionRanges::testing_new(
            sorted_vector_set! { VersionRange::begins_with(VersionNumber::new(0))}
        )
    );

    // verify that the instance we return and store is the same as the original instance
    assert_eq!(
        res.value()
            .downcast_maybe_transient::<InstanceEqual>()
            .unwrap()
            .instance_count,
        res2.value()
            .downcast_maybe_transient::<InstanceEqual>()
            .unwrap()
            .instance_count
    );

    Ok(())
}

#[tokio::test]
async fn spawn_with_no_previously_cancelled_task() {
    let dice = DiceModern::new(DiceData::new());

    let (shared_ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(0)).await;

    let is_ran = Arc::new(AtomicBool::new(false));
    let k = dice.key_index.index_key(IsRan(is_ran.dupe()));

    let extra = std::sync::Arc::new(UserComputationData::new());
    let eval = AsyncEvaluator {
        per_live_version_ctx: shared_ctx.dupe(),
        user_data: extra.dupe(),
        dice: dice.dupe(),
    };
    let cycles = UserCycleDetectorData::testing_new();
    let events_dispatcher = DiceEventDispatcher::new(std::sync::Arc::new(NoOpTracker), dice.dupe());
    let previously_cancelled_task = None;

    let task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval,
        cycles,
        events_dispatcher,
        previously_cancelled_task,
    );

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .await
            .is_ok()
    );

    assert!(is_ran.load(Ordering::SeqCst));
}

#[tokio::test]
async fn spawn_with_previously_cancelled_task_that_cancelled() {
    let dice = DiceModern::new(DiceData::new());

    let (shared_ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(0)).await;

    let extra = std::sync::Arc::new(UserComputationData::new());
    let eval = AsyncEvaluator {
        per_live_version_ctx: shared_ctx.dupe(),
        user_data: extra.dupe(),
        dice: dice.dupe(),
    };
    let cycles = UserCycleDetectorData::testing_new();
    let events_dispatcher = DiceEventDispatcher::new(std::sync::Arc::new(NoOpTracker), dice.dupe());

    #[derive(Allocative, Clone, Dupe, Debug, Display, PartialEq, Eq, Hash)]
    struct CancellableNeverFinish;

    #[async_trait]
    impl Key for CancellableNeverFinish {
        type Value = ();

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            futures::future::pending().await
        }

        fn equality(_: &Self::Value, _: &Self::Value) -> bool {
            unreachable!("test")
        }
    }

    let k = dice.key_index.index_key(CancellableNeverFinish);
    let previous_task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval.dupe(),
        cycles,
        events_dispatcher.dupe(),
        None,
    );

    previous_task.cancel();

    let previously_cancelled_task = Some(PreviouslyCancelledTask {
        previous: previous_task,
    });

    let is_ran = Arc::new(AtomicBool::new(false));
    let k = dice.key_index.index_key(IsRan(is_ran.dupe()));
    let cycles = UserCycleDetectorData::testing_new();
    let task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval,
        cycles,
        events_dispatcher,
        previously_cancelled_task,
    );

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .await
            .is_ok()
    );

    assert!(is_ran.load(Ordering::SeqCst));
}

#[tokio::test]
async fn spawn_with_previously_cancelled_task_that_finished() {
    let dice = DiceModern::new(DiceData::new());

    let (shared_ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(0)).await;

    let extra = std::sync::Arc::new(UserComputationData::new());
    let eval = AsyncEvaluator {
        per_live_version_ctx: shared_ctx.dupe(),
        user_data: extra.dupe(),
        dice: dice.dupe(),
    };
    let cycles = UserCycleDetectorData::testing_new();
    let events_dispatcher = DiceEventDispatcher::new(std::sync::Arc::new(NoOpTracker), dice.dupe());

    #[derive(Allocative, Clone, Dupe, Debug, Display, PartialEq, Eq, Hash)]
    struct Finish;

    #[async_trait]
    impl Key for Finish {
        type Value = ();

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
        }

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

    let k = dice.key_index.index_key(Finish);
    let previous_task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval.dupe(),
        cycles,
        events_dispatcher.dupe(),
        None,
    );
    // wait for it to finish then trigger cancel
    previous_task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap()
        .await
        .unwrap();
    previous_task.cancel();

    let previously_cancelled_task = Some(PreviouslyCancelledTask {
        previous: previous_task,
    });

    let is_ran = Arc::new(AtomicBool::new(false));
    let k = dice.key_index.index_key(IsRan(is_ran.dupe()));
    let cycles = UserCycleDetectorData::testing_new();
    let task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval,
        cycles,
        events_dispatcher,
        previously_cancelled_task,
    );

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .await
            .is_ok()
    );

    assert!(!is_ran.load(Ordering::SeqCst));
}

#[tokio::test]
async fn mismatch_epoch_results_in_cancelled_result() {
    let dice = DiceModern::new(DiceData::new());

    let (shared_ctx, guard) = dice.testing_shared_ctx(VersionNumber::new(0)).await;

    let extra = std::sync::Arc::new(UserComputationData::new());
    let eval = AsyncEvaluator {
        per_live_version_ctx: shared_ctx.dupe(),
        user_data: extra.dupe(),
        dice: dice.dupe(),
    };
    let cycles = UserCycleDetectorData::testing_new();
    let events_dispatcher = DiceEventDispatcher::new(std::sync::Arc::new(NoOpTracker), dice.dupe());

    // trigger dice to delete and update the epoch
    drop(guard);

    let k = dice.key_index.index_key(Finish);
    let task = IncrementalEngine::spawn_for_key(
        k,
        shared_ctx.testing_get_epoch(),
        eval.dupe(),
        cycles,
        events_dispatcher.dupe(),
        None,
    );
    // wait for it to finish then trigger cancel
    assert_matches!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .await,
        Err(_) => {}
    );
}

#[tokio::test]
async fn spawn_with_previously_cancelled_task_nested_cancelled() -> anyhow::Result<()> {
    let dice = DiceModern::new(DiceData::new());

    let (shared_ctx, _guard) = dice.testing_shared_ctx(VersionNumber::new(0)).await;

    let k = dice.key_index.index_key(Finish);

    let extra = std::sync::Arc::new(UserComputationData::new());
    let eval = AsyncEvaluator {
        per_live_version_ctx: shared_ctx.dupe(),
        user_data: extra.dupe(),
        dice: dice.dupe(),
    };
    let cycles = UserCycleDetectorData::testing_new();
    let events_dispatcher = DiceEventDispatcher::new(std::sync::Arc::new(NoOpTracker), dice.dupe());

    let first_task = unsafe { sync_dice_task(DiceKey { index: 0 }) };
    let first_promise = first_task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap();

    let second_task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval.dupe(),
        cycles,
        events_dispatcher.dupe(),
        Some(PreviouslyCancelledTask {
            previous: first_task,
        }),
    );

    second_task.cancel();

    let cycles = UserCycleDetectorData::testing_new();
    let third_task = IncrementalEngine::spawn_for_key(
        k,
        VersionEpoch::testing_new(0),
        eval,
        cycles,
        events_dispatcher,
        Some(PreviouslyCancelledTask {
            previous: second_task,
        }),
    );

    let promise = third_task
        .depended_on_by(ParentKey::None)
        .not_cancelled()
        .unwrap();

    pin_mut!(promise);

    assert_matches!(poll!(&mut promise), Poll::Pending);
    assert_matches!(poll!(&mut promise), Poll::Pending);
    assert_matches!(poll!(&mut promise), Poll::Pending);
    assert_matches!(poll!(&mut promise), Poll::Pending);
    assert_matches!(poll!(&mut promise), Poll::Pending);

    first_promise.get_or_complete(|| {
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(DiceValidValue::testing_new(
                DiceKeyValue::<Finish>::new(()),
            )),
            Arc::new(CellHistory::empty()),
        )
    })?;

    let _ignored = promise.await?;

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

//! A task stored by Dice that is shared for all transactions at the same version
use std::cell::UnsafeCell;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering;
use std::task::Context;
use std::task::Poll;

use allocative::Allocative;
use allocative::Visitor;
use dupe::Dupe;
use dupe::OptionDupedExt;
use futures::future::Shared;
use futures::task::AtomicWaker;
use futures::FutureExt;
use more_futures::cancellation::future::CancellationHandle;
use parking_lot::Mutex;
use parking_lot::MutexGuard;
use slab::Slab;
use tokio::sync::oneshot;

use crate::arc::Arc;
use crate::impls::key::DiceKey;
use crate::impls::key::ParentKey;
use crate::impls::task::handle::TaskState;
use crate::impls::task::promise::DicePromise;
use crate::impls::task::state::AtomicDiceTaskState;
use crate::impls::value::DiceComputedValue;
use crate::result::CancellableResult;
use crate::result::Cancelled;

///
/// 'DiceTask' is approximately a copy of Shared and Weak from std, but with some custom special
/// record keeping to allow us to track the waiters as DiceKeys.
///
/// 'std::future::Weak' is akin to 'DiceTask', and each 'DicePromise' is a strong reference to it
/// akin to a 'std::future::Shared'.
///
/// The DiceTask is always completed by a thread whose future is the 'JoinHandle'. The thread
/// reports updates to the state of the future via 'DiceTaskHandle'. Simplifying the future
/// implementation in that no poll will ever be doing real work. No Wakers sleeping will be awoken
/// unless the task is ready.
/// The task is not the "standard states" of Pending, Polling, etc as stored by Shared future,
/// but instead we store the Dice specific states so that its useful when we dump the state.
/// Wakers are tracked with their corresponding DiceKeys, allowing us to track the rdeps and see
/// which key is waiting on what
///
/// We can explicitly track cancellations by tracking the Waker drops.
///
/// Memory size difference:
/// DiceTask <-> Weak: DiceTask holds an extra JoinHandle which is a single ptr.
/// DiceTask now holds a 'triomphe::Arc' instead of 'std::Arc' which is slightly more efficient as it
/// doesn't require weak ptr handling. This is just so that we have the JoinHandle so we can abort
/// when canceled, but we could choose to change the implementation by moving cancellation
/// notification into the DiceTaskInternal
#[derive(Allocative, Clone, Dupe)]
pub(crate) struct DiceTask {
    pub(crate) internal: Arc<DiceTaskInternal>,
    /// Handle to cancel the spawned task
    #[allocative(skip)]
    pub(super) cancellations: Cancellations,
}

pub(crate) struct DiceTaskInternal {
    pub(crate) key: DiceKey,
    /// The internal progress state of the task
    pub(super) state: AtomicDiceTaskState,

    /// Internals that require mutex
    pub(super) critical: Mutex<DiceTaskInternalCritical>,
    /// The value if finished computing
    maybe_value: UnsafeCell<Option<CancellableResult<DiceComputedValue>>>,
}

pub(super) struct DiceTaskInternalCritical {
    /// Other DiceTasks that are awaiting the completion of this task.
    ///
    /// We hold a pair DiceKey and Waker.
    /// Compared to 'Shared', which just holds a standard 'Waker', the Waker itself is now an
    /// AtomicWaker, which is an extra AtomicUsize, so this is marginally larger than the standard
    /// Shared future.
    pub(super) dependants: Option<Slab<(ParentKey, Arc<AtomicWaker>)>>,
    termination_waiter: Shared<oneshot::Receiver<()>>,
    termination_sender: Option<oneshot::Sender<()>>,
}

impl Allocative for DiceTaskInternal {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(allocative::Key::new("critical"), &self.critical);
        if self.state.is_ready(Ordering::Acquire) {
            visitor.visit_field(allocative::Key::new("maybe_value"), unsafe {
                &*self.maybe_value.get()
            });
        }
        visitor.exit();
    }
}

impl Allocative for DiceTaskInternalCritical {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(allocative::Key::new("dependants"), &self.dependants);
        visitor.exit();
    }
}

pub(crate) enum MaybeCancelled {
    Ok(DicePromise),
    Cancelled,
}

impl MaybeCancelled {
    pub(crate) fn not_cancelled(self) -> Option<DicePromise> {
        match self {
            MaybeCancelled::Ok(promise) => Some(promise),
            MaybeCancelled::Cancelled => None,
        }
    }
}

/// Future when resolves when task is finished or cancelled and terminated.
pub(crate) struct TerminationObserver(TerminationObserverInternal);

enum TerminationObserverInternal {
    Waiting {
        internals: Arc<DiceTaskInternal>,
        waiter: Shared<oneshot::Receiver<()>>,
    },
    Done,
}

impl TerminationObserver {
    pub(crate) fn is_terminated(&self) -> bool {
        match &self.0 {
            TerminationObserverInternal::Waiting { internals, .. } => !internals.is_pending(),
            TerminationObserverInternal::Done => true,
        }
    }
}

impl Future for TerminationObserver {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match &mut self.0 {
            TerminationObserverInternal::Waiting { internals, waiter } => {
                if !internals.is_pending() {
                    Poll::Ready(())
                } else {
                    waiter
                        .poll_unpin(cx)
                        .map(|res| res.expect("task dropped without notifying"))
                }
            }
            TerminationObserverInternal::Done => Poll::Ready(()),
        }
    }
}

impl DiceTask {
    /// `k` depends on this task, returning a `DicePromise` that will complete when this task
    /// completes
    pub(crate) fn depended_on_by(&self, k: ParentKey) -> MaybeCancelled {
        if let Some(result) = self.internal.read_value() {
            match result {
                Ok(result) => MaybeCancelled::Ok(DicePromise::ready(result)),
                Err(_err) => MaybeCancelled::Cancelled,
            }
        } else {
            let mut critical = self.internal.critical.lock();
            if self.cancellations.is_cancelled(&critical) {
                return MaybeCancelled::Cancelled;
            }
            match &mut critical.dependants {
                None => {
                    match self
                        .internal
                        .read_value()
                        .expect("invalid state where deps are taken before state is ready")
                    {
                        Ok(result) => MaybeCancelled::Ok(DicePromise::ready(result)),
                        Err(_err) => MaybeCancelled::Cancelled,
                    }
                }
                Some(ref mut wakers) => {
                    let waker = Arc::new(AtomicWaker::new());
                    let id = wakers.insert((k, waker.dupe()));

                    MaybeCancelled::Ok(DicePromise::pending(
                        id,
                        self.internal.dupe(),
                        waker,
                        self.cancellations.dupe(),
                    ))
                }
            }
        }
    }

    pub(crate) fn get_finished_value(&self) -> Option<CancellableResult<DiceComputedValue>> {
        self.internal.read_value()
    }

    /// true if this task is not yet complete and not yet canceled.
    pub(crate) fn is_pending(&self) -> bool {
        self.internal.is_pending()
    }

    #[allow(unused)] // future introspection functions
    pub(crate) fn inspect_waiters(&self) -> Option<Vec<ParentKey>> {
        self.internal
            .critical
            .lock()
            .dependants
            .as_ref()
            .map(|deps| deps.iter().map(|(_, (k, _))| *k).collect())
    }

    pub(crate) fn cancel(&self) {
        let lock = self.internal.critical.lock();
        self.cancellations.cancel(&lock);
    }

    pub(crate) fn await_termination(&self) -> TerminationObserver {
        if let Some(_result) = self.internal.read_value() {
            TerminationObserver(TerminationObserverInternal::Done)
        } else {
            let mut critical = self.internal.critical.lock();
            match &mut critical.dependants {
                None => {
                    let _result = self
                        .internal
                        .read_value()
                        .expect("invalid state where deps are taken before state is ready");
                    TerminationObserver(TerminationObserverInternal::Done)
                }
                Some(_wakers) => TerminationObserver(TerminationObserverInternal::Waiting {
                    internals: self.internal.dupe(),
                    waiter: critical.termination_waiter.clone(),
                }),
            }
        }
    }
}

impl DiceTaskInternal {
    pub(super) fn drop_waiter(&self, slab: usize, cancellations: &Cancellations) {
        let mut critical = self.critical.lock();
        match critical.dependants {
            None => {}
            Some(ref mut deps) => {
                deps.remove(slab);

                if deps.is_empty() {
                    cancellations.cancel(&critical);
                }
            }
        }
    }

    pub(super) fn new(key: DiceKey) -> Arc<Self> {
        let (tx, rx) = oneshot::channel();
        Arc::new(Self {
            key,
            state: AtomicDiceTaskState::default(),
            maybe_value: UnsafeCell::new(None),
            critical: Mutex::new(DiceTaskInternalCritical {
                dependants: Some(Slab::new()),
                termination_waiter: rx.shared(),
                termination_sender: Some(tx),
            }),
        })
    }

    pub(crate) fn read_value(&self) -> Option<CancellableResult<DiceComputedValue>> {
        if self.state.is_ready(Ordering::Acquire) || self.state.is_terminated(Ordering::Acquire) {
            Some(
                unsafe {
                    // SAFETY: main thread only writes this before setting state to `READY`
                    &*self.maybe_value.get()
                }
                .as_ref()
                .duped()
                .expect("result should be present"),
            )
        } else {
            None
        }
    }

    pub(super) fn set_value(
        &self,
        value: DiceComputedValue,
    ) -> CancellableResult<DiceComputedValue> {
        match self.state.sync() {
            TaskState::Continue => {}
            TaskState::Finished => {
                return self
                    .read_value()
                    .expect("task finished must mean result is ready");
            }
        };

        let prev_exist = unsafe {
            // SAFETY: no tasks read the value unless state is converted to `READY`
            &mut *self.maybe_value.get()
        }
        .replace(Ok(value.dupe()))
        .is_some();
        assert!(
            !prev_exist,
            "invalid state where somehow value was already written"
        );

        self.state.report_ready();
        self.wake_dependents();

        Ok(value)
    }

    pub(super) fn wake_dependents(&self) {
        let mut critical = self.critical.lock();
        let mut deps = critical
            .dependants
            .take()
            .expect("Invalid state where deps where taken already");

        deps.drain().for_each(|(_k, waker)| waker.wake());

        critical
            .termination_sender
            .take()
            .expect("terminated twice")
            .send(())
            .expect("receiver held by self");
    }

    /// report the task as terminated. This should only be called once. No effect if called affect
    /// task is already ready
    pub(super) fn report_terminated(&self) {
        match self.state.sync() {
            TaskState::Continue => {}
            TaskState::Finished => {
                return;
            }
        };

        let prev_exist = unsafe {
            // SAFETY: no tasks read the value unless state is converted to `READY`
            &mut *self.maybe_value.get()
        }
        .replace(Err(Cancelled))
        .is_some();
        assert!(
            !prev_exist,
            "invalid state where somehow value was already written"
        );

        self.state.report_terminated();
        self.wake_dependents();
    }

    /// true if this task is not yet complete and not yet canceled.
    pub(crate) fn is_pending(&self) -> bool {
        !(self.state.is_ready(Ordering::Acquire) || self.state.is_terminated(Ordering::Acquire))
    }
}

// our use of `UnsafeCell` is okay to be send and sync.
// Each unsafe block around its access has comments explaining the invariants.
unsafe impl Send for DiceTaskInternal {}
unsafe impl Sync for DiceTaskInternal {}

/// Stores either task cancellation handle which can be used to cancel the task
/// or termination observers if task is being cancelled.
#[derive(Clone, Dupe)]
pub(super) struct Cancellations {
    /// `UnsafeCell` access is guarded by `DiceTaskInternal.critical` mutex.
    /// `None` means task is not cancellable.
    internal: Option<Arc<UnsafeCell<CancellationsInternal>>>,
}

enum CancellationsInternal {
    NotCancelled(CancellationHandle),
    Cancelled,
}

impl Cancellations {
    pub(super) fn new(cancellation_handle: CancellationHandle) -> Self {
        Self {
            internal: Some(Arc::new(UnsafeCell::new(
                CancellationsInternal::NotCancelled(cancellation_handle),
            ))),
        }
    }

    pub(super) fn not_cancellable() -> Self {
        Self { internal: None }
    }

    pub(super) fn cancel(&self, _lock: &MutexGuard<DiceTaskInternalCritical>) {
        if let Some(internal) = self.internal.as_ref() {
            take_mut::take(
                unsafe {
                    // SAFETY: locked by the MutexGuard of Slab
                    &mut *internal.get()
                },
                |internal| match internal {
                    CancellationsInternal::NotCancelled(handle) => {
                        handle.cancel();
                        CancellationsInternal::Cancelled
                    }
                    cancelled => cancelled,
                },
            )
        };
    }

    pub(super) fn is_cancelled(&self, _lock: &MutexGuard<DiceTaskInternalCritical>) -> bool {
        self.internal.as_ref().map_or(false, |internal| {
            match unsafe {
                // SAFETY: locked by the MutexGuard of Slab
                &*internal.get()
            } {
                CancellationsInternal::NotCancelled(_) => false,
                CancellationsInternal::Cancelled => true,
            }
        })
    }
}

// our use of `UnsafeCell` is okay to be send and sync.
// Each unsafe block around its access has comments explaining the invariants.
unsafe impl Send for Cancellations {}
unsafe impl Sync for Cancellations {}

pub(crate) mod introspection {
    use crate::impls::task::dice::DiceTask;
    use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;

    impl DiceTask {
        pub(crate) fn introspect_state(&self) -> DiceTaskStateForDebugging {
            self.internal.state.introspect_state()
        }
    }
}


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

//! Handle to the DiceTask as seen by the thread responsible for completing the task

use more_futures::cancellation::ExplicitCancellationContext;

use crate::arc::Arc;
use crate::impls::task::dice::DiceTaskInternal;
use crate::impls::value::DiceComputedValue;

/// The handle to the 'DiceTask' owned by the spawned thread that is responsible for completing
/// the task.
pub(crate) struct DiceTaskHandle<'a> {
    pub(super) internal: Arc<DiceTaskInternal>,
    pub(super) cancellations: &'a ExplicitCancellationContext,
}

/// After reporting that we are about to transition to a state, should we continue processing or
/// should we terminate
pub(crate) enum TaskState {
    /// continue processing as normal
    Continue,
    /// task was finished already
    Finished,
}

impl<'a> DiceTaskHandle<'a> {
    pub(crate) fn report_initial_lookup(&self) -> TaskState {
        self.internal.state.report_initial_lookup()
    }

    pub(crate) fn checking_deps(&self) -> TaskState {
        self.internal.state.report_checking_deps()
    }

    pub(crate) fn computing(&self) -> TaskState {
        self.internal.state.report_computing()
    }

    pub(crate) fn finished(self, value: DiceComputedValue) {
        let _ignore = self.internal.set_value(value);
    }

    pub(crate) fn cancellation_ctx(&self) -> &ExplicitCancellationContext {
        &self.cancellations
    }

    #[cfg(test)]
    pub(crate) fn testing_new() -> &'static DiceTaskHandle<'static> {
        static TEST: once_cell::sync::Lazy<DiceTaskHandle> =
            once_cell::sync::Lazy::new(|| DiceTaskHandle::<'static> {
                internal: DiceTaskInternal::new(crate::impls::key::DiceKey { index: 99999 }),
                cancellations: ExplicitCancellationContext::testing(),
            });

        &TEST
    }
}

impl<'a> Drop for DiceTaskHandle<'a> {
    fn drop(&mut self) {
        if self.internal.read_value().is_none() {
            debug!("task is terminated");

            // This is only owned by the main worker task. If this was dropped, and no result was
            // ever recorded, then we must have been terminated.
            self.internal.report_terminated()
        }
    }
}

unsafe impl<'a> Send for DiceTaskHandle<'a> {}


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

use dupe::Dupe;
use futures::future::BoxFuture;
use more_futures::spawn::spawn_cancellable;
use more_futures::spawn::FutureAndCancellationHandle;
use more_futures::spawner::Spawner;

use crate::impls::key::DiceKey;
use crate::impls::task::dice::Cancellations;
use crate::impls::task::dice::DiceTask;
use crate::impls::task::dice::DiceTaskInternal;
use crate::impls::task::handle::DiceTaskHandle;

pub(crate) mod dice;
pub(crate) mod handle;
pub(crate) mod promise;
mod state;

#[cfg(test)]
mod tests;

pub(crate) fn spawn_dice_task<S>(
    key: DiceKey,
    spawner: &dyn Spawner<S>,
    ctx: &S,
    f: impl for<'a> FnOnce(DiceTaskHandle<'a>) -> BoxFuture<'a, Box<dyn Any + Send>> + Send,
) -> DiceTask {
    let internal = DiceTaskInternal::new(key);

    // since the spawn is alive until cancelled via the handle, we can drop the spawn future itself
    let FutureAndCancellationHandle {
        cancellation_handle,
        ..
    } = spawn_cancellable(
        {
            let internal = internal.dupe();
            |cancellations| {
                let handle = DiceTaskHandle {
                    internal,
                    cancellations,
                };

                f(handle)
            }
        },
        spawner,
        ctx,
    );

    DiceTask {
        internal,
        cancellations: Cancellations::new(cancellation_handle),
    }
}

/// Unsafe as this creates a Task that must be completed explicitly otherwise polling will never
/// complete.
pub(crate) unsafe fn sync_dice_task(key: DiceKey) -> DiceTask {
    let internal = DiceTaskInternal::new(key);

    DiceTask {
        internal,
        cancellations: Cancellations::not_cancellable(),
    }
}

pub(crate) struct PreviouslyCancelledTask {
    pub(crate) previous: DiceTask,
}


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

//! Promise is a handle to a DiceTask that will be completed

use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use dupe::Dupe;
use futures::task::AtomicWaker;

use crate::arc::Arc;
use crate::impls::task::dice::Cancellations;
use crate::impls::task::dice::DiceTaskInternal;
use crate::impls::task::handle::TaskState;
use crate::impls::value::DiceComputedValue;
use crate::result::CancellableResult;

/// A string reference to a 'DiceTask' that is pollable as a future.
/// This is only awoken when the result is ready, as none of the pollers are responsible for
/// running the task to completion
///
/// Memory size:
/// DicePromise <-> Shared: DicePromise is 2 triomphe::Arc and a usize, whereas Shared is a usize
/// and 1 std::Arc, so we hold one extra Arc.
pub(crate) struct DicePromise(pub(super) DicePromiseInternal);

pub(super) enum DicePromiseInternal {
    Ready {
        result: DiceComputedValue,
    },
    Pending {
        slab: usize,
        task_internal: Arc<DiceTaskInternal>,
        waker: Arc<AtomicWaker>,
        cancellations: Cancellations,
    },
    Done,
}

impl DicePromise {
    pub(crate) fn ready(result: DiceComputedValue) -> Self {
        Self(DicePromiseInternal::Ready { result })
    }

    pub(super) fn pending(
        slab: usize,
        internal: Arc<DiceTaskInternal>,
        waker: Arc<AtomicWaker>,
        cancellations: Cancellations,
    ) -> Self {
        Self(DicePromiseInternal::Pending {
            slab,
            task_internal: internal,
            waker,
            cancellations,
        })
    }

    /// Get the value if already complete, or complete it. Note that `f` may run even if the result
    /// is not used.
    pub(crate) fn get_or_complete(
        self,
        f: impl FnOnce() -> DiceComputedValue,
    ) -> CancellableResult<DiceComputedValue> {
        match &self.0 {
            DicePromiseInternal::Ready { result } => Ok(result.dupe()),
            DicePromiseInternal::Pending { task_internal, .. } => {
                if let Some(res) = task_internal.read_value() {
                    res
                } else {
                    match task_internal.state.report_project() {
                        TaskState::Continue => {}
                        TaskState::Finished => {
                            return task_internal
                                .read_value()
                                .expect("task finished must mean result is ready");
                        }
                    }

                    let value = f();

                    task_internal.set_value(value)
                }
            }
            DicePromiseInternal::Done => panic!("poll after ready"),
        }
    }
}

impl Drop for DicePromise {
    fn drop(&mut self) {
        match &self.0 {
            DicePromiseInternal::Ready { .. } => {}
            DicePromiseInternal::Done => {}
            DicePromiseInternal::Pending {
                slab,
                task_internal,
                cancellations,
                ..
            } => task_internal.drop_waiter(*slab, cancellations),
        }
    }
}

impl Future for DicePromise {
    type Output = CancellableResult<DiceComputedValue>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match &mut self.0 {
            x @ DicePromiseInternal::Ready { .. } => {
                match mem::replace(x, DicePromiseInternal::Done) {
                    DicePromiseInternal::Ready { result } => Poll::Ready(Ok(result)),
                    _ => unreachable!(),
                }
            }
            DicePromiseInternal::Pending {
                task_internal,
                waker,
                ..
            } => {
                waker.register(cx.waker());
                if let Some(res) = task_internal.read_value() {
                    Poll::Ready(res)
                } else {
                    Poll::Pending
                }
            }
            DicePromiseInternal::Done => panic!("poll after ready"),
        }
    }
}


/*
 * 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::atomic::AtomicU8;
use std::sync::atomic::Ordering;

use allocative::Allocative;

use crate::impls::task::handle::TaskState;

/// The state of the DiceTask about what stage of evaluation we are in.
/// The state is an `u8` consisting of states in `DiceTaskState`.
/// Only certain state transitions are allowed. `INITIAL_LOOKUP` can transition to all other states.
/// Only `SYNC` can transition to `READY`. Transition from `READY` to any other state is a noop.
/// `TERMINATED` is used to indicate that no more status updates will be available despite not
/// being `READY`.
///
/// The state `SYNC` is also special in that it acts as a spinlock, in that once in that state, all
/// other state transitions except to `READY` are blocked.
#[derive(Default, Allocative)]
pub(super) struct AtomicDiceTaskState(AtomicU8);

impl AtomicDiceTaskState {
    pub(super) fn is_ready(&self, ordering: Ordering) -> bool {
        match DiceTaskState::from_u8_state(self.0.load(ordering)) {
            DiceTaskState::Ready => true,
            _ => false,
        }
    }

    pub(super) fn is_terminated(&self, ordering: Ordering) -> bool {
        match DiceTaskState::from_u8_state(self.0.load(ordering)) {
            DiceTaskState::Terminated => true,
            _ => false,
        }
    }

    fn transition(
        &self,
        maybe_transition: impl Fn(DiceTaskState) -> Option<DiceTaskState>,
    ) -> TaskState {
        loop {
            match self
                .0
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |old| {
                    let unpacked = DiceTaskState::from_u8_state(old);
                    maybe_transition(unpacked).map(DiceTaskState::into_u8_state)
                }) {
                Ok(_) => {
                    return TaskState::Continue;
                }
                Err(old) => {
                    let unpacked = DiceTaskState::from_u8_state(old);
                    if unpacked == DiceTaskState::Sync {
                        std::hint::spin_loop();
                        continue;
                    } else if unpacked == DiceTaskState::Ready {
                        return TaskState::Finished;
                    } else {
                        unreachable!("handled above")
                    }
                }
            }
        }
    }

    pub(super) fn report_checking_deps(&self) -> TaskState {
        self.transition(|s| s.transition(TargetState::CheckingDeps))
    }

    pub(super) fn report_initial_lookup(&self) -> TaskState {
        self.transition(|s| s.transition(TargetState::InitialLookup))
    }

    pub(super) fn report_computing(&self) -> TaskState {
        self.transition(|s| s.transition(TargetState::Computing))
    }

    pub(super) fn report_project(&self) -> TaskState {
        self.transition(|s| s.project())
    }

    pub(super) fn sync(&self) -> TaskState {
        self.transition(|s| s.transition(TargetState::Sync))
    }

    pub(super) fn report_ready(&self) {
        self.transition(|s| s.transition(TargetState::Ready));
    }

    pub(super) fn report_terminated(&self) {
        self.transition(|s| s.transition(TargetState::Terminated));
    }
}

pub(crate) mod introspection {
    use std::sync::atomic::Ordering;

    use crate::impls::task::state::AtomicDiceTaskState;
    use crate::impls::task::state::DiceTaskState;
    use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;

    impl AtomicDiceTaskState {
        pub(crate) fn introspect_state(&self) -> DiceTaskStateForDebugging {
            match DiceTaskState::from_u8_state(self.0.load(Ordering::Acquire)) {
                DiceTaskState::InitialLookup(_) => DiceTaskStateForDebugging::AsyncInProgress,
                DiceTaskState::AwaitingPrevious(_) => DiceTaskStateForDebugging::AsyncInProgress,
                DiceTaskState::CheckingDeps(_) => DiceTaskStateForDebugging::AsyncInProgress,
                DiceTaskState::Computing(_) => DiceTaskStateForDebugging::AsyncInProgress,
                DiceTaskState::Sync => DiceTaskStateForDebugging::SyncInProgress,
                DiceTaskState::Ready => DiceTaskStateForDebugging::AsyncReady,
                DiceTaskState::Terminated => DiceTaskStateForDebugging::AsyncDropped,
            }
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
enum DiceTaskState {
    /// When waiting for the initial lookup of the cache
    InitialLookup(IsProjecting),
    /// When we are waiting for our dependencies to see if the value can be reused
    CheckingDeps(IsProjecting),
    /// When we are actively computing the value by running the key's compute
    Computing(IsProjecting),
    /// When we are waiting for the previous version of this task to cancel
    AwaitingPrevious(IsProjecting),
    /// When we are synchronizing over the updating of the value
    Sync,
    /// When the value is ready to be used
    Ready,
    /// When this future will never become Ready
    Terminated,
}

impl DiceTaskState {
    fn from_u8_state(state: u8) -> Self {
        // The 4th bit is a boolean to indicate if its projecting
        // We pack the state is a 3 bits corresponding to integers 0 to 6.
        match state & 0b111 {
            0 => Self::AwaitingPrevious(IsProjecting::unpack(state)),
            1 => Self::InitialLookup(IsProjecting::unpack(state)),
            2 => Self::CheckingDeps(IsProjecting::unpack(state)),
            3 => Self::Computing(IsProjecting::unpack(state)),
            4 => Self::Sync,
            5 => Self::Ready,
            6 => Self::Terminated,
            _ => unreachable!("invalid state `{}`", state),
        }
    }

    fn into_u8_state(self) -> u8 {
        // We pack the state is a 3 bits corresponding to integers 0 to 6.
        // The 4th bit is a boolean to indicate if its projecting
        match self {
            DiceTaskState::AwaitingPrevious(proj) => proj.pack(),
            DiceTaskState::InitialLookup(proj) => 1 | proj.pack(),
            DiceTaskState::CheckingDeps(proj) => 2 | proj.pack(),
            DiceTaskState::Computing(proj) => 3 | proj.pack(),
            DiceTaskState::Sync => 4,
            DiceTaskState::Ready => 5,
            DiceTaskState::Terminated => 6,
        }
    }

    fn transition(self, target: TargetState) -> Option<Self> {
        match self {
            DiceTaskState::AwaitingPrevious(proj) => match target {
                target @ (TargetState::InitialLookup
                | TargetState::Sync
                | TargetState::Terminated) => Some(target.into_dice_task_state_with_proj(proj)),
                target => {
                    panic!(
                        "invalid state transition `{:?}` -> `{:?}`",
                        DiceTaskState::Computing(proj),
                        target
                    )
                }
            },
            DiceTaskState::InitialLookup(proj) => match target {
                TargetState::Ready => {
                    panic!(
                        "invalid state transition `{:?}` -> `{:?}`",
                        DiceTaskState::InitialLookup(proj),
                        target
                    )
                }
                target => Some(target.into_dice_task_state_with_proj(proj)),
            },
            DiceTaskState::CheckingDeps(proj) => match target {
                target @ (TargetState::Computing | TargetState::Sync | TargetState::Terminated) => {
                    Some(target.into_dice_task_state_with_proj(proj))
                }
                target => {
                    panic!(
                        "invalid state transition `{:?}` -> `{:?}`",
                        DiceTaskState::CheckingDeps(proj),
                        target
                    )
                }
            },
            DiceTaskState::Computing(proj) => match target {
                target @ (TargetState::Sync | TargetState::Terminated) => {
                    Some(target.into_dice_task_state_with_proj(proj))
                }
                target => {
                    panic!(
                        "invalid state transition `{:?}` -> `{:?}`",
                        DiceTaskState::Computing(proj),
                        target
                    )
                }
            },
            DiceTaskState::Sync => match target {
                TargetState::Ready => Some(target.into_dice_task_state()),
                TargetState::Terminated => Some(target.into_dice_task_state()),
                _ => None,
            },
            DiceTaskState::Ready => None,
            DiceTaskState::Terminated => {
                panic!(
                    "invalid state transition `{:?}` -> `{:?}`",
                    DiceTaskState::Terminated,
                    target
                )
            }
        }
    }

    fn project(self) -> Option<Self> {
        match self {
            DiceTaskState::AwaitingPrevious(_) => {
                Some(DiceTaskState::AwaitingPrevious(IsProjecting::Projecting))
            }
            DiceTaskState::InitialLookup(_) => {
                Some(DiceTaskState::InitialLookup(IsProjecting::Projecting))
            }
            DiceTaskState::CheckingDeps(_) => {
                Some(DiceTaskState::CheckingDeps(IsProjecting::Projecting))
            }
            DiceTaskState::Computing(_) => Some(DiceTaskState::Computing(IsProjecting::Projecting)),
            DiceTaskState::Sync => None,
            DiceTaskState::Ready => None,
            DiceTaskState::Terminated => {
                panic!(
                    "invalid projection when state is `{:?}`",
                    DiceTaskState::Terminated,
                )
            }
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
enum TargetState {
    /// When we are looking up this node from cache
    InitialLookup,
    /// When we are waiting for our dependencies to see if the value can be reused
    CheckingDeps,
    /// When we are actively computing the value by running the key's compute
    Computing,
    /// When we are synchronizing over the updating of the value
    Sync,
    /// When the value is ready to be used
    Ready,
    /// When this future will never become Ready
    Terminated,
}

impl TargetState {
    fn into_dice_task_state_with_proj(&self, proj: IsProjecting) -> DiceTaskState {
        match self {
            TargetState::InitialLookup => DiceTaskState::InitialLookup(proj),
            TargetState::CheckingDeps => DiceTaskState::CheckingDeps(proj),
            TargetState::Computing => DiceTaskState::Computing(proj),
            TargetState::Sync => DiceTaskState::Sync,
            TargetState::Ready => DiceTaskState::Ready,
            TargetState::Terminated => DiceTaskState::Terminated,
        }
    }

    fn into_dice_task_state(&self) -> DiceTaskState {
        match self {
            TargetState::Sync => DiceTaskState::Sync,
            TargetState::Ready => DiceTaskState::Ready,
            TargetState::Terminated => DiceTaskState::Terminated,
            _ => panic!("requires projection state"),
        }
    }
}

/// When we are actively projecting the value by running the key's compute synchronously.
/// This can occur simultaneously with other states except `Sync` and `Ready`.
#[derive(Debug, PartialEq, Eq)]
enum IsProjecting {
    Projecting,
    NotProjecting,
}

impl IsProjecting {
    /// We pack the DiceTaskState as a 3 bits corresponding to integers 0 to 4.
    /// The 4th bit is a boolean to indicate if its projecting, so this all operates on the 4th bit.
    fn unpack(state: u8) -> Self {
        if (state & (1 << 3)) != 0 {
            IsProjecting::Projecting
        } else {
            IsProjecting::NotProjecting
        }
    }

    fn pack(&self) -> u8 {
        match self {
            IsProjecting::Projecting => 1 << 3,
            IsProjecting::NotProjecting => 0,
        }
    }
}


/*
 * 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::Any;
use std::sync::atomic::Ordering;
use std::task::Poll;

use allocative::Allocative;
use assert_matches::assert_matches;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use futures::poll;
use futures::FutureExt;
use more_futures::cancellation::CancellationContext;
use more_futures::spawner::TokioSpawner;
use tokio::sync::Barrier;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;

use crate::api::computations::DiceComputations;
use crate::api::key::Key;
use crate::arc::Arc;
use crate::impls::core::graph::history::CellHistory;
use crate::impls::key::DiceKey;
use crate::impls::key::ParentKey;
use crate::impls::task::dice::MaybeCancelled;
use crate::impls::task::spawn_dice_task;
use crate::impls::task::sync_dice_task;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceKeyValue;
use crate::impls::value::DiceValidValue;
use crate::impls::value::MaybeValidDiceValue;

#[derive(Allocative, Clone, Debug, Display, Eq, PartialEq, Hash)]
struct K;

#[async_trait]
impl Key for K {
    type Value = usize;

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        unimplemented!("test")
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        x == y
    }
}

#[tokio::test]
async fn simple_task() -> anyhow::Result<()> {
    let lock = Arc::new(Mutex::new(()));

    let lock_dupe = lock.dupe();
    let locked = lock_dupe.lock().await;

    let task = spawn_dice_task(DiceKey { index: 10 }, &TokioSpawner, &(), |handle| {
        async move {
            // wait for the lock too
            let _lock = lock.lock().await;

            handle.finished(DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(2))),
                Arc::new(CellHistory::empty()),
            ));

            Box::new(()) as Box<dyn Any + Send + 'static>
        }
        .boxed()
    });

    assert!(task.is_pending());

    let mut promise = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    assert_eq!(
        task.inspect_waiters(),
        Some(vec![ParentKey::Some(DiceKey { index: 1 })])
    );

    let polled = futures::poll!(&mut promise);
    assert!(
        !polled.is_ready(),
        "lock should be held so shouldn't be ready"
    );

    drop(locked);

    // since test runtime is one thread, dropping the lock shouldn't cause the task to complete yet
    let polled = futures::poll!(&mut promise);
    assert!(
        !polled.is_ready(),
        "lock should be held so shouldn't be ready"
    );

    // now await on the task, and see that we wake up and complete

    let v = promise.await;
    assert!(
        v?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );

    assert!(!task.is_pending());

    Ok(())
}

#[tokio::test]
async fn never_ready_results_in_terminated() -> anyhow::Result<()> {
    let task = spawn_dice_task(DiceKey { index: 500 }, &TokioSpawner, &(), |handle| {
        async move {
            let _handle = handle;
            // never report ready

            Box::new(()) as Box<dyn Any + Send + 'static>
        }
        .boxed()
    });

    let promise = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    assert_eq!(
        task.inspect_waiters(),
        Some(vec![ParentKey::Some(DiceKey { index: 1 })])
    );

    let v = promise.await;
    assert!(v.is_err());

    assert!(!task.internal.state.is_ready(Ordering::SeqCst));
    assert!(task.internal.state.is_terminated(Ordering::SeqCst));
    assert!(!task.is_pending());

    Ok(())
}

#[tokio::test]
async fn multiple_promises_all_completes() -> anyhow::Result<()> {
    let task = spawn_dice_task(DiceKey { index: 20 }, &TokioSpawner, &(), |handle| {
        async move {
            // wait for the lock too
            handle.finished(DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(2))),
                Arc::new(CellHistory::empty()),
            ));

            Box::new(()) as Box<dyn Any + Send + 'static>
        }
        .boxed()
    });

    let promise1 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();
    let promise2 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 2 }))
        .not_cancelled()
        .unwrap();
    let promise3 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 3 }))
        .not_cancelled()
        .unwrap();
    let promise4 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 4 }))
        .not_cancelled()
        .unwrap();
    let promise5 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 5 }))
        .not_cancelled()
        .unwrap();

    assert_eq!(
        task.inspect_waiters(),
        Some(vec![
            ParentKey::Some(DiceKey { index: 1 }),
            ParentKey::Some(DiceKey { index: 2 }),
            ParentKey::Some(DiceKey { index: 3 }),
            ParentKey::Some(DiceKey { index: 4 }),
            ParentKey::Some(DiceKey { index: 5 }),
        ])
    );

    let (v1, v2, v3, v4, v5) =
        futures::future::join5(promise1, promise2, promise3, promise4, promise5).await;
    assert!(
        v1?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );
    assert!(
        v2?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );
    assert!(
        v3?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );
    assert!(
        v4?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );
    assert!(
        v5?.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );

    Ok(())
}

#[tokio::test]
async fn sync_complete_task_completes_promises() -> anyhow::Result<()> {
    let task = unsafe {
        // SAFETY: completed below later
        sync_dice_task(DiceKey { index: 100 })
    };

    let mut promise_before = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    assert!(poll!(&mut promise_before).is_pending());

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .get_or_complete(|| DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(2))),
                Arc::new(CellHistory::empty())
            ))?
            .value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );

    let promise_after = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    let polled = futures::poll!(promise_before);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    let polled = futures::poll!(promise_after);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    Ok(())
}

#[tokio::test]
async fn sync_complete_task_wakes_waiters() -> anyhow::Result<()> {
    let task = unsafe {
        // SAFETY: completed below later
        sync_dice_task(DiceKey { index: 88 })
    };

    let mut promise1 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    let mut promise2 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 2 }))
        .not_cancelled()
        .unwrap();
    let mut promise3 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 3 }))
        .not_cancelled()
        .unwrap();

    assert_eq!(
        task.inspect_waiters(),
        Some(vec![
            ParentKey::Some(DiceKey { index: 1 }),
            ParentKey::Some(DiceKey { index: 2 }),
            ParentKey::Some(DiceKey { index: 3 }),
        ])
    );

    let barrier = Arc::new(Barrier::new(4));

    let fut1 = tokio::spawn({
        let barrier = barrier.dupe();
        async move {
            assert!(poll!(&mut promise1).is_pending());
            barrier.wait().await;

            promise1.await
        }
    });
    let fut2 = tokio::spawn({
        let barrier = barrier.dupe();
        async move {
            assert!(poll!(&mut promise2).is_pending());
            barrier.wait().await;

            promise2.await
        }
    });
    let fut3 = tokio::spawn({
        let barrier = barrier.dupe();
        async move {
            assert!(poll!(&mut promise3).is_pending());
            barrier.wait().await;

            promise3.await
        }
    });

    barrier.wait().await;

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .get_or_complete(|| DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
                Arc::new(CellHistory::empty())
            ))?
            .value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
    );

    let (v1, v2, v3) = futures::future::join3(fut1, fut2, fut3).await;
    assert!(
        v1??.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
    );
    assert!(
        v2??.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
    );
    assert!(
        v3??.value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
    );

    Ok(())
}

#[tokio::test]
async fn sync_complete_unfinished_spawned_task() -> anyhow::Result<()> {
    let lock = Arc::new(Mutex::new(()));

    let g = lock.lock().await;

    let task = spawn_dice_task(DiceKey { index: 88 }, &TokioSpawner, &(), {
        let lock = lock.dupe();
        |handle| {
            async move {
                let _g = lock.lock().await;
                // wait for the lock too
                handle.finished(DiceComputedValue::new(
                    MaybeValidDiceValue::valid(DiceValidValue::testing_new(
                        DiceKeyValue::<K>::new(2),
                    )),
                    Arc::new(CellHistory::empty()),
                ));

                Box::new(()) as Box<dyn Any + Send + 'static>
            }
            .boxed()
        }
    });

    let promise_before = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .get_or_complete(|| DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
                Arc::new(CellHistory::empty())
            ))?
            .value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
    );

    drop(g);

    let promise_after = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    let polled = futures::poll!(promise_before);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    let polled = futures::poll!(promise_after);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    Ok(())
}

#[tokio::test]
async fn sync_complete_finished_spawned_task() -> anyhow::Result<()> {
    let sem = Arc::new(Semaphore::new(0));

    let task = spawn_dice_task(DiceKey { index: 44 }, &TokioSpawner, &(), {
        let sem = sem.dupe();
        |handle| {
            async move {
                // wait for the lock too
                handle.finished(DiceComputedValue::new(
                    MaybeValidDiceValue::valid(DiceValidValue::testing_new(
                        DiceKeyValue::<K>::new(2),
                    )),
                    Arc::new(CellHistory::empty()),
                ));

                sem.add_permits(1);

                Box::new(()) as Box<dyn Any + Send + 'static>
            }
            .boxed()
        }
    });

    let promise_before = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    let _g = sem.acquire().await.unwrap();

    // actually completes with `2` from the spawn
    assert!(
        task.depended_on_by(ParentKey::None)
            .not_cancelled()
            .unwrap()
            .get_or_complete(|| DiceComputedValue::new(
                MaybeValidDiceValue::valid(DiceValidValue::testing_new(DiceKeyValue::<K>::new(1))),
                Arc::new(CellHistory::empty())
            ))?
            .value()
            .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
    );

    let promise_after = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 1 }))
        .not_cancelled()
        .unwrap();

    let polled = futures::poll!(promise_before);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    let polled = futures::poll!(promise_after);

    match polled {
        Poll::Ready(v) => {
            assert!(
                v?.value()
                    .equality(&DiceValidValue::testing_new(DiceKeyValue::<K>::new(2)))
            )
        }
        Poll::Pending => panic!("Promise should be ready immediately"),
    }

    Ok(())
}

#[tokio::test]
async fn dropping_all_waiters_cancels_task() {
    let barrier = Arc::new(Barrier::new(2));

    let task = spawn_dice_task(DiceKey { index: 600 }, &TokioSpawner, &(), {
        let barrier = barrier.dupe();
        |handle| {
            async move {
                let _handle = handle;
                // wait for the lock too
                barrier.wait().await;
                futures::future::pending().await
            }
            .boxed()
        }
    });

    barrier.wait().await;

    let await_termination = task.await_termination();
    futures::pin_mut!(await_termination);

    assert!(!task.internal.state.is_ready(Ordering::SeqCst));
    assert!(!task.internal.state.is_terminated(Ordering::SeqCst));
    assert!(task.is_pending());
    assert_matches!(futures::poll!(&mut await_termination), Poll::Pending);

    let promise1 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    assert!(task.is_pending());
    assert_matches!(futures::poll!(&mut await_termination), Poll::Pending);

    let promise2 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    assert!(task.is_pending());
    assert_matches!(futures::poll!(&mut await_termination), Poll::Pending);

    drop(promise1);

    assert!(task.is_pending());
    assert_matches!(futures::poll!(&mut await_termination), Poll::Pending);

    let promise3 = task
        .depended_on_by(ParentKey::Some(DiceKey { index: 0 }))
        .not_cancelled()
        .unwrap();

    drop(promise2);
    drop(promise3);

    match task.depended_on_by(ParentKey::None) {
        MaybeCancelled::Ok(_) => {
            panic!("should be cancelled")
        }
        MaybeCancelled::Cancelled => {}
    }

    task.await_termination().await;

    assert!(!task.internal.state.is_ready(Ordering::SeqCst));
    assert!(task.internal.state.is_terminated(Ordering::SeqCst));
    assert!(!task.is_pending());
}

#[tokio::test]
async fn task_that_already_cancelled_returns_cancelled() {
    let task = spawn_dice_task(DiceKey { index: 777 }, &TokioSpawner, &(), {
        |_handle| async move { futures::future::pending().await }.boxed()
    });

    task.cancel();
    task.await_termination().await;

    match task.depended_on_by(ParentKey::None) {
        MaybeCancelled::Ok(_) => {
            panic!("should be cancelled")
        }
        MaybeCancelled::Cancelled => {}
    }
}


/*
 * 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::Any;
use std::sync::Arc;
use std::sync::Mutex;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::dice::Dice;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::ActivationData;
use crate::ActivationTracker;
use crate::DiceDataBuilder;
use crate::InjectedKey;

#[derive(Default, Allocative)]
struct Tracker {
    /// Key, deps, data, reused
    state: Mutex<Vec<(Kind, Vec<Kind>, Option<Data>, bool)>>,
}

impl Tracker {
    fn new() -> Self {
        Self {
            state: Mutex::new(Vec::new()),
        }
    }
}

impl ActivationTracker for Tracker {
    fn key_activated(
        &self,
        key: &dyn Any,
        deps: &mut dyn Iterator<Item = &dyn Any>,
        activation_data: ActivationData,
    ) {
        let (data, reused) = match activation_data {
            ActivationData::Evaluated(d) => (d.map(|d| *d.downcast::<Data>().unwrap()), false),
            ActivationData::Reused => (None, true),
        };

        self.state.lock().unwrap().push((
            Kind::from_any(key),
            deps.into_iter().map(Kind::from_any).collect(),
            data,
            reused,
        ));
    }
}

#[derive(PartialEq, Eq, Debug, Dupe, Clone, Allocative)]
enum Kind {
    Injected,
    Stage0,
    Stage1,
}

impl Kind {
    fn from_any(key: &dyn Any) -> Self {
        if key.is::<Injected>() {
            return Self::Injected;
        }

        if key.is::<Stage0>() {
            return Self::Stage0;
        }

        if key.is::<Stage1>() {
            return Self::Stage1;
        }

        panic!("Unexpected key: {:?}", key)
    }
}

#[derive(PartialEq, Eq, Debug, Dupe, Clone, Allocative)]
struct Data;

#[derive(Clone, Dupe, Debug, Display, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "{:?}", self)]
struct Injected;

#[async_trait]
impl InjectedKey for Injected {
    type Value = i32;

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        x == y
    }
}

#[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "{:?}", self)]
struct Stage0;

#[async_trait]
impl Key for Stage0 {
    type Value = ();

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.store_evaluation_data(Data).unwrap();
        ctx.compute(&Injected).await.unwrap();
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        true
    }
}

#[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "{:?}", self)]
struct Stage1;

#[async_trait]
impl Key for Stage1 {
    type Value = ();

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.store_evaluation_data(Data).unwrap();
        ctx.compute(&Stage0).await.unwrap()
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        true
    }
}

async fn test_events_impl(builder: DiceDataBuilder) -> anyhow::Result<()> {
    let dice = builder.build(DetectCycles::Enabled);

    {
        let activation_tracker = Arc::new(Tracker::new());

        let data = UserComputationData {
            activation_tracker: Some(activation_tracker.dupe()),
            ..Default::default()
        };

        let mut updater = dice.updater_with_data(data);
        updater.changed_to(vec![(Injected, 123)])?;

        let transaction = updater.commit().await;
        transaction.compute(&Stage1).await?;

        assert_eq!(
            &*activation_tracker.state.lock().unwrap(),
            &[
                (Kind::Stage0, vec![Kind::Injected], Some(Data), false),
                (Kind::Stage1, vec![Kind::Stage0], Some(Data), false),
            ]
        );
    }

    {
        let activation_tracker = Arc::new(Tracker::default());

        let data = UserComputationData {
            activation_tracker: Some(activation_tracker.dupe()),
            ..Default::default()
        };

        // Change the value.
        let mut updater = dice.updater_with_data(data);
        updater.changed_to(vec![(Injected, 456)])?;

        let transaction = updater.commit().await;
        transaction.compute(&Stage1).await?;

        assert_eq!(
            &*activation_tracker.state.lock().unwrap(),
            &[
                (Kind::Stage0, vec![Kind::Injected], Some(Data), false),
                (Kind::Stage1, vec![Kind::Stage0], None, true),
            ]
        );
    }

    Ok(())
}

#[tokio::test]
async fn test_events_legacy() -> anyhow::Result<()> {
    test_events_impl(Dice::builder()).await
}

#[tokio::test]
async fn test_events_modern() -> anyhow::Result<()> {
    test_events_impl(Dice::modern()).await
}


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

/////////////////////////////////////////////////////////////////////
// DEMO USAGE
// Reading a file is based on the value in that file, plus the file encoding

use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use tempfile::NamedTempFile;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::injected::InjectedKey;
use crate::api::key::Key;
use crate::api::transaction::DiceTransactionUpdater;
use crate::Dice;

#[derive(Debug, Clone, Dupe, PartialEq, Allocative)]
enum Encoding {
    Utf8,
    Ascii,
}

#[derive(Clone, Dupe, Debug, Display, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "{:?}", self)]
struct EncodingConfig();

impl InjectedKey for EncodingConfig {
    type Value = Encoding;

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        x == y
    }
}

struct Encodings<'c>(&'c DiceComputations);

impl<'c> Encodings<'c> {
    async fn get(&self) -> Result<Encoding, Arc<anyhow::Error>> {
        self.0
            .compute(&EncodingConfig())
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))
    }
}

trait HasEncodings {
    fn encodings(&self) -> Encodings;
}

impl HasEncodings for DiceComputations {
    fn encodings(&self) -> Encodings {
        Encodings(self)
    }
}

trait SetEncodings {
    fn set_encodings(&mut self, enc: Encoding) -> anyhow::Result<()>;
}

impl SetEncodings for DiceTransactionUpdater {
    fn set_encodings(&mut self, enc: Encoding) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(EncodingConfig(), enc)])?)
    }
}

struct Filesystem<'c>(&'c DiceComputations);

#[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "File({})", "_0.display()")]
struct File(PathBuf);

impl<'c> Filesystem<'c> {
    async fn read_file(&self, file: &Path) -> Result<Arc<String>, Arc<anyhow::Error>> {
        #[async_trait]
        impl Key for File {
            type Value = Result<Arc<String>, Arc<anyhow::Error>>;
            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellations: &CancellationContext,
            ) -> Self::Value {
                let encoding = ctx.encodings().get().await?;

                let s = fs::read_to_string(&self.0).unwrap();

                Ok(Arc::new(match encoding {
                    Encoding::Utf8 => s,
                    Encoding::Ascii => s.replace(":-)", "smile"),
                }))
            }

            fn equality(x: &Self::Value, y: &Self::Value) -> bool {
                match (x, y) {
                    (Ok(x), Ok(y)) => x == y,
                    _ => false,
                }
            }
        }

        self.0
            .compute(&File(file.to_path_buf()))
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?
    }
}

trait HasFilesystem<'c> {
    fn filesystem(&'c self) -> Arc<Filesystem<'c>>;
}

impl<'c> HasFilesystem<'c> for DiceComputations {
    fn filesystem(&'c self) -> Arc<Filesystem<'c>> {
        Arc::new(Filesystem(self))
    }
}

trait SetFilesystem {
    fn filesystem_changed(&mut self, file: &Path) -> anyhow::Result<()>;
}

impl SetFilesystem for DiceTransactionUpdater {
    fn filesystem_changed(&mut self, file: &Path) -> anyhow::Result<()> {
        Ok(self.changed(vec![File(file.to_path_buf())])?)
    }
}

#[tokio::test]
async fn demo() -> anyhow::Result<()> {
    let temp = NamedTempFile::new().unwrap();
    let f = PathBuf::from(temp.path());

    let dice = Dice::modern().build(DetectCycles::Enabled);

    let mut ctx = dice.updater();
    ctx.set_encodings(Encoding::Utf8)?;
    ctx.commit().await;

    let set = |x: &str| fs::write(&f, x).unwrap();

    async fn get(x: &str, dice: &Arc<Dice>, f: &Path) {
        let contents = dice
            .updater()
            .commit()
            .await
            .filesystem()
            .read_file(f)
            .await
            .unwrap();
        assert_eq!(*contents, x)
    }

    set(":-)");

    get(":-)", &dice, &f).await;

    // doesn't change because I didn't dirty it
    set("hello :-)");

    get(":-)", &dice, &f).await;

    let mut ctx = dice.updater();
    ctx.filesystem_changed(&f)?;
    ctx.commit().await;
    get("hello :-)", &dice, &f).await;

    let mut ctx = dice.updater();
    ctx.set_encodings(Encoding::Ascii)?;
    ctx.commit().await;
    get("hello smile", &dice, &f).await;

    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::sync::Arc;
use std::sync::Mutex;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::dice::Dice;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::DiceDataBuilder;
use crate::DiceEvent;
use crate::DiceEventListener;
use crate::InjectedKey;

#[derive(Default, Allocative)]
struct Tracker {
    state: Mutex<Vec<DiceEvent>>,
}

impl DiceEventListener for Tracker {
    fn event(&self, event: DiceEvent) {
        self.state.lock().unwrap().push(event);
    }
}

#[derive(Clone, Dupe, Debug, Display, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "{:?}", self)]
struct Injected;

#[async_trait]
impl InjectedKey for Injected {
    type Value = i32;

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        x == y
    }
}

#[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "{:?}", self)]
struct Stage0;

#[async_trait]
impl Key for Stage0 {
    type Value = ();

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.compute(&Injected).await.unwrap();
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        true
    }
}

#[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "{:?}", self)]
struct Stage1;

#[async_trait]
impl Key for Stage1 {
    type Value = ();

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.compute(&Stage0).await.unwrap()
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        true
    }
}

async fn test_events_impl(builder: DiceDataBuilder) -> anyhow::Result<()> {
    let dice = builder.build(DetectCycles::Enabled);

    {
        let tracker = Arc::new(Tracker::default());

        let data = UserComputationData {
            tracker: tracker.dupe(),
            ..Default::default()
        };

        let mut updater = dice.updater_with_data(data);
        updater.changed_to(vec![(Injected, 123)])?;

        let transaction = updater.commit().await;
        transaction.compute(&Stage1).await?;

        assert_eq!(
            &*tracker.state.lock().unwrap(),
            &[
                DiceEvent::Started { key_type: "Stage1" },
                DiceEvent::Started { key_type: "Stage0" },
                DiceEvent::Finished { key_type: "Stage0" },
                DiceEvent::Finished { key_type: "Stage1" },
            ]
        );
    }

    {
        let tracker = Arc::new(Tracker::default());

        let data = UserComputationData {
            tracker: tracker.dupe(),
            ..Default::default()
        };

        // Change the value.
        let mut updater = dice.updater_with_data(data);
        updater.changed_to(vec![(Injected, 456)])?;

        let transaction = updater.commit().await;
        transaction.compute(&Stage1).await?;

        assert_eq!(
            &*tracker.state.lock().unwrap(),
            &[
                DiceEvent::CheckDepsStarted { key_type: "Stage1" },
                DiceEvent::CheckDepsStarted { key_type: "Stage0" },
                DiceEvent::CheckDepsFinished { key_type: "Stage0" },
                DiceEvent::Started { key_type: "Stage0" },
                DiceEvent::Finished { key_type: "Stage0" },
                DiceEvent::CheckDepsFinished { key_type: "Stage1" }
            ]
        );
    }

    Ok(())
}

#[tokio::test]
async fn test_events_legacy() -> anyhow::Result<()> {
    test_events_impl(Dice::builder()).await
}

#[tokio::test]
async fn test_events_modern() -> anyhow::Result<()> {
    test_events_impl(Dice::modern()).await
}