/*
 * 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::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Barrier;
use std::sync::Mutex;

use allocative::Allocative;
use assert_matches::assert_matches;
use async_trait::async_trait;
use derivative::Derivative;
use derive_more::Display;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use tokio::sync::oneshot;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::error::DiceErrorImpl;
use crate::api::injected::InjectedKey;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::impls::dice::DiceModern;
use crate::versions::VersionNumber;
use crate::Dice;
use crate::DiceData;
use crate::UserCycleDetector;
use crate::UserCycleDetectorGuard;

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

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

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

#[derive(Clone, Dupe, Debug, Derivative, Allocative, Display)]
#[derivative(PartialEq, Eq, Hash)]
#[display(fmt = "{:?}", self)]
#[allocative(skip)]
struct KeyThatRuns {
    #[derivative(Hash = "ignore", PartialEq = "ignore")]
    barrier1: Arc<tokio::sync::Semaphore>,
    #[derivative(Hash = "ignore", PartialEq = "ignore")]
    barrier2: Arc<tokio::sync::Semaphore>,
    #[derivative(Hash = "ignore", PartialEq = "ignore")]
    is_ran: Arc<AtomicBool>,
}

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

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        self.barrier1.add_permits(1);
        let _s = self.barrier2.acquire().await.unwrap();
        self.is_ran.store(true, Ordering::SeqCst);
    }

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

#[tokio::test]
async fn set_injected_multiple_times_per_commit() -> anyhow::Result<()> {
    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    {
        let mut ctx = dice.updater();
        ctx.changed_to(vec![(Foo(0), 0)])?;
        ctx.changed_to(vec![(Foo(1), 1)])?;

        let ctx = ctx.commit().await;
        assert_eq!(ctx.compute(&Foo(0)).await?, 0);
        assert_eq!(ctx.compute(&Foo(1)).await?, 1);
    }

    {
        let mut ctx = dice.updater();
        ctx.changed_to(vec![(Foo(0), 0)])?;

        assert_matches!(
            ctx.changed_to(vec![(Foo(0), 1)]),
            Err(err) => assert_matches!(&*err.0, DiceErrorImpl::DuplicateChange(_))
        );
    }

    Ok(())
}

#[tokio::test]
async fn set_injected_with_no_change_no_new_ctx() -> anyhow::Result<()> {
    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    {
        let mut ctx = dice.updater();
        ctx.changed_to(vec![(Foo(0), 0)])?;

        let ctx = ctx.commit().await;

        assert_eq!(ctx.get_version(), VersionNumber::new(1));
    }

    {
        let mut ctx = dice.updater();
        ctx.changed_to(vec![(Foo(0), 0)])?;

        let ctx = ctx.commit().await;
        assert_eq!(ctx.get_version(), VersionNumber::new(1));
    }

    Ok(())
}

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

#[async_trait]
impl Key for K {
    type Value = Result<K, Arc<anyhow::Error>>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        let mut sum = self.0;
        for i in 0..self.0 {
            sum += ctx
                .compute(&K(i))
                .await
                .map_err(|e| Arc::new(anyhow::anyhow!(e)))??
                .0;
        }
        Ok(K(sum))
    }

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

#[test]
fn dice_computations_are_parallel() {
    let n_thread = 10;

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(n_thread)
        .max_blocking_threads(n_thread)
        .build()
        .unwrap();
    let barrier = Arc::new(Barrier::new(n_thread));

    #[derive(Clone, Debug, Display, Derivative, Allocative)]
    #[derivative(Hash, PartialEq, Eq)]
    #[display(fmt = "{:?}", self)]
    struct Blocking {
        index: usize,
        #[derivative(PartialEq = "ignore", Hash = "ignore")]
        #[allocative(skip)]
        barrier: Arc<Barrier>,
    }

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

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            self.barrier.wait();
            1
        }

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

    rt.block_on(async move {
        let dice = DiceModern::builder().build(DetectCycles::Disabled);
        let mut sum = 0;

        let dice = &dice;
        let barrier = &barrier;

        let futs = (0..n_thread)
            .map(|i| async move {
                dice.updater()
                    .commit()
                    .await
                    .compute(&Blocking {
                        index: i,
                        barrier: barrier.dupe(),
                    })
                    .await
                    .unwrap()
            })
            .collect::<Vec<_>>();

        futures::future::join_all(futs)
            .await
            .iter()
            .for_each(|res| sum += res);

        assert_eq!(sum, n_thread);
    })
}

#[tokio::test]
async fn different_data_per_compute_ctx() {
    struct U(usize);

    #[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
    #[display(fmt = "{:?}", self)]
    struct DataRequest(u8);
    #[async_trait]
    impl Key for DataRequest {
        type Value = usize;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            ctx.per_transaction_data().data.get::<U>().unwrap().0
        }

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

    let dice = DiceModern::builder().build(DetectCycles::Disabled);
    let per_cmd_data0 = {
        let mut d = UserComputationData::new();
        d.data.set(U(0));
        d
    };
    let per_cmd_data1 = {
        let mut d = UserComputationData::new();
        d.data.set(U(1));
        d
    };

    let ctx0 = dice.updater_with_data(per_cmd_data0).commit().await;

    let ctx1 = dice.updater_with_data(per_cmd_data1).commit().await;

    let request0 = ctx0.compute(&DataRequest(0));
    let request1 = ctx1.compute(&DataRequest(1));

    assert_eq!(request0.await.unwrap(), 0);
    assert_eq!(request1.await.unwrap(), 1);
}

#[test]
fn invalid_update() {
    #[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
    struct Invalid;

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

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

        fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
            unimplemented!("not needed for test")
        }

        fn validity(_x: &Self::Value) -> bool {
            false
        }
    }

    let dice = DiceModern::new(DiceData::new());
    let mut updater = dice.updater();

    assert!(updater.changed_to([(Invalid, ())]).is_err());
}

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

#[async_trait]
impl Key for Fib {
    type Value = Result<u64, Arc<anyhow::Error>>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        if self.0 > 93 {
            return Err(Arc::new(anyhow::anyhow!("that's too big")));
        }
        if self.0 < 2 {
            return Ok(self.0 as u64);
        }
        let (a, b) =
            futures::future::join(ctx.compute(&Fib(self.0 - 2)), ctx.compute(&Fib(self.0 - 1)))
                .await;
        match (a, b) {
            (Ok(a), Ok(b)) => Ok(a? + b?),
            _ => Err(Arc::new(anyhow::anyhow!("some dice error"))),
        }
    }

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

#[derive(Debug, Eq, PartialEq)]
enum CycleDetectorEvents {
    Start(Fib),
    Finish(Fib),
    Edge(Fib, Fib),
}

#[derive(Debug, Clone)]
struct CycleDetector {
    events: Arc<Mutex<Vec<CycleDetectorEvents>>>,
}

struct CycleDetectorGuard {
    key: Fib,
    events: Arc<Mutex<Vec<CycleDetectorEvents>>>,
}

impl UserCycleDetector for CycleDetector {
    fn start_computing_key(
        &self,
        key: &dyn std::any::Any,
    ) -> Option<Box<dyn UserCycleDetectorGuard>> {
        let f = key.downcast_ref::<Fib>().unwrap();
        self.events
            .lock()
            .unwrap()
            .push(CycleDetectorEvents::Start(*f));
        Some(Box::new(CycleDetectorGuard {
            key: *f,
            events: self.events.dupe(),
        }))
    }

    fn finished_computing_key(&self, key: &dyn std::any::Any) {
        let f = key.downcast_ref::<Fib>().unwrap();
        self.events
            .lock()
            .unwrap()
            .push(CycleDetectorEvents::Finish(*f));
    }
}

impl UserCycleDetectorGuard for CycleDetectorGuard {
    fn add_edge(&self, key: &dyn std::any::Any) {
        let f = key.downcast_ref::<Fib>().unwrap();
        self.events
            .lock()
            .unwrap()
            .push(CycleDetectorEvents::Edge(self.key, *f))
    }

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

    fn type_name(&self) -> &'static str {
        std::any::type_name::<Self>()
    }
}

#[test]
fn user_cycle_detector_receives_events() -> anyhow::Result<()> {
    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let events = Arc::new(Mutex::new(Vec::new()));
        let user_data = UserComputationData {
            cycle_detector: Some(Arc::new(CycleDetector {
                events: events.dupe(),
            })),
            ..Default::default()
        };
        let ctx = dice.updater_with_data(user_data).commit().await;
        let res = ctx.compute(&Fib(20)).await?.expect("should succeed");
        assert_eq!(res, 6765);

        let check_events = move |i, expected_edges| {
            let mut started = false;
            let mut finished = false;
            let mut edges = Vec::new();
            for ev in events.lock().unwrap().iter() {
                match ev {
                    CycleDetectorEvents::Start(Fib(v)) if i == *v => {
                        started = true;
                        assert!(!finished);
                    }
                    CycleDetectorEvents::Finish(Fib(v)) if i == *v => {
                        assert!(!finished);
                        finished = true;
                        assert!(started);
                    }
                    CycleDetectorEvents::Edge(Fib(v), Fib(j)) if i == *v => {
                        assert!(started);
                        assert!(!finished);
                        edges.push(*j)
                    }
                    _ => {
                        // ignore
                    }
                }
            }
            assert!(finished);
            edges.sort();
            assert_eq!(edges, expected_edges);
        };

        check_events(0, vec![]);
        check_events(1, vec![]);
        for i in 2..=20 {
            check_events(i, vec![i - 2, i - 1]);
        }

        Ok(())
    })
}

/// Test that dropping request cancel execution. Test scenario is:
/// The test has 1 barrier and 1 semaphore:
/// - barrier is used by the task to signal that the task is started, and that the task doesn't progress
///   to complete.
/// - semaphore ensures that the task doesn't run until the main thread drops the request
/// 1. we spawn a task
/// 2. assert task is spawned
/// 3. drop the request future
/// 4. ensure the task is cancelled since the task will panic if it proceeds beyond the semaphore
#[tokio::test]
async fn dropping_request_future_cancels_execution() {
    #[derive(Debug)]
    struct DropSignal(Option<oneshot::Sender<()>>);

    impl Drop for DropSignal {
        fn drop(&mut self) {
            self.0.take().unwrap().send(()).unwrap();
        }
    }

    #[derive(Clone, Dupe, Debug, Derivative, Allocative, Display)]
    #[derivative(PartialEq, Eq, Hash)]
    #[display(fmt = "{:?}", self)]
    #[allocative(skip)]
    struct KeyThatShouldntRun {
        #[derivative(Hash = "ignore", PartialEq = "ignore")]
        barrier1: Arc<tokio::sync::Barrier>,
        #[derivative(Hash = "ignore", PartialEq = "ignore")]
        barrier2: Arc<tokio::sync::Semaphore>,
        #[derivative(Hash = "ignore", PartialEq = "ignore")]
        drop_signal: Arc<Mutex<Option<DropSignal>>>,
    }

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

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            let _drop_signal = self.drop_signal.lock().unwrap().take();

            self.barrier1.wait().await;
            let _guard = self.barrier2.acquire().await.unwrap();

            panic!("shouldn't run")
        }

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

    let barrier1 = Arc::new(tokio::sync::Barrier::new(2));
    let barrier2 = Arc::new(tokio::sync::Semaphore::new(0));

    let (tx, rx) = oneshot::channel();
    let drop_signal = DropSignal(Some(tx));

    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    let ctx = dice.updater().commit().await;

    let key = KeyThatShouldntRun {
        barrier1: barrier1.dupe(),
        barrier2: barrier2.dupe(),
        drop_signal: Arc::new(Mutex::new(Some(drop_signal))),
    };
    let req = ctx.compute(&key);

    // ensure that the key starts computing
    barrier1.wait().await;

    drop(req);

    barrier2.add_permits(1);

    rx.await.unwrap();
}

/// Test that dropping request cancel execution. Test scenario is:
/// The test has 2 barriers:
/// - first is used by the task to signal that the task is started, and that the task doesn't progress
///   to complete.
/// - second ensures that the task doesn't run until the main thread drops the one of the request
/// 1. we spawn a task, but have two requests attached to it
/// 2. assert task is spawned
/// 3. drop one of the request future
/// 4. ensure the task is still finishes and runs by checking the `is-ran` atomic
#[tokio::test]
async fn dropping_request_future_doesnt_cancel_if_multiple_requests_active() {
    let barrier1 = Arc::new(tokio::sync::Semaphore::new(0));
    let barrier2 = Arc::new(tokio::sync::Semaphore::new(0));
    let is_ran = Arc::new(AtomicBool::new(false));

    let key = KeyThatRuns {
        barrier1: barrier1.dupe(),
        barrier2: barrier2.dupe(),
        is_ran: is_ran.dupe(),
    };

    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    let ctx = dice.updater().commit().await;
    let req1 = ctx.compute(&key);
    let req2 = ctx.compute(&key);

    // ensure that the key starts computing
    let _b = barrier1.acquire().await;

    drop(req1);

    barrier2.add_permits(1);

    // req2 still succeed
    req2.await.unwrap();

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

#[tokio::test]
async fn user_cycle_detector_is_present_legacy() -> anyhow::Result<()> {
    user_cycle_detector_is_present(Dice::builder().build(DetectCycles::Disabled)).await
}

#[tokio::test]
async fn user_cycle_detector_is_present_modern() -> anyhow::Result<()> {
    user_cycle_detector_is_present(Dice::modern().build(DetectCycles::Disabled)).await
}

async fn user_cycle_detector_is_present(dice: Arc<Dice>) -> anyhow::Result<()> {
    #[derive(Clone, Copy, Dupe, Display, Debug, Eq, PartialEq, Hash, Allocative)]
    #[display(fmt = "{:?}", self)]
    struct AccessCycleGuardKey;

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

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            assert!(
                ctx.cycle_guard::<AccessCycleDetectorGuard>()
                    .unwrap()
                    .is_some()
            );
        }

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

    #[derive(Debug, Clone)]
    struct AccessCycleDetector;

    struct AccessCycleDetectorGuard;

    impl UserCycleDetector for AccessCycleDetector {
        fn start_computing_key(
            &self,
            _key: &dyn std::any::Any,
        ) -> Option<Box<dyn UserCycleDetectorGuard>> {
            Some(Box::new(AccessCycleDetectorGuard))
        }

        fn finished_computing_key(&self, _key: &dyn std::any::Any) {}
    }

    impl UserCycleDetectorGuard for AccessCycleDetectorGuard {
        fn add_edge(&self, _key: &dyn std::any::Any) {}

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

        fn type_name(&self) -> &'static str {
            std::any::type_name::<Self>()
        }
    }

    let user_data = UserComputationData {
        cycle_detector: Some(Arc::new(AccessCycleDetector)),
        ..Default::default()
    };
    let ctx = dice.updater_with_data(user_data).commit().await;
    Ok(ctx.compute(&AccessCycleGuardKey).await?)
}

#[tokio::test]
async fn test_dice_usable_after_cancellations() {
    let dice = DiceModern::builder().build(DetectCycles::Disabled);

    let ctx = dice.updater().commit().await;

    let barrier1 = Arc::new(tokio::sync::Semaphore::new(0));
    let barrier2 = Arc::new(tokio::sync::Semaphore::new(0));
    let is_ran = Arc::new(AtomicBool::new(false));

    let key = KeyThatRuns {
        barrier1: barrier1.dupe(),
        barrier2: barrier2.dupe(),
        is_ran: is_ran.dupe(),
    };

    let req1 = ctx.compute(&key);

    // ensure that the key starts computing
    let _b = barrier1.acquire().await;

    // cancel and await for cancellation
    drop(req1);
    drop(ctx);

    dice.wait_for_idle().await;

    assert!(!is_ran.load(Ordering::Acquire));

    let ctx = dice.updater().commit().await;

    // req2 still succeed. Note that due to dice caching, even if we make a new key, the same
    // instance would be used, so just use the same one.
    let req2 = ctx.compute(&key);

    barrier2.add_permits(1);

    req2.await.unwrap();

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


/*
 * 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::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use futures::future::join3;
use more_futures::cancellation::CancellationContext;
use tokio::sync::Mutex;

use crate::api::computations::DiceComputations;
use crate::api::data::DiceData;
use crate::api::key::Key;
use crate::impls::dice::DiceModern;

#[tokio::test]
async fn concurrent_identical_requests_are_deduped() -> anyhow::Result<()> {
    #[derive(Allocative, Clone, Debug, Display)]
    #[display(fmt = "{:?}", self)]
    struct ComputeOnce(#[allocative(skip)] Arc<Mutex<u8>>);

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

    impl Eq for ComputeOnce {}

    impl Hash for ComputeOnce {
        fn hash<H: Hasher>(&self, _state: &mut H) {}
    }

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

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            let mut count = self.0.lock().await;
            *count += 1;
        }

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

    let dice = DiceModern::new(DiceData::new());

    let guard = Arc::new(Mutex::new(0));
    let _g = guard.lock().await;

    let ctx = dice.updater().commit().await;

    let k = ComputeOnce(guard.dupe());

    let compute1 = ctx.compute(&k);
    let compute2 = ctx.compute(&k);
    let compute3 = ctx.compute(&k);

    drop(_g);

    let (r1, r2, r3) = join3(compute1, compute2, compute3).await;
    r1?;
    r2?;
    r3?;

    assert_eq!(1, *guard.lock().await);

    Ok(())
}

#[test]
fn different_requests_are_spawned_in_parallel() -> anyhow::Result<()> {
    let n_thread = 10usize;

    #[derive(Allocative, Clone, Debug, Display)]
    #[display(fmt = "{:?}", self)]
    // purposely use a sync barrier to see that our compute is spawned
    struct ComputeParallel(#[allocative(skip)] Arc<std::sync::Barrier>);

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

    impl Eq for ComputeParallel {}

    impl Hash for ComputeParallel {
        fn hash<H: Hasher>(&self, _state: &mut H) {}
    }

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

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            self.0.wait();

            1
        }

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

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(n_thread)
        .build()
        .unwrap();

    // use barrier to ensure that n_threads in parallel are spawned by the engine
    let barrier = Arc::new(std::sync::Barrier::new(n_thread));

    rt.block_on(async move {
        let dice = DiceModern::new(DiceData::new());

        let ctx = dice.updater().commit().await;
        let k = ComputeParallel(barrier.dupe());

        let futs = (0..n_thread).map(|_| ctx.compute(&k)).collect::<Vec<_>>();

        let mut sum = 0;
        futures::future::join_all(futs)
            .await
            .into_iter()
            .try_for_each(|res| {
                sum += res?;
                anyhow::Ok(())
            })?;

        assert_eq!(sum, n_thread);

        anyhow::Ok::<()>(())
    })?;

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

mod activation_tracker;
mod demo;
mod events;
mod general;
mod keys;
mod spawner;
mod transients;
mod user_data;


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

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dupe::Dupe;
use futures::future::BoxFuture;
use more_futures::cancellation::CancellationContext;
use more_futures::spawner::Spawner;
use tokio::task::JoinHandle;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::impls::dice::DiceModern;

struct MySpawner(AtomicUsize);

impl<S> Spawner<S> for MySpawner {
    fn spawn(
        &self,
        _ctx: &S,
        fut: BoxFuture<'static, Box<dyn Any + Send + 'static>>,
    ) -> JoinHandle<Box<dyn Any + Send + 'static>> {
        self.0.fetch_add(1, Ordering::SeqCst);

        tokio::spawn(fut)
    }
}

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

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

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

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

#[tokio::test]
async fn uses_custom_spawner() {
    let dice = DiceModern::builder().build(DetectCycles::Disabled);
    let spawner = Arc::new(MySpawner(AtomicUsize::new(0)));

    let mut data = UserComputationData::new();
    data.spawner = spawner.dupe();
    let updater = dice.updater_with_data(data);

    let ctx = updater.commit().await;

    ctx.compute(&K).await.unwrap();

    assert_eq!(spawner.0.load(Ordering::SeqCst), 1);
}


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

use allocative::Allocative;
use async_trait::async_trait;
use derivative::Derivative;
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::key::Key;
use crate::impls::dice::DiceModern;

#[tokio::test]
async fn invalid_results_are_not_cached() -> anyhow::Result<()> {
    #[derive(Clone, Dupe, Debug, Display, Derivative, Allocative)]
    #[derivative(Hash, PartialEq, Eq)]
    #[display(fmt = "{:?}", self)]
    struct AlwaysTransient(#[derivative(PartialEq = "ignore", Hash = "ignore")] Arc<AtomicBool>);

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

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

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

        fn validity(_x: &Self::Value) -> bool {
            false
        }
    }

    let dice = DiceModern::builder().build(DetectCycles::Enabled);
    let is_ran = Arc::new(AtomicBool::new(false));
    {
        let ctx = dice.updater().commit().await;
        ctx.compute(&AlwaysTransient(is_ran.dupe())).await?;
        assert!(is_ran.load(Ordering::SeqCst));

        // same ctx, so should reuse the result and
        is_ran.store(false, Ordering::SeqCst);
        ctx.compute(&AlwaysTransient(is_ran.dupe())).await?;
        assert!(!is_ran.load(Ordering::SeqCst));

        // simultaneously ctx should also re-use the result
        let ctx1 = dice.updater().commit().await;
        is_ran.store(false, Ordering::SeqCst);
        ctx1.compute(&AlwaysTransient(is_ran.dupe())).await?;
        assert!(!is_ran.load(Ordering::SeqCst));
    }

    {
        // new context should re-run
        let ctx = dice.updater().commit().await;
        is_ran.store(false, Ordering::SeqCst);
        ctx.compute(&AlwaysTransient(is_ran.dupe())).await?;
        assert!(is_ran.load(Ordering::SeqCst));

        // same ctx, so should reuse the result and
        is_ran.store(false, Ordering::SeqCst);
        ctx.compute(&AlwaysTransient(is_ran.dupe())).await?;
        assert!(!is_ran.load(Ordering::SeqCst));
    }

    Ok(())
}

#[tokio::test]
async fn demo_with_transient() -> anyhow::Result<()> {
    #[derive(Clone, Dupe, Debug, Display, Derivative, Allocative)]
    #[derivative(Hash, PartialEq, Eq)]
    #[display(fmt = "{:?}", self)]
    struct MaybeTransient(
        usize,
        #[derivative(PartialEq = "ignore", Hash = "ignore")] Arc<AtomicBool>,
    );

    #[async_trait]
    impl Key for MaybeTransient {
        type Value = Result<usize, bool>;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            if self.0 == 0 {
                if !self.1.load(Ordering::SeqCst) {
                    Err(true)
                } else {
                    Ok(1)
                }
            } else {
                let mut sum = 0;
                for i in 0..self.0 {
                    if let Ok(v) = ctx
                        .compute(&MaybeTransient(i, self.1.dupe()))
                        .await
                        .unwrap()
                    {
                        sum += v;
                    } else {
                        return Err(false);
                    }
                }
                Ok(sum)
            }
        }

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

        fn validity(x: &Self::Value) -> bool {
            // intermediate nodes won't be directly invalid, but rely on the children to
            // propagate transient-ness
            if let Err(x) = x { !*x } else { true }
        }
    }

    let dice = DiceModern::builder().build(DetectCycles::Enabled);

    let ctx = dice.updater().commit().await;
    let validity = Arc::new(AtomicBool::new(false));

    assert!(
        ctx.compute(&MaybeTransient(10, validity.dupe()))
            .await?
            .is_err(),
    );

    validity.store(true, Ordering::SeqCst);
    assert!(
        ctx.compute(&MaybeTransient(10, validity.dupe()))
            .await?
            .is_err(),
    );

    drop(ctx);

    let ctx = dice.updater().commit().await;
    assert_eq!(
        ctx.compute(&MaybeTransient(10, validity.dupe())).await?,
        Ok(512)
    );

    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 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::key::Key;
use crate::api::user_data::UserComputationData;
use crate::impls::dice::DiceModern;

#[tokio::test]
async fn different_data_per_compute_ctx() {
    struct U(usize);

    #[derive(Clone, Dupe, Debug, Display, PartialEq, Eq, Hash, Allocative)]
    #[display(fmt = "{:?}", self)]
    struct DataRequest(u8);
    #[async_trait]
    impl Key for DataRequest {
        type Value = usize;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            ctx.per_transaction_data().data.get::<U>().unwrap().0
        }

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

    let dice = DiceModern::builder().build(DetectCycles::Enabled);
    let per_cmd_data0 = {
        let mut d = UserComputationData::new();
        d.data.set(U(0));
        d
    };
    let per_cmd_data1 = {
        let mut d = UserComputationData::new();
        d.data.set(U(1));
        d
    };

    let ctx0 = dice.updater_with_data(per_cmd_data0).commit().await;

    let ctx1 = dice.updater_with_data(per_cmd_data1).commit().await;

    let request0 = ctx0.compute(&DataRequest(0));
    let request1 = ctx1.compute(&DataRequest(1));

    assert_eq!(request0.await.unwrap(), 0);
    assert_eq!(request1.await.unwrap(), 1);
}


/*
 * 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 main worker thread for the dice task

use std::any::Any;

use dupe::Dupe;
use futures::FutureExt;
use tracing::Instrument;

use crate::impls::evaluator::AsyncEvaluator;
use crate::impls::events::DiceEventDispatcher;
use crate::impls::incremental::IncrementalEngine;
use crate::impls::key::DiceKey;
use crate::impls::task::dice::DiceTask;
use crate::impls::task::handle::DiceTaskHandle;
use crate::impls::task::spawn_dice_task;
use crate::impls::task::PreviouslyCancelledTask;
use crate::impls::user_cycle::UserCycleDetectorData;
use crate::impls::worker::state::DiceWorkerStateAwaitingPrevious;
use crate::impls::worker::state::DiceWorkerStateFinishedAndCached;
use crate::result::CancellableResult;
use crate::result::Cancelled;

pub(crate) mod state;

/// The worker on the spawned dice task
pub(crate) struct DiceTaskWorker {
    k: DiceKey,
    eval: AsyncEvaluator,
    cycles: UserCycleDetectorData,
    events_dispatcher: DiceEventDispatcher,
    previously_cancelled_task: Option<PreviouslyCancelledTask>,
    incremental: IncrementalEngine,
}

impl DiceTaskWorker {
    pub(crate) fn spawn(
        k: DiceKey,
        eval: AsyncEvaluator,
        cycles: UserCycleDetectorData,
        events_dispatcher: DiceEventDispatcher,
        previously_cancelled_task: Option<PreviouslyCancelledTask>,
        incremental: IncrementalEngine,
    ) -> DiceTask {
        let span = debug_span!(parent: None, "spawned_dice_task", k = ?k, v = %eval.per_live_version_ctx.get_version(), v_epoch = %incremental.version_epoch);

        let spawner = eval.user_data.spawner.dupe();
        let spawner_ctx = eval.user_data.dupe();

        let worker = DiceTaskWorker::new(
            k,
            eval,
            cycles,
            events_dispatcher,
            previously_cancelled_task,
            incremental,
        );

        spawn_dice_task(k, &*spawner, &spawner_ctx, move |handle| {
            async move {
                // we hold onto the handle and drop it last after consuming the `worker`. This
                // ensures any data being held for the actual evaluation is dropped before we
                // notify the future as done.
                match worker.do_work(&handle).await {
                    Ok(res) => {
                        debug!("{:?} finished. Notifying result", k);
                        handle.finished(res.value);
                    }
                    Err(Cancelled) => {
                        // we drop the current handle, leaving the original `DiceTask` as terminated
                        // state
                        debug!("{:?} cancelled. Notifying cancellation", k);
                        drop(handle)
                    }
                }

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

    fn new(
        k: DiceKey,
        eval: AsyncEvaluator,
        cycles: UserCycleDetectorData,
        events_dispatcher: DiceEventDispatcher,
        previously_cancelled_task: Option<PreviouslyCancelledTask>,
        incremental: IncrementalEngine,
    ) -> Self {
        Self {
            k,
            eval,
            cycles,
            events_dispatcher,
            previously_cancelled_task,
            incremental,
        }
    }

    pub(crate) async fn do_work<'a>(
        self,
        handle: &'a DiceTaskHandle<'a>,
    ) -> CancellableResult<DiceWorkerStateFinishedAndCached> {
        let state = DiceWorkerStateAwaitingPrevious::new(self.k, self.cycles, handle);

        let state = if let Some(previous) = self.previously_cancelled_task {
            handle
                .cancellation_ctx()
                .critical_section(|| previous.previous.await_termination())
                .await;

            // old task actually finished, so just use that result if it wasn't
            // cancelled

            match previous
                .previous
                .get_finished_value()
                .expect("Terminated task must have finished value")
            {
                Ok(res) => {
                    return Ok(state.previously_finished(res));
                }
                Err(Cancelled) => {
                    // actually was cancelled, so just continue re-evaluating
                }
            }

            state.previously_cancelled()
        } else {
            state.no_previous_task()
        };

        self.incremental
            .eval_entry_versioned(self.k, &self.eval, self.events_dispatcher, state)
            .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.
 */

