/*
 * 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 gazebo::variants::VariantName;

use crate::impls::core::internals::CoreState;
use crate::impls::core::state::CoreStateHandle;
use crate::impls::core::state::StateRequest;
use crate::impls::ctx::SharedLiveTransactionCtx;

pub(super) struct StateProcessor {
    state: CoreState,
    rx: tokio::sync::mpsc::UnboundedReceiver<StateRequest>,
}

impl StateProcessor {
    pub(super) fn spawn() -> CoreStateHandle {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let state = CoreState::new();

        std::thread::spawn(move || StateProcessor { state, rx }.event_loop());
        CoreStateHandle::new(tx)
    }

    fn event_loop(mut self) {
        loop {
            // Skip tokio scheduling.
            while let Ok(message) = self.rx.try_recv() {
                self.iteration(message);
            }
            if let Some(message) = self.rx.blocking_recv() {
                self.iteration(message);
            } else {
                break;
            }
        }
        debug!("Processor terminated");
    }

    #[instrument(skip_all, fields(kind = %message.variant_name()))]
    fn iteration(&mut self, message: StateRequest) {
        match message {
            StateRequest::UpdateState { changes, resp } => {
                // ignore error if the requester dropped it.
                let _ = resp.send(self.state.update_state(changes));
            }
            StateRequest::CtxAtVersion {
                version,
                guard,
                resp,
            } => {
                let (version_epoch, cache) = self.state.ctx_at_version(version);

                let ctx = SharedLiveTransactionCtx::new(version, version_epoch, cache);
                let _ignored = resp.send((ctx, guard));
            }
            StateRequest::DropCtxAtVersion { version } => self.state.drop_ctx_at_version(version),
            StateRequest::CurrentVersion { resp } => {
                // ignore error if the requester dropped it.
                let _ = resp.send(self.state.current_version());
            }
            StateRequest::LookupKey { key, resp } => drop(resp.send(self.state.lookup_key(key))),
            StateRequest::UpdateComputed {
                key,
                epoch,
                storage,
                value,
                deps,
                resp,
                ..
            } => {
                // ignore error if the requester dropped it.
                drop(resp.send(self.state.update_computed(key, epoch, storage, value, deps)));
            }
            StateRequest::GetTasksPendingCancellation { resp } => {
                let _ignored = resp.send(self.state.get_tasks_pending_cancellation());
            }
            StateRequest::UnstableDropEverything => self.state.unstable_drop_everything(),
            StateRequest::Metrics { resp } => {
                let _ignored = resp.send(self.state.metrics());
            }
            StateRequest::Introspection { resp, key_map } => {
                let _ignored = resp.send(self.state.introspection(key_map));
            }
        }
    }
}


/*
 * 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 derivative::Derivative;
use dupe::Dupe;
use gazebo::variants::VariantName;
use tokio::sync::oneshot::Sender;

use crate::api::storage_type::StorageType;
use crate::arc::Arc;
use crate::impls::core::graph::types::VersionedGraphKey;
use crate::impls::core::graph::types::VersionedGraphResult;
use crate::impls::core::processor::StateProcessor;
use crate::impls::core::versions::VersionEpoch;
use crate::impls::ctx::SharedLiveTransactionCtx;
use crate::impls::key::DiceKey;
use crate::impls::task::dice::TerminationObserver;
use crate::impls::transaction::ActiveTransactionGuard;
use crate::impls::transaction::ChangeType;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceValidValue;
use crate::introspection::graph::AnyKey;
use crate::introspection::graph::GraphIntrospectable;
use crate::metrics::Metrics;
use crate::result::CancellableResult;
use crate::versions::VersionNumber;
use crate::HashMap;

/// Core state is accessed via message passing to a single threaded processor
#[derive(Derivative, VariantName)]
#[derivative(Debug)]
pub(crate) enum StateRequest {
    /// Updates the core state with the given set of changes. The new VersionNumber that should be
    /// used is sent back via the channel provided
    UpdateState {
        changes: Vec<(DiceKey, ChangeType)>,
        resp: Sender<VersionNumber>,
    },
    /// Gets the current version number
    CurrentVersion { resp: Sender<VersionNumber> },
    /// Obtains the shared state ctx at the given version
    CtxAtVersion {
        version: VersionNumber,
        guard: ActiveTransactionGuard,
        resp: Sender<(SharedLiveTransactionCtx, ActiveTransactionGuard)>,
    },
    /// Report that a computation context at a version has been dropped
    DropCtxAtVersion { version: VersionNumber },
    /// Lookup the state of a key
    LookupKey {
        key: VersionedGraphKey,
        resp: Sender<VersionedGraphResult>,
    },
    /// Report that a value has been computed
    UpdateComputed {
        key: VersionedGraphKey,
        epoch: VersionEpoch,
        /// The storage selection for the key,
        storage: StorageType,
        /// The newly computed value
        value: DiceValidValue,
        /// The deps accessed during the computation of newly computed value
        deps: Arc<Vec<DiceKey>>,
        /// Response of the new value to use. This could be a different instance that is `Eq` to the
        /// given computed value if the state already stores an instance of value that is equal.
        resp: Sender<CancellableResult<DiceComputedValue>>,
    },
    /// Get all the tasks pending cancellation
    GetTasksPendingCancellation {
        #[derivative(Debug = "ignore")]
        resp: Sender<Vec<TerminationObserver>>,
    },
    /// For unstable take
    UnstableDropEverything,
    /// Collect metrics
    Metrics { resp: Sender<Metrics> },
    /// Collects the introspectable dice state
    Introspection {
        resp: Sender<GraphIntrospectable>,
        #[derivative(Debug = "ignore")]
        key_map: HashMap<DiceKey, AnyKey>,
    },
}

/// A handle to the core state that allows sending requests
#[derive(Allocative, Clone)]
pub(crate) struct CoreStateHandle {
    #[allocative(skip)]
    tx: tokio::sync::mpsc::UnboundedSender<StateRequest>,
    // should this handle hold onto the thread and terminate it when all of Dice is dropped?
}

impl CoreStateHandle {
    pub(crate) fn new(tx: tokio::sync::mpsc::UnboundedSender<StateRequest>) -> Self {
        Self { tx }
    }

    pub(crate) fn request(&self, message: StateRequest) {
        self.tx.send(message).expect("dice runner died");
    }
}

impl Dupe for CoreStateHandle {}

/// Start processing state
pub(crate) fn init_state() -> CoreStateHandle {
    StateProcessor::spawn()
}


/*
 * 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 derivative::Derivative;
use derive_more::Display;
use dupe::Dupe;

use crate::impls::cache::SharedCache;
use crate::versions::VersionNumber;
use crate::HashMap;

/// Tracks the currently in-flight versions for updates and reads to ensure
/// values are up to date.
#[derive(Allocative)]
pub(crate) struct VersionTracker {
    current: VersionNumber,
    /// Tracks the currently active versions and how many contexts are holding each of them.
    active_versions: HashMap<VersionNumber, ActiveVersionData>,
    epoch_tracker: VersionEpochTracker,
}

// an epoch counter that increases each time we create a new instance of state at a version. This
// will increase the first time we create a version AND when we reuse a version that has been idle
#[derive(Allocative)]
struct VersionEpochTracker(usize);

impl VersionEpochTracker {
    fn next(&mut self) -> VersionEpoch {
        let e = VersionEpoch(self.0);
        self.0 += 1;

        e
    }

    fn new() -> Self {
        Self(0)
    }
}

#[derive(Copy, Clone, Eq, Debug, Display, Dupe, PartialEq, Allocative)]
#[display(fmt = "v{}", "_0")]
pub(crate) struct VersionEpoch(usize);

impl VersionEpoch {
    #[cfg(test)]
    pub(crate) fn testing_new(e: usize) -> Self {
        Self(e)
    }
}

#[derive(Derivative, Allocative)]
#[derivative(Debug)]
struct ActiveVersionData {
    #[derivative(Debug = "ignore")]
    per_transaction_data: SharedCache,
    ref_count: usize,
    version_epoch: VersionEpoch,
}

impl VersionTracker {
    pub(crate) fn new() -> Self {
        VersionTracker {
            current: VersionNumber::ZERO,
            active_versions: HashMap::default(),
            epoch_tracker: VersionEpochTracker::new(),
        }
    }

    pub(crate) fn currently_active(&self) -> impl Iterator<Item = (usize, &SharedCache)> {
        self.active_versions
            .values()
            .map(|data| (data.ref_count, &data.per_transaction_data))
    }

    /// hands out the current "latest" committed version's associated transaction context
    pub(crate) fn current(&self) -> VersionNumber {
        self.current
    }

    pub(crate) fn at(&mut self, v: VersionNumber) -> (VersionEpoch, SharedCache) {
        let entry = self.active_versions.entry(v).or_insert_with(|| {
            let version_epoch = self.epoch_tracker.next();

            debug!(
                msg = "Creating new shared state",
                v = %v,
                v_epoch = %version_epoch
            );

            ActiveVersionData {
                // TODO properly create the PerLiveTransactionCtx
                per_transaction_data: SharedCache::new(),
                ref_count: 0,
                version_epoch,
            }
        });

        entry.ref_count += 1;

        (entry.version_epoch, entry.per_transaction_data.dupe())
    }

    /// check if the given version and epoch are "relevant", that is the current active version's
    /// epoch matches the given epoch
    pub(crate) fn is_relevant(&self, v: VersionNumber, epoch: VersionEpoch) -> bool {
        self.active_versions
            .get(&v)
            .map_or(false, |active| active.version_epoch == epoch)
    }

    /// Drops reference to a VersionNumber given the token
    pub(crate) fn drop_at_version(&mut self, v: VersionNumber) -> Option<SharedCache> {
        let ref_count = {
            let entry = self
                .active_versions
                .get_mut(&v)
                .expect("shouldn't be able to return version without obtaining one");

            entry.ref_count -= 1;
            entry.ref_count
        };

        if ref_count == 0 {
            let removed = self.active_versions.remove(&v).expect("existed above");

            debug!(
                msg = "shared state removed",
                v = %v,
                v_epoch = %removed.version_epoch
            );
            Some(removed.per_transaction_data)
        } else {
            None
        }
    }

    /// Requests the 'WriteVersion' that is intended to be used for updates to
    /// the incremental computations
    pub(crate) fn write(&mut self) -> VersionForWrites {
        VersionForWrites { tracker: self }
    }
}

pub(crate) struct VersionForWrites<'a> {
    tracker: &'a mut VersionTracker,
}

impl<'a> VersionForWrites<'a> {
    /// Commits the version write and increases the global version number
    pub(crate) fn commit(self) -> VersionNumber {
        self.tracker.current.inc();
        self.tracker.current
    }

    pub(crate) fn version(&self) -> VersionNumber {
        let mut v = self.tracker.current;
        v.inc();

        v
    }

    /// Undo the pending write to version
    pub(crate) fn undo(self) -> VersionNumber {
        self.tracker.current
    }
}

pub(crate) mod introspection {
    use dupe::Dupe;

    use crate::impls::cache::SharedCache;
    use crate::impls::core::versions::VersionTracker;
    use crate::impls::key::DiceKey;
    use crate::introspection::graph::AnyKey;
    use crate::introspection::graph::VersionNumber;
    use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
    use crate::HashMap;

    pub(crate) struct VersionIntrospectable(Vec<(usize, SharedCache)>);

    impl VersionIntrospectable {
        pub(crate) fn versions_currently_running(&self) -> Vec<VersionNumber> {
            self.0.iter().map(|(v, _)| VersionNumber(*v)).collect()
        }

        pub(crate) fn keys_currently_running(
            &self,
            key_map: &HashMap<DiceKey, AnyKey>,
        ) -> Vec<(AnyKey, VersionNumber, DiceTaskStateForDebugging)> {
            self.0
                .iter()
                .flat_map(|(v, cache)| {
                    cache.iter_tasks().map(|(k, state)| {
                        (
                            key_map.get(&k).expect("key should exist").clone(),
                            VersionNumber(*v),
                            state,
                        )
                    })
                })
                .collect()
        }

        pub(crate) fn currently_running_key_count(&self) -> usize {
            self.0
                .iter()
                .flat_map(|(_, cache)| {
                    cache.iter_tasks().filter(|(_, state)| match state {
                        DiceTaskStateForDebugging::AsyncInProgress => true,
                        DiceTaskStateForDebugging::SyncInProgress => true,
                        _ => false,
                    })
                })
                .count()
        }
    }

    impl VersionTracker {
        pub(crate) fn introspect(&self) -> VersionIntrospectable {
            VersionIntrospectable(
                self.currently_active()
                    .map(|(v, cache)| (v, cache.dupe()))
                    .collect(),
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;

    use crate::impls::core::versions::VersionEpoch;
    use crate::impls::core::versions::VersionTracker;
    use crate::versions::VersionNumber;

    #[test]
    fn simple_version_increases() {
        let mut vt = VersionTracker::new();

        let _vg = vt.at(VersionNumber::new(0));
        assert_matches!(
            vt.active_versions.get(&VersionNumber::new(0)), Some(active) if active.ref_count == 1
        );

        let _vg = vt.at(VersionNumber::new(0));
        assert_matches!(
            vt.active_versions.get(&VersionNumber::new(0)), Some(active) if active.ref_count == 2
        );

        vt.drop_at_version(VersionNumber::new(0));
        assert_matches!(
            vt.active_versions.get(&VersionNumber::new(0)), Some(active) if active.ref_count == 1
        );

        vt.drop_at_version(VersionNumber::new(0));
        assert_matches!(vt.active_versions.get(&VersionNumber::new(0)), None);
    }

    #[test]
    fn version_epoch_relevant() {
        let mut vt = VersionTracker::new();

        let (epoch, _s) = vt.at(VersionNumber::new(0));
        assert!(vt.is_relevant(VersionNumber::new(0), epoch));
        assert!(!vt.is_relevant(VersionNumber::new(0), VersionEpoch::testing_new(9999)));

        let (epoch1, _s) = vt.at(VersionNumber::new(2));
        assert!(!vt.is_relevant(VersionNumber::new(0), epoch1));
        assert!(vt.is_relevant(VersionNumber::new(2), epoch1));
        assert!(!vt.is_relevant(VersionNumber::new(2), epoch));

        vt.drop_at_version(VersionNumber::new(2));
        let (epoch2, _s) = vt.at(VersionNumber::new(2));
        assert!(!vt.is_relevant(VersionNumber::new(0), epoch1));
        assert!(!vt.is_relevant(VersionNumber::new(2), epoch1));
        assert!(vt.is_relevant(VersionNumber::new(2), epoch2));
    }

    #[test]
    fn write_version_commits_and_undo() {
        let mut vt = VersionTracker::new();

        {
            let v1 = vt.write();
            assert_eq!(v1.version(), VersionNumber::new(1));
            assert_eq!(v1.version(), VersionNumber::new(1));

            assert_eq!(v1.commit(), VersionNumber::new(1));
            assert_eq!(vt.current(), VersionNumber::new(1));
        }

        {
            let v2 = vt.write();
            assert_eq!(v2.version(), VersionNumber::new(2));

            assert_eq!(v2.undo(), VersionNumber::new(1));
            assert_eq!(vt.current(), VersionNumber::new(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.
 */

