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

//!
//! This library contains utilities for tracking a global version number. The
//! global version number is used for tagging computed values so that we can
//! track when a value needs to be updated because its version number is out of
//! date.

use std::cmp;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::ops::Bound;
use std::ops::RangeBounds;
use std::ops::Sub;

use allocative::Allocative;
use derive_more::Display;
use dupe::Dupe;
use sorted_vector_map::SortedVectorSet;

/// The incrementing Version number associated with all the cache entries
#[derive(Copy, Eq, Debug, Display, Dupe)]
// split this due to formatters not agreeing
#[derive(PartialEq, Hash, Clone, Ord, PartialOrd, Allocative)]
#[display(fmt = "v{}", "_0")]
pub struct VersionNumber(pub(crate) usize);

impl VersionNumber {
    /// First transaction has version number zero.
    pub(crate) const ZERO: VersionNumber = VersionNumber(0);

    pub(crate) fn new(num: usize) -> Self {
        VersionNumber(num)
    }

    pub(crate) fn inc(&mut self) {
        self.0 += 1;
    }

    pub(crate) fn dec(&mut self) {
        self.0 = self.0.checked_sub(1).expect("shouldn't underflow");
    }
}

impl Sub for VersionNumber {
    type Output = isize;

    fn sub(self, rhs: Self) -> Self::Output {
        self.0 as isize - rhs.0 as isize
    }
}

mod introspection {
    use crate::versions::VersionNumber;

    impl VersionNumber {
        pub fn to_introspectable(&self) -> crate::introspection::graph::VersionNumber {
            crate::introspection::graph::VersionNumber(self.0)
        }
    }
}

/// Represents a range of versions. This range must have a start that is inclusive, but may be
/// unbounded towards the end. The end, if present, is exclusive.
#[derive(Allocative, Eq, Debug, Dupe, PartialEq, Hash, Clone)]
pub(crate) struct VersionRange {
    begin: VersionNumber,
    end: Option<VersionNumber>,
}

impl Display for VersionRange {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "[{}, ", self.begin)?;
        match self.end {
            None => {
                write!(f, "Unbounded")?;
            }
            Some(end) => {
                write!(f, "{}", end)?;
            }
        }
        write!(f, ")")
    }
}

impl VersionRange {
    fn new(begin: VersionNumber, end: Option<VersionNumber>) -> Self {
        if let Some(end) = end {
            assert!(begin < end);
        }
        VersionRange { begin, end }
    }

    pub(crate) fn bounded(begin: VersionNumber, end: VersionNumber) -> Self {
        VersionRange::new(begin, Some(end))
    }

    pub(crate) fn begins_with(begin: VersionNumber) -> Self {
        VersionRange::new(begin, None)
    }

    pub(crate) fn intersect(&self, other: &VersionRange) -> Option<Self> {
        // we exclude the end bound, because intervals [1,2) and [2,3) do not intersect
        fn contains_end_exclusive(
            v: VersionNumber,
            begin: VersionNumber,
            end: Option<VersionNumber>,
        ) -> bool {
            v >= begin && end.map_or(true, |end| v < end)
        }

        if contains_end_exclusive(self.begin, other.begin, other.end)
            || contains_end_exclusive(other.begin, self.begin, self.end)
        {
            let begin = cmp::max(self.begin, other.begin);
            let end = match (self.end, other.end) {
                (Some(x), Some(y)) => Some(cmp::min(x, y)),
                (None, x) => x,
                (y, None) => y,
            };

            Some(VersionRange::new(begin, end))
        } else {
            None
        }
    }

    /// splits the current range into two ranges at the given version, updating the current range
    /// and returning the new range that occurs after this range. If the given version is
    /// outside the range, then this does nothing and returns None.
    #[allow(unused)] // useful function. Probably will be needed at some point in the future
    pub(crate) fn split(&mut self, at: VersionNumber) -> Option<Self> {
        if self.begin < at {
            match self.end {
                Some(end) => {
                    if at < end {
                        self.end = Some(at);
                        Some(Self::bounded(at, end))
                    } else {
                        None
                    }
                }
                None => {
                    self.end = Some(at);
                    Some(Self::begins_with(at))
                }
            }
        } else {
            None
        }
    }

    #[allow(unused)] // useful function
    pub(crate) fn begin(&self) -> &VersionNumber {
        &self.begin
    }

    /// Merges this range with the given range if they overlap, otherwise return `None`
    pub(crate) fn merge(&self, other: &VersionRange) -> Option<Self> {
        // when merging, we include the end point, because intervals of the form [1,2) and [2,3)
        // should be merged into [1,3)
        fn is_between_end_inclusive(
            v: VersionNumber,
            begin: VersionNumber,
            end: Option<VersionNumber>,
        ) -> bool {
            v >= begin && end.map_or(true, |end| v <= end)
        }

        if is_between_end_inclusive(self.begin, other.begin, other.end)
            || is_between_end_inclusive(other.begin, self.begin, self.end)
        {
            Some(VersionRange::new(
                cmp::min(self.begin, other.begin),
                match (self.end, other.end) {
                    (None, _) => None,
                    (_, None) => None,
                    (Some(e1), Some(e2)) => Some(cmp::max(e1, e2)),
                },
            ))
        } else {
            None
        }
    }
}

impl RangeBounds<VersionNumber> for VersionRange {
    fn start_bound(&self) -> Bound<&VersionNumber> {
        Bound::Included(&self.begin)
    }

    fn end_bound(&self) -> Bound<&VersionNumber> {
        self.end.as_ref().map_or(Bound::Unbounded, Bound::Included)
    }
}

impl PartialOrd for VersionRange {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

/// We form a total ordering for the range by first comparing the beginning version, followed by
/// the end version if the beginning is tied. Since `None` indicates unbounded, that would be
/// larger than any version.
impl Ord for VersionRange {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        cmp_chain! {
            self.begin.cmp(&other.begin),
            match (self.end, other.end) {
                (None, None) => cmp::Ordering::Equal,
                (None, Some(_)) => cmp::Ordering::Greater,
                (Some(_), None) => cmp::Ordering::Less,
                (Some(end), Some(other_end)) => end.cmp(&other_end)
            }
        }
    }
}

/// Represents a sequence of `VersionRange`s. Each range contained in this sorted set is disjoint.
/// Any operations that causes ranges to overlap will result in the ranges being merged.
/// A version is contained in the `VersionRanges` iff there exists a `VersionRange` in this
/// sequence of ranges that contains the version. That is, consider the sequence `[[1,3), [5,6)]`,
/// 3, and 4 would not be in the sequence of ranges, but 1, 2, 5, would be. This is essentially
/// a list of numerical end-exclusive intervals.
#[derive(Allocative, Eq, Debug, PartialEq, Hash, Clone, PartialOrd, Ord)]
pub(crate) struct VersionRanges(SortedVectorSet<VersionRange>);

impl Display for VersionRanges {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{{")?;
        for (i, range) in self.0.iter().enumerate() {
            if i != 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}", range)?;
        }
        write!(f, "}}")
    }
}

impl VersionRanges {
    pub(crate) fn new() -> Self {
        Self(Default::default())
    }

    /// inserts a single range, merging different ranges if necessary
    pub(crate) fn insert(&mut self, mut range: VersionRange) {
        if let Some(smaller) = self
            .0
            .range((Bound::Unbounded, Bound::Included(range.dupe())))
            .max()
        {
            if let Some(merged) = smaller.merge(&range) {
                let to_remove = smaller.dupe();
                let removed = self.0.remove(&to_remove);
                assert!(removed);
                range = merged;
            }
        }

        if let Some(larger) = self
            .0
            .range((Bound::Included(range.dupe()), Bound::Unbounded))
            .min()
        {
            if let Some(merged) = larger.merge(&range) {
                let to_remove = larger.dupe();
                let removed = self.0.remove(&to_remove);
                assert!(removed);
                range = merged;
            }
        }

        let inserted = self.0.insert(range);
        assert!(inserted);
    }