//! The main worker thread for the dice task

use more_futures::cancellable_future::DisableCancellationGuard;
use more_futures::cancellation::ExplicitCancellationContext;

use crate::impls::evaluator::AsyncEvaluator;
use crate::impls::evaluator::KeyEvaluationResult;
use crate::impls::key::DiceKey;
use crate::impls::task::handle::DiceTaskHandle;
use crate::impls::user_cycle::KeyComputingUserCycleDetectorData;
use crate::impls::user_cycle::UserCycleDetectorData;
use crate::impls::value::DiceComputedValue;
use crate::result::CancellableResult;
use crate::result::Cancelled;

/// Represents when we are in a spawned dice task worker and are currently waiting for the previous
/// cancelled instance of this task to finish cancelling.
pub(crate) struct DiceWorkerStateAwaitingPrevious<'a> {
    k: DiceKey,
    cycles: UserCycleDetectorData,
    internals: &'a DiceTaskHandle<'a>,
}

impl<'a> DiceWorkerStateAwaitingPrevious<'a> {
    pub(crate) fn new(
        k: DiceKey,
        cycles: UserCycleDetectorData,
        handle: &'a DiceTaskHandle<'a>,
    ) -> Self {
        debug!(msg = "Task started. Waiting for previously cancelled task if any");
        Self {
            k,
            internals: handle,
            cycles,
        }
    }