//! Represents the forward and backward dependencies of the computation graph

use std::collections::hash_map::Entry;

use allocative::Allocative;
use dupe::Dupe;

use crate::arc::Arc;
use crate::impls::key::DiceKey;
use crate::versions::VersionNumber;
use crate::HashMap;

#[derive(Allocative, Clone)]
pub(crate) struct VersionedDependencies {
    /// once the deps at a particular version is written, it is final and never modified
    /// We only store the dependencies relevant to the most recent result
    recorded_at: VersionNumber,
    deps: Arc<Vec<DiceKey>>,
}

impl Dupe for VersionedDependencies {
    // triomphe is dupe
}

impl VersionedDependencies {
    pub(crate) fn new(recorded_at: VersionNumber, deps: Arc<Vec<DiceKey>>) -> Self {
        Self { recorded_at, deps }
    }

    pub(crate) fn deps(&self) -> Arc<Vec<DiceKey>> {
        self.deps.dupe()
    }

    pub(crate) fn replace_deps(&mut self, v: VersionNumber, deps: Arc<Vec<DiceKey>>) {
        if self.recorded_at < v {
            // we only ever write the newest version of the dependencies of this node for simplicity
            // That way, if we are ever dirtied, we just check if the latest version of the deps
            // have changed at the dirtied version which only requires spawning one set of deps.
            // It might cause us to falsely fail to reuse some nodes, but this is less memory
            // and less work per node when in incremental cases.
            self.deps = deps;
            self.recorded_at = v;
        }
    }
}

// the set of reverse dependencies of a node
#[derive(Allocative, Clone)] // TODO(bobyf) remove need to clone
pub(crate) struct VersionedRevDependencies {
    rdeps: HashMap<DiceKey, VersionNumber>,
}

impl VersionedRevDependencies {
    pub(crate) fn new() -> Self {
        Self {
            rdeps: Default::default(),
        }
    }

    pub(crate) fn add_rdep(&mut self, dependent: DiceKey, current_version: VersionNumber) {
        match self.rdeps.entry(dependent) {
            Entry::Occupied(entry) => {
                if *entry.get() < current_version {
                    entry.replace_entry(current_version);
                }
            }
            Entry::Vacant(v) => {
                v.insert(current_version);
            }
        }
    }

    pub(crate) fn rdeps(&self) -> &HashMap<DiceKey, VersionNumber> {
        &self.rdeps
    }
}