    /// Computes the union of this set of ranges and another
    #[allow(unused)] // useful function
    pub(crate) fn union(&self, other: &VersionRanges) -> VersionRanges {
        let mut this = self.0.iter().peekable();
        let mut other = other.0.iter().peekable();

        let mut out = SortedVectorSet::new();
        let mut pending: Option<VersionRange> = None;
        loop {
            let smaller = match (this.peek(), other.peek()) {
                (Some(this_range), Some(other_range)) => {
                    if this_range < other_range {
                        this.next().expect("just peeked")
                    } else {
                        other.next().expect("just peeked")
                    }
                }
                (Some(_), None) => this.next().expect("just peeked"),
                (None, Some(_)) => other.next().expect("just peeked"),
                (None, None) => break,
            };

            pending = Some(pending.map_or_else(
                || smaller.dupe(),
                |last| {
                    if let Some(merged) = last.merge(smaller) {
                        merged
                    } else {
                        out.insert(last);
                        smaller.dupe()
                    }
                },
            ));
        }

        if let Some(last) = pending {
            out.insert(last);
        }

        VersionRanges(out)
    }

    /// Computes the intersection of this set of ranges and another
    pub(crate) fn intersect(&self, other: &VersionRanges) -> VersionRanges {
        let mut this = self.0.iter().peekable();
        let mut other = other.0.iter().peekable();

        let mut out = SortedVectorSet::new();
        // Pending is the last range we saw that has the largest end point, which is not the
        // standard sorting of intervals.
        // We want the largest end point interval to handle cases where there is one large interval
        // in a list that overlaps with several distinct intervals in the other list.
        let mut pending: Option<VersionRange> = None;

        loop {
            let smaller = match (this.peek(), other.peek()) {
                (Some(this_range), Some(other_range)) => {
                    if this_range < other_range {
                        this.next().expect("just peeked")
                    } else {
                        other.next().expect("just peeked")
                    }
                }
                (Some(_), None) => this.next().expect("just peeked"),
                (None, Some(_)) => other.next().expect("just peeked"),
                (None, None) => break,
            };

            if let Some(r) = pending {
                if let Some(intersect) = r.intersect(smaller) {
                    // we know that within an VersionRange, there are no overlaps, so as soon as we
                    // have an intersection, it can be pushed to the result and no other ranges
                    // will overlap with the intersection
                    out.insert(intersect);

                    // get the largest ending range
                    pending = Some(cmp::max_by(r, smaller.dupe(), |r1, r2| {
                        match (r1.end, r2.end) {
                            (None, None) => cmp::Ordering::Equal,
                            (None, Some(_)) => cmp::Ordering::Greater,
                            (Some(_), None) => cmp::Ordering::Less,
                            (Some(end), Some(other_end)) => end.cmp(&other_end),
                        }
                    }));
                } else {
                    // if there's no overlap, the current interval must have a larger ending point
                    // than `pending`, since `pending` < `smaller` and they don't intersect
                    pending = Some(smaller.dupe());
                }
            } else {
                pending = Some(smaller.dupe());
            }
        }

        VersionRanges(out)
    }

    /// Computes the intersection of this set of ranges and a range.
    #[allow(unused)] // useful function
    pub(crate) fn intersect_range(&self, range: &VersionRange) -> VersionRanges {
        let mut ranges = VersionRanges::new();
        ranges.insert(range.dupe());
        self.intersect(&ranges)
    }

    pub(crate) fn ranges(&self) -> &SortedVectorSet<VersionRange> {
        &self.0
    }

    pub(crate) fn is_empty(&self) -> bool {
        // Ranges in the set are not empty, so self is not empty if the ranges set is not empty.
        self.ranges().is_empty()
    }
}

#[cfg(test)]
pub(crate) mod testing {
    use sorted_vector_map::SortedVectorSet;

    use crate::versions::VersionRange;
    use crate::versions::VersionRanges;

    pub(crate) trait VersionRangesExt {
        fn testing_new(ranges: SortedVectorSet<VersionRange>) -> Self;
    }

    impl VersionRangesExt for VersionRanges {
        fn testing_new(ranges: SortedVectorSet<VersionRange>) -> Self {
            Self(ranges)
        }
    }
}

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

    use crate::versions::VersionNumber;
    use crate::versions::VersionRange;
    use crate::versions::VersionRanges;

    #[test]
    fn version_range_intersects() {
        let r1 = VersionRange::bounded(VersionNumber::new(0), VersionNumber::new(4));
        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(5));
        assert_eq!(
            r1.intersect(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(4)
            ))
        );

        let r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(5));
        assert_eq!(
            r1.intersect(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(4)
            ))
        );

        let r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        let r2 = VersionRange::begins_with(VersionNumber::new(0));
        assert_eq!(
            r1.intersect(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(4)
            ))
        );

        let r1 = VersionRange::begins_with(VersionNumber::new(2));
        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(5));
        assert_eq!(
            r1.intersect(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(2),
                VersionNumber::new(5)
            ))
        );

        let r1 = VersionRange::begins_with(VersionNumber::new(2));
        let r2 = VersionRange::begins_with(VersionNumber::new(1));
        assert_eq!(
            r1.intersect(&r2),
            Some(VersionRange::begins_with(VersionNumber::new(2)))
        );

        let r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(2));
        let r2 = VersionRange::bounded(VersionNumber::new(3), VersionNumber::new(4));
        assert_eq!(r1.intersect(&r2), None,);
    }

    #[test]
    fn version_range_splits() {
        let mut r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        assert_eq!(r1.split(VersionNumber::new(0)), None);
        assert_eq!(
            r1,
            VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4))
        );

        let mut r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        assert_eq!(r1.split(VersionNumber::new(5)), None,);
        assert_eq!(
            r1,
            VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4))
        );

        let mut r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        assert_eq!(
            r1.split(VersionNumber::new(3)),
            Some(VersionRange::bounded(
                VersionNumber::new(3),
                VersionNumber::new(4)
            )),
        );
        assert_eq!(
            r1,
            VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3))
        );

        let mut r1 = VersionRange::begins_with(VersionNumber::new(2));
        assert_eq!(r1.split(VersionNumber::new(1)), None);
        assert_eq!(r1, VersionRange::begins_with(VersionNumber::new(2)));

        let mut r1 = VersionRange::begins_with(VersionNumber::new(2));
        assert_eq!(
            r1.split(VersionNumber::new(4)),
            Some(VersionRange::begins_with(VersionNumber::new(4)))
        );
        assert_eq!(
            r1,
            VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(4))
        );
    }

    #[test]
    fn version_range_ops() {
        let r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));

        assert_eq!(r1 == r2, true);
        assert_eq!(r1 < r2, false);
        assert_eq!(r1 > r2, false);

        let r2 = VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(5));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, true);
        assert_eq!(r1 > r2, false);

        let r2 = VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(3));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, true);
        assert_eq!(r1 > r2, false);

        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, false);
        assert_eq!(r1 > r2, true);

        let r2 = VersionRange::begins_with(VersionNumber::new(2));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, true);
        assert_eq!(r1 > r2, false);

        let r2 = VersionRange::begins_with(VersionNumber::new(0));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, false);
        assert_eq!(r1 > r2, true);

        let r1 = VersionRange::begins_with(VersionNumber::new(1));
        let r2 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, false);
        assert_eq!(r1 > r2, true);

        let r1 = VersionRange::begins_with(VersionNumber::new(1));
        let r2 = VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(4));
        assert_eq!(r1 == r2, false);
        assert_eq!(r1 < r2, true);
        assert_eq!(r1 > r2, false);

        let r1 = VersionRange::begins_with(VersionNumber::new(1));
        let r2 = VersionRange::begins_with(VersionNumber::new(1));
        assert_eq!(r1 == r2, true);
        assert_eq!(r1 < r2, false);
        assert_eq!(r1 > r2, false);
    }

    #[test]
    fn version_range_merge() {
        let r1 = VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(4));
        let r2 = VersionRange::bounded(VersionNumber::new(5), VersionNumber::new(6));

        assert_eq!(r1.merge(&r2), None);

        let r2 = VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(3));
        assert_eq!(
            r1.merge(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(4)
            ))
        );

        let r2 = VersionRange::bounded(VersionNumber::new(2), VersionNumber::new(6));
        assert_eq!(
            r1.merge(&r2),
            Some(VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(6)
            ))
        );

        let r2 = VersionRange::begins_with(VersionNumber::new(6));
        assert_eq!(r1.merge(&r2), None);

        let r2 = VersionRange::begins_with(VersionNumber::new(2));
        assert_eq!(
            r1.merge(&r2),
            Some(VersionRange::begins_with(VersionNumber::new(1),))
        );
    }

    #[test]
    fn version_ranges_union() {
        let r1 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3)),
            VersionRange::bounded(VersionNumber::new(6), VersionNumber::new(8)),
            VersionRange::bounded(VersionNumber::new(10), VersionNumber::new(11))
        ]);

        let r2 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber::new(0), VersionNumber::new(1)),
            VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
            VersionRange::bounded(VersionNumber::new(8), VersionNumber::new(10)),
            VersionRange::bounded(VersionNumber::new(11), VersionNumber::new(12)),
            VersionRange::begins_with(VersionNumber::new(13)),
        ]);

        assert_eq!(
            r1.union(&r2),
            VersionRanges(sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(0), VersionNumber::new(3)),
                VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
                VersionRange::bounded(VersionNumber::new(6), VersionNumber::new(12)),
                VersionRange::begins_with(VersionNumber::new(13)),
            ])
        );
    }

    #[test]
    fn version_ranges_intersect() {
        let r1 = VersionRanges(sorted_vector_set![VersionRange::bounded(
            VersionNumber::new(1),
            VersionNumber::new(3)
        ),]);

        let r2 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
            VersionRange::bounded(VersionNumber::new(11), VersionNumber::new(12)),
            VersionRange::begins_with(VersionNumber::new(13)),
        ]);

        assert_eq!(r1.intersect(&r2), VersionRanges::new());

        let r1 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3)),
            VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
            VersionRange::bounded(VersionNumber::new(6), VersionNumber::new(9)),
            VersionRange::bounded(VersionNumber::new(10), VersionNumber::new(14)),
            VersionRange::begins_with(VersionNumber::new(15)),
        ]);

        let r2 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber::new(0), VersionNumber::new(1)),
            VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
            VersionRange::bounded(VersionNumber::new(8), VersionNumber::new(10)),
            VersionRange::bounded(VersionNumber::new(11), VersionNumber::new(12)),
            VersionRange::begins_with(VersionNumber::new(13)),
        ]);

        assert_eq!(
            r1.intersect(&r2),
            VersionRanges(sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(5)),
                VersionRange::bounded(VersionNumber::new(8), VersionNumber::new(9)),
                VersionRange::bounded(VersionNumber::new(11), VersionNumber::new(12)),
                VersionRange::bounded(VersionNumber::new(13), VersionNumber::new(14)),
                VersionRange::begins_with(VersionNumber::new(15)),
            ])
        );
    }

    #[test]
    fn version_ranges_intersects_range() {
        let r1 = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber(10), VersionNumber(20)),
            VersionRange::bounded(VersionNumber(30), VersionNumber(40)),
        ]);
        let r2 = VersionRange::bounded(VersionNumber(15), VersionNumber(35));
        let expected = VersionRanges(sorted_vector_set![
            VersionRange::bounded(VersionNumber(15), VersionNumber(20)),
            VersionRange::bounded(VersionNumber(30), VersionNumber(35)),
        ]);
        assert_eq!(expected, r1.intersect_range(&r2));
    }

    #[test]
    fn version_ranges_insert() {
        let mut r = VersionRanges::new();

        r.insert(VersionRange::bounded(
            VersionNumber::new(1),
            VersionNumber::new(3),
        ));
        assert_eq!(
            r.ranges(),
            &sorted_vector_set![VersionRange::bounded(
                VersionNumber::new(1),
                VersionNumber::new(3)
            )]
        );

        r.insert(VersionRange::bounded(
            VersionNumber::new(4),
            VersionNumber::new(6),
        ));
        assert_eq!(
            r.ranges(),
            &sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3)),
                VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(6))
            ]
        );

        r.insert(VersionRange::bounded(
            VersionNumber::new(5),
            VersionNumber::new(7),
        ));
        assert_eq!(
            r.ranges(),
            &sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(1), VersionNumber::new(3)),
                VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(7))
            ]
        );

        r.insert(VersionRange::bounded(
            VersionNumber::new(0),
            VersionNumber::new(1),
        ));
        assert_eq!(
            r.ranges(),
            &sorted_vector_set![
                VersionRange::bounded(VersionNumber::new(0), VersionNumber::new(3)),
                VersionRange::bounded(VersionNumber::new(4), VersionNumber::new(7))
            ]
        );
    }
}


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