    pub(crate) fn previously_finished(
        self,
        value: DiceComputedValue,
    ) -> DiceWorkerStateFinishedAndCached {
        debug!(msg = "previously cancelled task actually finished");

        DiceWorkerStateFinishedAndCached { value }
    }

    pub(crate) fn previously_cancelled(self) -> DiceWorkerStateLookupNode<'a> {
        debug!(msg = "previously cancelled task was cancelled");

        self.internals.report_initial_lookup();

        DiceWorkerStateLookupNode {
            k: self.k,
            internals: self.internals,
            cycles: self.cycles,
        }
    }

    pub(crate) fn no_previous_task(self) -> DiceWorkerStateLookupNode<'a> {
        debug!(msg = "no previous task to wait for");

        self.internals.report_initial_lookup();

        DiceWorkerStateLookupNode {
            k: self.k,
            internals: self.internals,
            cycles: self.cycles,
        }
    }
}

/// Represents when we are currently looking up the current requested key from the core state, and
/// are waiting for it to respond.
pub(crate) struct DiceWorkerStateLookupNode<'a> {
    k: DiceKey,
    cycles: UserCycleDetectorData,
    internals: &'a DiceTaskHandle<'a>,
}

impl<'a> DiceWorkerStateLookupNode<'a> {
    pub(crate) fn checking_deps(self, eval: &AsyncEvaluator) -> DiceWorkerStateCheckingDeps<'a> {
        debug!(msg = "found existing entry with mismatching version. checking if deps changed.",);

        self.internals.checking_deps();

        let cycles = self.cycles.start_computing_key(
            self.k,
            &eval.dice.key_index,
            eval.user_data.cycle_detector.as_ref(),
        );