/*
 * 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 gazebo::prelude::SliceExt;

use crate::arc::Arc;
use crate::impls::core::graph::nodes::VersionedGraphNode;
use crate::impls::core::graph::storage::VersionedGraph;
use crate::impls::key::DiceKey;
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::introspection::graph::VersionNumber;
use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;
use crate::HashMap;
use crate::HashSet;

pub struct VersionedGraphIntrospectable {
    nodes: HashMap<AnyKey, SerializedGraphNodesForKey>,
    edges: HashMap<AnyKey, Vec<AnyKey>>,
}

impl VersionedGraphIntrospectable {
    pub(crate) fn keys<'a>(&'a self) -> impl Iterator<Item = AnyKey> + 'a {
        self.nodes.keys().cloned()
    }
    pub(crate) fn edges<'a>(&'a self) -> impl Iterator<Item = (AnyKey, Vec<AnyKey>)> + 'a {
        self.edges.iter().map(|(k, v)| (k.clone(), v.clone()))
    }
    pub(crate) fn nodes<'a>(&'a self) -> impl Iterator<Item = SerializedGraphNodesForKey> + 'a {
        self.nodes.values().cloned()
    }
    pub(crate) fn len_for_introspection(&self) -> usize {
        self.nodes.len()
    }
}

impl VersionedGraph {
    pub(crate) fn introspect(
        &self,
        key_map: HashMap<DiceKey, AnyKey>,
    ) -> VersionedGraphIntrospectable {
        let mut edges = HashMap::default();
        let mut nodes = HashMap::default();

        fn visit_node(key: DiceKey, node: &VersionedGraphNode) -> Option<SerializedGraphNode> {
            match node {
                VersionedGraphNode::Occupied(o) => Some(SerializedGraphNode {
                    node_id: NodeID(key.index as usize),
                    kind: GraphNodeKind::Occupied,
                    history: o.metadata().hist.to_introspectable(),
                    deps: Some(visit_deps(o.metadata().deps.deps())),
                    rdeps: Some(visit_rdeps(o.metadata().rdeps.rdeps())),
                }),
                VersionedGraphNode::Vacant(_) => {
                    // TODO(bobyf) should probably write the metadata of vacant
                    None
                }
            }
        }

        for (k, versioned_nodes) in self.last_n.iter() {
            let dyn_k = key_map.get(k).expect("should be present");

            nodes.insert(
                key_map.get(k).expect("key should exist").clone(),
                SerializedGraphNodesForKey {
                    id: KeyID(k.index as usize),
                    key: dyn_k.to_string(),
                    type_name: dyn_k.short_type_name().to_owned(),
                    nodes: versioned_nodes
                        .iter()
                        .map(|(v, node)| (v.to_introspectable(), visit_node(*k, node)))
                        .collect(),
                },
            );

            if let Some(last) = versioned_nodes.iter().last() {
                edges.insert(
                    dyn_k.clone(),
                    last.1
                        .unpack_occupied()
                        .map(|node| {
                            node.metadata()
                                .deps
                                .deps()
                                .map(|k| key_map.get(k).expect("key should exist").clone())
                        })
                        .unwrap_or_default(),
                );
            }
        }

        fn visit_deps(deps: Arc<Vec<DiceKey>>) -> HashSet<KeyID> {
            deps.iter().map(|d| d.introspect()).collect()
        }

        fn visit_rdeps(
            rdeps: &HashMap<DiceKey, crate::versions::VersionNumber>,
        ) -> BTreeMap<VersionNumber, Vec<NodeID>> {
            let mut res = BTreeMap::<VersionNumber, Vec<NodeID>>::new();

            for (rdep, v) in rdeps {
                res.entry(v.to_introspectable())
                    .or_default()
                    .push(NodeID(rdep.index as usize))
            }

            res
        }
        VersionedGraphIntrospectable { nodes, edges }
    }
}


/*
 * 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 versioned dice graph of dependencies
mod dependencies;
pub(crate) mod history;
#[allow(unused)]
pub(crate) mod introspection;
mod nodes;
pub(crate) mod storage;
pub(crate) mod types;


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

//!
//! A cache that deals with versions
//!
//! This is responsible for performing incremental caching and invalidations
//! with multiple versions in-flight at the same time.
//!
//! The 'VersionedCache' will track dependency edges and use computed version
//! number for each cache entry and a global version counter to determine
//! up-to-date-ness of cache entries.
//!

use allocative::Allocative;
use dupe::Dupe;
use gazebo::variants::UnpackVariants;

use crate::arc::Arc;
use crate::impls::core::graph::dependencies::VersionedDependencies;
use crate::impls::core::graph::dependencies::VersionedRevDependencies;
use crate::impls::core::graph::history::CellHistory;
use crate::impls::key::DiceKey;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceValidValue;
use crate::impls::value::MaybeValidDiceValue;
use crate::versions::VersionNumber;

/// actual entries as seen when querying the cache
/// The placeholder will be used to indicate known dirty entries.
#[derive(UnpackVariants, Allocative)]
pub(crate) enum VersionedGraphNode {
    Occupied(OccupiedGraphNode),
    Vacant(VacantGraphNode),
}

impl VersionedGraphNode {
    pub(crate) fn force_dirty(&mut self, v: VersionNumber) -> bool {
        match self {
            VersionedGraphNode::Occupied(e) => e.metadata.hist.force_dirty(v),
            VersionedGraphNode::Vacant(e) => e.hist.force_dirty(v),
        }
    }

    pub(crate) fn mark_invalidated(&mut self, v: VersionNumber) -> bool {
        match self {
            VersionedGraphNode::Occupied(e) => e.metadata.hist.mark_invalidated(v),
            VersionedGraphNode::Vacant(e) => e.hist.mark_invalidated(v),
        }
    }

    pub(crate) fn history(&self) -> &CellHistory {
        match self {
            VersionedGraphNode::Occupied(o) => &o.metadata().hist,
            VersionedGraphNode::Vacant(v) => &v.hist,
        }
    }
}

/// The stored entry of the cache
#[derive(Allocative, Clone)] // TODO(bobyf) remove need to clone
pub(crate) struct OccupiedGraphNode {
    key: DiceKey,
    res: DiceValidValue,
    metadata: NodeMetadata,
}

/// Meta data about a DICE node, which are its edges and history information
#[derive(Allocative, Clone)] // TODO(bobyf) remove need to clone
pub(crate) struct NodeMetadata {
    pub(crate) deps: VersionedDependencies,
    pub(crate) rdeps: VersionedRevDependencies,
    pub(crate) hist: CellHistory,
}

impl OccupiedGraphNode {
    pub(crate) fn new(
        key: DiceKey,
        res: DiceValidValue,
        deps: VersionedDependencies,
        hist: CellHistory,
    ) -> Self {
        Self {
            key,
            res,
            metadata: NodeMetadata {
                hist,
                deps,
                rdeps: VersionedRevDependencies::new(),
            },
        }
    }

    pub(crate) fn metadata(&self) -> &NodeMetadata {
        &self.metadata
    }

    pub(crate) fn metadata_mut(&mut self) -> &mut NodeMetadata {
        &mut self.metadata
    }

    pub(crate) fn mark_unchanged(
        &mut self,
        v: VersionNumber,
        latest_dep_verified: Option<VersionNumber>,
        first_dep_dirtied: Option<VersionNumber>,
        deps: Arc<Vec<DiceKey>>,
    ) -> VersionNumber {
        // Marking a node as unchanged ALWAYS requires the dependencies for which we used to deem
        // that the node is unchanged.
        //
        // Consider a node n2 that depends on n1 at version v0.
        // We then dirty versions v1, v2, v3 at n1. We'd defer dirtying v2, v3 on n2 due
        // to the fact that it's possible that at v2 and v3, n2 no longer depends on n1
        // and we rely on deferred propagation of dirtiness. However, if at v2, we recompute
        // and find that the values are equal to that at v0, then we can resurrect v0's n2.
        // However, at this point, we will need to deferred propagate the dirty at v3.
        let changed_since =
            self.metadata
                .hist
                .mark_verified_modern(v, latest_dep_verified, first_dep_dirtied);
        self.metadata.deps.replace_deps(changed_since, deps);

        changed_since
    }

    pub(crate) fn val(&self) -> &DiceValidValue {
        &self.res
    }

    pub(crate) fn computed_val(&self) -> DiceComputedValue {
        DiceComputedValue::new(
            MaybeValidDiceValue::valid(self.res.dupe()),
            Arc::new(self.metadata.hist.clone()),
        )
    }
}

/// An entry in the graph that has no computation value associated. This is used to store the
/// history information that is known.
/// This will be replaced by `OccupiedGraphNode` when a computed value is associated with
/// this node. There is no guarantees of when, or even if that will occur since users may never
/// need the associated value at this node.
#[derive(Allocative)]
pub(crate) struct VacantGraphNode {
    pub(crate) key: DiceKey,
    pub(crate) hist: CellHistory,
}

#[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::key::Key;
    use crate::arc::Arc;
    use crate::impls::core::graph::dependencies::VersionedDependencies;
    use crate::impls::core::graph::history::testing::CellHistoryExt;
    use crate::impls::core::graph::history::testing::HistoryExt;
    use crate::impls::core::graph::history::CellHistory;
    use crate::impls::core::graph::nodes::OccupiedGraphNode;
    use crate::impls::key::DiceKey;
    use crate::impls::value::DiceKeyValue;
    use crate::impls::value::DiceValidValue;
    use crate::versions::VersionNumber;

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

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

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

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

    #[test]
    fn update_versioned_graph_entry_tracks_versions_and_deps() {
        let deps0: Arc<Vec<DiceKey>> = Arc::new(vec![DiceKey { index: 5 }]);
        let mut entry = OccupiedGraphNode::new(
            DiceKey { index: 1335 },
            DiceValidValue::testing_new(DiceKeyValue::<K>::new(1)),
            VersionedDependencies::new(VersionNumber::new(0), deps0.dupe()),
            CellHistory::testing_new(
                &[VersionNumber::new(0)],
                &[VersionNumber::new(1), VersionNumber::new(2)],
            ),
        );

        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(0))
            .assert_verified();
        assert_eq!(entry.metadata().deps.deps(), deps0);

        entry.mark_unchanged(VersionNumber::new(1), None, None, Arc::new(vec![]));
        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(0))
            .assert_verified();
        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(1))
            .assert_verified();
        assert_eq!(entry.metadata().deps.deps(), Arc::new(Vec::new()));

        let deps1 = Arc::new(vec![DiceKey { index: 7 }]);
        entry.mark_unchanged(
            VersionNumber::new(2),
            Some(VersionNumber::new(1)),
            None,
            deps1.dupe(),
        );

        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(0))
            .assert_verified();
        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(1))
            .assert_verified();
        entry
            .metadata()
            .hist
            .get_history(&VersionNumber::new(2))
            .assert_verified();

        assert_eq!(entry.metadata().deps.deps(), deps1);
    }
}


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

//!
//! A cache that deals with versions
//!
//! This is responsible for performing incremental caching and invalidations
//! with multiple versions in-flight at the same time.
//!
//! The 'VersionedCache' will track dependency edges and use computed version
//! number for each cache entry and a global version counter to determine
//! up-to-date-ness of cache entries.
//!

use std::cmp;
use std::ops::Bound;

use allocative::Allocative;
use dupe::Dupe;
use sorted_vector_map::SortedVectorMap;

use crate::api::storage_type::StorageType;
use crate::arc::Arc;
use crate::impls::core::graph::dependencies::VersionedDependencies;
use crate::impls::core::graph::history::CellHistory;
use crate::impls::core::graph::history::HistoryState;
use crate::impls::core::graph::nodes::OccupiedGraphNode;
use crate::impls::core::graph::nodes::VacantGraphNode;
use crate::impls::core::graph::nodes::VersionedGraphNode;
use crate::impls::core::graph::types::VersionedGraphKey;
use crate::impls::core::graph::types::VersionedGraphResult;
use crate::impls::core::graph::types::VersionedGraphResultMismatch;
use crate::impls::key::DiceKey;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceValidValue;
use crate::versions::VersionNumber;
use crate::HashMap;

/// The actual incremental cache that checks versions and dependency's versions
/// to maintain correct caching based on versions and the versions of its
/// dependencies.
///
/// TODO refactor this so that the storage doesn't handle multi version, instead the nodes do, and
/// support different node types of different storage persistency.
#[derive(Allocative)]
pub(crate) struct VersionedGraph {
    /// storage that stores every version forever
    /// This storage is implemented so that the map keys are composed of the versions for which
    /// the node changes. Corresponding to each key is a node storing the values and the history.
    /// VacantGraphEntries can only be present when no other entries are present for the key at
    /// any version.
    pub(crate) last_n: HashMap<DiceKey, SortedVectorMap<VersionNumber, VersionedGraphNode>>,
}

impl VersionedGraph {
    pub(crate) fn new() -> Self {
        Self {
            last_n: Default::default(),
        }
    }

    /// gets the cache entry corresponding to the cache entry if up to date.
    /// returns 'None' if entry is missing or versions are out of date.
    pub(crate) fn get(&self, key: VersionedGraphKey) -> VersionedGraphResult {
        fn handle_occupied(
            key: VersionedGraphKey,
            entry: &OccupiedGraphNode,
        ) -> VersionedGraphResult {
            match entry.metadata().hist.get_history(&key.v) {
                HistoryState::Verified => VersionedGraphResult::Match(entry.computed_val()),
                HistoryState::Unknown(verified_versions) => {
                    VersionedGraphResult::CheckDeps(VersionedGraphResultMismatch {
                        entry: entry.val().dupe(),
                        verified_versions,
                        deps_to_validate: entry.metadata().deps.deps(),
                    })
                }
                HistoryState::Dirty => VersionedGraphResult::Compute,
            }
        }

        fn handle_vacant() -> VersionedGraphResult {
            // vacant entries only occur if no other graph entries are
            // present, so we know this has to be dirty
            VersionedGraphResult::Compute
        }

        if let Some(versioned) = self.last_n.get(&key.k) {
            let mut potential = versioned.range((
                Bound::Included(VersionNumber::new(0)),
                Bound::Included(key.v),
            ));
            if let Some(found) = potential.next_back().map(|e| match e.1 {
                VersionedGraphNode::Occupied(entry) => handle_occupied(key, entry),
                VersionedGraphNode::Vacant(_) => handle_vacant(),
            }) {
                found
            } else {
                // this branch takes care of an ongoing computation that is operating on an older
                // version than anything stored currently. However, it has a problem where it's nodes
                // would fail to share work due to nothing going into the cache if its evaluating
                // to a different result. TODO add some per ctx result caching for old versions
                versioned
                    .range((Bound::Included(key.v), Bound::Unbounded))
                    .find_map(|(v, e)| match e {
                        VersionedGraphNode::Occupied(e) => Some((v, e)),
                        VersionedGraphNode::Vacant(_) => None,
                    })
                    .map_or_else(
                        || VersionedGraphResult::Compute,
                        |(_, entry)| {
                            VersionedGraphResult::CheckDeps(VersionedGraphResultMismatch {
                                entry: entry.val().dupe(),
                                verified_versions: entry.metadata().hist.get_verified_ranges(),
                                deps_to_validate: entry.metadata().deps.deps(),
                            })
                        },
                    )
            }
        } else {
            VersionedGraphResult::Compute
        }
    }

    /// gets the cache entry corresponding to the cache entry if up to date.
    /// returns 'None' if entry is missing or versions are out of date.
    fn get_internal<'a>(
        &'a mut self,
        key: VersionedGraphKey,
    ) -> Option<&'a mut VersionedGraphNode> {
        if let Some(versioned) = self.last_n.get_mut(&key.k) {
            let v = versioned
                .range((
                    Bound::Included(VersionNumber::new(0)),
                    Bound::Included(key.v),
                ))
                .next_back()
                .map(|e| *e.0);

            match v {
                None => {
                    // this branch takes care of an ongoing computation that is operating on an older
                    // version than anything stored currently. However, it has a problem where it's nodes
                    // would fail to share work due to nothing going into the cache if its evaluating
                    // to a different result. TODO add some per ctx result caching for old versions
                    versioned
                        .range_mut((Bound::Included(key.v), Bound::Unbounded))
                        .map(|(_v, e)| e)
                        .next()
                }
                Some(v) => versioned.get_mut(&v),
            }
        } else {
            None
        }
    }

    /// updates the cached value based on the given key and versions. The value
    /// is only updated if the version of the new value is of a newer
    /// version than what is stored.
    /// Returns the new entry, and an optional old entry that was invalidated due to this update
    pub(crate) fn update(
        &mut self,
        key: VersionedGraphKey,
        value: DiceValidValue,
        deps: Arc<Vec<DiceKey>>,
        storage_type: StorageType,
    ) -> (DiceComputedValue, bool) {
        let StorageType::LastN(num_to_keep) = storage_type;
        // persistent keys, if any changes, are committed at the moment when the version
        // is increased. therefore, it must be the case that the current update for the
        // persistent key is the largest/newest version. it's also the case that they are
        // never updated to the cache more than once per version.
        // TODO refactor this to be less error prone.

        // we pick the nearest entry because the closest version number to the current key would
        // have the least number of changes recorded in dice, which we assume naively to mean
        // most likely to reuse a node. We could implement this to check for reuse against both
        // the previous and the next version, but that complexity is likely not worth the benefit
        // of trying to reuse a node. Maybe this is worth revisiting at some point.
        let nearest = {
            let versioned_map = self.last_n.entry(key.k).or_default();
            Self::nearest_entry(&key, versioned_map)
        };

        let mut latest_dep_verified = None;
        let mut first_dep_dirtied = None;
        for dep in deps.iter() {
            match self.get_internal(VersionedGraphKey::new(key.v, dep.dupe())) {
                None => {
                    unreachable!("dependency should exist")
                }
                Some(node) => match node {
                    VersionedGraphNode::Occupied(occ) => {
                        if let Some(dep_v) = occ.metadata().hist.latest_verified_before(key.v) {
                            latest_dep_verified = cmp::max(latest_dep_verified, Some(dep_v));

                            let dep_d_v = occ.metadata().hist.first_dirty_after(key.v);
                            first_dep_dirtied = cmp::min(first_dep_dirtied.or(dep_d_v), dep_d_v);

                            occ.metadata_mut().rdeps.add_rdep(key.k, key.v);
                        } else {
                            let dep_d_v = occ.metadata().hist.first_verified_after(key.v);
                            first_dep_dirtied = cmp::min(first_dep_dirtied.or(dep_d_v), dep_d_v);
                        }
                    }
                    VersionedGraphNode::Vacant(_) => {
                        unreachable!("dependency should exist")
                    }
                },
            }
        }

        if let Some(key_of_e) = nearest {
            self.update_entry(
                key_of_e,
                key,
                value,
                first_dep_dirtied,
                latest_dep_verified,
                deps,
                num_to_keep,
            )
        } else {
            (
                self.update_empty(
                    key.k,
                    key.v,
                    value,
                    first_dep_dirtied,
                    latest_dep_verified,
                    deps,
                ),
                true,
            )
        }
    }

    /// find the nearest entry to the given key, preferring the smaller version number when tied
    fn nearest_entry<'a>(
        key: &VersionedGraphKey,
        versioned_map: &'a SortedVectorMap<VersionNumber, VersionedGraphNode>,
    ) -> Option<VersionNumber> {
        let newest_previous = versioned_map
            .range((
                Bound::Included(VersionNumber::new(0)),
                Bound::Included(key.v),
            ))
            .next_back()
            .map(|(v, _e)| *v);
        let oldest_newer = versioned_map
            .range((Bound::Included(key.v), Bound::Unbounded))
            .next()
            .map(|(v, _e)| *v);

        match (newest_previous, oldest_newer) {
            (Some(prev_v), Some(next_v)) => {
                if next_v - key.v < prev_v - key.v {
                    Some(next_v)
                } else {
                    Some(prev_v)
                }
            }
            (Some(x), None) => Some(x),
            (None, Some(x)) => Some(x),
            (None, None) => None,
        }
    }

    fn update_empty(
        &mut self,
        key: DiceKey,
        v: VersionNumber,
        value: DiceValidValue,
        first_dep_dirtied: Option<VersionNumber>,
        latest_dep_verified: Option<VersionNumber>,

        deps: Arc<Vec<DiceKey>>,
    ) -> DiceComputedValue {
        let since = latest_dep_verified.unwrap_or(v);
        let mut hist = CellHistory::verified(since);
        hist.propagate_from_deps_version(since, first_dep_dirtied);
        let entry =
            OccupiedGraphNode::new(key, value, VersionedDependencies::new(since, deps), hist);

        let res = entry.computed_val();

        self.last_n
            .get_mut(&key)
            .unwrap()
            .insert(v, VersionedGraphNode::Occupied(entry));

        res
    }

    /// Returns the newly updated value for the key, and whether or not any state changed.
    fn update_entry(
        &mut self,
        key_of_e: VersionNumber,
        key: VersionedGraphKey,
        value: DiceValidValue,
        first_dep_dirtied: Option<VersionNumber>,
        latest_dep_verified: Option<VersionNumber>,
        deps: Arc<Vec<DiceKey>>,
        num_to_keep: usize,
    ) -> (DiceComputedValue, bool) {
        let versioned_map = self.last_n.get_mut(&key.k).unwrap();
        let (ret, map_fixup) = match versioned_map.get_mut(&key_of_e).unwrap() {
            VersionedGraphNode::Occupied(entry) if value.equality(entry.val()) => {
                let since =
                    entry.mark_unchanged(key.v, latest_dep_verified, first_dep_dirtied, deps);

                let ret = entry.computed_val();

                (ret, MapFixup::Reused { since, key_of_e })
            }
            entry => {
                let (since, end, mut hist) = entry
                    .history()
                    .make_new_verified_history(key.v, latest_dep_verified);

                hist.propagate_from_deps_version(key.v, first_dep_dirtied);

                let new = OccupiedGraphNode::new(
                    key.k,
                    value,
                    VersionedDependencies::new(since, deps),
                    hist,
                );

                let ret = new.computed_val();

                (
                    ret,
                    MapFixup::NewEntry {
                        since,
                        end,
                        new,
                        key_of_e,
                        num_to_keep,
                    },
                )
            }
        };

        let any_invalidated = map_fixup.fixup(versioned_map);

        (ret, any_invalidated)
    }

    /// Invalidates an entry and its transitive rdeps. Returning true if this caused any type of
    /// change
    pub(crate) fn invalidate(
        &mut self,
        key: VersionedGraphKey,
        invalidate: InvalidateKind,
    ) -> bool {
        let rdeps = {
            match invalidate {
                invalidate @ (InvalidateKind::ForceDirty | InvalidateKind::Invalidate) => {
                    let versioned_map = self
                        .last_n
                        .entry(key.k)
                        .or_insert_with(SortedVectorMap::new);
                    if let Some(e) = versioned_map
                        .range_mut((Bound::Unbounded, Bound::Included(key.v)))
                        .next_back()
                        .map(|(_, e)| e)
                    {
                        let dirtied = match invalidate {
                            InvalidateKind::ForceDirty => e.force_dirty(key.v),
                            InvalidateKind::Invalidate => e.mark_invalidated(key.v),
                            _ => unreachable!("handled elsewhere"),
                        };

                        if dirtied {
                            if let Some(e) = e.unpack_occupied() {
                                let queue = {
                                    let metadata = e.metadata();
                                    let rdeps = metadata.rdeps.rdeps();

                                    rdeps
                                        .iter()
                                        .map(|(r, v)| (r.dupe(), *v))
                                        .collect::<Vec<_>>()
                                };

                                queue
                            } else {
                                return true;
                            }
                        } else {
                            return false;
                        }
                    } else {
                        let mut entry = VersionedGraphNode::Vacant(VacantGraphNode {
                            key: key.k,
                            hist: CellHistory::empty(),
                        });

                        entry.mark_invalidated(key.v);

                        versioned_map.insert(key.v, entry);

                        return true;
                    }
                }
                InvalidateKind::Update(value, StorageType::LastN(num_to_keep)) => {
                    let rdeps = {
                        let entry = self.last_n.get(&key.k).and_then(|versioned_map| {
                            versioned_map
                                .range((Bound::Unbounded, Bound::Included(key.v)))
                                .next_back()
                                .map(|(_, e)| e)
                        });

                        match entry {
                            Some(VersionedGraphNode::Occupied(occ)) => {
                                if !occ.val().equality(&value) {
                                    occ.metadata()
                                        .rdeps
                                        .rdeps()
                                        .iter()
                                        .map(|(r, v)| (r.dupe(), *v))
                                        .collect::<Vec<_>>()
                                } else {
                                    return false;
                                }
                            }
                            _ => vec![],
                        }
                    };

                    let versioned_map = self.last_n.entry(key.k).or_default();
                    let fixup = if let Some((key_of_e, entry)) = versioned_map
                        .range_mut((Bound::Unbounded, Bound::Included(key.v)))
                        .next_back()
                    {
                        let (since, end, mut hist) =
                            entry.history().make_new_verified_history(key.v, None);

                        hist.propagate_from_deps_version(key.v, None);

                        let new = OccupiedGraphNode::new(
                            key.k,
                            value,
                            VersionedDependencies::new(since, Arc::new(vec![])),
                            hist,
                        );

                        MapFixup::NewEntry {
                            since,
                            end,
                            new,
                            key_of_e: *key_of_e,
                            num_to_keep,
                        }
                    } else {
                        let entry = VersionedGraphNode::Occupied(OccupiedGraphNode::new(
                            key.k,
                            value,
                            VersionedDependencies::new(key.v, Arc::new(vec![])),
                            CellHistory::verified(key.v),
                        ));

                        versioned_map.insert(key.v, entry);

                        return true;
                    };

                    fixup.fixup(versioned_map);

                    rdeps
                }
            }
        };

        self.invalidate_rdeps(key.v, rdeps);
        true
    }

    fn invalidate_rdeps(
        &mut self,
        version: VersionNumber,
        mut queue: Vec<(DiceKey, VersionNumber)>,
    ) {
        while let Some((rdep, relevant_version)) = queue.pop() {
            if let Some(node) = self.get_internal(VersionedGraphKey::new(relevant_version, rdep)) {
                if node.mark_invalidated(version) {
                    // since dirty always occurs in increasing order, it must be the case that if
                    // the history was already dirtied, it was by a version number less than the
                    // current version number.
                    // furthermore, if the rdep was dirtied, at any future versions larger than
                    // the version it was dirtied at, it may no longer depend on the current node
                    // so we skip marking it as dirty, and rely on delayed propagation of dirty

                    if let Some(node) = node.unpack_occupied() {
                        queue.extend({
                            let rdeps = node.metadata().rdeps.rdeps();

                            rdeps
                                .iter()
                                .map(|(r, v)| (r.dupe(), *v))
                                .collect::<Vec<_>>()
                        })
                    }
                }
            }
        }
    }
}

pub(crate) enum InvalidateKind {
    ForceDirty,
    #[allow(unused)] // constructed for tests
    Invalidate,
    Update(DiceValidValue, StorageType),
}

// because of rust mut lifetimes, we have to fixup the entries in our map after we drop the
// references to the nodes we modify. This forces us to express things via these enums
// that approximately acts as lambdas
enum MapFixup {
    Reused {
        since: VersionNumber,
        key_of_e: VersionNumber,
    },
    NewEntry {
        since: VersionNumber,
        end: Option<VersionNumber>,
        new: OccupiedGraphNode,
        key_of_e: VersionNumber,
        num_to_keep: usize,
    },
}

impl MapFixup {
    fn fixup(self, versioned_map: &mut SortedVectorMap<VersionNumber, VersionedGraphNode>) -> bool {
        match self {
            MapFixup::Reused { since, key_of_e } => {
                if since < key_of_e {
                    let old = versioned_map.remove(&key_of_e).unwrap();
                    versioned_map.insert(since, old);
                }
                false
            }
            MapFixup::NewEntry {
                since,
                end,
                new,
                key_of_e,
                num_to_keep,
            } => {
                match versioned_map.get(&key_of_e).unwrap() {
                    VersionedGraphNode::Occupied(occ) => {
                        if let Some(end) = end {
                            // if there is newer data, we also need to store that at a newer
                            // key to make it reachable.
                            // TODO(bobyf): we probably want a custom versioned map here to
                            // better represent this and reduce complexity

                            if versioned_map.len() == num_to_keep {
                                // if we are already at max entries to store, then we should
                                // just skip doing this entirely, as the most up to date
                                // entry we will store will be the entry at "end", which
                                // is no different than the original entry.
                                // We also don't need to store rdeps since this node will be discarded
                                return true;
                            }

                            // TODO change storage so that we don't clone
                            versioned_map.insert(end, VersionedGraphNode::Occupied(occ.clone()));
                        }

                        if versioned_map.len() == num_to_keep {
                            let min_version_stored = *versioned_map.iter().next().expect("should be at least one entry if there is more entries than what we want to keep").0;

                            if since < min_version_stored {
                                return true;
                            }

                            versioned_map.remove(&min_version_stored);
                        }

                        versioned_map.insert(since, VersionedGraphNode::Occupied(new));

                        true
                    }
                    VersionedGraphNode::Vacant(_) => {
                        // remove the vacant entry since we now have an actual entry
                        versioned_map.remove(&key_of_e);
                        versioned_map.insert(since, VersionedGraphNode::Occupied(new));

                        true
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub(crate) mod testing {

    use gazebo::variants::VariantName;

    use crate::impls::core::graph::storage::VersionedGraphResult;
    use crate::impls::core::graph::storage::VersionedGraphResultMismatch;
    use crate::impls::value::DiceComputedValue;

    pub(crate) trait VersionedCacheResultAssertsExt {
        fn assert_compute(&self);

        fn assert_match(&self) -> &DiceComputedValue;

        fn assert_check_deps(&self) -> &VersionedGraphResultMismatch;
    }

    impl VersionedCacheResultAssertsExt for VersionedGraphResult {
        fn assert_compute(&self) {
            self.unpack_compute()
                .unwrap_or_else(|| panic!("expected Compute, but was {}", self.variant_name()))
        }

        fn assert_match(&self) -> &DiceComputedValue {
            self.unpack_match()
                .unwrap_or_else(|| panic!("expected Match, but was {}", self.variant_name()))
        }

        fn assert_check_deps(&self) -> &VersionedGraphResultMismatch {
            self.unpack_check_deps()
                .unwrap_or_else(|| panic!("expected Mismatch, but was {}", self.variant_name()))
        }
    }
}

#[cfg(test)]
mod tests {
    use std::fmt::Debug;
    use std::hash::Hash;

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

    use crate::api::computations::DiceComputations;
    use crate::api::key::Key;
    use crate::arc::Arc;
    use crate::impls::core::graph::storage::testing::VersionedCacheResultAssertsExt;
    use crate::impls::core::graph::storage::InvalidateKind;
    use crate::impls::core::graph::storage::StorageType;
    use crate::impls::core::graph::storage::VersionedGraph;
    use crate::impls::core::graph::types::VersionedGraphKey;
    use crate::impls::key::DiceKey;
    use crate::impls::value::DiceKeyValue;
    use crate::impls::value::DiceValidValue;
    use crate::versions::testing::VersionRangesExt;
    use crate::versions::VersionNumber;
    use crate::versions::VersionRange;
    use crate::versions::VersionRanges;

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

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

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

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

    #[test]
    fn latest_only_stores_latest_only() {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));
        let key = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });

        // first, empty cache gives none
        cache.get(key).assert_compute();

        assert!(
            cache
                .update(
                    key.dupe(),
                    res.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );

        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        let res2 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(200));

        let key2 = VersionedGraphKey::new(VersionNumber::new(2), DiceKey { index: 0 });
        assert!(cache.invalidate(key2.dupe(), InvalidateKind::Invalidate));
        assert!(
            cache
                .update(
                    key2.dupe(),
                    res2.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );

        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        // old version is gone
        let entry = cache.get(key.dupe());
        let mismatch = entry.assert_check_deps();
        assert!(mismatch.entry.equality(&res2));
        assert_eq!(
            mismatch.verified_versions,
            VersionRanges::testing_new(sorted_vector_set![VersionRange::begins_with(
                VersionNumber::new(2),
            )])
        );

        // if the value is the same, then versions are shared
        let res3 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(200));
        let key3 = VersionedGraphKey::new(VersionNumber::new(5), DiceKey { index: 0 });
        let key4 = VersionedGraphKey::new(VersionNumber::new(4), DiceKey { index: 0 });
        assert!(cache.invalidate(key4.dupe(), InvalidateKind::Invalidate));
        assert!(cache.invalidate(key3.dupe(), InvalidateKind::Invalidate));
        assert!(
            !cache
                .update(
                    key3.dupe(),
                    res3.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );

        assert!(
            cache
                .get(key3.dupe(),)
                .assert_match()
                .value()
                .equality(&res2)
        );
        assert!(
            cache
                .get(key2.dupe(),)
                .assert_match()
                .value()
                .equality(&res2)
        );
        // the first result is gone still
        let entry = cache.get(key.dupe());
        let mismatch = entry.assert_check_deps();
        assert!(mismatch.entry.equality(&res2));
        assert_eq!(
            mismatch.verified_versions,
            VersionRanges::testing_new(sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(4)),
                VersionRange::begins_with(VersionNumber::new(5))
            ])
        );

        // smaller version numbers don't get cached
        let res4 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(400));
        assert!(
            cache
                .update(
                    key4.dupe(),
                    res4.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1),
                )
                .1
        );
        let entry = cache.get(key4.dupe());
        let mismatch = entry.assert_check_deps();
        assert!(mismatch.entry.equality(&res2));
        assert_eq!(
            mismatch.verified_versions,
            VersionRanges::testing_new(sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(4)),
                VersionRange::begins_with(VersionNumber::new(5))
            ])
        );

        assert!(
            cache
                .get(key3.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        // the first result is gone still
        let entry = cache.get(key.dupe());
        let mismatch = entry.assert_check_deps();
        assert!(mismatch.entry.equality(&res2));
        assert_eq!(
            mismatch.verified_versions,
            VersionRanges::testing_new(sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(4)),
                VersionRange::begins_with(VersionNumber::new(5))
            ])
        );
        // different key is miss
        cache
            .get(VersionedGraphKey::new(
                VersionNumber::new(5),
                DiceKey { index: 1000 },
            ))
            .assert_compute();

        let key5 = VersionedGraphKey::new(VersionNumber::new(7), DiceKey { index: 0 });
        assert!(cache.invalidate(key5, InvalidateKind::ForceDirty));
        cache.get(key5.dupe()).assert_compute()
    }

    #[test]
    fn last_n_max_usize_stores_everything() {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));
        let key = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 0 });

        assert!(
            cache
                .update(
                    key,
                    res.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(usize::MAX)
                )
                .1
        );

        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        let res2 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(200));
        let key2 = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });
        assert!(cache.invalidate(key2.dupe(), InvalidateKind::Invalidate));
        assert!(
            cache
                .update(
                    key2.dupe(),
                    res2.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(usize::MAX)
                )
                .1
        );

        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        // skip a few versions
        let res3 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(300));
        let key3 = VersionedGraphKey::new(VersionNumber::new(5), DiceKey { index: 0 });
        let key2 = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });
        assert!(cache.invalidate(key3.dupe(), InvalidateKind::Invalidate));
        assert!(
            cache
                .update(
                    key3,
                    res3.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(usize::MAX)
                )
                .1
        );

        assert!(
            cache
                .get(key3.dupe())
                .assert_match()
                .value()
                .equality(&res3)
        );
        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        // keys goes to the largest version that's smaller than it
        let key4 = VersionedGraphKey::new(VersionNumber::new(4), DiceKey { index: 0 });
        assert!(
            cache
                .get(key4.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );

        let key5 = VersionedGraphKey::new(VersionNumber::new(6), DiceKey { index: 0 });
        assert!(
            cache
                .get(key5.dupe())
                .assert_match()
                .value()
                .equality(&res3)
        );

        // different key is none
        let key6 = VersionedGraphKey::new(VersionNumber::new(6), DiceKey { index: 100 });
        cache.get(key6.dupe()).assert_compute();

        let key7 = VersionedGraphKey::new(VersionNumber::new(7), DiceKey { index: 0 });
        assert!(cache.invalidate(key7.dupe(), InvalidateKind::ForceDirty));
        cache.get(key7.dupe()).assert_compute()
    }

    #[tokio::test]
    async fn last_2_stores_last_2() {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));
        let key = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 0 });

        assert!(
            cache
                .update(
                    key.dupe(),
                    res.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(2)
                )
                .1
        );

        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        let res2 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(200));
        let key2 = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });
        assert!(cache.invalidate(key2.dupe(), InvalidateKind::Invalidate));
        assert!(
            cache
                .update(
                    key2.dupe(),
                    res2.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(2)
                )
                .1
        );

        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        // skip a few versions
        let res3 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(300));
        let key3 = VersionedGraphKey::new(VersionNumber::new(5), DiceKey { index: 0 });
        let key2 = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });
        assert!(cache.invalidate(key3.dupe(), InvalidateKind::Invalidate));
        assert!(
            cache
                .update(
                    key3.dupe(),
                    res3.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(2)
                )
                .1
        );

        assert!(
            cache
                .get(key3.dupe())
                .assert_match()
                .value()
                .equality(&res3)
        );
        assert!(
            cache
                .get(key2.dupe())
                .assert_match()
                .value()
                .equality(&res2)
        );
        // the oldest entry should be evicted because we don't store more than 2
        let entry = cache.get(key.dupe());
        let mismatch = entry.assert_check_deps();
        assert!(mismatch.entry.equality(&res2));
        assert_eq!(
            mismatch.verified_versions,
            VersionRanges::testing_new(sorted_vector_set![VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(5)
            )])
        )
    }

    #[test]
    fn test_dirty_for_persistent_storage() {
        fn key(v: usize) -> VersionedGraphKey {
            VersionedGraphKey::new(VersionNumber::new(v), DiceKey { index: 0 })
        }

        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));

        let existing = cache.invalidate(key(0), InvalidateKind::Invalidate);
        assert!(existing);

        cache.get(key(0).dupe()).assert_compute();
        cache.get(key(1).dupe()).assert_compute();

        let existing = cache.invalidate(key(2), InvalidateKind::Invalidate);
        assert!(existing);

        cache.get(key(0).dupe()).assert_compute();
        cache.get(key(1).dupe()).assert_compute();
        cache.get(key(2).dupe()).assert_compute();

        cache.update(
            key(0),
            res.dupe(),
            Arc::new(vec![]),
            StorageType::LastN(usize::MAX),
        );
        assert!(
            cache
                .get(key(0).dupe())
                .assert_match()
                .value()
                .equality(&res)
        );
        assert!(
            cache
                .get(key(1).dupe())
                .assert_match()
                .value()
                .equality(&res)
        );
        cache.get(key(2)).assert_check_deps();
    }

    #[test]
    fn test_dirty_for_nonpersistent_storage() {
        fn key(v: usize) -> VersionedGraphKey {
            VersionedGraphKey::new(VersionNumber::new(v), DiceKey { index: 1 })
        }

        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));

        let existing = cache.invalidate(key(0), InvalidateKind::Invalidate);
        assert!(existing);

        cache.get(key(0).dupe()).assert_compute();
        cache.get(key(1).dupe()).assert_compute();

        let existing = cache.invalidate(key(2), InvalidateKind::Invalidate);
        assert!(existing);

        cache.get(key(0).dupe()).assert_compute();
        cache.get(key(1).dupe()).assert_compute();
        cache.get(key(2).dupe()).assert_compute();

        cache.update(key(0), res.dupe(), Arc::new(vec![]), StorageType::LastN(1));
        assert!(
            cache
                .get(key(0).dupe())
                .assert_match()
                .value()
                .equality(&res)
        );
        assert!(
            cache
                .get(key(1).dupe())
                .assert_match()
                .value()
                .equality(&res)
        );
        cache.get(key(2).dupe()).assert_check_deps();
    }

    #[test]
    fn reuse_inserts_into_cache() {
        // This tests a very specific condition of resurrecting a value.
        // Consider a node n at version v0 that was dirtied at v1, v2.
        // It was evaluated at v1, resulting in a different value, but at v2, it results in the same
        // value as v0.
        // It is possible that we attempt to resurrect the entry from v0 and v2, which actually
        // requires actually requires insertion of a new entry at v2, rather than simply marking
        // v0 as reusable.

        let mut cache = VersionedGraph::new();

        let key1 = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 0 });
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(1));

        let value = cache.update(key1, res.dupe(), Arc::new(vec![]), StorageType::LastN(1));

        let key2 = VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 });
        let res2 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(2));

        cache.update(key2, res2.dupe(), Arc::new(vec![]), StorageType::LastN(1));

        let key3 = VersionedGraphKey::new(VersionNumber::new(2), DiceKey { index: 0 });
        let res3 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(1));
        let value3 = cache.update(
            key3.dupe(),
            res3.dupe(),
            Arc::new(vec![]),
            StorageType::LastN(1),
        );

        // should have created a new entry because of key2
        #[allow(clippy::vtable_address_comparisons)] // this should be same exact ptr copy
        let is_same_ptr = !std::sync::Arc::ptr_eq(
            value.0.value().testing_value(),
            value3.0.value().testing_value(),
        );
        assert!(is_same_ptr);
        // should actually be cached though
        cache.get(key3).assert_match();
    }

    #[test]
    fn update_prior_version_reuses_nodes_correctly() {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));

        let key = VersionedGraphKey::new(VersionNumber::new(5), DiceKey { index: 0 });

        // first, empty cache gives none
        cache.get(key.dupe()).assert_compute();

        assert!(
            cache
                .update(
                    key.dupe(),
                    res.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );

        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));

        // now insert a new value of a older version, this shouldn't evict anything.
        let res2 = DiceValidValue::testing_new(DiceKeyValue::<K>::new(200));
        let key2 = VersionedGraphKey::new(VersionNumber::new(4), DiceKey { index: 0 });
        assert!(
            cache
                .update(
                    key2.dupe(),
                    res2.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );
        cache.get(key2.dupe()).assert_check_deps();
        // the newer version should still be there
        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));
        // there should be size 1
        assert_eq!(cache.last_n.get(&DiceKey { index: 0 }).unwrap().len(), 1);

        // now insert the same value of a older version, this shouldn't evict anything but reuses
        // the existing node.
        let key3 = VersionedGraphKey::new(VersionNumber::new(3), DiceKey { index: 0 });
        assert!(
            !cache
                .update(
                    key3.dupe(),
                    res.dupe(),
                    Arc::new(vec![]),
                    StorageType::LastN(1)
                )
                .1
        );

        assert!(cache.get(key.dupe()).assert_match().value().equality(&res));
        assert!(cache.get(key3.dupe()).assert_match().value().equality(&res));
    }

    #[test]
    fn dirty_invalidates_rdeps() -> anyhow::Result<()> {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));

        let key = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 0 });
        cache.update(key, res.dupe(), Arc::new(vec![]), StorageType::LastN(1));

        let key1 = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 1 });
        cache.update(
            key1,
            res.dupe(),
            Arc::new(vec![DiceKey { index: 0 }]),
            StorageType::LastN(1),
        );

        let key2 = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 2 });
        cache.update(
            key2,
            res.dupe(),
            Arc::new(vec![DiceKey { index: 0 }]),
            StorageType::LastN(1),
        );

        assert!(cache.invalidate(
            VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 }),
            InvalidateKind::ForceDirty
        ));

        assert_eq!(
            cache
                .get(VersionedGraphKey::new(
                    VersionNumber::new(1),
                    DiceKey { index: 1 }
                ))
                .assert_check_deps()
                .deps_to_validate,
            Arc::new(vec![DiceKey { index: 0 }])
        );
        assert_eq!(
            cache
                .get(VersionedGraphKey::new(
                    VersionNumber::new(1),
                    DiceKey { index: 2 }
                ))
                .assert_check_deps()
                .deps_to_validate,
            Arc::new(vec![DiceKey { index: 0 }])
        );

        Ok(())
    }

    #[test]
    fn dirty_same_nodes() -> anyhow::Result<()> {
        let mut cache = VersionedGraph::new();
        let res = DiceValidValue::testing_new(DiceKeyValue::<K>::new(100));

        let key = VersionedGraphKey::new(VersionNumber::new(0), DiceKey { index: 0 });
        cache.update(key, res.dupe(), Arc::new(vec![]), StorageType::LastN(1));

        assert!(!cache.invalidate(
            VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 }),
            InvalidateKind::Update(res.dupe(), StorageType::LastN(1))
        ));

        let key = VersionedGraphKey::new(VersionNumber::new(2), DiceKey { index: 0 });
        cache.update(key, res.dupe(), Arc::new(vec![]), StorageType::LastN(1));

        assert!(cache.invalidate(
            VersionedGraphKey::new(VersionNumber::new(1), DiceKey { index: 0 }),
            InvalidateKind::Update(
                DiceValidValue::testing_new(DiceKeyValue::<K>::new(30)),
                StorageType::LastN(1)
            )
        ));

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

//! Type definitions for objects relating to the Versioned Graph

use dupe::Dupe;
use gazebo::variants::UnpackVariants;
use gazebo::variants::VariantName;

use crate::arc::Arc;
use crate::impls::key::DiceKey;
use crate::impls::value::DiceComputedValue;
use crate::impls::value::DiceValidValue;
use crate::versions::VersionNumber;
use crate::versions::VersionRanges;

/// The Key for a Versioned, incremental computation
#[derive(Copy, Clone, Dupe, Debug)]
pub(crate) struct VersionedGraphKey {
    pub(crate) v: VersionNumber,
    pub(crate) k: DiceKey,
}

impl VersionedGraphKey {
    pub(crate) fn new(v: VersionNumber, k: DiceKey) -> Self {
        VersionedGraphKey { v, k }
    }
}

#[derive(Debug)]
pub(crate) struct VersionedGraphResultMismatch {
    /// Last known value for the key.
    pub(crate) entry: DiceValidValue,
    /// Versions at which the value for given key is valid.
    pub(crate) verified_versions: VersionRanges,
    pub(crate) deps_to_validate: Arc<Vec<DiceKey>>,
}

#[derive(Debug, VariantName, UnpackVariants)]
pub(crate) enum VersionedGraphResult {
    /// when the version cache has the exact matching entry via versions
    Match(DiceComputedValue),
    /// when the version cache found an entry, but the versions were mismatching. The existing entry
    /// is returned, along with the last known version
    CheckDeps(VersionedGraphResultMismatch),
    /// An entry that is known to require re-evaluation because it was marked as dirty at the
    /// requested version or that it was missing
    Compute,
}

#[cfg(test)]
pub(crate) mod testing {
    use crate::gazebo::variants::VariantName;
    use crate::impls::core::graph::types::VersionedGraphResult;
    use crate::impls::core::graph::types::VersionedGraphResultMismatch;
    use crate::impls::value::DiceComputedValue;

    pub(crate) trait VersionedCacheResultAssertsExt {
        fn assert_compute(&self);

        fn assert_match(&self) -> &DiceComputedValue;

        fn assert_mismatch(&self) -> &VersionedGraphResultMismatch;
    }

    impl VersionedCacheResultAssertsExt for VersionedGraphResult {
        fn assert_compute(&self) {
            self.unpack_compute()
                .unwrap_or_else(|| panic!("expected Compute, but was {}", self.variant_name()))
        }
        fn assert_match(&self) -> &DiceComputedValue {
            self.unpack_match()
                .unwrap_or_else(|| panic!("expected Match, but was {}", self.variant_name()))
        }

        fn assert_mismatch(&self) -> &VersionedGraphResultMismatch {
            self.unpack_check_deps()
                .unwrap_or_else(|| panic!("expected Mismatch, but was {}", self.variant_name()))
        }
    }
}


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

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

use std::borrow::Cow;
use std::fmt::Debug;

use allocative::Allocative;
use dupe::Dupe;
use futures::stream::FuturesUnordered;
use futures::FutureExt;
use futures::StreamExt;
use tokio::sync::oneshot;

use crate::api::activation_tracker::ActivationData;
use crate::arc::Arc;
use crate::impls::core::graph::history::CellHistory;
use crate::impls::core::graph::types::VersionedGraphKey;
use crate::impls::core::graph::types::VersionedGraphResult;
use crate::impls::core::state::CoreStateHandle;
use crate::impls::core::state::StateRequest;
use crate::impls::core::versions::VersionEpoch;
use crate::impls::evaluator::AsyncEvaluator;
use crate::impls::evaluator::SyncEvaluator;
use crate::impls::events::DiceEventDispatcher;
use crate::impls::key::DiceKey;
use crate::impls::key::ParentKey;
use crate::impls::key_index::DiceKeyIndex;
use crate::impls::task::dice::DiceTask;
use crate::impls::task::promise::DicePromise;
use crate::impls::task::PreviouslyCancelledTask;
use crate::impls::user_cycle::UserCycleDetectorData;
use crate::impls::value::DiceComputedValue;
use crate::impls::worker::state::DiceWorkerStateCheckingDeps;
use crate::impls::worker::state::DiceWorkerStateComputing;
use crate::impls::worker::state::DiceWorkerStateFinishedAndCached;
use crate::impls::worker::state::DiceWorkerStateLookupNode;
use crate::impls::worker::DiceTaskWorker;
use crate::result::CancellableResult;
use crate::result::Cancelled;
use crate::versions::VersionNumber;
use crate::versions::VersionRanges;
use crate::ActivationTracker;

#[cfg(test)]
mod tests;

/// The incremental engine that manages all the handling of the results of a
/// specific key, performing the recomputation if necessary
///
/// The computation of an identical request (same key and version) is
/// automatically deduplicated, so that identical requests share the same set of
/// work. It is guaranteed that there is at most one computation in flight at a
/// time if they share the same key and version.
#[derive(Allocative)]
pub(crate) struct IncrementalEngine {
    state: CoreStateHandle,
    pub(crate) version_epoch: VersionEpoch,
}

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

impl IncrementalEngine {
    fn new(state: CoreStateHandle, version_epoch: VersionEpoch) -> Self {
        Self {
            state,
            version_epoch,
        }
    }

    pub(crate) fn spawn_for_key(
        k: DiceKey,
        version_epoch: VersionEpoch,
        eval: AsyncEvaluator,
        cycles: UserCycleDetectorData,
        events_dispatcher: DiceEventDispatcher,
        previously_cancelled_task: Option<PreviouslyCancelledTask>,
    ) -> DiceTask {
        let state_handle = eval.dice.state_handle.dupe();
        DiceTaskWorker::spawn(
            k,
            eval,
            cycles,
            events_dispatcher,
            previously_cancelled_task,
            IncrementalEngine::new(state_handle, version_epoch),
        )
    }

    #[instrument(
        level = "debug",
        skip(state, promise, eval, event_dispatcher),
        fields(k = ?k, version = %v),
    )]
    pub(crate) fn project_for_key(
        state: CoreStateHandle,
        promise: DicePromise,
        k: DiceKey,
        v: VersionNumber,
        version_epoch: VersionEpoch,
        eval: SyncEvaluator,
        event_dispatcher: DiceEventDispatcher,
    ) -> CancellableResult<DiceComputedValue> {
        promise.get_or_complete(|| {
            event_dispatcher.started(k);

            debug!(msg = "running projection");

            let eval_result = eval.evaluate(k);

            debug!(msg = "projection finished. updating caches");

            let res = {
                // send the update but don't wait for it
                match eval_result.value.dupe().into_valid_value() {
                    Ok(value) => {
                        let (tx, _rx) = tokio::sync::oneshot::channel();
                        state.request(StateRequest::UpdateComputed {
                            key: VersionedGraphKey::new(v, k),
                            epoch: version_epoch,
                            storage: eval_result.storage,
                            value,
                            deps: Arc::new(eval_result.deps.into_iter().collect()),
                            resp: tx,
                        });
                        // TODO(bobyf) consider if we want to block and wait for the cache
                    }
                    Err(_) => {}
                }

                eval_result.value
            };

            debug!(msg = "update future completed");
            event_dispatcher.finished(k);

            DiceComputedValue::new(res, Arc::new(CellHistory::verified(v)))
        })
    }

    pub(crate) async fn eval_entry_versioned(
        &self,
        k: DiceKey,
        eval: &AsyncEvaluator,
        events_dispatcher: DiceEventDispatcher,
        task_state: DiceWorkerStateLookupNode<'_>,
    ) -> CancellableResult<DiceWorkerStateFinishedAndCached> {
        let v = eval.per_live_version_ctx.get_version();
        let (tx, rx) = oneshot::channel();
        self.state.request(StateRequest::LookupKey {
            key: VersionedGraphKey::new(v, k),
            resp: tx,
        });

        let state_result = rx.await.unwrap();

        match state_result {
            VersionedGraphResult::Match(entry) => Ok(task_state.lookup_matches(entry)),
            VersionedGraphResult::Compute => {
                self.compute(k, eval, &events_dispatcher, task_state.lookup_dirtied(eval))
                    .await
            }

            VersionedGraphResult::CheckDeps(mismatch) => {
                let task_state = task_state.checking_deps(eval);

                let deps_changed = {
                    events_dispatcher.check_deps_started(k);
                    scopeguard::defer! {
                        events_dispatcher.check_deps_finished(k);
                    }

                    self.compute_whether_dependencies_changed(
                        ParentKey::Some(k), // the computing of deps is triggered by this key as the parent
                        eval.dupe(),
                        &mismatch.verified_versions,
                        mismatch.deps_to_validate,
                        &task_state,
                    )
                    .await?
                };

                match deps_changed {
                    DidDepsChange::Changed | DidDepsChange::NoDeps => {
                        self.compute(k, eval, &events_dispatcher, task_state.deps_not_match())
                            .await
                    }
                    DidDepsChange::NoChange(deps) => {
                        report_key_activation(
                            &eval.dice.key_index,
                            eval.user_data.activation_tracker.as_deref(),
                            k,
                            deps.iter().copied(),
                            ActivationData::Reused,
                        );

                        let task_state = task_state.deps_match()?;

                        // report reuse
                        let (tx, rx) = tokio::sync::oneshot::channel();
                        self.state.request(StateRequest::UpdateComputed {
                            key: VersionedGraphKey::new(v, k),
                            epoch: self.version_epoch,
                            storage: eval.storage_type(k),
                            value: mismatch.entry,
                            deps,
                            resp: tx,
                        });

                        rx.await.unwrap().map(|r| task_state.cached(r))
                    }
                }
            }
        }
    }

    async fn compute(
        &self,
        k: DiceKey,
        eval: &AsyncEvaluator,
        event_dispatcher: &DiceEventDispatcher,
        task_state: DiceWorkerStateComputing<'_>,
    ) -> CancellableResult<DiceWorkerStateFinishedAndCached> {
        event_dispatcher.started(k);
        scopeguard::defer! {
            event_dispatcher.finished(k);
        };

        let v = eval.per_live_version_ctx.get_version();

        // TODO(bobyf) these also make good locations where we want to perform instrumentation
        debug!(msg = "running evaluator");

        let eval_result_state = eval.evaluate(k, task_state).await?;
        let eval_result = eval_result_state.result;

        let res = {
            report_key_activation(
                &eval.dice.key_index,
                eval.user_data.activation_tracker.as_deref(),
                k,
                eval_result.deps.iter().copied(),
                eval_result.evaluation_data.into_activation_data(),
            );

            match eval_result.value.into_valid_value() {
                Ok(value) => {
                    let (tx, rx) = tokio::sync::oneshot::channel();
                    self.state.request(StateRequest::UpdateComputed {
                        key: VersionedGraphKey::new(v, k),
                        epoch: self.version_epoch,
                        storage: eval_result.storage,
                        value,
                        deps: Arc::new(eval_result.deps.into_iter().collect()),
                        resp: tx,
                    });

                    rx.await.unwrap()
                }
                Err(value) => CancellableResult::Ok(DiceComputedValue::new(
                    value,
                    Arc::new(CellHistory::verified(v)),
                )),
            }
        };

        res.map(|res| eval_result_state.state.cached(res))
    }

    /// determines if the given 'Dependency' has changed between versions 'last_version' and
    /// 'target_version'
    #[instrument(
        level = "debug",
        skip(self, eval, deps, check_deps_state),
        fields(version = %eval.per_live_version_ctx.get_version(), verified_versions = %verified_versions)
    )]
    async fn compute_whether_dependencies_changed(
        &self,
        parent_key: ParentKey,
        eval: AsyncEvaluator,
        verified_versions: &VersionRanges,
        deps: Arc<Vec<DiceKey>>,
        check_deps_state: &DiceWorkerStateCheckingDeps<'_>,
    ) -> CancellableResult<DidDepsChange> {
        trace!(deps = ?deps);

        if deps.is_empty() {
            return Ok(DidDepsChange::NoDeps);
        }

        let mut fs: FuturesUnordered<_> = deps
            .iter()
            .map(|dep| {
                eval.per_live_version_ctx
                    .compute_opaque(
                        dep.dupe(),
                        parent_key,
                        &eval,
                        check_deps_state.cycles_for_dep(*dep, &eval),
                    )
                    .map(|r| r.map(|v| v.history().get_verified_ranges()))
            })
            .collect();

        let mut verified_versions = Cow::Borrowed(verified_versions);

        while let Some(dep_result) = fs.next().await {
            match dep_result {
                Ok(dep_version_ranges) => {
                    verified_versions =
                        Cow::Owned(verified_versions.intersect(&dep_version_ranges));
                    if verified_versions.is_empty() {
                        return Ok(DidDepsChange::Changed);
                    }
                }
                Err(Cancelled) => {
                    return Err(Cancelled);
                }
            }
        }

        Ok(DidDepsChange::NoChange(deps))
    }
}

enum DidDepsChange {
    Changed,
    /// These deps did not change
    NoChange(Arc<Vec<DiceKey>>),
    NoDeps,
}

fn report_key_activation(
    key_index: &DiceKeyIndex,
    activation_tracker: Option<&dyn ActivationTracker>,
    key: DiceKey,
    deps: impl Iterator<Item = DiceKey>,
    activation_data: ActivationData,
) {
    if let Some(activation_tracker) = &activation_tracker {
        let key = key_index.get(key).as_any();
        let mut iter = deps.map(|dep| key_index.get(dep).as_any());
        activation_tracker.key_activated(key, &mut iter, activation_data);
    }
}

#[cfg(test)]
pub(crate) mod testing {

    use crate::impls::incremental::DidDepsChange;

    pub(crate) trait DidDepsChangeExt {
        fn is_changed(&self) -> bool;
    }

    impl DidDepsChangeExt for DidDepsChange {
        fn is_changed(&self) -> bool {
            match self {
                DidDepsChange::Changed => true,
                DidDepsChange::NoChange(..) => false,
                DidDepsChange::NoDeps => false,
            }
        }
    }
}