use std::any::Any;

/// An ActivationTracker can be used to identify which keys were either reused or computed during a
/// transaction.
pub trait ActivationTracker: Send + Sync + 'static {
    /// Receives when a key was activated (computed, or reused). The caller will want to downcast
    /// the key and deps to types they care about. The caller also receives whatever the key passed
    /// to `store_evaluation_data` (if any).
    fn key_activated(
        &self,
        key: &dyn Any,
        deps: &mut dyn Iterator<Item = &dyn Any>,
        activation_data: ActivationData,
    );
}

/// Describes the kind of activation, and possibly carries data passed by the key's evaluation.
pub enum ActivationData {
    /// This key was evaluated. Evaluation data will be passed if the key's evaluation set any.
    Evaluated(Option<Box<dyn Any + Send + Sync + 'static>>),

    /// This key was reused. No data is passed.
    Reused,
}


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

use std::future::Future;

use allocative::Allocative;

use crate::api::data::DiceData;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::opaque::OpaqueValue;
use crate::api::user_data::UserComputationData;
use crate::ctx::DiceComputationsImpl;
use crate::UserCycleDetectorGuard;

/// The context for computations to register themselves, and request for additional dependencies.
/// The dependencies accessed are tracked for caching via the `DiceCtx`.
///
/// The computations are registered onto `DiceComputations` via implementing traits for the
/// `DiceComputation`.
///
/// The context is valid only for the duration of the computation of a single key, and cannot be
/// owned.
#[derive(Allocative)]
#[repr(transparent)]
pub struct DiceComputations(pub(crate) DiceComputationsImpl);

fn _test_computations_sync_send() {
    fn _assert_sync_send<T: Sync + Send>() {}
    _assert_sync_send::<DiceComputations>();
}

impl DiceComputations {
    /// Gets all the result of of the given computation key.
    /// recorded as dependencies of the current computation for which this
    /// context is for.
    pub fn compute<'a, K>(
        &'a self,
        key: &K,
    ) -> impl Future<Output = DiceResult<<K as Key>::Value>> + 'a
    where
        K: Key,
    {
        self.0.compute(key)
    }

    /// Compute "opaque" value where the value is only accessible via projections.
    /// Projections allow accessing derived results from the "opaque" value,
    /// where the dependency of reading a projection is the projection value rather
    /// than the entire opaque value.
    pub fn compute_opaque<'a, K>(
        &'a self,
        key: &K,
    ) -> impl Future<Output = DiceResult<OpaqueValue<'a, K>>> + 'a
    where
        K: Key,
    {
        self.0.compute_opaque(key)
    }

    /// Data that is static per the entire lifetime of Dice. These data are initialized at the
    /// time that Dice is initialized via the constructor.
    pub fn global_data(&self) -> &DiceData {
        self.0.global_data()
    }

    /// Data that is static for the lifetime of the current request context. This lifetime is
    /// the lifetime of the top-level `DiceComputation` used for all requests.
    /// The data is also specific to each request context, so multiple concurrent requests can
    /// each have their own individual data.
    pub fn per_transaction_data(&self) -> &UserComputationData {
        self.0.per_transaction_data()
    }

    /// Gets the current cycle guard if its set. If it's set but a different type, an error will be returned.
    pub fn cycle_guard<T: UserCycleDetectorGuard>(&self) -> DiceResult<Option<&T>> {
        self.0.cycle_guard()
    }

    /// Store some extra data that the ActivationTracker will receive if / when this key finishes
    /// executing.
    pub fn store_evaluation_data<T: Send + Sync + 'static>(&self, value: T) -> DiceResult<()> {
        self.0.store_evaluation_data(value)
    }
}

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

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

    use crate::api::cycles::DetectCycles;
    use crate::api::error::DiceErrorImpl;
    use crate::api::storage_type::StorageType;
    use crate::api::user_data::UserComputationData;
    use crate::legacy::ctx::ComputationData;
    use crate::legacy::cycles::RequestedKey;
    use crate::legacy::incremental::graph::storage_properties::StorageProperties;

    #[derive(Clone, Dupe, Display, Debug, PartialEq, Eq, Hash, Allocative)]
    struct K(usize);
    impl StorageProperties for K {
        type Key = K;

        type Value = ();

        fn key_type_name() -> &'static str {
            unreachable!()
        }

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

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

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

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

    #[test]
    fn cycle_detection_when_no_cycles() -> anyhow::Result<()> {
        let ctx = ComputationData::new(UserComputationData::new(), DetectCycles::Enabled);
        let ctx = ctx.subrequest::<K>(&K(1))?;
        let ctx = ctx.subrequest::<K>(&K(2))?;
        let ctx = ctx.subrequest::<K>(&K(3))?;
        let _ctx = ctx.subrequest::<K>(&K(4))?;

        Ok(())
    }

    #[test]
    fn cycle_detection_when_cycles() -> anyhow::Result<()> {
        let ctx = ComputationData::new(UserComputationData::new(), DetectCycles::Enabled);
        let ctx = ctx.subrequest::<K>(&K(1))?;
        let ctx = ctx.subrequest::<K>(&K(2))?;
        let ctx = ctx.subrequest::<K>(&K(3))?;
        let ctx = ctx.subrequest::<K>(&K(4))?;
        match ctx.subrequest::<K>(&K(1)) {
            Ok(_) => {
                panic!("should have cycle error")
            }
            Err(e) => match &*e.0 {
                DiceErrorImpl::Cycle {
                    trigger,
                    cyclic_keys,
                } => {
                    assert!(
                        (**trigger).get_key_equality() == K(1).get_key_equality(),
                        "expected trigger key to be `{}` but was `{}`",
                        K(1),
                        trigger
                    );
                    assert_eq!(
                        cyclic_keys,
                        &indexset![
                            Arc::new(K(1)) as Arc<dyn RequestedKey>,
                            Arc::new(K(2)) as Arc<dyn RequestedKey>,
                            Arc::new(K(3)) as Arc<dyn RequestedKey>,
                            Arc::new(K(4)) as Arc<dyn RequestedKey>
                        ]
                    )
                }
                _ => {
                    panic!("wrong error type")
                }
            },
        }

        Ok(())
    }
}


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