        DiceWorkerStateCheckingDeps {
            cycles,
            internals: self.internals,
        }
    }

    pub(crate) fn lookup_dirtied(self, eval: &AsyncEvaluator) -> DiceWorkerStateComputing<'a> {
        debug!(msg = "lookup requires recompute.");

        self.internals.computing();

        let cycles = self.cycles.start_computing_key(
            self.k,
            &eval.dice.key_index,
            eval.user_data.cycle_detector.as_ref(),
        );

        DiceWorkerStateComputing {
            cycles,
            internals: self.internals,
        }
    }

    pub(crate) fn lookup_matches(
        self,
        value: DiceComputedValue,
    ) -> DiceWorkerStateFinishedAndCached {
        debug!(msg = "found existing entry with matching version in cache. reusing result.",);

        DiceWorkerStateFinishedAndCached { value }
    }
}

/// When the spawned dice task worker is checking if the dependencies have changed since the last
/// time this node was verified, and are waiting for the results of the dependency re-computation.
pub(crate) struct DiceWorkerStateCheckingDeps<'a> {
    cycles: KeyComputingUserCycleDetectorData,
    internals: &'a DiceTaskHandle<'a>,
}

impl<'a> DiceWorkerStateCheckingDeps<'a> {
    pub(crate) fn cycles_for_dep(
        &self,
        dep: DiceKey,
        eval: &AsyncEvaluator,
    ) -> UserCycleDetectorData {
        self.cycles.subrequest(dep, &eval.dice.key_index)
    }

    pub(crate) fn deps_not_match(self) -> DiceWorkerStateComputing<'a> {
        debug!(msg = "deps changed");
        self.internals.computing();

        DiceWorkerStateComputing {
            cycles: self.cycles,
            internals: self.internals,
        }
    }

    pub(crate) fn deps_match(self) -> CancellableResult<DiceWorkerStateFinished> {
        debug!(msg = "reusing previous value because deps didn't change. Updating caches");

        let guard = match self
            .internals
            .cancellation_ctx()
            .try_to_keep_going_on_cancellation()
        {
            Some(g) => g,
            None => {
                debug!("evaluation cancelled, skipping cache updates");
                return Err(Cancelled);
            }
        };

        Ok(DiceWorkerStateFinished {
            _prevent_cancellation: guard,
        })
    }

    #[cfg(test)]
    pub(crate) fn testing() -> Self {
        DiceWorkerStateCheckingDeps {
            cycles: KeyComputingUserCycleDetectorData::Untracked,
            internals: DiceTaskHandle::testing_new(),
        }
    }
}

/// When the spawned dice worker is currently computing the requested Key.
pub(crate) struct DiceWorkerStateComputing<'a> {
    cycles: KeyComputingUserCycleDetectorData,
    internals: &'a DiceTaskHandle<'a>,
}

impl<'a> DiceWorkerStateComputing<'a> {
    pub(crate) fn evaluating(
        self,
    ) -> (
        KeyComputingUserCycleDetectorData,
        DiceWorkerStateEvaluating<'a>,
    ) {
        (
            self.cycles,
            DiceWorkerStateEvaluating {
                internals: self.internals,
            },
        )
    }
}

/// When the spawned dice worker is currently actively evaluating the `Key::compute` function
pub(crate) struct DiceWorkerStateEvaluating<'a> {
    internals: &'a DiceTaskHandle<'a>,
}

impl<'a> DiceWorkerStateEvaluating<'a> {
    pub(crate) fn cancellation_ctx(&self) -> &ExplicitCancellationContext {
        self.internals.cancellation_ctx()
    }

    pub(crate) fn finished(
        self,
        cycles: KeyComputingUserCycleDetectorData,
        result: KeyEvaluationResult,
    ) -> CancellableResult<DiceWorkerStateFinishedEvaluating> {
        debug!(msg = "evaluation finished. updating caches");

        let guard = match self
            .internals
            .cancellation_ctx()
            .try_to_keep_going_on_cancellation()
        {
            Some(g) => g,
            None => {
                debug!("evaluation cancelled, skipping cache updates");
                return Err(Cancelled);
            }
        };

        drop(cycles);

        Ok(DiceWorkerStateFinishedEvaluating {
            state: DiceWorkerStateFinished {
                _prevent_cancellation: guard,
            },
            result,
        })
    }
}

/// When the spawned dice worker has just finished evaluating the `Key::compute` function
pub(crate) struct DiceWorkerStateFinishedEvaluating {
    pub(crate) state: DiceWorkerStateFinished,
    pub(crate) result: KeyEvaluationResult,
}

/// When the spawned dice worker is finished checking dependencies or finished computing the key.
/// At this point, the value of the node is known. We are just waiting for core state to finish
/// updating the caches and return the correct instance of the value.
pub(crate) struct DiceWorkerStateFinished {
    _prevent_cancellation: DisableCancellationGuard,
}

impl DiceWorkerStateFinished {
    pub(crate) fn cached(self, value: DiceComputedValue) -> DiceWorkerStateFinishedAndCached {
        debug!(msg = "Update caches complete");

        DiceWorkerStateFinishedAndCached { value }
    }
}

/// When the spawned dice worker is done computing and saving the value to core state cache.
/// The final value is known.
pub(crate) struct DiceWorkerStateFinishedAndCached {
    pub(crate) value: DiceComputedValue,
}


