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

use dupe::Dupe;

use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::projection::ProjectionKey;
use crate::legacy::ctx::DiceComputationsImplLegacy;
use crate::legacy::incremental::dep_trackers::BothDeps;
use crate::legacy::incremental::graph::GraphNode;
use crate::legacy::incremental::IncrementalEngine;
use crate::StoragePropertiesForKey;

/// Computed value which is not directly visible to user.
///
/// The value can be accessed only via "projection" operation,
/// so projection result is recorded as a dependency
/// of a computation which requested the opaqued value,
/// but the opaque value key is not.
pub(crate) struct OpaqueValueImplLegacy<'a, K: Key> {
    /// Computed value.
    pub(crate) value: GraphNode<StoragePropertiesForKey<K>>,
    /// Computations which requested this value, parent of K.
    pub(crate) parent_computations: &'a Arc<DiceComputationsImplLegacy>,
    incremental_engine: Arc<IncrementalEngine<StoragePropertiesForKey<K>>>,
}

impl<'a, K> Debug for OpaqueValueImplLegacy<'a, K>
where
    K: Key,
    K::Value: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("OpaqueValue")
            .field("key", self.value.key())
            .field("value", self.value.val())
            .finish_non_exhaustive()
    }
}

impl<'a, K: Key> OpaqueValueImplLegacy<'a, K> {
    pub(crate) fn new(
        value: GraphNode<StoragePropertiesForKey<K>>,
        parent_computations: &'a Arc<DiceComputationsImplLegacy>,
        incremental_engine: Arc<IncrementalEngine<StoragePropertiesForKey<K>>>,
    ) -> OpaqueValueImplLegacy<'a, K> {
        OpaqueValueImplLegacy {
            value,
            parent_computations,
            incremental_engine,
        }
    }

    pub(crate) fn key(&self) -> &K {
        self.value.key()
    }

    pub(crate) fn as_both_deps(&self) -> BothDeps {
        BothDeps::only_one_dep(
            self.parent_computations.transaction_ctx.get_version(),
            self.value.dupe(),
            &self.incremental_engine,
        )
    }

    /// Get a value and record parent computation dependency on `K`.
    pub(crate) fn into_value(self) -> K::Value {
        let value = self.value.val().dupe();

        // Track dependencies.
        self.parent_computations.dep_trackers.record(
            self.parent_computations.transaction_ctx.get_version(),
            self.incremental_engine,
            self.value,
        );

        value
    }

    pub(crate) fn projection<P>(&self, projection_key: &P) -> DiceResult<P::Value>
    where
        P: ProjectionKey<DeriveFromKey = K>,
    {
        self.parent_computations
            .compute_projection_sync(self, projection_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.
 */

//! "Projection" keys: synchronously computed keys from "opaque" values.

use std::fmt::Debug;
use std::hash::Hash;
use std::sync::Arc;
use std::sync::Weak;

use allocative::Allocative;
use async_trait::async_trait;

use crate::api::error::DiceResult;
use crate::api::projection::ProjectionKey;
use crate::api::storage_type::StorageType;
use crate::introspection::graph::short_type_name;
use crate::legacy::ctx::ComputationData;
use crate::legacy::dice_futures::sync_handle::SyncDiceTaskHandle;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::legacy::incremental::IncrementalComputeProperties;
use crate::legacy::incremental::IncrementalEngine;
use crate::legacy::DiceLegacy;
use crate::GraphNode;
use crate::TransactionCtx;

/// Actual key in the graph (well, actual key is `StorageKeyForKey<ProjectionKeyAsKey<P>>`).
/// `ProjectionKey` alone does not have information about what it is derived from.
/// So we record dependency on projection key, we depend on this key.
#[derive(Debug, derive_more::Display, Clone, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "({}, {})", derive_from_key, k)]
pub(crate) struct ProjectionKeyAsKey<P: ProjectionKey> {
    pub(crate) derive_from_key: P::DeriveFromKey,
    pub(crate) k: P,
}

#[derive(Debug, Allocative)]
#[allocative(bound = "")]
pub(crate) struct ProjectionKeyProperties<P: ProjectionKey> {
    _marker: std::marker::PhantomData<P>,
    pub(crate) dice: Weak<DiceLegacy>,
}

impl<P: ProjectionKey> ProjectionKeyProperties<P> {
    pub(crate) fn new(dice: &Arc<DiceLegacy>) -> Self {
        ProjectionKeyProperties {
            _marker: std::marker::PhantomData,
            dice: Arc::downgrade(dice),
        }
    }
}

impl<P: ProjectionKey> StorageProperties for ProjectionKeyProperties<P> {
    type Key = ProjectionKeyAsKey<P>;
    type Value = P::Value;

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

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

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

    /// Provides a short informative name for this projection type.
    fn key_type_name() -> &'static str {
        short_type_name(std::any::type_name::<Self>())
    }

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

#[async_trait]
impl<P: ProjectionKey> IncrementalComputeProperties for ProjectionKeyProperties<P> {
    type DiceTask = SyncDiceTaskHandle<Self>;

    async fn recompute(
        key: &Self::Key,
        engine: &Arc<IncrementalEngine<Self>>,
        transaction_ctx: &Arc<TransactionCtx>,
        extra: &ComputationData,
    ) -> DiceResult<GraphNode<Self>> {
        engine
            .recompute_projection(key, transaction_ctx, extra.subrequest::<Self>(key)?)
            .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.
 */

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

use futures::future::BoxFuture;
use futures::FutureExt;
use more_futures::instrumented_shared::SharedEventsFuture;
use more_futures::spawn::StrongJoinHandle;
use more_futures::spawn::WeakFutureError;

use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::result::CancellableResult;
use crate::GraphNode;

type DiceJoinHandle<S> = StrongJoinHandle<
    SharedEventsFuture<
        BoxFuture<'static, Result<CancellableResult<GraphNode<S>>, WeakFutureError>>,
    >,
>;

pub(crate) enum DiceFuture<S: StorageProperties> {
    /// Earlier computed value.
    Ready(Option<GraphNode<S>>),
    /// Current computation spawned the task.
    AsyncCancellableSpawned(DiceJoinHandle<S>),
    /// Other computation for current key spawned the task.
    AsyncCancellableJoining(DiceJoinHandle<S>),
}

impl<S> Future for DiceFuture<S>
where
    S: StorageProperties,
{
    type Output = GraphNode<S>;

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
        match self.get_mut() {
            DiceFuture::Ready(value) => Poll::Ready(value.take().expect("polled after ready")),
            DiceFuture::AsyncCancellableSpawned(fut) | DiceFuture::AsyncCancellableJoining(fut) => {
                Pin::new(&mut fut.map(|cancellable| match cancellable {
                    Ok(res) => res,
                    Err(_) => {
                        unreachable!("Strong Join Handle was cancelled while still polled")
                    }
                }))
                .poll(cx)
            }
        }
    }
}


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

#[derive(Debug, Clone, Copy, Dupe)]
pub(crate) enum DiceTaskStateForDebugging {
    /// Async task finished.
    AsyncReady,
    /// Async task in progress.
    AsyncInProgress,
    /// Weak handle is empty.
    AsyncDropped,
    /// Sync task finished.
    SyncReady,
    /// Sync in progress.
    SyncInProgress,
}

/// Marker trait for a task currently executed in `IncrementalEngine`.
pub(crate) trait DiceTask: Allocative + Send + Sync + 'static {
    fn state_for_debugging(&self) -> DiceTaskStateForDebugging;
}


/*
 * 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 future that is spawned and managed by DICE. This is a single computation unit that is
//! shareable across different computation units.
//!
use allocative::Allocative;
use more_futures::spawn::CompletionObserver;
use more_futures::spawn::WeakFutureError;
use more_futures::spawn::WeakJoinHandle;

use crate::legacy::dice_futures::dice_future::DiceFuture;
use crate::legacy::dice_futures::dice_task::DiceTask;
use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::result::CancellableResult;
use crate::GraphNode;

#[derive(Allocative)]
pub(crate) struct WeakDiceFutureHandle<S: StorageProperties> {
    #[allocative(skip)] // TODO(nga): value may be hiding in there.
    handle: WeakJoinHandle<Result<CancellableResult<GraphNode<S>>, WeakFutureError>>,
}

impl<S: StorageProperties> DiceTask for WeakDiceFutureHandle<S> {
    fn state_for_debugging(&self) -> DiceTaskStateForDebugging {
        match self.handle.pollable() {
            Some(p) => {
                if p.inner().inner().peek().is_some() {
                    DiceTaskStateForDebugging::AsyncReady
                } else {
                    DiceTaskStateForDebugging::AsyncInProgress
                }
            }
            None => DiceTaskStateForDebugging::AsyncDropped,
        }
    }
}

impl<S: StorageProperties> WeakDiceFutureHandle<S> {
    pub(crate) fn async_cancellable(
        handle: WeakJoinHandle<Result<CancellableResult<GraphNode<S>>, WeakFutureError>>,
    ) -> WeakDiceFutureHandle<S> {
        WeakDiceFutureHandle { handle }
    }

    pub(crate) fn pollable(&self) -> Option<DiceFuture<S>> {
        self.handle
            .pollable()
            .map(DiceFuture::AsyncCancellableJoining)
    }

    /// Turn this into a JoinHandle. The output is erased. This is used to observe this future
    /// exiting, but that's it.
    pub fn into_completion_observer(self) -> CompletionObserver<CancellableResult<GraphNode<S>>> {
        self.handle.into_completion_observer()
    }
}


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

pub mod dice_future;
pub mod dice_task;
pub mod future_handle;
pub mod sync_handle;


/*
 * 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 dupe::Clone_;
use dupe::Dupe;
use futures::future::Shared;

use crate::legacy::dice_futures::dice_task::DiceTask;
use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::GraphNode;

/// `IncrementalEngine` task type for projection key.
#[derive(Clone_, Allocative)]
pub(crate) struct SyncDiceTaskHandle<S: StorageProperties> {
    // Doesn't have peek, so we cannot check if the value is there.
    // Could measure the size of the channel though.
    #[allocative(skip)]
    pub(crate) rx: Shared<tokio::sync::oneshot::Receiver<GraphNode<S>>>,
}

impl<S: StorageProperties> Dupe for SyncDiceTaskHandle<S> {}

impl<S: StorageProperties> DiceTask for SyncDiceTaskHandle<S> {
    fn state_for_debugging(&self) -> DiceTaskStateForDebugging {
        if self.rx.peek().is_some() {
            DiceTaskStateForDebugging::SyncReady
        } else {
            DiceTaskStateForDebugging::SyncInProgress
        }
    }
}


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

//! Trackers that records dependencies and reverse dependencies during execution of requested nodes

use std::sync::Arc;

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

use crate::legacy::incremental::dep_trackers::internals::ComputedDep;
use crate::legacy::incremental::graph::dependencies::ComputedDependency;
use crate::legacy::incremental::graph::GraphNode;
use crate::legacy::incremental::graph::GraphNodeDyn;
use crate::legacy::incremental::IncrementalComputeProperties;
use crate::legacy::incremental::IncrementalEngine;
use crate::versions::VersionNumber;
use crate::HashSet;

/// The 'DepsTracker' is used to record dependencies of a particular compute node by calling
/// 'record' for each dependency, and then getting a list of 'Dependency's at the end by calling
/// 'collect_deps'.
#[derive(Allocative)]
struct RecordingDepsTracker {
    deps: HashSet<Box<dyn ComputedDependency>>,
}

impl RecordingDepsTracker {
    fn new() -> Self {
        Self {
            deps: HashSet::default(),
        }
    }

    /// records k with the given evaluator and engine
    fn record<K>(&mut self, v: VersionNumber, engine: Arc<IncrementalEngine<K>>, node: GraphNode<K>)
    where
        K: IncrementalComputeProperties,
    {
        self.deps.insert(Box::new(ComputedDep {
            engine: Arc::downgrade(&engine),
            version: v,
            node,
        }));
    }

    fn collect_deps(self) -> HashSet<Box<dyn ComputedDependency>> {
        self.deps
    }
}

#[derive(Allocative)]
struct RecordingRdepsTracker {
    rdeps: Vec<Arc<dyn GraphNodeDyn>>,
}

impl RecordingRdepsTracker {
    fn new() -> Self {
        Self { rdeps: Vec::new() }
    }

    fn record(&mut self, dep: Arc<dyn GraphNodeDyn>) {
        self.rdeps.push(dep)
    }

    fn collect_rdeps(self) -> Vec<Arc<dyn GraphNodeDyn>> {
        self.rdeps
    }
}

#[derive(Allocative)]
struct BothRecordingDepTrackers {
    deps: RecordingDepsTracker,
    rdeps: RecordingRdepsTracker,
}

#[derive(Default)]
pub(crate) struct BothDeps {
    pub(crate) deps: HashSet<Box<dyn ComputedDependency>>,
    pub(crate) rdeps: Vec<Arc<dyn GraphNodeDyn>>,
}

impl BothDeps {
    pub(crate) fn only_one_dep<S: IncrementalComputeProperties>(
        version: VersionNumber,
        node: GraphNode<S>,
        incremental_engine: &Arc<IncrementalEngine<S>>,
    ) -> BothDeps {
        let dep: Box<dyn ComputedDependency> = Box::new(ComputedDep::<S> {
            engine: Arc::downgrade(incremental_engine),
            version,
            node: node.dupe(),
        });
        BothDeps {
            deps: HashSet::from_iter([dep]),
            rdeps: Vec::from_iter([node.into_dyn()]),
        }
    }
}

#[derive(Allocative)]
enum BothDepTrackersImpl {
    Noop,
    Recording(Mutex<BothRecordingDepTrackers>),
}

#[derive(Allocative)]
pub(crate) struct BothDepTrackers(BothDepTrackersImpl);

/// There are two variants, a 'Recording' tracker and a 'Noop' tracker. The 'Noop' tracker never
/// tracks any dependencies such that 'collect_deps' is always empty. The 'Recording' tracker will
/// actually track the dependencies.
impl BothDepTrackers {
    pub(crate) fn noop() -> BothDepTrackers {
        BothDepTrackers(BothDepTrackersImpl::Noop)
    }

    pub(crate) fn recording() -> BothDepTrackers {
        BothDepTrackers(BothDepTrackersImpl::Recording(Mutex::new(
            BothRecordingDepTrackers {
                deps: RecordingDepsTracker::new(),
                rdeps: RecordingRdepsTracker::new(),
            },
        )))
    }

    /// records k with the given evaluator and engine
    pub(crate) fn record<K>(
        &self,
        v: VersionNumber,
        engine: Arc<IncrementalEngine<K>>,
        node: GraphNode<K>,
    ) where
        K: IncrementalComputeProperties,
    {
        match &self.0 {
            BothDepTrackersImpl::Noop => {}
            BothDepTrackersImpl::Recording(recording) => {
                let mut recording = recording.lock();
                let BothRecordingDepTrackers { deps, rdeps } = &mut *recording;
                deps.record(v, engine, node.dupe());
                rdeps.record(node.into_dyn());
            }
        }
    }

    pub(crate) fn collect_deps(self) -> BothDeps {
        match self.0 {
            BothDepTrackersImpl::Noop => BothDeps::default(),
            BothDepTrackersImpl::Recording(recording) => {
                let BothRecordingDepTrackers { deps, rdeps } = recording.into_inner();
                let deps = deps.collect_deps();
                let rdeps = rdeps.collect_rdeps();
                BothDeps { deps, rdeps }
            }
        }
    }
}

mod internals {
    use std::any::type_name;
    use std::any::Any;
    use std::fmt;
    use std::fmt::Debug;
    use std::fmt::Display;
    use std::fmt::Formatter;
    use std::hash::Hash;
    use std::hash::Hasher;
    use std::sync::Arc;
    use std::sync::Weak;

    use allocative::Allocative;
    use async_trait::async_trait;
    use cmp_any::PartialEqAny;
    use dupe::Dupe;

    use crate::api::error::DiceResult;
    use crate::introspection::graph::AnyKey;
    use crate::legacy::ctx::ComputationData;
    use crate::legacy::incremental::graph::GraphNode;
    use crate::legacy::incremental::graph::GraphNodeDyn;
    use crate::legacy::incremental::graph::ReadOnlyHistory;
    use crate::legacy::incremental::graph::VersionedGraphKeyRef;
    use crate::legacy::incremental::transaction_ctx::TransactionCtx;
    use crate::legacy::incremental::versions::MinorVersion;
    use crate::legacy::incremental::ComputedDependency;
    use crate::legacy::incremental::Dependency;
    use crate::legacy::incremental::IncrementalComputeProperties;
    use crate::legacy::incremental::IncrementalEngine;
    use crate::versions::VersionNumber;

    #[derive(Allocative)]
    pub(crate) struct ComputedDep<K: IncrementalComputeProperties> {
        pub(crate) engine: Weak<IncrementalEngine<K>>,
        pub(crate) version: VersionNumber,
        pub(crate) node: GraphNode<K>,
    }

    impl<K> ComputedDependency for ComputedDep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn get_history(&self) -> ReadOnlyHistory {
            self.node.get_history()
        }

        fn into_dependency(self: Box<Self>) -> Box<dyn Dependency> {
            Box::new(Dep {
                engine: self.engine,
                k: self.node.key().clone(),
            })
        }

        fn get_key_equality(&self) -> (PartialEqAny, VersionNumber) {
            (PartialEqAny::new(self.node.key()), self.version)
        }

        fn to_key_any(&self) -> &dyn Any {
            K::to_key_any(self.node.key())
        }

        fn hash(&self, mut state: &mut dyn Hasher) {
            self.node.key().hash(&mut state);
            self.version.hash(&mut state);
        }

        fn is_valid(&self) -> bool {
            self.node.is_valid()
        }
    }

    impl<K> Debug for ComputedDep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
            write!(
                f,
                "ComputedDependency(({:?}={:?}) -> {:?}, version={:?})",
                type_name::<K::Key>(),
                self.node.key(),
                type_name::<K::Value>(),
                self.version,
            )
        }
    }

    #[derive(Allocative)]
    pub(crate) struct Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        pub(crate) engine: Weak<IncrementalEngine<K>>,
        pub(crate) k: K::Key,
    }

    impl<K> Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        pub(crate) fn engine(&self) -> Arc<IncrementalEngine<K>> {
            self.engine.upgrade().expect(
                "IncrementalEngine should not be destroyed because IncrementalEngine owns Dep",
            )
        }
    }

    #[async_trait]
    impl<K> Dependency for Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        #[instrument(level = "info", skip(self, transaction_ctx, extra), fields(k = %self.k, version = %transaction_ctx.get_version()))]
        async fn recompute(
            &self,
            transaction_ctx: &Arc<TransactionCtx>,
            extra: &ComputationData,
        ) -> DiceResult<(Box<dyn ComputedDependency>, Arc<dyn GraphNodeDyn>)> {
            let res = K::recompute(&self.k, &self.engine(), transaction_ctx, extra).await?;

            Ok((
                Box::new(ComputedDep {
                    engine: self.engine.dupe(),
                    version: transaction_ctx.get_version(),
                    node: res.dupe(),
                }),
                res.into_dyn(),
            ))
        }

        fn lookup_node(&self, v: VersionNumber, mv: MinorVersion) -> Option<Arc<dyn GraphNodeDyn>> {
            if let Some(node) = self
                .engine()
                .versioned_cache
                .get(VersionedGraphKeyRef::new(v, &self.k), mv)
                .unpack_match()
            {
                Some(node.dupe().into_dyn())
            } else {
                None
            }
        }

        fn dirty(&self, v: VersionNumber) {
            self.engine().dirty(self.k.clone(), v, false)
        }

        fn get_key_equality(&self) -> PartialEqAny {
            PartialEqAny::new(&self.k)
        }

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

        fn introspect(&self) -> AnyKey {
            AnyKey::new(self.k.clone())
        }

        fn to_key_any(&self) -> &dyn std::any::Any {
            K::to_key_any(&self.k)
        }
    }

    impl<K> Debug for Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
            write!(
                f,
                "Dependency(({:?}={:?}) -> {:?})",
                type_name::<K::Key>(),
                self.k,
                type_name::<K::Value>()
            )
        }
    }

    impl<K> Display for Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
            write!(f, "Dependency({})", self.k)
        }
    }
}

#[cfg(test)]
pub(crate) mod testing {
    use std::sync::Arc;
    use std::sync::Weak;

    pub(crate) use crate::legacy::incremental::dep_trackers::internals::ComputedDep;
    pub(crate) use crate::legacy::incremental::dep_trackers::internals::Dep;
    use crate::legacy::incremental::graph::GraphNode;
    use crate::legacy::incremental::graph::OccupiedGraphNode;
    use crate::legacy::incremental::IncrementalComputeProperties;
    use crate::legacy::incremental::IncrementalEngine;
    use crate::versions::VersionNumber;

    pub(crate) trait DepExt<K: IncrementalComputeProperties> {
        fn testing_new(engine: Weak<IncrementalEngine<K>>, k: K::Key) -> Self;
    }

    impl<K> DepExt<K> for Dep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn testing_new(engine: Weak<IncrementalEngine<K>>, k: K::Key) -> Self {
            Dep { engine, k }
        }
    }

    pub(crate) trait ComputedDepExt<K: IncrementalComputeProperties> {
        fn testing_new(
            engine: Weak<IncrementalEngine<K>>,
            version: VersionNumber,
            node: Arc<OccupiedGraphNode<K>>,
        ) -> Self;
    }

    impl<K> ComputedDepExt<K> for ComputedDep<K>
    where
        K: IncrementalComputeProperties,
    {
        fn testing_new(
            engine: Weak<IncrementalEngine<K>>,
            version: VersionNumber,
            node: Arc<OccupiedGraphNode<K>>,
        ) -> Self {
            ComputedDep {
                engine,
                version,
                node: GraphNode::occupied(node),
            }
        }
    }
}

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

    use dupe::Dupe;
    use futures::FutureExt;

    use crate::impls::core::graph::history::CellHistory;
    use crate::legacy::ctx::testing::ComputationDataExt;
    use crate::legacy::ctx::ComputationData;
    use crate::legacy::incremental::dep_trackers::BothDeps;
    use crate::legacy::incremental::dep_trackers::RecordingDepsTracker;
    use crate::legacy::incremental::dep_trackers::RecordingRdepsTracker;
    use crate::legacy::incremental::evaluator::testing::EvaluatorFn;
    use crate::legacy::incremental::evaluator::testing::EvaluatorUnreachable;
    use crate::legacy::incremental::graph::OccupiedGraphNode;
    use crate::legacy::incremental::testing::ComputedDependencyExt;
    use crate::legacy::incremental::IncrementalEngine;
    use crate::legacy::incremental::TransactionCtx;
    use crate::versions::VersionNumber;
    use crate::HashSet;

    #[test]
    fn recording_rdeps_tracker_tracks_rdeps() {
        let mut rdeps_tracker = RecordingRdepsTracker::new();

        let node = Arc::new(OccupiedGraphNode::<EvaluatorFn<usize, usize>>::new(
            1337,
            2,
            CellHistory::verified(VersionNumber::new(0)),
        ));
        rdeps_tracker.record(node.dupe());
        let tracked = rdeps_tracker.collect_rdeps();

        assert_eq!(tracked.len(), 1);
    }

    #[tokio::test]
    async fn recording_deps_tracker_tracks_deps() -> anyhow::Result<()> {
        let mut deps_tracker = RecordingDepsTracker::new();
        // set up so that we have keys 2 and 3 with a history of VersionNumber(1)
        let fn_for_2_and_3 = |k| async move { (k, BothDeps::default()) }.boxed();

        let engine = IncrementalEngine::new(EvaluatorFn::new(move |k, _| fn_for_2_and_3(k)));

        let ctx = Arc::new(TransactionCtx::testing_new(VersionNumber::new(1)));

        let node1 = engine
            .eval_entry_versioned(&2, &ctx, ComputationData::testing_new())
            .await;
        let node2 = engine
            .eval_entry_versioned(&3, &ctx, ComputationData::testing_new())
            .await;

        deps_tracker.record(VersionNumber::new(1), engine.dupe(), node1);
        deps_tracker.record(VersionNumber::new(1), engine.dupe(), node2);

        let deps = deps_tracker.collect_deps();

        let expected = HashSet::from_iter([
            ComputedDependencyExt::<EvaluatorUnreachable<_, i32>>::testing_raw(
                2,
                VersionNumber::new(1),
                true,
            ),
            ComputedDependencyExt::<EvaluatorUnreachable<_, i32>>::testing_raw(
                3,
                VersionNumber::new(1),
                true,
            ),
        ]);
        assert_eq!(deps, expected);

        Ok(())
    }
}


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

use std::sync::Arc;

use async_trait::async_trait;
use more_futures::cancellation::CancellationContext;

use crate::legacy::ctx::ComputationData;
use crate::legacy::dice_futures::dice_task::DiceTask;
use crate::legacy::incremental::graph::GraphNode;
use crate::legacy::incremental::IncrementalComputeProperties;
use crate::legacy::incremental::IncrementalEngine;
use crate::legacy::EvaluationResult;
use crate::TransactionCtx;
use crate::WeakDiceFutureHandle;

#[async_trait]
pub(crate) trait Recompute: IncrementalComputeProperties {
    /// Dice task executed in the `IncrementalEngine` for compute or recompute.
    type DiceTask: DiceTask;

    async fn recompute(
        key: &Self::Key,
        engine: &Arc<IncrementalEngine<Self>>,
        transaction_ctx: &Arc<TransactionCtx>,
        extra: &ComputationData,
    ) -> GraphNode<Self>;
}

#[async_trait]
pub(crate) trait Evaluator:
    IncrementalComputeProperties<DiceTask = WeakDiceFutureHandle<Self>>
{
    async fn eval(
        &self,
        k: &Self::Key,
        transaction_ctx: Arc<TransactionCtx>,
        cancellations: &CancellationContext,
        extra: ComputationData,
    ) -> EvaluationResult<Self::Value>;
}

#[cfg(test)]
pub(crate) mod testing {
    use std::fmt;
    use std::fmt::Debug;
    use std::fmt::Display;
    use std::hash::Hash;
    use std::sync::Arc;

    use allocative::Allocative;
    use async_trait::async_trait;
    use dupe::Dupe;
    use futures::future::BoxFuture;
    use gazebo::prelude::*;
    use more_futures::cancellation::CancellationContext;

    use crate::api::error::DiceResult;
    use crate::api::storage_type::StorageType;
    use crate::legacy::ctx::ComputationData;
    use crate::legacy::incremental::evaluator::Evaluator;
    use crate::legacy::incremental::graph::storage_properties::StorageProperties;
    use crate::legacy::incremental::graph::GraphNode;
    use crate::legacy::incremental::Computable;
    use crate::legacy::incremental::IncrementalComputeProperties;
    use crate::legacy::incremental::IncrementalEngine;
    use crate::legacy::BothDeps;
    use crate::legacy::EvaluationResult;
    use crate::TransactionCtx;
    use crate::WeakDiceFutureHandle;

    /// Evaluator which panics on attempt to evaluate.
    #[derive(Default_, Allocative)]
    #[allocative(bound = "")]
    pub(crate) struct EvaluatorUnreachable<K, V> {
        _kv: std::marker::PhantomData<fn(K) -> V>,
    }

    impl<K, V> fmt::Debug for EvaluatorUnreachable<K, V> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("EvaluatorUnreachable")
                .finish_non_exhaustive()
        }
    }

    impl<
        K: Send + Sync + Hash + Eq + Display + Debug + Clone + Allocative + 'static,
        V: Send + Sync + PartialEq + Dupe + Allocative + 'static,
    > StorageProperties for EvaluatorUnreachable<K, V>
    {
        type Key = K;
        type Value = V;

        fn key_type_name() -> &'static str {
            "TestingUnreachable"
        }

        fn storage_type(&self) -> StorageType {
            unreachable!()
        }

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

        fn validity(&self, _x: &Self::Value) -> bool {
            true
        }

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

    #[async_trait]
    impl<
        K: Send + Sync + Hash + Eq + Display + Debug + Clone + Allocative + 'static,
        V: Send + Sync + PartialEq + Dupe + Allocative + 'static,
    > IncrementalComputeProperties for EvaluatorUnreachable<K, V>
    {
        type DiceTask = WeakDiceFutureHandle<Self>;

        async fn recompute(
            _key: &Self::Key,
            _engine: &Arc<IncrementalEngine<Self>>,
            _transaction_ctx: &Arc<TransactionCtx>,
            _extra: &ComputationData,
        ) -> DiceResult<GraphNode<Self>> {
            unreachable!()
        }
    }

    #[async_trait]
    impl<
        K: Send + Sync + Hash + Eq + Display + Debug + Clone + Allocative + 'static,
        V: Send + Sync + PartialEq + Dupe + Allocative + 'static,
    > Evaluator for EvaluatorUnreachable<K, V>
    {
        async fn eval(
            &self,
            _k: &Self::Key,
            _transaction_ctx: Arc<TransactionCtx>,
            _cancellations: &CancellationContext,
            _extra: ComputationData,
        ) -> EvaluationResult<Self::Value> {
            unreachable!()
        }
    }

    /// Evaluator which invokes boxed closure on evaluation.
    #[derive(Allocative)]
    #[allocative(bound = "")]
    pub(crate) struct EvaluatorFn<K, V> {
        #[allocative(skip)]
        f: Box<
            dyn for<'a> Fn(K, &'a CancellationContext) -> BoxFuture<'a, (V, BothDeps)>
                + Send
                + Sync
                + 'static,
        >,
    }

    impl<K, V> EvaluatorFn<K, V> {
        pub(crate) fn new<F>(f: F) -> Self
        where
            F: for<'a> FnOnce(K, &'a CancellationContext) -> BoxFuture<'a, (V, BothDeps)>
                + Clone
                + Sync
                + Send
                + 'static,
        {
            Self {
                f: Box::new(move |k, cancellations| {
                    let f = f.clone();
                    f(k, cancellations)
                }),
            }
        }
    }

    impl<K, V> Debug for EvaluatorFn<K, V> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("EvaluatorFn").finish_non_exhaustive()
        }
    }

    impl<
        K: Clone + Eq + Hash + Display + Debug + Send + Sync + Allocative + 'static,
        V: Dupe + PartialEq + Send + Sync + Allocative + 'static,
    > StorageProperties for EvaluatorFn<K, V>
    {
        type Key = K;
        type Value = V;

        fn key_type_name() -> &'static str {
            "EvaluatorFn"
        }

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

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

        fn validity(&self, _x: &Self::Value) -> bool {
            true
        }

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

    #[async_trait]
    impl<K, V> IncrementalComputeProperties for EvaluatorFn<K, V>
    where
        K: Computable + 'static,
        V: Dupe + PartialEq + Allocative + Send + Sync + 'static,
    {
        type DiceTask = WeakDiceFutureHandle<Self>;

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

    #[async_trait]
    impl<K, V> Evaluator for EvaluatorFn<K, V>
    where
        K: Clone + Eq + Hash + Display + Debug + Allocative + Send + Sync + 'static,
        V: Dupe + PartialEq + Allocative + Send + Sync + 'static,
    {
        async fn eval(
            &self,
            k: &K,
            _: Arc<TransactionCtx>,
            cancellations: &CancellationContext,
            extra: ComputationData,
        ) -> EvaluationResult<V> {
            let (value, both_deps) = (self.f)(k.clone(), cancellations).await;
            EvaluationResult {
                value,
                both_deps,
                extra,
            }
        }
    }
}


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

use dupe::Dupe;
use sorted_vector_map::SortedVectorMap;

use crate::introspection::graph::AnyKey;
use crate::introspection::graph::EngineForIntrospection;
use crate::introspection::graph::GraphNodeKind;
use crate::introspection::graph::KeyID;
use crate::introspection::graph::NodeID;
use crate::introspection::graph::SerializedGraphNode;
use crate::introspection::graph::SerializedGraphNodesForKey;
use crate::legacy::dice_futures::dice_task::DiceTask;
use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
use crate::legacy::incremental::graph::dependencies::VersionedDependencies;
use crate::legacy::incremental::graph::dependencies::VersionedRevDependencies;
use crate::legacy::incremental::graph::storage_properties::StorageProperties;
use crate::legacy::incremental::graph::GraphNodeDyn;
use crate::legacy::incremental::graph::VersionedGraphNodeInternal;
use crate::legacy::incremental::Dependency;
use crate::legacy::incremental::IncrementalComputeProperties;
use crate::legacy::incremental::IncrementalEngine;
use crate::versions::VersionNumber;
use crate::HashMap;
use crate::HashSet;

impl<K, T> EngineForIntrospection for IncrementalEngine<K>
where
    K: IncrementalComputeProperties<Value = T> + 'static,
    T: Dupe + Send + Sync + 'static,
{
    fn keys<'a>(&'a self) -> Box<dyn Iterator<Item = AnyKey> + 'a> {
        Box::new(
            self.versioned_cache
                .iter()
                .map(|e| AnyKey::new(e.key().clone())),
        )
    }

    fn edges<'a>(&'a self) -> Box<dyn Iterator<Item = (AnyKey, Vec<AnyKey>)> + 'a> {
        fn extract_deps<K: StorageProperties>(
            e: &SortedVectorMap<VersionNumber, VersionedGraphNodeInternal<K>>,
        ) -> Option<Arc<Vec<Box<dyn Dependency>>>> {
            e.iter()
                .last()?
                .1
                .unpack_occupied()?
                .read_meta()
                .deps
                .deps()
        }

        Box::new(self.versioned_cache.iter().map(|e| {
            let k = AnyKey::new(e.key().clone());
            let deps = match extract_deps(e.value()) {
                Some(deps) => deps.iter().map(|d| d.introspect()).collect(),
                None => Vec::new(),
            };
            (k, deps)
        }))
    }

    fn keys_currently_running<'a>(
        &'a self,
    ) -> Vec<(
        AnyKey,
        crate::introspection::graph::VersionNumber,
        DiceTaskStateForDebugging,
    )> {
        self.currently_running
            .read()
            .iter()
            .flat_map(|(v, es)| {
                es.iter()
                    .map(move |entry| {
                        let k = entry.key();
                        let e = entry.value();
                        (
                            AnyKey::new(k.clone()),
                            crate::introspection::graph::VersionNumber(v.0),
                            e.task.state_for_debugging(),
                        )
                    })
                    .collect::<Vec<_>>()
            })
            .collect()
    }

    fn versions_currently_running<'a>(&'a self) -> Vec<crate::introspection::graph::VersionNumber> {
        self.currently_running
            .read()
            .iter()
            .map(|e| e.0.to_introspectable())
            .collect()
    }

    fn nodes<'a>(
        &'a self,
        keys: &'a mut HashMap<AnyKey, KeyID>,
    ) -> Box<dyn Iterator<Item = SerializedGraphNodesForKey> + 'a> {
        let mut map_id = move |key: AnyKey| -> KeyID {
            let num_keys = keys.len();
            *keys.entry(key).or_insert_with(|| KeyID(num_keys))
        };
        fn visit_deps<M: FnMut(AnyKey) -> KeyID>(
            deps: &VersionedDependencies,
            map_id: &mut M,
        ) -> Option<HashSet<KeyID>> {
            deps.debug_deps().try_read().and_then(|deps| {
                deps.as_ref()
                    .map(|deps| deps.1.iter().map(|d| map_id(d.introspect())).collect())
            })
        }
        fn visit_rdeps(
            rdeps: &VersionedRevDependencies,
        ) -> BTreeMap<crate::introspection::graph::VersionNumber, Vec<NodeID>> {
            let mut res = BTreeMap::new();

            let rdeps = rdeps.rdeps();
            for rdep in rdeps.rdeps.iter() {
                if let Some(node) = rdep.0.0.upgrade() {
                    res.entry(rdep.1.to_introspectable())
                        .or_insert_with(Vec::new)
                        .push(NodeID(node.id()));
                }
            }

            res
        }
        fn visit_node<K: StorageProperties, M: FnMut(AnyKey) -> KeyID>(
            node: &VersionedGraphNodeInternal<K>,
            map_id: &mut M,
        ) -> Option<SerializedGraphNode> {
            node.unpack_graph_value().map(|graph_value| {
                let m = graph_value.try_read_meta();
                SerializedGraphNode {
                    node_id: match node {
                        VersionedGraphNodeInternal::Occupied(o) => NodeID(o.id()),
                        VersionedGraphNodeInternal::Transient(t) => NodeID(t.id()),
                        VersionedGraphNodeInternal::Vacant(_) => {
                            unreachable!("node was unpacked, can't be vacant")
                        }
                    },
                    kind: GraphNodeKind::from(node),
                    history: (*graph_value.get_history()).to_introspectable(),
                    deps: m.as_ref().and_then(|meta| visit_deps(&meta.deps, map_id)),
                    rdeps: m.map(|meta| visit_rdeps(&meta.rdeps)),
                }
            })
        }
        Box::new(self.versioned_cache.iter().map(move |e| {
            let k = AnyKey::new(e.key().clone());
            SerializedGraphNodesForKey {
                id: map_id(k.clone()),
                key: k.to_string(),
                type_name: k.short_type_name().to_owned(),
                nodes: e
                    .value()
                    .iter()
                    .map(|(v, node)| (v.to_introspectable(), visit_node(node, &mut map_id)))
                    .collect(),
            }
        }))
    }

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

    fn currently_running_key_count(&self) -> usize {
        self.currently_running
            .read()
            .iter()
            .map(|(_, e)| e.len())
            .sum()
    }
}