//! Cycle detection in DICE

use std::fmt::Debug;
use std::str::FromStr;

use allocative::Allocative;
use dupe::Dupe;
use gazebo::variants::VariantName;
use thiserror::Error;

#[derive(Clone, Dupe, Copy, Debug, VariantName, Allocative)]
pub enum DetectCycles {
    Enabled,
    Disabled,
}

#[derive(Error, Debug)]
#[error("Invalid type of DetectCycles: `{0}`")]
pub struct InvalidType(String);

impl FromStr for DetectCycles {
    type Err = InvalidType;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "ENABLED" => Ok(DetectCycles::Enabled),
            "DISABLED" => Ok(DetectCycles::Disabled),
            _ => Err(InvalidType(s.to_owned())),
        }
    }
}


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

//!
//! Dice allows users to attach a set of immutable data when constructing it. These data are for
//! users to attach context information for computations to use. These values are NOT tracked by
//! Dice. They are obtained synchronously.
//!
//! To attach data to Dice, you can implement a trait on DiceData, and use the methods to retrieve
//! values itself as the key for setting values.
//! e.g.
//! ```
//! use crate::dice::DiceData;
//!
//! pub trait HasData {
//!    fn my_data(&self) -> usize;
//!
//!    fn other_data(&self) -> &String;
//!
//!    fn set_multi(&mut self, i: usize, s: String);
//! }
//!
//! struct HasDataContainer(usize, String);
//!
//! impl HasData for DiceData {
//!     fn my_data(&self) -> usize {
//!         self.get::<HasDataContainer>().unwrap().0
//!     }
//!
//!     fn other_data(&self) -> &String {
//!         &self.get::<HasDataContainer>().unwrap().1
//!     }
//!
//!     fn set_multi(&mut self, i: usize, s: String) {
//!         self.set(HasDataContainer(i, s));
//!     }
//! }
//!
//! let mut data = DiceData::new();
//! data.set_multi(1, "foo".to_string());
//!
//! assert_eq!(data.other_data(), &"foo".to_string());
//! assert_eq!(data.my_data(), 1);
//!
//! ```
//!

use std::collections::BTreeSet;

use allocative::Allocative;
use anymap::any::Any;
use anymap::Map;
use itertools::Itertools;
use thiserror::Error;

#[derive(Error, Debug)]
#[error(
    "should store a value first before requesting a value for requested data key of type `{0}`. Known types are `{1}`"
)]
pub struct MissingData(&'static str, String);

#[derive(Allocative)]
pub struct DiceData(
    #[allocative(skip)] // TODO(nga): measure this.
    Map<dyn Any + Send + Sync>,
    BTreeSet<&'static str>,
);

impl DiceData {
    pub fn new() -> Self {
        Self(Map::new(), BTreeSet::new())
    }

    /// Stores the given data, overriding the previous value if any.
    pub fn set<K: Send + Sync + 'static>(&mut self, val: K) {
        if self.0.insert(val).is_none() {
            self.1.insert(std::any::type_name::<K>());
        }
    }

    pub fn get<K: Send + Sync + 'static>(&self) -> Result<&K, MissingData> {
        self.0
            .get::<K>()
            .ok_or_else(|| MissingData(std::any::type_name::<K>(), self.1.iter().join(", ")))
    }
}