/*
 * 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::collections::BTreeMap;
use std::collections::BTreeSet;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;
use std::iter;
use std::sync::Arc;

use cmp_any::PartialEqAny;
use derivative::Derivative;
use dupe::Dupe;
use itertools::Either;
use serde::de::Error;
use serde::de::Unexpected;
use serde::de::Visitor;
use serde::Deserialize;
use serde::Deserializer;
use serde::Serialize;
use serde::Serializer;

use crate::impls::core::graph::introspection::VersionedGraphIntrospectable;
use crate::impls::core::versions::introspection::VersionIntrospectable;
use crate::impls::key::DiceKey;
use crate::introspection::serialize_dense_graph;
use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
use crate::legacy::incremental::ErasedEngine;
use crate::HashMap;
use crate::HashSet;

#[derive(Derivative)]
#[derivative(Debug)]
pub enum GraphIntrospectable {
    Legacy {
        #[derivative(Debug = "ignore")]
        introspectables: LegacyIntrospectable,
    },
    Modern {
        #[derivative(Debug = "ignore")]
        introspection: ModernIntrospectable,
    },
}

pub struct LegacyIntrospectable(pub(crate) Vec<Arc<dyn ErasedEngine + Send + Sync + 'static>>);

impl GraphIntrospectable {
    pub(crate) fn introspectables(&self) -> impl Iterator<Item = &dyn EngineForIntrospection> {
        match self {
            GraphIntrospectable::Legacy { introspectables } => {
                Either::Left(introspectables.0.iter().map(|e| e.introspect()))
            }
            GraphIntrospectable::Modern { introspection } => {
                Either::Right(iter::once(introspection as _))
            }
        }
    }
}

pub struct ModernIntrospectable {
    pub(crate) graph: VersionedGraphIntrospectable,
    pub(crate) version_data: VersionIntrospectable,
    pub(crate) key_map: HashMap<DiceKey, AnyKey>,
}

impl EngineForIntrospection for ModernIntrospectable {
    fn keys<'a>(&'a self) -> Box<dyn Iterator<Item = AnyKey> + 'a> {
        Box::new(self.graph.keys())
    }

    fn edges<'a>(&'a self) -> Box<dyn Iterator<Item = (AnyKey, Vec<AnyKey>)> + 'a> {
        Box::new(self.graph.edges())
    }

    fn keys_currently_running<'a>(
        &'a self,
    ) -> Vec<(AnyKey, VersionNumber, DiceTaskStateForDebugging)> {
        self.version_data.keys_currently_running(&self.key_map)
    }

    fn versions_currently_running<'a>(&'a self) -> Vec<VersionNumber> {
        self.version_data.versions_currently_running()
    }

    fn nodes<'a>(
        &'a self,
        _keys: &'a mut HashMap<AnyKey, KeyID>,
    ) -> Box<dyn Iterator<Item = SerializedGraphNodesForKey> + 'a> {
        Box::new(self.graph.nodes())
    }

    fn len_for_introspection(&self) -> usize {
        self.graph.len_for_introspection()
    }

    fn currently_running_key_count(&self) -> usize {
        self.version_data.currently_running_key_count()
    }
}

impl Serialize for GraphIntrospectable {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serialize_dense_graph(self, serializer)
    }
}

#[derive(PartialEq, Eq, Hash, Serialize, Deserialize, Clone, Dupe, Copy)]
#[serde(transparent)]
pub struct KeyID(pub usize);

#[derive(PartialEq, Eq, Hash, Serialize, Deserialize, Clone, Dupe, Copy)]
#[serde(transparent)]
pub struct NodeID(pub usize);

#[derive(
    PartialEq,
    Eq,
    Hash,
    Clone,
    Dupe,
    Copy,
    Ord,
    PartialOrd,
    derive_more::Display
)]
pub struct VersionNumber(pub usize);

impl Serialize for VersionNumber {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&format!("v{}", self.0))
    }
}

struct VersionNumberVisitor;

impl<'de> Visitor<'de> for VersionNumberVisitor {
    type Value = VersionNumber;

    fn expecting(&self, formatter: &mut Formatter) -> fmt::Result {
        formatter.write_str("string of format `vX` where X is a usize, like `v2`")
    }

    fn visit_borrowed_str<E>(self, v: &'de str) -> Result<Self::Value, E>
    where
        E: Error,
    {
        let mut chars = v.chars();

        if chars.next() != Some('v') {
            Err(Error::invalid_value(Unexpected::Str(v), &self))
        } else {
            chars
                .as_str()
                .parse::<usize>()
                .map_err(|_| Error::invalid_value(Unexpected::Str(v), &self))
                .map(VersionNumber)
        }
    }
}

impl<'de> Deserialize<'de> for VersionNumber {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_str(VersionNumberVisitor)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub enum GraphNodeKind {
    Occupied,
    Transient,
    Vacant,
}

#[derive(Clone, Serialize, Deserialize)]
pub struct CellHistory {
    pub history: BTreeMap<VersionNumber, HistoryState>,
}

impl CellHistory {
    pub fn new(verified: BTreeSet<VersionNumber>, dirtied: BTreeMap<VersionNumber, bool>) -> Self {
        Self {
            history: verified
                .into_iter()
                .map(|v| (v, HistoryState::Verified))
                .chain(dirtied.into_iter().map(|(v, f)| {
                    (
                        v,
                        if f {
                            HistoryState::ForceDirty
                        } else {
                            HistoryState::Dirty
                        },
                    )
                }))
                .collect(),
        }
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub enum HistoryState {
    Verified,
    Dirty,
    ForceDirty,
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SerializedGraphNode {
    pub node_id: NodeID,
    pub kind: GraphNodeKind,
    pub history: CellHistory,
    /// Deps and Rdeps are behind read locks, and if dumping after a panic
    /// it's theoretically possible for those locks to be poisoned.
    /// Therefore, they're optional.
    pub deps: Option<HashSet<KeyID>>,
    pub rdeps: Option<BTreeMap<VersionNumber, Vec<NodeID>>>,
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SerializedGraphNodesForKey {
    pub id: KeyID,
    pub key: String,
    pub type_name: String,
    pub nodes: BTreeMap<VersionNumber, Option<SerializedGraphNode>>,
}

pub(crate) trait EngineForIntrospection {
    fn keys<'a>(&'a self) -> Box<dyn Iterator<Item = AnyKey> + 'a>;
    fn edges<'a>(&'a self) -> Box<dyn Iterator<Item = (AnyKey, Vec<AnyKey>)> + 'a>;
    fn keys_currently_running<'a>(
        &'a self,
    ) -> Vec<(AnyKey, VersionNumber, DiceTaskStateForDebugging)>;
    fn versions_currently_running<'a>(&'a self) -> Vec<VersionNumber>;
    fn nodes<'a>(
        &'a self,
        keys: &'a mut HashMap<AnyKey, KeyID>,
    ) -> Box<dyn Iterator<Item = SerializedGraphNodesForKey> + 'a>;
    fn len_for_introspection(&self) -> usize;
    fn currently_running_key_count(&self) -> usize;
}

pub(crate) trait KeyForIntrospection: Display + Send + 'static {
    fn get_key_equality(&self) -> PartialEqAny;

    fn hash(&self, state: &mut dyn Hasher);

    fn type_name(&self) -> &'static str {
        std::any::type_name::<Self>()
    }

    fn box_clone(&self) -> Box<dyn KeyForIntrospection>;
}

impl<K> KeyForIntrospection for K
where
    K: Clone + Display + Hash + Eq + Send + 'static,
{
    fn get_key_equality(&self) -> PartialEqAny {
        PartialEqAny::new(self)
    }

    fn hash(&self, mut state: &mut dyn Hasher) {
        K::hash(self, &mut state)
    }

    fn box_clone(&self) -> Box<dyn KeyForIntrospection> {
        Box::new(self.clone())
    }
}

pub(crate) struct AnyKey {
    pub inner: Box<dyn KeyForIntrospection>,
}

/// Shorten a type name like
/// ```ignore
/// <dice::ctx::DiceComputations as buck2_interpreter::extra::buckconfig::HasLegacyBuckConfigForStarlark>
///     ::get_legacy_buck_config_for_starlark::{{closure}}::LegacyBuckConfigForStarlarkKey
/// ```
/// to
/// `LegacyBuckConfigForStarlarkKey`.
pub fn short_type_name(type_name: &str) -> &str {
    type_name
        .rsplit("::")
        .next()
        .filter(|s| {
            // TODO(nga): shorten types like `Vec<String>`.
            //   Ignore them for now.
            !s.contains('>')
        })
        .unwrap_or(type_name)
}

impl AnyKey {
    pub fn new(k: impl KeyForIntrospection) -> Self {
        Self { inner: Box::new(k) }
    }

    pub fn type_name(&self) -> &'static str {
        self.inner.type_name()
    }

    pub fn short_type_name(&self) -> &'static str {
        short_type_name(self.type_name())
    }
}

impl Clone for AnyKey {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.box_clone(),
        }
    }
}

impl Display for AnyKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.inner.fmt(f)
    }
}

impl PartialEq for AnyKey {
    fn eq(&self, other: &Self) -> bool {
        self.inner.get_key_equality() == other.inner.get_key_equality()
    }
}

impl Eq for AnyKey {}

impl Hash for AnyKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.inner.hash(state)
    }
}

#[cfg(test)]
mod tests {
    use std::any::type_name;

    use crate::introspection::graph::short_type_name;

    #[test]
    fn test_short_type_name() {
        assert_eq!("String", short_type_name(type_name::<String>()));

        let closure = || {
            struct MyLocalTypeInClosure;
            short_type_name(type_name::<MyLocalTypeInClosure>())
        };
        assert_eq!("MyLocalTypeInClosure", closure());

        // We do not shorten generic types yet.
        assert_eq!(
            type_name::<Vec<String>>(),
            short_type_name(type_name::<Vec<String>>())
        );
    }
}


/*
 * 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::collections::hash_map::Entry;
use std::io::Write;

use anyhow::Context as _;
use serde::ser::SerializeSeq;
use serde::Serializer;

use crate::introspection::graph::GraphIntrospectable;
use crate::introspection::AnyKey;
use crate::HashMap;

pub fn serialize_graph(
    graph: &GraphIntrospectable,
    nodes: impl Write,
    mut edges: impl Write,
    mut nodes_currently_running: impl Write,
) -> anyhow::Result<()> {
    let mut reg = NodeRegistry::new();

    for engine in graph.introspectables() {
        for (k, vs) in engine.edges() {
            let k = reg.map(k);

            for v in vs.into_iter() {
                let v = reg.map(v);
                edges
                    .write_all(format!("{}\t{}\n", k, v).as_bytes())
                    .context("Failed to write edge")?;
            }
        }

        for (k, v, s) in engine.keys_currently_running() {
            let k_short_type_name = k.short_type_name();
            let k_str = k.to_string();
            let k_n = reg.map(k);
            writeln!(
                nodes_currently_running,
                "{k_n}\t{v}\t{s:?}\t{k_short_type_name}\t{k_str}",
            )?;
        }
    }

    reg.write(nodes)?;

    Ok(())
}

pub fn serialize_dense_graph<S>(graph: &GraphIntrospectable, writer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let mut reg = HashMap::default();

    let num_nodes = graph
        .introspectables()
        .map(|engine| engine.len_for_introspection())
        .sum();

    let mut seq = writer.serialize_seq(Some(num_nodes))?;
    for engine in graph.introspectables() {
        for node in engine.nodes(&mut reg) {
            seq.serialize_element(&node)?;
        }
    }
    seq.end()
}

struct NodeRegistry {
    keys: HashMap<AnyKey, u64>,
}

impl NodeRegistry {
    fn new() -> Self {
        Self {
            keys: HashMap::default(),
        }
    }

    fn map(&mut self, key: AnyKey) -> u64 {
        let next_idx = self.keys.len() as u64;
        match self.keys.entry(key) {
            Entry::Occupied(e) => *e.get(),
            Entry::Vacant(e) => *e.insert(next_idx),
        }
    }

    fn write(self, mut out: impl Write) -> anyhow::Result<()> {
        // NOTE: Writing out the node type every time is wasteful here, we could optimize, though
        // if we compress then that goes away, so optimizing *here* might not make the most sense.
        let mut keys = self
            .keys
            .into_iter()
            .map(|(key, idx)| (idx, key))
            .collect::<Vec<_>>();

        keys.sort_by_key(|(idx, _)| *idx);

        for (idx, key) in keys {
            out.write_all(format!("{}\t{}\t{}\n", idx, key.short_type_name(), key).as_bytes())
                .context("Failed to write node")?;
        }
        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.
 */

//!
//! Interfaces for introspection of the DICE graph

use crate::introspection::graph::AnyKey;
use crate::introspection::graph::GraphIntrospectable;
use crate::introspection::graph::LegacyIntrospectable;
use crate::Dice;
use crate::DiceImplementation;

pub mod graph;
pub(crate) mod introspect;

pub use crate::introspection::introspect::serialize_dense_graph;
pub use crate::introspection::introspect::serialize_graph;
use crate::legacy::DiceLegacy;

impl Dice {
    pub fn to_introspectable(&self) -> GraphIntrospectable {
        match &self.implementation {
            DiceImplementation::Legacy(dice) => dice.to_introspectable(),
            DiceImplementation::Modern(_) => {
                unimplemented!("todo")
            }
        }
    }
}

impl DiceLegacy {
    pub fn to_introspectable(&self) -> GraphIntrospectable {
        GraphIntrospectable::Legacy {
            introspectables: LegacyIntrospectable(self.map.read().engines().to_vec()),
        }
    }
}

#[cfg(test)]
mod tests {
    use allocative::Allocative;
    use anyhow::Context as _;
    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::key::Key;
    use crate::introspection::graph::SerializedGraphNodesForKey;
    use crate::introspection::serialize_graph;
    use crate::DiceLegacy;
    use crate::HashMap;
    use crate::WhichSpawner;

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

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

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            if self.0 > 0 {
                ctx.compute(&KeyA(self.0 - 1)).await.unwrap();
            } else {
                ctx.compute(&KeyB).await.unwrap();
            }
        }

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

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

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

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

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

    #[tokio::test]
    async fn test_serialization() -> anyhow::Result<()> {
        let dice =
            DiceLegacy::builder().build(DetectCycles::Disabled, WhichSpawner::ExplicitCancel);
        let ctx = dice.updater().commit().await;
        ctx.compute(&KeyA(3)).await?;

        let mut nodes = Vec::new();
        let mut edges = Vec::new();
        let mut nodes_currently_running = Vec::new();

        serialize_graph(
            &dice.to_introspectable(),
            &mut nodes,
            &mut edges,
            &mut nodes_currently_running,
        )
        .unwrap();
        let nodes = String::from_utf8(nodes)?;
        let edges = String::from_utf8(edges)?;

        let mut node_map = HashMap::<String, u64>::default();
        let mut edge_list = Vec::<(u64, u64)>::new();

        for line in nodes.lines() {
            let mut it = line.trim().split('\t');
            let idx = it.next().context("No idx")?.parse()?;
            let _key_type = it.next().context("No key type")?;
            let key = it.next().context("No key")?;
            node_map.insert(key.into(), idx);
        }

        for line in edges.lines() {
            let mut it = line.trim().split('\t');
            let from = it.next().context("No idx")?.parse()?;
            let to = it.next().context("No key")?.parse()?;
            edge_list.push((from, to));
        }

        let a3 = *node_map.get("KeyA(3)").context("Missing key")?;
        let a2 = *node_map.get("KeyA(2)").context("Missing key")?;
        let a1 = *node_map.get("KeyA(1)").context("Missing key")?;
        let a0 = *node_map.get("KeyA(0)").context("Missing key")?;
        let b = *node_map.get("KeyB").context("Missing key")?;

        let mut expected_edge_list = vec![(a3, a2), (a2, a1), (a1, a0), (a0, b)];
        expected_edge_list.sort_unstable();
        edge_list.sort_unstable();
        assert_eq!(expected_edge_list, edge_list);

        assert!(nodes_currently_running.is_empty());

        Ok(())
    }

    #[tokio::test]
    async fn test_serialization_dense() -> anyhow::Result<()> {
        let dice =
            DiceLegacy::builder().build(DetectCycles::Disabled, WhichSpawner::ExplicitCancel);
        let ctx = dice.updater().commit().await;
        ctx.compute(&KeyA(3)).await?;

        let node = bincode::serialize(&dice.to_introspectable())?;

        let _out: Vec<SerializedGraphNodesForKey> = bincode::deserialize(&node)?;
        Ok(())
    }
}


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