/*
 * 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 dynamic incremental caching computation engine that powers buckv2.
//!
//! The computation engine will output values corresponding to given `Key`s,
//! reusing previously computed values when possible. `Key`s computations are
//! allowed to request other `Key`s via a `ComputationContext`.
//!
//! Example:
//! ```
//! mod c {
//!
//!     /// declaring computations:
//!     use async_trait::async_trait;
//!     use derive_more::Display;
//!     use dice::{Key, InjectedKey, DiceComputations, DiceDataBuilder, DiceData, DiceTransactionUpdater};
//!     use std::sync::Arc;
//!     use allocative::Allocative;
//! use more_futures::cancellation::CancellationContext;
//!
//!     /// A configuration computation that consists of values that are pre-computed outside of DICE
//!     pub struct InjectConfigs<'compute>(&'compute DiceComputations);
//!
//!     impl<'compute> InjectConfigs<'compute> {
//!         pub async fn get_config(&self) -> usize {
//!             self.0.compute(&ConfigKey).await.unwrap()
//!         }
//!     }
//!
//!     #[derive(Clone, Debug, Display, Eq, Hash, PartialEq, Allocative)]
//!     #[display(fmt = "{:?}", self)]
//!     struct ConfigKey;
//!
//!     #[async_trait]
//!     impl InjectedKey for ConfigKey {
//!         type Value = usize;
//!
//!         fn equality(x: &Self::Value,y: &Self::Value) -> bool {
//!             x == y
//!         }
//!     }
//!
//!     pub struct MyComputation<'compute>(&'compute DiceComputations);
//!
//!     impl<'compute> MyComputation<'compute> {
//!         // declaring a computation function
//!         pub async fn compute_a(&self, a: usize, s: String) -> Arc<String> {
//!             #[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
//!             #[display(fmt = "{:?}", self)]
//!             struct ComputeA(usize, String);
//!
//!             #[async_trait]
//!             impl Key for ComputeA {
//!                 type Value = Arc<String>;
//!
//!                 async fn compute(&self, ctx: &DiceComputations, _cancellations: &CancellationContext) -> Self::Value {
//!                     // request for other computations on the self
//!                     let n = ctx.my_computation().compute_b(self.0).await;
//!                     Arc::new(self.1.repeat(n))
//!                 }
//!
//!                 fn equality(x: &Self::Value,y: &Self::Value) -> bool {
//!                     x == y
//!                 }
//!             }
//!
//!             self.0.compute(&ComputeA(a, s)).await.unwrap()
//!         }
//!
//!         // second computation function
//!         pub async fn compute_b(&self, a: usize) -> usize {
//!                 self.0.compute(&ComputeB(a)).await.unwrap()
//!         }
//!     }
//!
//!     #[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
//!     #[display(fmt = "{:?}", self)]
//!     struct ComputeB(usize);
//!
//!     #[async_trait]
//!     impl Key for ComputeB {
//!         type Value = usize;
//!
//!         async fn compute(&self, ctx: &DiceComputations, cancellations: &CancellationContext) -> Self::Value {
//!             self.0 + ctx.injected_configs().get_config().await + ctx.global_data().static_data().len()
//!         }
//!
//!         fn equality(x: &Self::Value,y: &Self::Value) -> bool {
//!             x == y
//!         }
//!     }
//!
//!     // trait to register the computation to DICE
//!     pub trait HasMyComputation {
//!         fn my_computation(&self) -> MyComputation;
//!     }
//!
//!     // attach the declared computation to DICE via the context
//!     impl HasMyComputation for DiceComputations {
//!         fn my_computation(&self) -> MyComputation {
//!             MyComputation(self)
//!         }
//!     }
//!
//!     // trait to register the precomputed configs to DICE
//!     pub trait HasInjectedConfig {
//!         fn injected_configs(&self) -> InjectConfigs;
//!     }
//!
//!     impl HasInjectedConfig for DiceComputations {
//!         fn injected_configs(&self) -> InjectConfigs {
//!             InjectConfigs(self)
//!         }
//!     }
//!
//!     pub trait SetInjectedConfig {
//!         fn inject_config(&mut self, i: usize);
//!     }
//!
//!     impl SetInjectedConfig for DiceTransactionUpdater {
//!         /// ways to inject the precomputed values to dice
//!         fn inject_config(&mut self, i: usize) {
//!             self.changed_to(vec![(ConfigKey, i)]).unwrap()
//!         }
//!     }
//!
//!     pub trait StaticData {
//!         fn static_data(&self) -> &String;
//!     }
//!
//!     impl StaticData for DiceData {
//!         fn static_data(&self) -> &String {
//!             self.get::<String>().unwrap()
//!         }
//!     }
//!
//!     pub trait SetStaticData {
//!         fn set_static(&mut self, s: String);
//!     }
//!
//!     impl SetStaticData for DiceDataBuilder {
//!         fn set_static(&mut self, s: String) {
//!             self.set(s);
//!         }
//!     }
//! }
//!
//! /// how to use computations
//! use dice::{Dice, DetectCycles};
//! use std::sync::Arc;
//! use c::*;
//!
//! let mut rt = tokio::runtime::Runtime::new().unwrap();
//! let mut builder = Dice::builder();
//! builder.set_static("len4".into());
//! let engine = builder.build(DetectCycles::Disabled);
//!
//! // inject config
//! let mut ctx = engine.updater();
//! ctx.inject_config(0);
//!
//! let ctx = rt.block_on(ctx.commit());
//!
//! // request the computation from DICE
//! rt.block_on(async {
//!     assert_eq!("aaaaaaaa", &*ctx.my_computation().compute_a(4, "a".into()).await);
//! });
//!
//! let mut ctx = engine.updater();
//! ctx.inject_config(2);
//!
//! let ctx = rt.block_on(ctx.commit());
//!
//! // request the computation from DICE
//! rt.block_on(async {
//!     assert_eq!("aaaaaaaaaa", &*ctx.my_computation().compute_a(4, "a".into()).await);
//! });
//! ```

use std::fmt::Debug;
use std::io::Write;
use std::sync::Arc;

use allocative::Allocative;
use futures::future::Future;
use serde::Serializer;

use crate::api::cycles::DetectCycles;
use crate::api::transaction::DiceTransactionUpdater;
use crate::api::user_data::UserComputationData;
use crate::api::which::WhichSpawner;
use crate::metrics::Metrics;
use crate::DiceDataBuilderImpl;
use crate::DiceImplementation;
use crate::WhichDice;

/// An incremental computation engine that executes arbitrary computations that
/// maps `Key`s to values.
#[derive(Allocative, Debug)]
pub struct Dice {
    pub(crate) implementation: DiceImplementation,
}

impl Dice {
    pub fn builder() -> DiceDataBuilder {
        DiceDataBuilder(DiceDataBuilderImpl::new_legacy())
    }

    pub fn modern() -> DiceDataBuilder {
        DiceDataBuilder(DiceDataBuilderImpl::new_modern())
    }

    pub(crate) fn new(implementation: DiceImplementation) -> Arc<Self> {
        Arc::new(Self { implementation })
    }

    pub fn updater(self: &Arc<Dice>) -> DiceTransactionUpdater {
        self.implementation.updater()
    }

    pub fn updater_with_data(
        self: &Arc<Dice>,
        extra: UserComputationData,
    ) -> DiceTransactionUpdater {
        self.implementation.updater_with_data(extra)
    }

    pub fn serialize_tsv(
        &self,
        nodes: impl Write,
        edges: impl Write,
        nodes_currently_running: impl Write,
    ) -> anyhow::Result<()> {
        self.implementation
            .serialize_tsv(nodes, edges, nodes_currently_running)
    }

    pub fn serialize_serde<S>(&self, serializer: S) -> Result<(), S::Error>
    where
        S: Serializer,
    {
        self.implementation.serialize_serde(serializer)
    }

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

    pub fn which_dice(&self) -> WhichDice {
        match self.implementation {
            DiceImplementation::Legacy(_) => WhichDice::Legacy,
            DiceImplementation::Modern(_) => WhichDice::Modern,
        }
    }

    pub fn metrics(&self) -> Metrics {
        self.implementation.metrics()
    }

    /// Wait until all active versions have exited.
    pub fn wait_for_idle(&self) -> impl Future<Output = ()> + 'static {
        self.implementation.wait_for_idle()
    }

    /// true when there are no active tasks nor transactions alive
    pub async fn is_idle(&self) -> bool {
        self.implementation.is_idle().await
    }
}

pub struct DiceDataBuilder(DiceDataBuilderImpl);

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

    pub fn build(self, detect_cycles: DetectCycles) -> Arc<Dice> {
        self.build_with_which_spawner(detect_cycles, WhichSpawner::ExplicitCancel)
    }

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

pub mod testing {
    use crate::api::cycles::DetectCycles;
    use crate::api::key::Key;
    use crate::api::transaction::DiceTransactionUpdater;
    use crate::api::user_data::UserComputationData;
    use crate::Dice;
    use crate::DiceDataBuilder;

    /// Testing utility that can be used to build a specific `DiceComputation` where certain keys
    /// of computation mocked to return a specific result.
    ///
    /// TODO(bobyf): ideally, we want something where we don't have to use the specific keys
    /// but rather the computation function, like `mock.expect(|c| c.other_compute(4), "4 res")`
    pub struct DiceBuilder {
        builder: DiceDataBuilder,
        mocked: Vec<Box<dyn FnOnce(&mut DiceTransactionUpdater) -> anyhow::Result<()>>>,
    }

    impl DiceBuilder {
        pub fn new() -> Self {
            let builder = Dice::builder();

            Self {
                builder,
                mocked: Vec::new(),
            }
        }

        pub fn set_data(mut self, setter: impl FnOnce(&mut DiceDataBuilder)) -> Self {
            setter(&mut self.builder);
            self
        }

        /// mocks the call of compute for the key `expected_k` so that it returns `expected_res`
        pub fn mock_and_return<K>(mut self, expected_k: K, expected_res: K::Value) -> Self
        where
            K: Key,
        {
            self.mocked.push(Box::new(move |ctx| {
                Ok(ctx.changed_to(vec![(expected_k, expected_res)])?)
            }));
            self
        }

        pub fn build(self, extra: UserComputationData) -> anyhow::Result<DiceTransactionUpdater> {
            let dice = self.builder.build(DetectCycles::Enabled);
            let mut ctx = dice.updater_with_data(extra);

            self.mocked.into_iter().try_for_each(|f| f(&mut ctx))?;
            Ok(ctx)
        }
    }
}


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

use std::sync::Arc;

use allocative::Allocative;
use dupe::Dupe;
use indexmap::IndexSet;
use itertools::Itertools;
use thiserror::Error;

use crate::legacy::cycles::RequestedKey;

#[derive(Clone, Dupe, Debug, Error, Allocative)]
#[error(transparent)]
pub struct DiceError(pub(crate) Arc<DiceErrorImpl>);

impl DiceError {
    pub fn cycle(
        trigger: Arc<dyn RequestedKey>,
        cyclic_keys: IndexSet<Arc<dyn RequestedKey>>,
    ) -> Self {
        DiceError(Arc::new(DiceErrorImpl::Cycle {
            trigger,
            cyclic_keys,
        }))
    }

    pub fn invalid_change(key: Arc<dyn RequestedKey>) -> Self {
        DiceError(Arc::new(DiceErrorImpl::ChangedToInvalid(key)))
    }

    pub fn duplicate(key: Arc<dyn RequestedKey>) -> Self {
        DiceError(Arc::new(DiceErrorImpl::DuplicateChange(key)))
    }

    pub fn cancelled() -> Self {
        DiceError(Arc::new(DiceErrorImpl::Cancelled))
    }

    pub fn duplicate_activation_data() -> Self {
        DiceError(Arc::new(DiceErrorImpl::DuplicateActivationData))
    }
}

#[derive(Debug, Error, Allocative)]
pub(crate) enum DiceErrorImpl {
    #[error("Cyclic computation detected when computing key `{}`, which forms a cycle in computation chain: `{}`", trigger, cyclic_keys.iter().join(" -> "))]
    Cycle {
        trigger: Arc<dyn RequestedKey>,
        cyclic_keys: IndexSet<Arc<dyn RequestedKey>>,
    },
    #[error("Key `{0}` was marked as changed multiple times on the same transaction.")]
    DuplicateChange(Arc<dyn RequestedKey>),
    #[error("Key `{0}` was reported as changed to an invalid value")]
    ChangedToInvalid(Arc<dyn RequestedKey>),
    /// NOTE: This isn't an error users normally see, since if the user is waiting on a result, the
    /// future doesn't get cancelled.
    #[error("The evaluation of this key was cancelled")]
    Cancelled,
    #[error(
        "Requested cycle_guard of type {}, but current guard has type {}",
        expected_type_name,
        actual_type_name
    )]
    UnexpectedCycleGuardType {
        expected_type_name: String,
        actual_type_name: String,
    },
    #[error("Activation data was already provided for this key")]
    DuplicateActivationData,
}

pub type DiceResult<T> = Result<T, DiceError>;


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

#[derive(Allocative, PartialEq, Eq, Debug)]
pub enum DiceEvent {
    /// Computation started.
    Started { key_type: &'static str },

    /// Computation finished.
    Finished { key_type: &'static str },

    /// Checking dependencies has started.
    CheckDepsStarted { key_type: &'static str },

    /// Checking dependencies has finished.
    CheckDepsFinished { key_type: &'static str },
}

pub trait DiceEventListener: Allocative + Send + Sync + 'static {
    fn event(&self, ev: DiceEvent);
}


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

use std::fmt::Debug;
use std::fmt::Display;
use std::hash::Hash;

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

use crate::api::computations::DiceComputations;
use crate::api::key::Key;
use crate::api::storage_type::StorageType;

/// Specialized version of `Key` above. This type of Key is never computed. It
/// should always be injected onto the graph before being requested via
/// `DiceCtx::changed_to(Key, Value)`. Therefore, the `compute` function on this
/// Key is never called.
///
/// TODO when rust gets smarter about negative trait bounds and specializations,
/// we can make injected and normal key types disjoint, and not have this clone issue nor
/// require a `panic!` implementation in `compute` function, both of which are
/// horrible and breaks semantics of traits.
pub trait InjectedKey:
    Allocative + Clone + Debug + Display + Send + Sync + Eq + Hash + 'static
{
    type Value: Allocative + Dupe + Send + Sync + 'static;

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

#[async_trait]
impl<K> Key for K
where
    K: InjectedKey,
{
    type Value = K::Value;

    async fn compute(
        &self,
        _ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        panic!(
            "Injected Keys must be injected onto the graph before being requested, hence \
            computes should never be called; however, `compute` on `{:?}` was called.",
            self
        )
    }

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

    fn storage_type() -> StorageType {
        // if we store more than usize max value, we are in trouble.
        StorageType::LastN(usize::max_value())
    }
}


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

use std::fmt::Debug;
use std::fmt::Display;
use std::hash::Hash;

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

use crate::api::computations::DiceComputations;
use crate::api::storage_type::StorageType;
use crate::introspection::graph::short_type_name;

/// The computation Key that maps to a value. The key will be used as an index
/// for caching the computed values.
///
/// ## Cloning
/// The Key will need to be cloned. This occurs whenever the key is recomputed
/// and needs to be stored. During normal lookups, only a reference is needed.
///
/// The result also needs to be cloned whenever it is requested from DICE.
/// Cloning therefore should be cheap.
#[async_trait]
pub trait Key: Allocative + Debug + Display + Clone + Eq + Hash + Send + Sync + 'static {
    type Value: Allocative + Dupe + Send + Sync + 'static;

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

    async fn compute(
        &self,
        ctx: &DiceComputations,
        cancellations: &CancellationContext,
    ) -> Self::Value;

    /// If computed value is equal to previously cached value,
    /// DICE won't invalidate graph nodes depending on this node.
    ///
    /// It is safe to return `false` from this function when values are equal,
    /// but returning `true` when values are not equal would result
    /// in inconsistent graph state.
    fn equality(x: &Self::Value, y: &Self::Value) -> bool;

    /// If the computed value is `false`, DICE will consider that result to be a transient value
    /// that won't be re-used on subsequent computations. It will, however, reuse that value for
    /// all on-going computations at the current version.
    ///
    /// The default here is true, but computations should override this if its expected that they
    /// may occasionally produce transient values.
    fn validity(_x: &Self::Value) -> bool {
        true
    }

    fn storage_type() -> StorageType {
        StorageType::LastN(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.
 */

//! Public DICE API

pub mod activation_tracker;
pub mod computations;
pub mod cycles;
pub mod data;
pub mod dice;
pub mod error;
pub mod events;
pub mod injected;
pub mod key;
pub mod opaque;
pub mod projection;
pub mod storage_type;
pub mod transaction;
pub mod user_data;
pub mod which;


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

use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;

use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::projection::ProjectionKey;
use crate::opaque::OpaqueValueImpl;

/// Computed value which is not directly visible to user.
///
/// The value can be accessed only via "projection" operation,
/// so projection result is recorded as a dependency
/// of a computation which requested the opaqued value,
/// but the opaque value key is not.
pub struct OpaqueValue<'a, K: Key> {
    pub(crate) implementation: OpaqueValueImpl<'a, K>,
}

impl<'a, K> Debug for OpaqueValue<'a, K>
where
    K: Key,
    K::Value: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.implementation.fmt(f)
    }
}

impl<'a, K: Key> OpaqueValue<'a, K> {
    pub(crate) fn new(implementation: OpaqueValueImpl<'a, K>) -> Self {
        Self { implementation }
    }

    pub fn projection<P>(&self, projection_key: &P) -> DiceResult<P::Value>
    where
        P: ProjectionKey<DeriveFromKey = K>,
    {
        self.implementation.projection(projection_key)
    }
}


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

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

use std::fmt::Debug;
use std::fmt::Display;
use std::hash::Hash;

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

use crate::api::data::DiceData;
use crate::api::key::Key;
use crate::api::storage_type::StorageType;
use crate::api::user_data::UserComputationData;
use crate::introspection::graph::short_type_name;

/// Synchronously computed key from an "opaque" value.
pub trait ProjectionKey:
    Allocative + Clone + PartialEq + Eq + Hash + Display + Debug + Send + Sync + 'static
{
    /// Key of the value that this projection key is computed from.
    type DeriveFromKey: Key;
    /// This projection key is mapped to this value.
    type Value: Allocative + Dupe + Send + Sync;

    /// Compute the projection key value.
    // Implementation note:
    // This function is called from two places:
    // * `OpaqueValue::projection`, where the projection key is computed synchronously.
    // * `ProjectionKeyAsKey`, where the projection key is computed asynchronously
    //   together with opaque key this key is derived from.
    //   This operation cannot be triggered manually by user,
    //   but it happens implicitly during dependency recompute.
    fn compute(
        &self,
        derive_from: &<<Self as ProjectionKey>::DeriveFromKey as Key>::Value,
        ctx: &DiceProjectionComputations,
    ) -> Self::Value;

    fn equality(x: &Self::Value, y: &Self::Value) -> bool;

    fn validity(x: &Self::Value) -> bool {
        let _ = x;
        true
    }

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

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

/// Context for projection key computation.
///
/// Only provide access to globals.
pub struct DiceProjectionComputations<'a> {
    pub(crate) data: &'a DiceData,
    pub(crate) user_data: &'a UserComputationData,
}

impl<'a> DiceProjectionComputations<'a> {
    /// Data that is static per the entire lifetime of Dice. These data are initialized at the
    /// time that Dice is initialized via the constructor.
    pub fn global_data(&self) -> &DiceData {
        self.data
    }

    /// Data that is static for the lifetime of the current request context. This lifetime is
    /// the lifetime of the top-level `DiceComputation` used for all requests.
    /// The data is also specific to each request context, so multiple concurrent requests can
    /// each have their own individual data.
    pub fn per_transaction_data(&self) -> &UserComputationData {
        self.user_data
    }
}


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

use std::any::Any;
use std::ops::ControlFlow;

use once_cell::sync::OnceCell;

use crate::legacy::incremental::Computable;
use crate::UserComputationData;


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

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

/// Storage type for a cached entry.
/// The oldest entry will be evicted once the cache stores more than N entries of the same key
/// request to compute them. TODO think about whether or not we can
/// optimize to delete injected keys when no more computation will request that version
#[derive(UnpackVariants, Debug, Clone, Copy, Dupe, Allocative)]
pub enum StorageType {
    LastN(usize),
}


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

use std::future::Future;
use std::ops::Deref;
use std::ops::DerefMut;

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