use std::any::Any;
use std::future::Future;
use std::sync::Arc;

use allocative::Allocative;
use dupe::Dupe;
use futures::FutureExt;
use parking_lot::Mutex;

use crate::api::activation_tracker::ActivationData;
use crate::api::cycles::DetectCycles;
use crate::api::data::DiceData;
use crate::api::error::DiceErrorImpl;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::projection::ProjectionKey;
use crate::api::user_data::UserComputationData;
use crate::api::user_data::UserCycleDetectorGuard;
use crate::legacy::cycles::CycleDetector;
use crate::legacy::incremental::dep_trackers::BothDepTrackers;
use crate::legacy::incremental::dep_trackers::BothDeps;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::legacy::incremental::transaction_ctx::ActiveTransactionCountGuard;
use crate::legacy::incremental::transaction_ctx::Changes;
use crate::legacy::incremental::transaction_ctx::TransactionCtx;
use crate::legacy::incremental::versions::VersionForWrites;
use crate::legacy::incremental::versions::VersionGuard;
use crate::legacy::key::StoragePropertiesForKey;
use crate::legacy::map::DiceMap;
use crate::legacy::opaque::OpaqueValueImplLegacy;
use crate::legacy::projection::ProjectionKeyAsKey;
use crate::legacy::projection::ProjectionKeyProperties;
use crate::legacy::DiceLegacy;
use crate::versions::VersionNumber;
use crate::DiceError;

/// A context for the duration of a top-level compute request.
///
/// This contains both user-visible and dice-internal computation-specific data.
#[derive(Allocative)]
pub(crate) struct ComputationData {
    pub(crate) user_data: Arc<UserComputationData>,
    cycle_detector: Option<Box<CycleDetector>>,
    // TODO(bobyf): this seems a natural place to gather some stats about the compute too
    #[allocative(skip)]
    pub(crate) user_cycle_detector_guard: Option<Box<dyn UserCycleDetectorGuard>>,
    /// Store extra data from provided by the key's evaluation, which will be passed to the
    /// user_data's ActivationTracker when the key evaluation finishes.
    #[allocative(skip)]
    pub(crate) evaluation_data: Mutex<Option<Box<dyn Any + Send + Sync + 'static>>>,
}

impl ComputationData {
    pub(crate) fn new(data: UserComputationData, detect_cycles: DetectCycles) -> Self {
        Self {
            user_data: Arc::new(data),
            cycle_detector: match detect_cycles {
                DetectCycles::Enabled => Some(Box::new(CycleDetector::new())),
                DetectCycles::Disabled => None,
            },
            user_cycle_detector_guard: None,
            evaluation_data: Mutex::new(None),
        }
    }

    /// records that we are entering the computation of another key as part of this main request
    /// i.e. computing key a, which during its evaluation requests key b, enters a new subrequest.
    pub(crate) fn subrequest<K>(&self, key: &K::Key) -> DiceResult<Self>
    where
        K: StorageProperties,
    {
        if let Some(v) = &self.user_cycle_detector_guard {
            v.add_edge(K::to_key_any(key));
        }
        Ok(Self {
            user_data: self.user_data.dupe(),
            cycle_detector: self
                .cycle_detector
                .as_ref()
                .map(|detector| Ok(Box::new(CycleDetector::visit(detector, key)?)))
                .transpose()?,
            user_cycle_detector_guard: None,
            evaluation_data: Mutex::new(None),
        })
    }

    pub(crate) fn start_computing_key<K: StorageProperties>(&mut self, k: &K::Key) {
        assert!(self.user_cycle_detector_guard.is_none());
        self.user_cycle_detector_guard = self
            .user_data
            .cycle_detector
            .as_ref()
            .and_then(|v| v.start_computing_key(K::to_key_any(k)));
    }

    pub(crate) fn finished_computing_key<K: StorageProperties>(
        self,
        k: &K::Key,
        deps: &BothDeps,
        reused: bool,
    ) {
        if let Some(v) = &self.user_data.cycle_detector {
            v.finished_computing_key(K::to_key_any(k))
        }

        if let Some(v) = &self.user_data.activation_tracker {
            let mut iter = deps.deps.iter().map(|d| d.to_key_any());

            let activation_data = if reused {
                ActivationData::Reused
            } else {
                ActivationData::Evaluated(self.evaluation_data.lock().take())
            };

            v.key_activated(K::to_key_any(k), &mut iter, activation_data);
        }
    }
}

/// A context for computations to request for additional dependencies. The
/// dependencies accessed are tracked for caching, if enabled based on
/// 'Strategy'.
///
/// The context is valid only for the duration of the computation of a single
/// key.
///
/// When marking values as changed on the ctx, the changes are part of the next
/// version. The next version is only committed when the current context is
/// dropped. context, which means that the "current" context will not see
/// updated values.
#[derive(Allocative)]
pub(crate) struct DiceComputationsImplLegacy {
    pub(crate) transaction_ctx: Arc<TransactionCtx>,
    pub(crate) dice: Arc<DiceLegacy>,
    pub(crate) dep_trackers: BothDepTrackers,
    pub(crate) extra: ComputationData,
}

impl DiceComputationsImplLegacy {
    pub(crate) fn new_transaction(
        dice: Arc<DiceLegacy>,
        version: VersionGuard,
        version_for_writes: VersionForWrites,
        extra: ComputationData,
    ) -> Self {
        Self {
            transaction_ctx: Arc::new(TransactionCtx::new(
                version,
                version_for_writes,
                Changes::new(),
                ActiveTransactionCountGuard::new(&dice),
            )),
            dep_trackers: BothDepTrackers::noop(),
            dice: dice.dupe(),
            extra,
        }
    }

    pub(crate) fn new_for_key_evaluation(
        dice: Arc<DiceLegacy>,
        transaction_ctx: Arc<TransactionCtx>,
        extra: ComputationData,
    ) -> Arc<Self> {
        // TODO(bobyf): for memory, handle cases where we don't want explicit tracking
        Arc::new(Self {
            transaction_ctx,
            dice: dice.dupe(),
            dep_trackers: BothDepTrackers::recording(),
            extra,
        })
    }

    pub(crate) fn finalize(self: Arc<Self>) -> (BothDeps, ComputationData) {
        // TODO express this via lifetimes
        let this = Arc::try_unwrap(self).map_err(|_| "The computation lifetime of the `ctx` has ended and there should be no further references to the `Arc`").unwrap();

        (this.dep_trackers.collect_deps(), this.extra)
    }

    pub(crate) fn compute_opaque<'a, K>(
        self: &'a Arc<Self>,
        key: &K,
    ) -> impl Future<Output = DiceResult<OpaqueValueImplLegacy<'a, K>>> + 'a
    where
        K: Key,
    {
        // This would be simpler with an `async fn/async move {}`, but we create these for every edge in the computation
        // and many of those may be live at a time, and so we need to take more care and ensure this is fairly small.
        let cache = self.dice.find_cache::<K>();
        let extra = self.extra.subrequest::<StoragePropertiesForKey<K>>(key);
        match extra {
            Ok(extra) => cache
                .eval_for_opaque(key, &self.transaction_ctx, extra)
                .map(move |value| Ok(OpaqueValueImplLegacy::new(value, self, cache)))
                .left_future(),
            Err(e) => futures::future::ready(Err(e)).right_future(),
        }
    }

    pub(crate) fn compute_projection_sync<P>(
        self: &Arc<Self>,
        derive_from: &OpaqueValueImplLegacy<P::DeriveFromKey>,
        projection_key: &P,
    ) -> DiceResult<P::Value>
    where
        P: ProjectionKey,
    {
        assert!(Arc::ptr_eq(self, derive_from.parent_computations));

        let cache = self.dice.find_projection_cache::<P>();

        let projection_key_as_key = ProjectionKeyAsKey {
            derive_from_key: derive_from.key().clone(),
            k: projection_key.clone(),
        };

        let extra = self
            .extra
            .subrequest::<ProjectionKeyProperties<P>>(&projection_key_as_key)?;

        Ok(cache.eval_projection(
            &projection_key_as_key,
            derive_from,
            &self.transaction_ctx,
            &extra,
        ))
    }

    pub(crate) fn changed<K, I>(&self, changed: I) -> DiceResult<()>
    where
        K: Key,
        I: IntoIterator<Item = K> + Send + Sync + 'static,
    {
        let mut changes = self.transaction_ctx.changes();

        changed.into_iter().try_for_each(|k| {
            let dice = self.dice.dupe();
            changes.change(
                k.clone(),
                Box::new(move |version| {
                    debug!(msg = "marking value as changed", version = %version, key = %k);
                    let cache = dice.find_cache::<K>();
                    cache.dirty(k, version, true);

                    true
                }),
            )
        })
    }

    pub(crate) fn changed_to<K, I>(&self, changed: I) -> DiceResult<()>
    where
        K: Key,
        I: IntoIterator<Item = (K, K::Value)> + Send + Sync + 'static,
    {
        let mut changes = self.transaction_ctx.changes();

        changed.into_iter().try_for_each(|(k, v)| {
            if !K::validity(&v) {
                return Err(DiceError::invalid_change(Arc::new(k)));
            }
            let dice = self.dice.dupe();
            changes.change(
                k.clone(),
                Box::new(move |version| {
                    let cache = dice.find_cache::<K>();
                    debug!(msg = "marking value as updated", version = %version, key = %k);
                    cache.update_injected_value(k, version, v)
                }),
            )
        })
    }

    /// Commit the changes registered via 'changed' and 'changed_to' to the current newest version.
    /// This can only be called when the this is the only node remaining in the computation graph
    pub(crate) fn commit(self: Arc<Self>) -> Arc<DiceComputationsImplLegacy> {
        // TODO need to clean up these ctxs so we have less runtime errors from Arc references
        let this = Arc::try_unwrap(self)
            .map_err(|_| "Error: tried to commit when there are more references")
            .unwrap();
        let eval = Arc::try_unwrap(this.transaction_ctx)
            .map_err(|_| "Error: tried to commit when there are more references")
            .unwrap();

        // hold onto the prev version until we get the new one below so we don't increment minor
        // version needlessly.
        let _prev_v = eval.commit();

        this.dice.make_ctx(this.extra)
    }

    /// Same as `commit`, but replacing the user data with the given
    pub(crate) fn commit_with_data(
        self: Arc<Self>,
        extra: UserComputationData,
    ) -> Arc<DiceComputationsImplLegacy> {
        // TODO need to clean up these ctxs so we have less runtime errors from Arc references
        let mut this = Arc::try_unwrap(self)
            .map_err(|_| "Error: tried to commit when there are more references")
            .unwrap();
        let eval = Arc::try_unwrap(this.transaction_ctx)
            .map_err(|_| "Error: tried to commit when there are more references")
            .unwrap();

        // hold onto the prev version until we get the new one below so we don't increment minor
        // version needlessly.
        let _prev_v = eval.commit();

        this.dice.make_ctx(ComputationData {
            user_data: Arc::new(extra),
            cycle_detector: this.extra.cycle_detector.take(),
            user_cycle_detector_guard: None,
            evaluation_data: Mutex::new(None),
        })
    }

    pub(crate) fn get_version(&self) -> VersionNumber {
        self.transaction_ctx.get_version()
    }

    pub(crate) fn unstable_take(self: &Arc<Self>) -> DiceMap {
        self.dice.unstable_take()
    }

    pub(crate) fn global_data(&self) -> &DiceData {
        &self.dice.data
    }

    pub(crate) fn per_transaction_data(&self) -> &UserComputationData {
        &self.extra.user_data
    }

    pub(crate) fn cycle_guard<T: UserCycleDetectorGuard>(&self) -> DiceResult<Option<&T>> {
        match &self.extra.user_cycle_detector_guard {
            None => Ok(None),
            Some(guard) => match guard.as_any().downcast_ref() {
                Some(guard) => Ok(Some(guard)),
                None => Err(DiceError(Arc::new(
                    DiceErrorImpl::UnexpectedCycleGuardType {
                        expected_type_name: std::any::type_name::<T>().to_owned(),
                        actual_type_name: guard.type_name().to_owned(),
                    },
                ))),
            },
        }
    }

    pub fn store_evaluation_data<T: Send + Sync + 'static>(&self, value: T) -> DiceResult<()> {
        let mut evaluation_data = self.extra.evaluation_data.lock();
        if evaluation_data.is_some() {
            return Err(DiceError::duplicate_activation_data());
        }
        *evaluation_data = Some(Box::new(value) as _);
        Ok(())
    }
}

#[cfg(test)]
pub(crate) mod testing {
    use crate::api::cycles::DetectCycles;
    use crate::api::user_data::UserComputationData;
    use crate::ctx::DiceComputationsImpl;
    use crate::legacy::ctx::ComputationData;
    use crate::legacy::incremental::versions::MinorVersion;
    use crate::transaction::DiceTransactionImpl;

    pub(crate) trait DiceCtxExt {
        fn get_minor_version(&self) -> MinorVersion;
    }

    impl DiceCtxExt for DiceComputationsImpl {
        fn get_minor_version(&self) -> MinorVersion {
            match self {
                DiceComputationsImpl::Legacy(delegate) => {
                    delegate.transaction_ctx.get_minor_version()
                }
                DiceComputationsImpl::Modern(_delegate) => {
                    unimplemented!("todo")
                }
            }
        }
    }

    impl DiceCtxExt for DiceTransactionImpl {
        fn get_minor_version(&self) -> MinorVersion {
            match self {
                DiceTransactionImpl::Legacy(delegate) => delegate.0.get_minor_version(),
                DiceTransactionImpl::Modern(_delegate) => {
                    unimplemented!("todo")
                }
            }
        }
    }

    pub(crate) trait ComputationDataExt {
        fn testing_new() -> Self;
    }

    impl ComputationDataExt for ComputationData {
        fn testing_new() -> Self {
            Self::new(UserComputationData::new(), DetectCycles::Enabled)
        }
    }
}


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

//! Cycle detection in DICE

use std::any::TypeId;
use std::fmt::Debug;
use std::fmt::Display;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;

use allocative::Allocative;
use cmp_any::PartialEqAny;
use dupe::IterDupedExt;
use indexmap::set::IndexSet;

use crate::api::error::DiceError;
use crate::api::error::DiceResult;
use crate::api::key::Key;

/// A `Key` that has been requested within Dice.
pub trait RequestedKey: Allocative + Display + Debug + Send + Sync {
    fn get_key_equality(&self) -> PartialEqAny;
    fn hash(&self, state: &mut dyn Hasher);
    fn type_id(&self) -> TypeId;
}

impl dyn RequestedKey {
    pub fn is_key<K: Key>(&self) -> bool {
        TypeId::of::<K>() == self.type_id()
    }
}

impl<T> RequestedKey for T
where
    T: Allocative + Display + Debug + Hash + Eq + Send + Sync + 'static,
{
    fn get_key_equality(&self) -> PartialEqAny {
        PartialEqAny::new(self)
    }

    fn hash(&self, mut state: &mut dyn Hasher) {
        self.hash(&mut state)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<T>()
    }
}

impl Hash for dyn RequestedKey {
    fn hash<H: Hasher>(&self, mut state: &mut H) {
        self.hash(&mut state)
    }
}

impl PartialEq for dyn RequestedKey {
    fn eq(&self, other: &Self) -> bool {
        self.get_key_equality() == other.get_key_equality()
    }
}

impl Eq for dyn RequestedKey {}

#[derive(Allocative)]
pub(crate) struct CycleDetector {
    stack: IndexSet<Arc<dyn RequestedKey>>,
}

impl CycleDetector {
    pub(crate) fn new() -> Self {
        Self {
            stack: IndexSet::new(),
        }
    }

    pub(crate) fn visit<K>(&self, key: &K) -> DiceResult<Self>
    where
        K: Allocative + Clone + Debug + Display + Eq + Hash + Send + Sync + 'static,
    {
        // quick and dirty cycle detection. we will have to make this more efficient
        // TODO(bobyf)
        let mut stack = self.stack.clone();
        if !stack.insert(Arc::new(key.clone())) {
            Err(DiceError::cycle(
                Arc::new(key.clone()),
                stack.iter().duped().collect(),
            ))
        } else {
            Ok(Self { stack })
        }
    }
}

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

    use allocative::Allocative;
    use derive_more::Display;
    use dupe::Dupe;
    use indexmap::indexset;

    use crate::api::error::DiceErrorImpl;
    use crate::legacy::cycles::CycleDetector;
    use crate::legacy::cycles::RequestedKey;

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

    #[test]
    fn cycle_detection_when_no_cycles() -> anyhow::Result<()> {
        let detector = CycleDetector::new();
        let detector1 = detector.visit(&K(1))?;
        let detector12 = detector1.visit(&K(2))?;
        let detector123 = detector12.visit(&K(3))?;
        let _detector1234 = detector123.visit(&K(4))?;

        let detector13 = detector1.visit(&K(3))?;
        let _detector132 = detector13.visit(&K(2))?;

        Ok(())
    }

    #[test]
    fn cycle_detection_when_cycles() -> anyhow::Result<()> {
        let detector = CycleDetector::new();
        let detector = detector.visit(&K(1))?;
        let detector = detector.visit(&K(2))?;
        let detector = detector.visit(&K(3))?;
        let detector = detector.visit(&K(4))?;

        match detector.visit(&K(1)) {
            Ok(_) => {
                panic!("should have cycle error")
            }
            Err(e) => match &*e.0 {
                DiceErrorImpl::Cycle {
                    cyclic_keys,
                    trigger,
                } => {
                    assert!(
                        (**trigger).get_key_equality() == K(1).get_key_equality(),
                        "expected trigger key to be `{}` but was `{}`",
                        K(1),
                        trigger
                    );
                    assert_eq!(
                        cyclic_keys,
                        &indexset![
                            Arc::new(K(1)) as Arc<dyn RequestedKey>,
                            Arc::new(K(2)) as Arc<dyn RequestedKey>,
                            Arc::new(K(3)) as Arc<dyn RequestedKey>,
                            Arc::new(K(4)) as Arc<dyn RequestedKey>
                        ]
                    )
                }
                _ => {
                    panic!("wrong error type")
                }
            },
        }

        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::fmt::Debug;
use std::sync::Arc;
use std::sync::Weak;

use allocative::Allocative;

use crate::api::key::Key;
use crate::api::storage_type::StorageType;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::legacy::DiceLegacy;

#[derive(Debug, Allocative)]
#[allocative(bound = "")]
pub(crate) struct StoragePropertiesForKey<K: Key> {
    _k: std::marker::PhantomData<K>,
    pub(crate) dice: Weak<DiceLegacy>,
}

impl<K: Key> StoragePropertiesForKey<K> {
    pub(crate) fn new(dice: &Arc<DiceLegacy>) -> StoragePropertiesForKey<K> {
        StoragePropertiesForKey {
            _k: std::marker::PhantomData,
            dice: Arc::downgrade(dice),
        }
    }
}

impl<K> StorageProperties for StoragePropertiesForKey<K>
where
    K: Key,
{
    type Key = K;
    type Value = K::Value;

    fn storage_type(&self) -> StorageType {
        K::storage_type()
    }

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

    fn validity(&self, x: &Self::Value) -> bool {
        K::validity(x)
    }

    fn key_type_name() -> &'static str {
        K::key_type_name()
    }

    fn to_key_any(key: &Self::Key) -> &dyn std::any::Any {
        key
    }
}


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

use std::sync::Arc;

use allocative::Allocative;
use anymap::any::Any;
use anymap::Map;
use dupe::Dupe;
use dupe::OptionDupedExt;

use crate::legacy::incremental::ErasedEngine;
use crate::legacy::incremental::IncrementalComputeProperties;
use crate::legacy::incremental::IncrementalEngine;

/// A dynamically typed Map for DICE to map computations to their key, value
/// cache maps.
#[derive(Allocative)]
pub(crate) struct DiceMap {
    #[allocative(skip)]
    typed: Map<dyn Any + Sync + Send>,
    erased: Vec<Arc<dyn ErasedEngine + Send + Sync + 'static>>,
}

impl DiceMap {
    pub(crate) fn new() -> Self {
        Self {
            typed: Map::new(),
            erased: Vec::new(),
        }
    }

    pub(crate) fn find_cache_opt<S>(&self) -> Option<Arc<IncrementalEngine<S>>>
    where
        S: IncrementalComputeProperties,
    {
        self.typed.get::<Arc<IncrementalEngine<S>>>().duped()
    }