use crate::api::computations::DiceComputations;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::transaction::DiceTransactionImpl;
use crate::transaction_update::DiceTransactionUpdaterImpl;
use crate::versions::VersionNumber;

/// The struct for which we build transactions. This is where changes are recorded, and committed
/// to DICE, which returns the Transaction where we spawn computations.
#[derive(Allocative)]
#[repr(transparent)]
pub struct DiceTransactionUpdater(pub(crate) DiceTransactionUpdaterImpl);

impl DiceTransactionUpdater {
    pub fn existing_state(&self) -> impl Future<Output = DiceTransaction> + '_ {
        self.0.existing_state()
    }

    /// Records a set of `Key`s as changed so that they, and any dependents will
    /// be recomputed on the next set of requests at the next version.
    pub fn changed<K, I>(&mut self, changed: I) -> DiceResult<()>
    where
        K: Key,
        I: IntoIterator<Item = K> + Send + Sync + 'static,
    {
        self.0.changed(changed)
    }

    /// Records a set of `Key`s as changed to a particular value so that any
    /// dependents will be recomputed on the next set of requests. The
    /// `Key`s themselves will be update to the new value such that they
    /// will not need to be recomputed as long as they aren't recorded to be
    /// `changed` again (or invalidated by other means). Calling this method
    /// does not in anyway alter the types of the key such that they
    /// permanently becomes a special "inject value only" key.
    pub fn changed_to<K, I>(&mut self, changed: I) -> DiceResult<()>
    where
        K: Key,
        I: IntoIterator<Item = (K, K::Value)> + Send + Sync + 'static,
    {
        self.0.changed_to(changed)
    }

    /// Commit the changes registered via 'changed' and 'changed_to' to the current newest version.
    pub fn commit(self) -> impl Future<Output = DiceTransaction> {
        self.0.commit()
    }

    /// Commit the changes registered via 'changed' and 'changed_to' to the current newest version,
    /// replacing the user data with the given set
    pub fn commit_with_data(
        self,
        extra: UserComputationData,
    ) -> impl Future<Output = DiceTransaction> {
        self.0.commit_with_data(extra)
    }

    pub fn unstable_take(self) -> Self {
        Self(self.0.unstable_take())
    }
}

/// The base struct for which all computations start. This is clonable, and dupe, and can be
/// moved to different runtimes to start computations.
/// All computations on this transaction will see only changes at the most-up-to-date version at
/// the time of creation of this transaction.
///
/// This SHOULD NOT be ever stored by computations, or any results of computations.
#[derive(Allocative, Clone, Dupe)]
pub struct DiceTransaction(pub(crate) DiceTransactionImpl);

impl DiceTransaction {
    /// Returns whether the `DiceTransaction` is equivalent. Equivalent is defined as whether the
    /// two Transactions are based off the same underlying set of key states. That is, all
    /// injected keys are the same, and the same compute keys are dirtied, and that any computations
    /// that occur between the two transactions can be shared.
    pub fn equivalent<E>(&self, other: &E) -> bool
    where
        E: DiceEquivalent,
    {
        self.version_for_equivalence() == other.version_for_equivalence()
    }

    pub fn equality_token(&self) -> DiceEquality {
        DiceEquality(self.0.get_version())
    }

    /// Creates an Updater to record changes to DICE that upon committing, creates a new transaction
    /// that keeps the same set of user data. This is equivalent to `Dice::updater_with_user_data(data)`
    /// where the `data` is taken from the current Transaction.
    pub fn into_updater(self) -> DiceTransactionUpdater {
        self.0.into_updater()
    }
}

#[derive(Allocative, Eq, PartialEq, Copy, Clone, derive_more::Display)]
#[repr(transparent)]
pub struct DiceEquality(VersionNumber);

mod private {
    use super::*;

    pub trait Sealed {}

    impl Sealed for DiceTransaction {}

    impl Sealed for DiceEquality {}
}

pub trait DiceEquivalent: private::Sealed {
    fn version_for_equivalence(&self) -> DiceEquality;
}

impl DiceEquivalent for DiceTransaction {
    fn version_for_equivalence(&self) -> DiceEquality {
        self.equality_token()
    }
}

impl DiceEquivalent for DiceEquality {
    fn version_for_equivalence(&self) -> DiceEquality {
        *self
    }
}

impl Deref for DiceTransaction {
    type Target = DiceComputations;

    fn deref(&self) -> &Self::Target {
        self.0.as_computations()
    }
}

impl DerefMut for DiceTransaction {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.0.as_computations_mut()
    }
}


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

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

use allocative::Allocative;
use more_futures::spawner::Spawner;
use more_futures::spawner::TokioSpawner;

use crate::api::activation_tracker::ActivationTracker;
use crate::api::data::DiceData;
use crate::api::events::DiceEvent;
use crate::api::events::DiceEventListener;

/// Includes all user related computation-specific data.
#[derive(Allocative)]
pub struct UserComputationData {
    /// The DiceData provides a spot for users to attach whatever extra things they want.
    ///
    /// This can contain arbitrary data from users that will not be part of the dice graph.
    /// As an example, users may want to inject some form of event dispatcher to send events from their computations.
    pub data: DiceData,
    pub tracker: Arc<dyn DiceEventListener>,
    #[allocative(skip)]
    pub spawner: Arc<dyn Spawner<Self>>,

    #[allocative(skip)]
    pub cycle_detector: Option<Arc<dyn UserCycleDetector>>,

    #[allocative(skip)]
    pub activation_tracker: Option<Arc<dyn ActivationTracker>>,

    /// We require that UserComputationData always be constructed with `..Default::default()`
    pub _requires_default: RequireDefault,
}

/// A UserCycleDetector can be used for custom cycle detection in the DICE computation.
pub trait UserCycleDetector: Send + Sync + 'static {
    /// Called by DICE when it starts computing a key. `key` will be a user Key type (and so user can reliably downcast it to known types).
    fn start_computing_key(&self, key: &dyn Any) -> Option<Box<dyn UserCycleDetectorGuard>>;

    /// Called by DICE when the key finished computing.
    fn finished_computing_key(&self, key: &dyn Any);
}

/// A UserCycleDetectorGuard is used to track the currently computing key. User code can access this through
/// ComputationData::cycle_guard() (and then downcast it with as_any to potentially access custom cycle behavior).
pub trait UserCycleDetectorGuard: Send + Sync + 'static {
    /// Called by dice when a dependency edge is encountered.
    fn add_edge(&self, key: &dyn Any);

    /// This is used to allow user code to get at the concrete guard instance.
    fn as_any(&self) -> &dyn Any;

    /// Used in error messages.
    fn type_name(&self) -> &'static str;
}

#[derive(Allocative)]
pub struct RequireDefault(());

impl UserComputationData {
    pub fn new() -> Self {
        Self::default()
    }
}

impl Default for UserComputationData {
    fn default() -> Self {
        Self {
            data: DiceData::new(),
            tracker: Arc::new(NoOpTracker),
            spawner: Arc::new(TokioSpawner),
            cycle_detector: None,
            activation_tracker: None,
            _requires_default: RequireDefault(()),
        }
    }
}

#[derive(Allocative)]
pub(crate) struct NoOpTracker;

impl DiceEventListener for NoOpTracker {
    fn event(&self, _ev: DiceEvent) {}
}


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

use allocative::Allocative;
use dupe::Dupe;
use gazebo::variants::VariantName;
use thiserror::Error;

/// which dice impl to use
#[derive(Clone, Dupe, Copy, Debug, VariantName, Allocative)]
pub enum WhichDice {
    Legacy,
    Modern,
}

#[derive(Error, Debug)]
#[error("Invalid type of WhichDice: `{0}`")]
pub struct InvalidType(String);

impl FromStr for WhichDice {
    type Err = InvalidType;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "MODERN" => Ok(WhichDice::Modern),
            "LEGACY" => Ok(WhichDice::Legacy),
            _ => Err(InvalidType(s.to_owned())),
        }
    }
}

#[derive(Clone, Dupe, Copy, Debug, VariantName, Allocative)]
pub enum WhichSpawner {
    DropCancel,
    ExplicitCancel,
}

#[derive(Error, Debug)]
#[error("Invalid type of WhichSpawner: `{0}`")]
pub struct InvalidWhichSpawner(String);

impl FromStr for WhichSpawner {
    type Err = InvalidType;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "DROPCANCEL" => Ok(WhichSpawner::DropCancel),
            "EXPLICITCANCEL" => Ok(WhichSpawner::ExplicitCancel),
            _ => Err(InvalidType(s.to_owned())),
        }
    }
}


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