    /// finds the computation cache for the given computation type
    pub(crate) fn find_cache<S>(
        &mut self,
        new: impl FnOnce() -> Arc<IncrementalEngine<S>>,
    ) -> Arc<IncrementalEngine<S>>
    where
        S: IncrementalComputeProperties,
    {
        if let Some(cache) = self.typed.get::<Arc<IncrementalEngine<S>>>() {
            cache.dupe()
        } else {
            let cache = new();
            self.typed.insert::<Arc<IncrementalEngine<S>>>(cache.dupe());
            self.erased
                .push(cache.dupe() as Arc<dyn ErasedEngine + Send + Sync + 'static>);
            cache
        }
    }

    pub(crate) fn engines(&self) -> &[Arc<dyn ErasedEngine + Send + Sync + 'static>] {
        self.erased.as_slice()
    }

    pub(crate) fn key_count(&self) -> usize {
        self.erased
            .iter()
            .map(|e| e.introspect().len_for_introspection())
            .sum()
    }

    pub(crate) fn currently_running_key_count(&self) -> usize {
        self.erased
            .iter()
            .map(|e| e.introspect().currently_running_key_count())
            .sum()
    }
}

#[cfg(test)]
mod tests {
    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::key::Key;
    use crate::legacy::incremental::testing::IncrementalEngineExt;
    use crate::legacy::incremental::versions::MinorVersion;
    use crate::legacy::incremental::IncrementalEngine;
    use crate::legacy::map::DiceMap;
    use crate::legacy::DiceLegacy;
    use crate::versions::VersionNumber;
    use crate::StoragePropertiesForKey;
    use crate::WhichSpawner;

    #[tokio::test]
    async fn test_find_caches() {
        #[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
        #[display(fmt = "{:?}", self)]
        struct MyKey;
        #[derive(Clone, Dupe, Display, Debug, Eq, PartialEq, Allocative)]
        #[display(fmt = "{:?}", self)]
        struct Bar;

        #[async_trait]
        impl Key for MyKey {
            type Value = Bar;

            async fn compute(
                &self,
                _ctx: &DiceComputations,
                _cancellations: &CancellationContext,
            ) -> Self::Value {
                panic!("value should be cached, not evaluated")
            }

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

        let mut map = DiceMap::new();
        let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);
        let cache = map.find_cache(|| IncrementalEngine::new(StoragePropertiesForKey::new(&dice)));
        {
            cache.update_injected_value(MyKey, VersionNumber::new(0), Bar);
            assert_eq!(
                cache
                    .get_cached(MyKey, VersionNumber::new(0), MinorVersion::testing_new(0))
                    .val(),
                &Bar
            )
        }
    }
}


/*
 * 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::fmt::Debug;
use std::sync::atomic::AtomicU32;
use std::sync::Arc;
use std::sync::Weak;

use allocative::Allocative;
use async_trait::async_trait;
use dice_futures::future_handle::WeakDiceFutureHandle;
use dupe::Dupe;
use futures::future::Future;
use futures::StreamExt;
use gazebo::prelude::*;
use incremental::evaluator::Evaluator;
use incremental::graph::GraphNode;
use incremental::transaction_ctx::TransactionCtx;
use incremental::versions::VersionTracker;
use incremental::IncrementalComputeProperties;
use incremental::IncrementalEngine;
use key::StoragePropertiesForKey;
use map::DiceMap;
use more_futures::cancellation::CancellationContext;
use parking_lot::RwLock;
use projection::ProjectionKeyProperties;
use tokio::sync::watch;

use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::data::DiceData;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::projection::ProjectionKey;
use crate::api::transaction::DiceTransactionUpdater;
use crate::api::user_data::UserComputationData;
use crate::api::which::WhichSpawner;
use crate::ctx::DiceComputationsImpl;
use crate::legacy::ctx::ComputationData;
use crate::legacy::ctx::DiceComputationsImplLegacy;
use crate::legacy::incremental::dep_trackers::BothDeps;
use crate::metrics::Metrics;
use crate::transaction_update::DiceTransactionUpdaterImpl;

pub(crate) mod ctx;
pub(crate) mod cycles;
pub(crate) mod dice_futures;
pub(crate) mod key;
pub(crate) mod map;
pub(crate) mod opaque;
pub(crate) mod projection;

pub mod incremental;
#[cfg(test)]
mod tests;

/// An incremental computation engine that executes arbitrary computations that
/// maps `Key`s to values.
#[derive(Allocative)]
pub(crate) struct DiceLegacy {
    pub(crate) data: DiceData,
    pub(crate) map: Arc<RwLock<DiceMap>>,
    pub(crate) global_versions: Arc<VersionTracker>,
    detect_cycles: DetectCycles,
    /// Number of active transactions.
    /// Or more precisely, the number of alive transaction context objects.
    pub(crate) active_transaction_count: AtomicU32,
    #[allocative(skip)]
    active_versions_observer: watch::Receiver<usize>,
    which_spawner: WhichSpawner,
}

impl Debug for DiceLegacy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Dice")
            .field("detect_cycles", &self.detect_cycles)
            .finish_non_exhaustive()
    }
}

pub(crate) struct DiceLegacyDataBuilder(DiceData);

impl DiceLegacyDataBuilder {
    pub(crate) fn new() -> Self {
        Self(DiceData::new())
    }

    pub fn set<K: Send + Sync + 'static>(&mut self, val: K) {
        self.0.set(val);
    }

    pub fn build(
        self,
        detect_cycles: DetectCycles,
        which_spawner: WhichSpawner,
    ) -> Arc<DiceLegacy> {
        DiceLegacy::new(self.0, detect_cycles, which_spawner)
    }
}

impl DiceLegacy {
    #[cfg(test)]
    pub(crate) fn builder() -> DiceLegacyDataBuilder {
        DiceLegacyDataBuilder::new()
    }

    pub(crate) fn new(
        data: DiceData,
        detect_cycles: DetectCycles,
        which_spawner: WhichSpawner,
    ) -> Arc<Self> {
        let map = Arc::new(RwLock::new(DiceMap::new()));
        let weak_map = Arc::downgrade(&map);
        let (active_versions_sender, active_versions_observer) = watch::channel(0);

        Arc::new(DiceLegacy {
            data,
            map,
            global_versions: VersionTracker::new(Box::new(move |update| {
                tracing::debug!("VersionTracker update: {:?}", update);

                if let Some(deleted) = update.deleted_version() {
                    if let Some(engines) = weak_map.upgrade() {
                        engines
                            .read()
                            .engines()
                            .map(|engine| engine.gc_version(deleted));
                    }
                }

                // If the corresponding Dice has been dropped, then so be it, ignore the error.
                active_versions_sender.send_replace(update.active_version_count());
            })),
            detect_cycles,
            which_spawner,
            active_transaction_count: AtomicU32::new(0),
            active_versions_observer,
        })
    }

    pub fn updater(self: &Arc<DiceLegacy>) -> DiceTransactionUpdater {
        self.updater_with_data(UserComputationData::new())
    }

    pub fn updater_with_data(
        self: &Arc<DiceLegacy>,
        extra: UserComputationData,
    ) -> DiceTransactionUpdater {
        let ctx = self.make_ctx(ComputationData::new(extra, self.detect_cycles));
        DiceTransactionUpdater(DiceTransactionUpdaterImpl::Legacy(ctx))
    }

    pub(crate) fn make_ctx(
        self: &Arc<DiceLegacy>,
        extra: ComputationData,
    ) -> Arc<DiceComputationsImplLegacy> {
        Arc::new(DiceComputationsImplLegacy::new_transaction(
            self.dupe(),
            self.global_versions.current(),
            self.global_versions.write(),
            extra,
        ))
    }

    /// finds the computation index for the given key
    pub(crate) fn find_cache<K>(
        self: &Arc<DiceLegacy>,
    ) -> Arc<IncrementalEngine<StoragePropertiesForKey<K>>>
    where
        K: Key,
    {
        if let Some(cache) = self
            .map
            .read()
            .find_cache_opt::<StoragePropertiesForKey<K>>()
        {
            return cache;
        }

        self.map
            .write()
            .find_cache(|| IncrementalEngine::new(StoragePropertiesForKey::<K>::new(self)))
    }

    pub(crate) fn find_projection_cache<P: ProjectionKey>(
        self: &Arc<DiceLegacy>,
    ) -> Arc<IncrementalEngine<ProjectionKeyProperties<P>>>
    where
        P: ProjectionKey,
    {
        if let Some(cache) = self
            .map
            .read()
            .find_cache_opt::<ProjectionKeyProperties<P>>()
        {
            return cache;
        }

        self.map
            .write()
            .find_cache(|| IncrementalEngine::new(ProjectionKeyProperties::<P>::new(self)))
    }

    pub(crate) fn unstable_take(self: &Arc<DiceLegacy>) -> DiceMap {
        debug!(msg = "clearing all Dice state");
        let mut map = self.map.write();
        std::mem::replace(&mut map, DiceMap::new())
    }

    pub fn detect_cycles(&self) -> &DetectCycles {
        &self.detect_cycles
    }

    pub fn metrics(&self) -> Metrics {
        let dice_map = self.map.read();
        Metrics {
            key_count: dice_map.key_count(),
            currently_active_key_count: dice_map.currently_running_key_count(),
            active_transaction_count: self
                .active_transaction_count
                .load(std::sync::atomic::Ordering::SeqCst),
        }
    }

    /// Wait until all active versions have exited.
    pub fn wait_for_idle(&self) -> impl Future<Output = ()> + 'static {
        let obs = self.active_versions_observer.clone();
        let mut obs = tokio_stream::wrappers::WatchStream::new(obs);

        async move {
            while let Some(v) = obs.next().await {
                if v == 0 {
                    break;
                }
            }
        }
    }

    pub fn is_idle(&self) -> bool {
        *self.active_versions_observer.borrow() == 0
    }
}

#[derive(Clone, Dupe)]
struct Eval(Weak<DiceLegacy>);

#[async_trait]
impl<K: Key> IncrementalComputeProperties for StoragePropertiesForKey<K> {
    type DiceTask = WeakDiceFutureHandle<Self>;

    async fn recompute(
        key: &Self::Key,
        engine: &Arc<IncrementalEngine<Self>>,
        transaction_ctx: &Arc<TransactionCtx>,
        extra: &ComputationData,
    ) -> DiceResult<GraphNode<StoragePropertiesForKey<K>>> {
        Ok(engine
            .eval_entry_versioned(key, transaction_ctx, extra.subrequest::<Self>(key)?)
            .await)
    }
}

#[async_trait]
impl<K: Key> Evaluator for StoragePropertiesForKey<K> {
    async fn eval(
        &self,
        k: &K,
        transaction_ctx: Arc<TransactionCtx>,
        cancellations: &CancellationContext,
        extra: ComputationData,
    ) -> EvaluationResult<K::Value> {
        let ctx = DiceComputationsImplLegacy::new_for_key_evaluation(
            self.dice
                .upgrade()
                .expect("Dice holds DiceMap so it should still be alive here"),
            transaction_ctx,
            extra,
        );

        let value = k
            .compute(
                &DiceComputations(DiceComputationsImpl::Legacy(ctx.dupe())),
                cancellations,
            )
            .await;

        let (both_deps, extra) = ctx.finalize();

        EvaluationResult {
            value,
            both_deps,
            extra,
        }
    }
}

/// Result of evaluation computation.
pub(crate) struct EvaluationResult<T> {
    pub(crate) value: T,
    pub(crate) both_deps: BothDeps,
    pub(crate) extra: ComputationData,
}