//! Shared, concurrent dice task cache that is shared between computations at the same version

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use dashmap::DashMap;
use dupe::Dupe;
use fnv::FnvBuildHasher;
use lock_free_hashtable::sharded::ShardedLockFreeRawTable;

use crate::arc::Arc;
use crate::impls::key::DiceKey;
use crate::impls::task::dice::DiceTask;
use crate::impls::task::dice::DiceTaskInternal;
use crate::impls::value::DiceComputedValue;
use crate::result::CancellableResult;

#[derive(Allocative)]
struct Data {
    completed: ShardedLockFreeRawTable<Arc<DiceTaskInternal>, 64>,
    /// Completed tasks lazily moved into `completed` from this map.
    storage: DashMap<DiceKey, DiceTask, FnvBuildHasher>,
    is_cancelled: AtomicBool,
}

#[derive(Allocative, Clone, Dupe)]
pub(crate) struct SharedCache {
    data: Arc<Data>,
}

/// Reference to the task in the cache.
pub(crate) enum DiceTaskRef<'a> {
    Computed(DiceComputedValue),
    Occupied(dashmap::mapref::entry::OccupiedEntry<'a, DiceKey, DiceTask, FnvBuildHasher>),
    Vacant(dashmap::mapref::entry::VacantEntry<'a, DiceKey, DiceTask, FnvBuildHasher>),
    TransactionCancelled,
}

impl<'a> DiceTaskRef<'a> {
    #[cfg(test)]
    pub(crate) fn testing_insert(self, task: DiceTask) {
        if let Self::Vacant(e) = self {
            e.insert(task);
        } else {
            panic!("inserting into non-vacant entry");
        }
    }
}

impl SharedCache {
    fn key_hash(key: DiceKey) -> u64 {
        (key.index as u64).wrapping_mul(0x9e3779b97f4a7c15)
    }

    fn try_get_computed(&self, key: DiceKey) -> Option<DiceComputedValue> {
        let hash = Self::key_hash(key);
        self.data
            .completed
            .lookup(hash, |task| task.key == key)
            .map(|task| {
                task.read_value()
                    .expect("must be computed")
                    .expect("must not be cancelled")
            })
    }

    pub(crate) fn get(&self, key: DiceKey) -> DiceTaskRef {
        if let Some(computed) = self.try_get_computed(key) {
            if self.data.is_cancelled.load(Ordering::Acquire) {
                return DiceTaskRef::TransactionCancelled;
            }

            return DiceTaskRef::Computed(computed);
        }

        let entry = self.data.storage.entry(key);

        // Not we acquired the lock, check computed map again.
        let computed = self.try_get_computed(key);

        if self.data.is_cancelled.load(Ordering::Acquire) {
            return DiceTaskRef::TransactionCancelled;
        }

        if let Some(computed) = computed {
            return DiceTaskRef::Computed(computed);
        }

        match entry {
            dashmap::mapref::entry::Entry::Occupied(e) => {
                if let Some(CancellableResult::Ok(result)) = e.get().internal.read_value() {
                    // Promote entry to computed.
                    // So lookup will be faster next time.

                    // TODO(nga): insert unique unchecked,
                    //   which `LockFreeRawTable` does not support yet.
                    let (_ignore, original) = self.data.completed.insert(
                        Self::key_hash(key),
                        e.get().internal.dupe(),
                        |a, b| a.key == b.key,
                        |task| Self::key_hash(task.key),
                    );
                    assert!(original.is_none());

                    // Must remove from dashmap after inserting into completed.
                    e.remove();
                    return DiceTaskRef::Computed(result);
                }
                DiceTaskRef::Occupied(e)
            }
            dashmap::mapref::entry::Entry::Vacant(e) => DiceTaskRef::Vacant(e),
        }
    }

    pub(crate) fn new() -> Self {
        SharedCache {
            data: Arc::new(Data {
                storage: DashMap::default(),
                completed: ShardedLockFreeRawTable::new(),
                is_cancelled: AtomicBool::new(false),
            }),
        }
    }

    pub(crate) fn active_tasks_count(&self) -> usize {
        self.data.storage.len() + self.data.completed.len()
    }

    /// This function gets the termination observer for all running tasks when transaction is
    /// cancelled and prevents further tasks from being added
    pub(crate) fn cancel_pending_tasks(self) -> Vec<DiceTask> {
        self.data.is_cancelled.store(true, Ordering::Release);
        self.data
            .storage
            .iter()
            .filter_map(|entry| {
                if entry.value().is_pending() {
                    entry.value().cancel();
                    Some(entry.value().clone())
                } else {
                    None
                }
            })
            .collect()
    }
}

#[cfg(test)]
impl SharedCache {
    pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.data, &other.data)
    }
}

pub(crate) mod introspection {
    use crate::impls::cache::SharedCache;
    use crate::impls::key::DiceKey;
    use crate::legacy::dice_futures::dice_task::DiceTaskStateForDebugging;

    impl SharedCache {
        pub(crate) fn iter_tasks(
            &self,
        ) -> impl Iterator<Item = (DiceKey, DiceTaskStateForDebugging)> + '_ {
            self.data
                .storage
                .iter()
                .map(|entry| (*entry.key(), entry.value().introspect_state()))
        }
    }
}

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

    use allocative::Allocative;
    use async_trait::async_trait;
    use derive_more::Display;
    use dupe::Dupe;
    use futures::FutureExt;
    use more_futures::cancellation::CancellationContext;
    use more_futures::spawner::TokioSpawner;

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

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

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

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

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

    async fn make_completed_task(key: DiceKey, val: usize) -> DiceTask {
        let task = spawn_dice_task(key, &TokioSpawner, &(), |handle| {
            async move {
                handle.finished(DiceComputedValue::new(
                    MaybeValidDiceValue::valid(DiceValidValue::testing_new(
                        DiceKeyValue::<K>::new(val),
                    )),
                    Arc::new(CellHistory::empty()),
                ));

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

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

        task
    }

    async fn make_finished_cancelling_task(key: DiceKey) -> DiceTask {
        let finished_cancelling_tasks = spawn_dice_task(key, &TokioSpawner, &(), |handle| {
            async move {
                let _handle = handle;
                futures::future::pending().await
            }
            .boxed()
        });
        finished_cancelling_tasks.cancel();

        finished_cancelling_tasks.await_termination().await;

        finished_cancelling_tasks
    }

    fn make_never_finish_yet_to_cancel_task(key: DiceKey) -> DiceTask {
        spawn_dice_task(key, &TokioSpawner, &(), |handle| {
            async move {
                let _handle = handle;
                futures::future::pending().await
            }
            .boxed()
        })
    }

    #[tokio::test]
    async fn test_drain_task() {
        let cache = SharedCache::new();

        let completed_task1 = make_completed_task(DiceKey { index: 10 }, 1).await;
        let completed_task2 = make_completed_task(DiceKey { index: 20 }, 2).await;

        let finished_cancelling_tasks1 = make_finished_cancelling_task(DiceKey { index: 30 }).await;
        let finished_cancelling_tasks2 = make_finished_cancelling_task(DiceKey { index: 40 }).await;

        let yet_to_cancel_tasks1 = make_never_finish_yet_to_cancel_task(DiceKey { index: 50 });
        let yet_to_cancel_tasks2 = make_never_finish_yet_to_cancel_task(DiceKey { index: 60 });
        let yet_to_cancel_tasks3 = make_never_finish_yet_to_cancel_task(DiceKey { index: 70 });

        cache
            .get(DiceKey { index: 1 })
            .testing_insert(completed_task1);
        cache
            .get(DiceKey { index: 2 })
            .testing_insert(completed_task2);
        cache
            .get(DiceKey { index: 3 })
            .testing_insert(finished_cancelling_tasks1);
        cache
            .get(DiceKey { index: 4 })
            .testing_insert(finished_cancelling_tasks2);
        cache
            .get(DiceKey { index: 5 })
            .testing_insert(yet_to_cancel_tasks1);
        cache
            .get(DiceKey { index: 6 })
            .testing_insert(yet_to_cancel_tasks2);
        cache
            .get(DiceKey { index: 7 })
            .testing_insert(yet_to_cancel_tasks3);

        let pending_tasks = cache.dupe().cancel_pending_tasks();

        assert_eq!(pending_tasks.len(), 3);

        assert!(matches!(
            cache.get(DiceKey { index: 999 }),
            DiceTaskRef::TransactionCancelled
        ));
    }
}