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

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Barrier;
use std::sync::Mutex;
use std::time::Duration;

use assert_matches::assert_matches;
use async_trait::async_trait;
use derivative::Derivative;
use derive_more::Display;
use dupe::Dupe;
use futures::future::FutureExt;
use futures::future::Shared;
use more_futures::cancellation::CancellationContext;
use tokio::sync::oneshot;
use tokio::time::timeout;

use super::*;
use crate::api::computations::DiceComputations;
use crate::api::cycles::DetectCycles;
use crate::api::error::DiceErrorImpl;
use crate::api::injected::InjectedKey;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::legacy::ctx::testing::DiceCtxExt;
use crate::legacy::incremental::evaluator::testing::EvaluatorUnreachable;
use crate::legacy::incremental::testing::DependencyExt;
use crate::legacy::incremental::testing::IncrementalEngineExt;
use crate::legacy::incremental::testing::VersionedCacheResultAssertsExt;
use crate::legacy::incremental::versions::MinorVersion;
use crate::versions::VersionNumber;
use crate::HashSet;
use crate::UserCycleDetector;
use crate::UserCycleDetectorGuard;

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

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

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

#[tokio::test]
async fn set_injected_multiple_times_per_commit() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

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

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

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

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

    Ok(())
}

#[tokio::test]
async fn set_injected_with_no_change_no_new_ctx() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

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

        let ctx = ctx.commit().await;

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

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

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

    Ok(())
}

#[tokio::test]
async fn updates_caches_only_on_ctx_finalize_in_order() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

    {
        let mut ctx = dice.updater();

        // now we write something and commit
        ctx.changed_to(vec![(Foo(1), 1)])?;
        let vg = dice.global_versions.current();
        dice.find_cache::<Foo>()
            .get_maybe_cached(Foo(1), vg.version, *vg.minor_version_guard)
            .assert_none();

        ctx.commit().await;

        // committing the context commits the value
        let vg = dice.global_versions.current();
        assert_eq!(
            *dice
                .find_cache::<Foo>()
                .get_maybe_cached(Foo(1), vg.version, *vg.minor_version_guard)
                .assert_match()
                .val(),
            1
        );
    }

    {
        let mut ctx = dice.updater();
        let mut ctx1 = dice.updater();
        // even if we do a change on this ctx first.
        ctx.changed_to(vec![(Foo(2), 2)])?;
        ctx1.changed_to(vec![(Foo(3), 3)])?;

        // as long as we commit ctx1 first, it's values are committed first, in linear
        // history
        ctx1.commit().await;

        let vg = dice.global_versions.current();
        assert_eq!(
            *dice
                .find_cache::<Foo>()
                .get_maybe_cached(Foo(3), vg.version, *vg.minor_version_guard)
                .assert_match()
                .val(),
            3
        );

        dice.find_cache::<Foo>()
            .get_maybe_cached(Foo(2), vg.version, *vg.minor_version_guard)
            .assert_none();

        ctx.commit().await;

        // only now is 'ctx' committed
        let vg = dice.global_versions.current();
        assert_eq!(
            *dice
                .find_cache::<Foo>()
                .get_maybe_cached(Foo(2), vg.version, *vg.minor_version_guard)
                .assert_match()
                .val(),
            2
        );
    }

    Ok(())
}

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

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

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

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

#[test]
fn ctx_tracks_deps_properly() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        let ctx = dice.updater().commit().await;
        let res = ctx
            .compute(&K(5))
            .await?
            .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
        assert_eq!(res, K(31));

        // introspect the caches for dependency info
        fn assert_cached_deps(dice: &Arc<DiceLegacy>, k: i32) {
            let vg = dice.global_versions.current();
            let cached = dice
                .find_cache::<K>()
                .get_cached(K(k), vg.version, *vg.minor_version_guard)
                .dupe();

            let expected_deps = (0..k)
                .map(K)
                .map(DependencyExt::<EvaluatorUnreachable<_, K>>::testing_raw)
                .collect::<Vec<_>>();

            // TODO(bobyf) better assert the versions stored in deps
            let meta = cached.read_meta();
            let deps = meta.deps.debug_deps().read();
            let (version, deps) = deps.as_ref().expect("No deps");

            assert_eq!(*version, VersionNumber::new(0));

            let deps = deps.iter().collect::<HashSet<_>>();
            let expected_deps = expected_deps.iter().collect::<HashSet<_>>();

            assert_eq!(deps, expected_deps);
        }

        assert_cached_deps(&dice, 5);
        assert_cached_deps(&dice, 4);
        assert_cached_deps(&dice, 3);
        assert_cached_deps(&dice, 2);
        assert_cached_deps(&dice, 1);
        assert_cached_deps(&dice, 0);

        Ok(())
    })
}

#[test]
fn ctx_tracks_rdeps_properly() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(1)
        .max_blocking_threads(1)
        .build()
        .unwrap();
    rt.block_on(async {
        let ctx = dice.updater().commit().await;
        let res = ctx
            .compute(&K(5))
            .await?
            .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
        assert_eq!(res, K(31));

        // introspect the caches for dependency info
        fn assert_cached_rdeps(dice: &Arc<DiceLegacy>, k: i32) {
            let vg = dice.global_versions.current();

            let cached = dice
                .find_cache::<K>()
                .get_cached(K(k), vg.version, *vg.minor_version_guard)
                .dupe();

            let mut expected_deps = ((k + 1)..6)
                .map(K)
                .map(|k| {
                    Arc::as_ptr(
                        &dice
                            .find_cache::<K>()
                            .get_cached(k, VersionNumber::new(0), *vg.minor_version_guard)
                            .into_dyn(),
                    )
                })
                .collect::<HashSet<_>>();

            for rdep in cached.read_meta().rdeps.rdeps().rdeps.iter() {
                assert!(
                    expected_deps.remove(&Arc::as_ptr(&rdep.0.0.upgrade().unwrap())),
                    "Extra rdeps"
                )
            }
            assert!(
                expected_deps.is_empty(),
                "Missing {} rdeps",
                expected_deps.len()
            )
        }

        assert_cached_rdeps(&dice, 0);
        assert_cached_rdeps(&dice, 4);
        assert_cached_rdeps(&dice, 3);
        assert_cached_rdeps(&dice, 2);
        assert_cached_rdeps(&dice, 1);
        assert_cached_rdeps(&dice, 0);

        Ok(())
    })
}

// ignore this for now. Need to change ctx to better represent lifetimes and ownership
// to support this
//
// #[test]
// fn compute_that_requests_changes() {
//     let dice = Dice::new();
//
//     #[derive(Clone, Debug, Eq, PartialEq, Hash)]
//     struct InvalidatingKey;
//
//     #[async_trait]
//     impl Key for InvalidatingKey {
//         type Value = i32;
//
//         async fn compute(&self, ctx: &Arc<DiceCtx>) -> Self::Value {
//             ctx.changed_to(vec![(InvalidatingKey, 1)]);
//             ctx.commit();
//             0
//         }
//     }
//
//     let mut rt = tokio::runtime::Runtime::new().unwrap();
//     rt.block_on(async {
//         assert_eq!(dice.ctx().compute(&InvalidatingKey).await, 0);
//         assert_eq!(dice.ctx().compute(&InvalidatingKey).await, 1);
//     });
// }

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

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

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

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

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

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

    rt.block_on(async move {
        let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);
        let mut sum = 0;

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

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

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

        assert_eq!(sum, n_thread);
    })
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    Ok(())
}

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

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

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

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

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

    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

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

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

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

    drop(ctx);

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

    Ok(())
}

#[tokio::test]
async fn test_wait_for_idle() -> anyhow::Result<()> {
    #[derive(Clone, Debug, Display, Derivative, Allocative)]
    #[derivative(Hash, PartialEq, Eq)]
    #[display(fmt = "{:?}", self)]
    struct TestKey {
        id: usize,

        #[allocative(skip)]
        #[derivative(PartialEq = "ignore", Hash = "ignore")]
        channel: Shared<oneshot::Receiver<()>>,
    }

    impl Dupe for TestKey {}

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

        async fn compute(
            &self,
            _ctx: &DiceComputations,
            cancellations: &CancellationContext,
        ) -> Self::Value {
            cancellations
                .critical_section(|| self.channel.clone())
                .await
                .unwrap()
        }

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

    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

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

    let (tx, rx) = oneshot::channel();
    let rx = rx.shared();

    let key = TestKey { id: 1, channel: rx };
    let handle = ctx.compute(&key);

    let idle = dice.wait_for_idle();
    futures::pin_mut!(idle);

    assert_matches!(timeout(Duration::from_secs(1), &mut idle).await, Err(..));
    assert!(!dice.is_idle());

    drop(handle);
    drop(ctx);
    assert_matches!(timeout(Duration::from_secs(1), &mut idle).await, Err(..));
    assert!(!dice.is_idle());

    tx.send(()).unwrap();
    assert_matches!(timeout(Duration::from_secs(1), &mut idle).await, Ok(..));
    assert!(dice.is_idle());

    // Still idle.
    let stays_idle = async {
        dice.wait_for_idle().await;
        dice.wait_for_idle().await;
    };
    assert_matches!(timeout(Duration::from_secs(1), stays_idle).await, Ok(..));

    Ok(())
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        Ok(())
    })
}

#[tokio::test]
async fn compute_and_update_uses_proper_version_numbers() -> anyhow::Result<()> {
    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

    {
        let ctx = dice.updater().commit().await;
        assert_eq!(ctx.0.get_version(), VersionNumber::new(0));
        assert_eq!(ctx.0.get_minor_version(), MinorVersion::testing_new(0));
    }

    {
        // second context that didn't have any writes should still be the same version
        let ctx = dice.updater().commit().await;
        assert_eq!(ctx.0.get_version(), VersionNumber::new(0));
        assert_eq!(ctx.0.get_minor_version(), MinorVersion::testing_new(1));

        // now we write something and commit
        let mut ctx = dice.updater();
        ctx.changed_to(vec![(Foo(1), 1)])?;
        // current version shouldn't be updated
        assert_eq!(
            ctx.existing_state().await.0.get_version(),
            VersionNumber::new(0)
        );
        assert_eq!(
            ctx.existing_state().await.0.get_minor_version(),
            MinorVersion::testing_new(1)
        );

        let mut ctx1 = dice.updater();
        // previous ctx isn't dropped, so versions shouldn't be committed yet.
        assert_eq!(
            ctx1.existing_state().await.0.get_version(),
            VersionNumber::new(0)
        );
        assert_eq!(
            ctx1.existing_state().await.0.get_minor_version(),
            MinorVersion::testing_new(1)
        );

        // if we update on the new context, nothing committed
        ctx1.changed_to(vec![(Foo(2), 2)])?;
        assert_eq!(
            ctx1.existing_state().await.0.get_version(),
            VersionNumber::new(0)
        );
        assert_eq!(
            ctx1.existing_state().await.0.get_minor_version(),
            MinorVersion::testing_new(1)
        );

        // drop a context
        ctx1.commit().await;
        // we should only have committed once, and in increasing order
        let vg = dice.global_versions.current();
        assert_eq!(
            (vg.version, *vg.minor_version_guard),
            (VersionNumber::new(1), MinorVersion::testing_new(1))
        );

        ctx.commit().await;
        // both versions finalized.
        let vg = dice.global_versions.current();
        assert_eq!(
            (vg.version, *vg.minor_version_guard),
            (VersionNumber::new(2), MinorVersion::testing_new(1))
        );
        assert!(dice.map.read().engines().iter().all(|engine| {
            engine
                .introspect()
                .versions_currently_running()
                .first()
                .is_none()
        }));
    }

    {
        let mut ctx = dice.updater();
        assert_eq!(
            ctx.existing_state().await.0.get_version(),
            VersionNumber::new(2)
        );
        assert_eq!(
            ctx.existing_state().await.0.get_minor_version(),
            MinorVersion::testing_new(2)
        );

        ctx.changed_to(vec![(Foo(3), 3)])?;
        assert_eq!(
            ctx.existing_state().await.0.get_version(),
            VersionNumber::new(2)
        );
        assert_eq!(
            ctx.existing_state().await.0.get_minor_version(),
            MinorVersion::testing_new(2)
        );

        ctx.commit().await;
        let vg = dice.global_versions.current();
        assert_eq!(
            (vg.version, *vg.minor_version_guard),
            (VersionNumber::new(3), MinorVersion::testing_new(1))
        );
        assert!(dice.map.read().engines().iter().all(|engine| {
            engine
                .introspect()
                .versions_currently_running()
                .first()
                .is_none()
        }));
    }

    Ok(())
}

#[test]
fn test_active_transaction_count() {
    let dice = Arc::new(DiceLegacy::new(
        DiceData::new(),
        DetectCycles::Enabled,
        WhichSpawner::ExplicitCancel,
    ));
    assert_eq!(0, dice.metrics().active_transaction_count);
    let ctx = dice.updater().commit();
    assert_eq!(1, dice.metrics().active_transaction_count);
    drop(ctx);
    assert_eq!(0, dice.metrics().active_transaction_count);
}

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

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

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

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

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

    let dice = DiceLegacy::builder().build(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);
    let mut updater = dice.updater();

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


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

pub mod math_computation;
pub mod supply_chain;


/*
 * 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 math computation built on dice.
//! The math computation will calculate a series of addition equations resolving variables.
//! e.g. `x = 1; y = 2, a = x + y; b = a + a; eval(b)`;

use std::convert::Infallible;
use std::str::FromStr;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dice::Key;
use dupe::Dupe;
use futures::future;
use futures::FutureExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;

#[derive(Clone, Dupe, PartialEq, Eq, Hash, Display, Debug, Allocative)]
#[display(fmt = "Var({})", _0)]
pub struct Var(pub Arc<String>);

#[derive(Clone, PartialEq, Eq, Hash, Debug, Allocative)]
pub enum Unit {
    Var(Var),
    Literal(i64),
}

#[derive(PartialEq, Eq, Clone, Debug, Allocative)]
pub enum Equation {
    Add(Vec<Unit>),
    Unit(Unit),
}

impl FromStr for Unit {
    type Err = Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(match s.parse::<i64>() {
            Ok(i) => Unit::Literal(i),
            Err(_) => Unit::Var(Var(Arc::new(s.into()))),
        })
    }
}

impl FromStr for Equation {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        Ok(if s.contains('+') {
            let units = s
                .split_terminator('+')
                .map(|unit| unit.parse::<Unit>())
                .collect::<Result<Vec<Unit>, _>>()
                .unwrap();

            if units.len() < 2 {
                return Err(anyhow::anyhow!("+ must have left and right"));
            }

            Equation::Add(units)
        } else {
            Equation::Unit(s.parse().unwrap())
        })
    }
}

pub fn parse_math_equations<'a>(
    math: impl IntoIterator<Item = &'a str>,
) -> anyhow::Result<Vec<(Var, Equation)>> {
    math.into_iter().map(parse_math_equation).collect()
}

pub fn parse_math_equation(math: &str) -> anyhow::Result<(Var, Equation)> {
    let (l, r) = math.split1("=");
    if l.is_empty() || r.is_empty() {
        return Err(anyhow::anyhow!("= must have left and right"));
    }

    Ok((Var(Arc::new(l.to_owned())), r.parse()?))
}

pub trait MathEquations {
    fn set_equation(&mut self, var: Var, equation: Equation) -> anyhow::Result<()>;
    fn set_equations(
        &mut self,
        equations: impl IntoIterator<Item = (Var, Equation)>,
    ) -> anyhow::Result<()>;
}

#[async_trait]
pub trait Math {
    async fn eval(&self, var: Var) -> Result<i64, Arc<anyhow::Error>>;
}

impl MathEquations for DiceTransactionUpdater {
    fn set_equation(&mut self, var: Var, equation: Equation) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(LookupVar(var), Arc::new(equation))])?)
    }
    fn set_equations(
        &mut self,
        equations: impl IntoIterator<Item = (Var, Equation)>,
    ) -> anyhow::Result<()> {
        Ok(self.changed_to(
            equations
                .into_iter()
                .map(|(var, eq)| (LookupVar(var), Arc::new(eq)))
                .collect::<Vec<_>>(),
        )?)
    }
}

#[derive(Clone, Display, Debug, Dupe, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "Eval({})", _0)]
pub struct EvalVar(pub Var);
#[async_trait]
impl Key for EvalVar {
    type Value = Result<i64, Arc<anyhow::Error>>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Result<i64, Arc<anyhow::Error>> {
        let equation = lookup_unit(ctx, &self.0).await.map_err(Arc::new)?;
        Ok(match &*equation {
            Equation::Add(adds) => resolve_units(ctx, &adds[..]).await?.iter().sum(),
            Equation::Unit(unit) => resolve_units(ctx, &[unit.clone()]).await?[0],
        })
    }

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

#[async_trait]
impl Math for DiceComputations {
    async fn eval(&self, var: Var) -> Result<i64, Arc<anyhow::Error>> {
        Ok(self
            .compute(&EvalVar(var))
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))??)
    }
}

async fn resolve_units(
    ctx: &DiceComputations,
    units: &[Unit],
) -> Result<Vec<i64>, Arc<anyhow::Error>> {
    let futs = units.map(|unit| match unit {
        Unit::Var(var) => ctx.eval(var.clone()),
        Unit::Literal(lit) => async move { Ok(*lit) }.boxed(),
    });

    future::join_all(futs)
        .await
        .into_iter()
        .collect::<Result<_, _>>()
}

async fn lookup_unit(ctx: &DiceComputations, var: &Var) -> anyhow::Result<Arc<Equation>> {
    Ok(ctx.compute(&LookupVar(var.clone())).await?)
}

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "Lookup({})", _0)]
struct LookupVar(Var);
impl InjectedKey for LookupVar {
    type Value = Arc<Equation>;

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

#[cfg(test)]
mod parser_tests {
    use super::*;

    fn var(name: &str) -> Var {
        Var(Arc::new(name.to_owned()))
    }

    #[test]
    fn test_parser_unit_var() {
        let eq = "a=bbab";
        assert_eq!(
            vec![(var("a"), Equation::Unit(Unit::Var(var("bbab"))))],
            parse_math_equations(vec![eq]).unwrap()
        );
    }

    #[test]
    fn test_parser_unit_int() {
        let eq = "hello world=25";
        assert_eq!(
            vec![(var("hello world"), Equation::Unit(Unit::Literal(25)))],
            parse_math_equations(vec![eq]).unwrap()
        );
    }

    #[test]
    fn test_parser_expr() {
        assert_eq!(
            vec![(
                var("a"),
                Equation::Add(vec![
                    Unit::Var(var("bd")),
                    Unit::Literal(25),
                    Unit::Var(var("b"))
                ])
            )],
            parse_math_equations(vec!["a=bd+25+b"]).unwrap()
        )
    }

    #[test]
    #[allow(clippy::many_single_char_names)]
    fn test_parser_compound() {
        let eq = vec!["x=1", "y=2", "a=x+y", "b=a+c"];

        let x = var("x");
        let y = var("y");
        let a = var("a");
        let b = var("b");
        let c = var("c");

        assert_eq!(
            vec![
                (x.clone(), Equation::Unit(Unit::Literal(1))),
                (y.clone(), Equation::Unit(Unit::Literal(2))),
                (a.clone(), Equation::Add(vec![Unit::Var(x), Unit::Var(y)])),
                (b, Equation::Add(vec![Unit::Var(a), Unit::Var(c)]))
            ],
            parse_math_equations(eq).unwrap()
        );
    }

    #[test]
    #[should_panic(expected = "+ must have left and right")]
    fn test_parser_invalid() {
        parse_math_equations(vec!["a=a+"]).unwrap();
    }
}


/*
 * 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 supply chain simulation built on top of DICE.
//! * Each company has a name and manufactures a list of resources.
//! * Each resource has a recipe associated with it.
//!
//! When a company is asked to produce a resource:
//! * It looks up the recipe of sub-resources.
//! * For each sub-resource, it finds the cheapest manufacturer.
//! * It then return a quote of the cost.
//!
//! The cost of a resource is the sum of each item cost + a company specific flat fee.

use std::collections::HashMap;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dice::Key;
use dupe::Dupe;
use futures::future::join_all;
use futures::stream::FuturesUnordered;
use futures::FutureExt;
use futures::StreamExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;
use ref_cast::RefCast;

#[derive(Display, Debug, Hash, Eq, Clone, PartialEq, Dupe, Allocative)]
pub enum Resource {
    Wood,
    Plank,
    Stick,
    Cobblestone,
    Pickaxe,
    CraftingTable,
}

pub struct Recipe {
    /// The number of items this recipe produces (currently unused)
    pub produces: u8,
    /// A list of the resources (and how many) are required.
    pub ingredients: &'static [(u8, Resource)],
}

impl Recipe {
    /// Singleton for a primitive resource.
    const PRIMITIVE: Self = Self {
        produces: 1,
        ingredients: &[],
    };
    pub const WOOD: Self = Self::PRIMITIVE;
    pub const COBBLESTONE: Self = Self::PRIMITIVE;
    pub const PLANK: Self = Self {
        produces: 4,
        ingredients: &[(1, Resource::Wood)],
    };
    pub const STICK: Self = Self {
        produces: 4,
        ingredients: &[(2, Resource::Plank)],
    };
    pub const PICKAXE: Self = Self {
        produces: 1,
        ingredients: &[(3, Resource::Cobblestone), (2, Resource::Stick)],
    };
    pub const CRAFTING_TABLE: Self = Self {
        produces: 1,
        ingredients: &[(4, Resource::Plank)],
    };
}

impl Resource {
    pub const RESOURCES: [Resource; 6] = [
        Self::Wood,
        Self::Plank,
        Self::Stick,
        Self::Cobblestone,
        Self::Pickaxe,
        Self::CraftingTable,
    ];
    pub fn recipe(&self) -> Recipe {
        match self {
            Self::Wood => Recipe::WOOD,
            Self::Plank => Recipe::PLANK,
            Self::Stick => Recipe::STICK,
            Self::Cobblestone => Recipe::COBBLESTONE,
            Self::Pickaxe => Recipe::PICKAXE,
            Self::CraftingTable => Recipe::CRAFTING_TABLE,
        }
    }
}

/// Each company manufactures a set of resources for a given upcharge
#[derive(Eq, Clone, PartialEq, Allocative)]
pub struct Company {
    pub name: Arc<String>,
    /// A mapping between resource the company makes and its flat upcharge
    pub makes: HashMap<Resource, u16>,
}

#[async_trait]
pub trait Setup {
    /// Must be called before any companies are added.
    /// Sets the state of the resource map.
    fn init_state(&mut self) -> anyhow::Result<()>;
    /// Adds a list of companies and maps them to their resources.
    async fn add_companies(&mut self, companies: Vec<Company>) -> anyhow::Result<()>;
}

#[async_trait]
impl Setup for DiceTransactionUpdater {
    fn init_state(&mut self) -> anyhow::Result<()> {
        self.changed_to(
            Resource::RESOURCES
                .iter()
                .map(|resource| (LookupResource(resource.dupe()), Arc::new(vec![]))),
        )?;

        Ok(())
    }

    async fn add_companies(&mut self, companies: Vec<Company>) -> anyhow::Result<()> {
        let mut resource_to_company_local = HashMap::new();

        // convert company to insertion ready format
        let insertion_ready_companies = companies.into_map(|company| {
            let lookup = LookupCompany(company.name.dupe());

            // construct a resource => company lookups mapping across all companies
            for resource in company.makes.keys() {
                resource_to_company_local
                    .entry(resource.dupe())
                    .or_insert_with(Vec::new)
                    .push(lookup.clone());
            }

            (lookup, Arc::new(company))
        });

        self.changed_to(insertion_ready_companies)?;

        // split iterators to talk about resources and companies separately
        let (resources, companies): (Vec<_>, Vec<_>) = resource_to_company_local
            .into_iter()
            .map(|(k, v)| (LookupResource(k), v))
            .unzip();

        // get the remote resources => company mapping
        let state = self.existing_state().await;
        let remote_resources = join_all(resources.iter().map(|res| state.compute(res))).await;

        // combine remote company list with local company list for reach resource
        let joined: Vec<_> = resources
            .into_iter()
            .zip(companies.into_iter())
            .zip(remote_resources.into_iter())
            .map(|((resource, mut local_companies), remote_companies)| {
                local_companies.append(&mut (*remote_companies?).clone());

                Ok((resource, Arc::new(local_companies)))
            })
            .collect::<anyhow::Result<_>>()?;

        self.changed_to(joined)?;

        Ok(())
    }
}

#[async_trait]
pub trait Cost {
    /// Find the cheapest manufacturing cost for a resource
    async fn resource_cost(&self, resource: &Resource) -> Result<Option<u16>, Arc<anyhow::Error>>;
}

#[async_trait]
pub trait CostUpdater {
    /// Change the upcharge of a company for a resource.
    async fn change_company_resource_cost(
        &mut self,
        company: &str,
        resource: &Resource,
        new_price: u16,
    ) -> anyhow::Result<()>;
}

async fn lookup_company_resource_cost(
    ctx: &DiceComputations,
    company: &LookupCompany,
    resource: &Resource,
) -> Result<Option<u16>, Arc<anyhow::Error>> {
    #[derive(Display, Debug, Hash, Eq, Clone, Dupe, PartialEq, Allocative)]
    #[display(fmt = "{:?}", self)]
    struct LookupCompanyResourceCost(LookupCompany, Resource);
    #[async_trait]
    impl Key for LookupCompanyResourceCost {
        type Value = Result<Option<u16>, Arc<anyhow::Error>>;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellations: &CancellationContext,
        ) -> Self::Value {
            let company = ctx
                .compute(&self.0)
                .await
                .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;
            let recipe = self.1.recipe();

            let upcharge = company.makes.get(&self.1);
            if upcharge.is_none() {
                return Ok(None);
            }

            // get the unit cost for each resource needed to make item
            let mut futs: FuturesUnordered<_> = recipe
                .ingredients
                .iter()
                .map(|(required, resource)| {
                    ctx.resource_cost(resource)
                        .map(|res| Ok::<_, Arc<anyhow::Error>>(res?.map(|x| x * *required as u16)))
                })
                .collect();

            let mut sum = 0;
            while let Some(x) = futs.next().await {
                if let Some(x) = x? {
                    sum += x;
                } else {
                    return Ok(None);
                }
            }

            Ok(Some(sum + upcharge.unwrap()))
        }

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

    ctx.compute(&LookupCompanyResourceCost(company.clone(), resource.dupe()))
        .await
        .map_err(|e| Arc::new(anyhow::anyhow!(e)))?
}

#[async_trait]
impl Cost for DiceComputations {
    async fn resource_cost(&self, resource: &Resource) -> Result<Option<u16>, Arc<anyhow::Error>> {
        #[derive(Display, Debug, Hash, Eq, Dupe, Clone, PartialEq, RefCast, Allocative)]
        #[repr(transparent)]
        struct LookupResourceCost(Resource);
        #[async_trait]
        impl Key for LookupResourceCost {
            type Value = Result<Option<u16>, Arc<anyhow::Error>>;

            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellations: &CancellationContext,
            ) -> Self::Value {
                let companies = ctx
                    .compute(LookupResource::ref_cast(&self.0))
                    .await
                    .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;

                let costs = join_all(
                    companies
                        .iter()
                        .map(|company| lookup_company_resource_cost(ctx, company, &self.0)),
                )
                .await;

                Ok(costs
                    .into_iter()
                    .collect::<Result<Vec<_>, Arc<_>>>()?
                    .into_iter()
                    .min()
                    .flatten())
            }

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

        self.compute(LookupResourceCost::ref_cast(resource))
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?
    }
}

#[async_trait]
impl CostUpdater for DiceTransactionUpdater {
    async fn change_company_resource_cost(
        &mut self,
        company: &str,
        resource: &Resource,
        new_price: u16,
    ) -> anyhow::Result<()> {
        let company_lookup = LookupCompany(Arc::new(company.to_owned()));
        let old_company = self.existing_state().await.compute(&company_lookup).await?;
        let mut new_company = (*old_company).clone();
        let old_price = new_company.makes.get_mut(resource).ok_or_else(|| {
            anyhow::anyhow!("Tried to update cost for a resource company does not make")
        })?;
        *old_price = new_price;

        self.changed_to(vec![(company_lookup, Arc::new(new_company))])?;

        Ok(())
    }
}

#[derive(Display, Debug, Hash, Eq, Clone, Dupe, PartialEq, RefCast, Allocative)]
#[repr(transparent)]
pub struct LookupCompany(pub Arc<String>);
impl InjectedKey for LookupCompany {
    type Value = Arc<Company>;

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

#[derive(Display, Debug, Hash, Eq, Dupe, Clone, PartialEq, RefCast, Allocative)]
#[repr(transparent)]
pub struct LookupResource(pub Resource);
impl InjectedKey for LookupResource {
    type Value = Arc<Vec<LookupCompany>>;

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


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

#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))]
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

pub mod math_computation;
pub mod supply_chain;


/*
 * 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 dice::DetectCycles;
use dice::Dice;
use dice_examples::math_computation::parse_math_equation;
use dice_examples::math_computation::parse_math_equations;
use dice_examples::math_computation::Equation;
use dice_examples::math_computation::Math;
use dice_examples::math_computation::MathEquations;
use dice_examples::math_computation::Unit;
use dice_examples::math_computation::Var;
use dupe::Dupe;

fn var(name: &str) -> Var {
    Var(Arc::new(name.to_owned()))
}

#[tokio::test]
async fn test_literal() -> Result<(), Arc<anyhow::Error>> {
    let dice = Dice::builder().build(DetectCycles::Enabled);
    let mut ctx = dice.updater();
    let (var, eq) = parse_math_equation("a=5").unwrap();
    ctx.set_equation(var.dupe(), eq)?;
    let ctx = ctx.commit().await;

    assert_eq!(5, ctx.eval(var).await?);

    Ok(())
}

#[tokio::test]
async fn test_var() -> Result<(), Arc<anyhow::Error>> {
    let dice = Dice::builder().build(DetectCycles::Enabled);
    let mut ctx = dice.updater();

    let eqs = parse_math_equations(vec!["b=a", "a=3"]).unwrap();

    ctx.set_equations(eqs)?;
    let ctx = ctx.commit().await;

    assert_eq!(3, ctx.eval(var("b")).await?);
    Ok(())
}

#[tokio::test]
async fn test_compound() -> Result<(), Arc<anyhow::Error>> {
    let dice = Dice::builder().build(DetectCycles::Enabled);
    let mut ctx = dice.updater();

    let eq = vec!["x=1", "y=2", "a=x+y", "b=a+a"];
    let eq = parse_math_equations(eq).unwrap();

    ctx.set_equations(eq)?;
    let ctx = ctx.commit().await;

    assert_eq!(3, ctx.eval(var("a")).await?);
    assert_eq!(6, ctx.eval(var("b")).await?);

    Ok(())
}

#[tokio::test]
async fn test_changed_eq() -> Result<(), Arc<anyhow::Error>> {
    let dice = Dice::builder().build(DetectCycles::Enabled);
    let mut ctx = dice.updater();

    let eq = vec!["x=1", "y=2", "a=x+y", "b=a+a"];
    let eq = parse_math_equations(eq).unwrap();

    ctx.set_equations(eq)?;
    let ctx = ctx.commit().await;

    assert_eq!(6, ctx.eval(var("b")).await?);

    let mut ctx = dice.updater();
    ctx.set_equation(var("a"), Equation::Unit(Unit::Literal(4)))?;
    let ctx = ctx.commit().await;

    assert_eq!(8, ctx.eval(var("b")).await?);

    Ok(())
}


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

use std::collections::HashMap;
use std::sync::Arc;

use dice::DetectCycles;
use dice::Dice;
use dice::DiceTransaction;
use dice_examples::supply_chain::Company;
use dice_examples::supply_chain::Cost;
use dice_examples::supply_chain::CostUpdater;
use dice_examples::supply_chain::Resource;
use dice_examples::supply_chain::Setup;

async fn setup(companies: Vec<Company>) -> anyhow::Result<(Arc<Dice>, DiceTransaction)> {
    let dice = Dice::builder().build(DetectCycles::Enabled);
    let mut ctx = dice.updater();
    ctx.init_state()?;
    ctx.commit().await;

    let mut ctx = dice.updater();
    ctx.add_companies(companies).await?;

    Ok((dice, ctx.commit().await))
}

#[tokio::test]
async fn test_no_resources() -> Result<(), Arc<anyhow::Error>> {
    let (dice, ctx) = setup(vec![Company {
        name: Arc::new("hello world".to_owned()),
        makes: HashMap::new(),
    }])
    .await?;

    assert_eq!(None, ctx.resource_cost(&Resource::Wood).await?);
    let mut ctx = dice.updater();
    let success = ctx
        .change_company_resource_cost("hello world", &Resource::Stick, 5)
        .await;
    assert!(success.is_err());
    Ok(())
}

#[tokio::test]
async fn test_other_resource() -> Result<(), Arc<anyhow::Error>> {
    let (_, ctx) = setup(vec![Company {
        name: Arc::new("hello world".to_owned()),
        makes: [(Resource::Wood, 2)].iter().cloned().collect(),
    }])
    .await?;

    assert_eq!(None, ctx.resource_cost(&Resource::Stick).await?);
    assert_eq!(Some(2), ctx.resource_cost(&Resource::Wood).await?);
    Ok(())
}

#[tokio::test]
async fn test_simple() -> Result<(), Arc<anyhow::Error>> {
    let (_, ctx) = setup(vec![
        Company {
            name: Arc::new("Steve".to_owned()),
            makes: [
                (Resource::Plank, 2),
                (Resource::Wood, 1),
                (Resource::Stick, 4),
                (Resource::Cobblestone, 7),
            ]
            .iter()
            .cloned()
            .collect(),
        },
        Company {
            name: Arc::new("Alex".to_owned()),
            makes: [(Resource::Pickaxe, 10), (Resource::CraftingTable, 5)]
                .iter()
                .cloned()
                .collect(),
        },
    ])
    .await?;

    let expected = [
        (Resource::Wood, 1),
        (Resource::Cobblestone, 7),
        (Resource::Plank, 1 + 2),
        (Resource::Stick, 2 * (1 + 2) + 4),
        (Resource::Pickaxe, 3 * 7 + 2 * (2 * (1 + 2) + 4) + 10),
        (Resource::CraftingTable, 4 * (1 + 2) + 5),
    ];

    for (resource, cost) in &expected {
        assert_eq!(
            Some(*cost),
            ctx.resource_cost(resource).await?,
            "Testing {}",
            resource
        );
    }
    Ok(())
}

#[tokio::test]
async fn test_complex() -> Result<(), Arc<anyhow::Error>> {
    let (_, ctx) = setup(vec![
        Company {
            name: Arc::new("Steve".to_owned()),
            makes: [
                (Resource::Plank, 2),
                (Resource::Wood, 1),
                (Resource::Stick, 4),
                (Resource::Cobblestone, 7),
            ]
            .iter()
            .cloned()
            .collect(),
        },
        Company {
            name: Arc::new("Alex".to_owned()),
            makes: [(Resource::Pickaxe, 10), (Resource::CraftingTable, 5)]
                .iter()
                .cloned()
                .collect(),
        },
        Company {
            name: Arc::new("Bob".to_owned()),
            makes: [
                (Resource::Pickaxe, 5),
                (Resource::Plank, 1),
                (Resource::Stick, 2),
            ]
            .iter()
            .cloned()
            .collect(),
        },
        Company {
            name: Arc::new("Raj".to_owned()),
            makes: [(Resource::Stick, 1)].iter().cloned().collect(),
        },
    ])
    .await?;

    let expected = [
        (Resource::Wood, 1),
        (Resource::Cobblestone, 7),
        (Resource::Plank, 1 + 1),
        (Resource::Stick, 2 * (1 + 1) + 1),
        (Resource::Pickaxe, 3 * 7 + 2 * (2 * (1 + 1) + 1) + 5),
        (Resource::CraftingTable, 4 * (1 + 1) + 5),
    ];

    for (resource, cost) in &expected {
        assert_eq!(
            Some(*cost),
            ctx.resource_cost(resource).await?,
            "Testing {}",
            resource
        );
    }
    Ok(())
}

#[tokio::test]
async fn test_change_cost() -> Result<(), Arc<anyhow::Error>> {
    let (dice, _) = setup(vec![
        Company {
            name: Arc::new("Steve".to_owned()),
            makes: [
                (Resource::Plank, 2),
                (Resource::Wood, 1),
                (Resource::Stick, 4),
                (Resource::Cobblestone, 7),
            ]
            .iter()
            .cloned()
            .collect(),
        },
        Company {
            name: Arc::new("Alex".to_owned()),
            makes: [(Resource::Pickaxe, 10), (Resource::CraftingTable, 5)]
                .iter()
                .cloned()
                .collect(),
        },
    ])
    .await?;

    let mut ctx = dice.updater();
    ctx.change_company_resource_cost("Steve", &Resource::Stick, 2)
        .await
        .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;

    let ctx = ctx.commit().await;

    assert_eq!(
        Some(3 * 7 + 2 * (2 * (1 + 2) + 2) + 10),
        ctx.resource_cost(&Resource::Pickaxe).await?
    );

    let mut ctx = dice.updater();
    let update_success = ctx
        .change_company_resource_cost("Steve", &Resource::Pickaxe, 5)
        .await;

    assert!(update_success.is_err());

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

#![cfg(test)]

mod legacy;
mod opaque;
mod projection;


/*
 * 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 async_trait::async_trait;
use dice::DetectCycles;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceData;
use dice::Key;
use dice::UserComputationData;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use parking_lot::Mutex;

#[derive(Debug, PartialEq)]
enum KeyType {
    DoesNotReadOpaque,
    IsOpaque,
}

/// Key used in `compute_opaque`.
#[derive(
    Debug,
    derive_more::Display,
    Copy,
    Clone,
    Dupe,
    Eq,
    PartialEq,
    Hash,
    Allocative
)]
struct IsOpaque;
/// Key which computes the opaque value, but does not read it.
#[derive(
    Debug,
    derive_more::Display,
    Copy,
    Clone,
    Dupe,
    Eq,
    PartialEq,
    Hash,
    Allocative
)]
struct DoesNotReadOpaque;

/// Record computations, for test.
struct ComputationsTracker {
    computations: Vec<KeyType>,
}

#[async_trait]
impl Key for IsOpaque {
    type Value = Arc<String>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.global_data()
            .get::<Arc<Mutex<ComputationsTracker>>>()
            .unwrap()
            .lock()
            .computations
            .push(KeyType::IsOpaque);

        let number = ctx.per_transaction_data().data.get::<i32>().unwrap();

        Arc::new(format!("{}", number))
    }

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

#[async_trait]
impl Key for DoesNotReadOpaque {
    type Value = Arc<String>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        ctx.global_data()
            .get::<Arc<Mutex<ComputationsTracker>>>()
            .unwrap()
            .lock()
            .computations
            .push(KeyType::DoesNotReadOpaque);

        // Compute the key but ignore it.
        let _opaque = ctx.compute_opaque(&IsOpaque).await;

        Arc::new("Aaa".to_owned())
    }

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

#[tokio::test]
async fn key_does_not_read_opaque() -> anyhow::Result<()> {
    let tracker = Arc::new(Mutex::new(ComputationsTracker {
        computations: Vec::new(),
    }));

    let dice = {
        let mut builder = Dice::builder();
        builder.set(tracker.dupe());
        builder.build(DetectCycles::Enabled)
    };

    // Part 1: compute key which requests on opaque key, but does not use it.

    let mut dice_data = DiceData::new();
    dice_data.set(17);
    let mut ctx = dice.updater_with_data(UserComputationData {
        data: dice_data,
        ..Default::default()
    });

    // Both keys are computed.
    assert_eq!(
        "Aaa",
        &*ctx
            .existing_state()
            .await
            .compute(&DoesNotReadOpaque)
            .await?
    );
    assert_eq!(
        vec![KeyType::DoesNotReadOpaque, KeyType::IsOpaque],
        tracker.lock().computations
    );

    // Part 2: invalidate opaque key, and "compute" the key again.

    ctx.changed([IsOpaque])?;
    let ctx = ctx.commit().await;

    tracker.lock().computations.clear();

    // No keys are recomputed, because `Aaa` does not depend on `KeyWhichIsOpaque`.
    assert_eq!("Aaa", &*ctx.compute(&DoesNotReadOpaque).await?);
    assert_eq!(Vec::<KeyType>::new(), tracker.lock().computations);

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

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

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

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dice::DetectCycles;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dice::Key;
use dice::WhichSpawner;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use tempfile::NamedTempFile;

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

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

impl InjectedKey for EncodingConfig {
    type Value = Encoding;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    let dice = Dice::builder()
        .build_with_which_spawner(DetectCycles::Enabled, WhichSpawner::ExplicitCancel);

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

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

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

    set(":-)");

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

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

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

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

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

    Ok(())
}


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

mod demo;
mod projection;


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

#![cfg(test)]

use std::collections::HashMap;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use dice::DetectCycles;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceData;
use dice::DiceProjectionComputations;
use dice::Key;
use dice::ProjectionKey;
use dice::UserComputationData;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use parking_lot::Mutex;

/// We have three keys in this test:
/// * key for a global "configuration"
/// * key for a configuration "property"
/// * key for a "file" which reads a "property" from a "configuration"
///
/// This enum describes types of these keys.
#[derive(PartialEq, Debug)]
enum Computation {
    File,
    Config,
    ConfigProperty,
}

/// Produce side effects during computation of each key.
/// Of course, users should not do that, but we are testing internals.
struct RecordedComputations {
    computations: Vec<Computation>,
}

/// This is what "configuration" key reads from the outside world.
struct GlobalConfig {
    config: HashMap<String, String>,
}

/// "Evaluate" a file.
#[derive(Debug, derive_more::Display, Clone, Hash, PartialEq, Eq, Allocative)]
#[display(fmt = "{}", name)]
struct FileKey {
    name: String,
}

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

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        // Read "config".
        let config = ctx
            .compute_opaque(&ConfigKey)
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;
        // But use only one "property" of the "config",
        // which is the result of file evaluation.
        // We are testing that file evaluation is not invalidated
        // if unrelated configurations changed.
        let value = config
            .projection(&ConfigPropertyKey {
                key: "x".to_owned(),
            })
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;
        // Record we executed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::File);
        Ok(Arc::new(format!("<{}>", value)))
    }

    fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
        unreachable!("not used in test")
    }
}

/// Global "configuration".
#[derive(
    Debug,
    derive_more::Display,
    Clone,
    Dupe,
    Hash,
    PartialEq,
    Eq,
    Allocative
)]
#[display(fmt = "{:?}", self)]
struct ConfigKey;

#[async_trait]
impl Key for ConfigKey {
    type Value = Arc<HashMap<String, String>>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Arc<HashMap<String, String>> {
        // Record we performed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::Config);
        // And produce a value fetched from the outside world.
        Arc::new(
            ctx.per_transaction_data()
                .data
                .get::<GlobalConfig>()
                .unwrap()
                .config
                .clone(),
        )
    }

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

/// One "property" of the "configuration".
#[derive(Debug, derive_more::Display, Clone, Hash, PartialEq, Eq, Allocative)]
#[display(fmt = "{}", key)]
struct ConfigPropertyKey {
    key: String,
}

impl ProjectionKey for ConfigPropertyKey {
    /// We read a property from the config.
    type DeriveFromKey = ConfigKey;
    /// And produce a string.
    type Value = Arc<String>;

    fn compute(
        &self,
        derive_from: &Arc<HashMap<String, String>>,
        ctx: &DiceProjectionComputations,
    ) -> Arc<String> {
        // Record we performed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::ConfigProperty);
        // Fetch the config property.
        let value = derive_from
            .get(&self.key)
            .map_or_else(|| "NO".to_owned(), |x| x.to_owned());
        Arc::new(value)
    }

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

#[tokio::test]
async fn smoke() -> anyhow::Result<()> {
    let tracker = Arc::new(Mutex::new(RecordedComputations {
        computations: Vec::new(),
    }));

    let mut dice = Dice::builder();
    dice.set(tracker.dupe());
    let dice = dice.build(DetectCycles::Enabled);

    // Part 1: full evaluation. We request a file,
    // and dice evaluates: config -> config property -> file.

    let mut data = DiceData::new();
    data.set(GlobalConfig {
        config: HashMap::from_iter([("x".to_owned(), "X".to_owned())]),
    });
    let ctx = dice.updater_with_data(UserComputationData {
        data,
        ..Default::default()
    });

    let ctx = ctx.commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [
            Computation::Config,
            Computation::ConfigProperty,
            Computation::File
        ]
        .as_slice(),
        tracker.lock().computations.as_slice()
    );

    let mut ctx = ctx.into_updater();
    ctx.changed([ConfigKey])?;
    ctx.commit().await;
    tracker.lock().computations.clear();

    // Part 2: we update the config with the identical config.
    // Dice performs only "config" computation,
    // and the rest remains cached.

    let mut data = UserComputationData::new();
    data.data.set(GlobalConfig {
        config: HashMap::from_iter([("x".to_owned(), "X".to_owned())]),
    });
    let ctx = dice.updater_with_data(data).commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [Computation::Config].as_slice(),
        tracker.lock().computations.as_slice()
    );

    let mut ctx = ctx.into_updater();
    ctx.changed([ConfigKey])?;
    ctx.commit().await;
    tracker.lock().computations.clear();

    // Part 3: we update the config with a different config,
    // which however preserves the config property we are interested in.
    // So dice performs "config" and "config property" computations,
    // but since "config property" result is unchanged, "file" is not reevaluated.

    let mut data = UserComputationData::new();
    data.data.set(GlobalConfig {
        config: HashMap::from_iter([
            ("x".to_owned(), "X".to_owned()),
            ("y".to_owned(), "Y".to_owned()),
        ]),
    });
    let ctx = dice.updater_with_data(data).commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [Computation::Config, Computation::ConfigProperty].as_slice(),
        tracker.lock().computations.as_slice()
    );

    Ok(())
}


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

use std::hash::Hash;
use std::hash::Hasher;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use derive_more::Display;
use dice::DetectCycles;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceProjectionComputations;
use dice::Key;
use dice::ProjectionKey;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;

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

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

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

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

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

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

    impl Eq for ComputeOnce {}

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

    impl ProjectionKey for ComputeOnce {
        type DeriveFromKey = BaseK;
        type Value = ();

        fn compute(
            &self,
            _derive_from: &<<Self as ProjectionKey>::DeriveFromKey as Key>::Value,
            _ctx: &DiceProjectionComputations,
        ) -> Self::Value {
            self.0.fetch_add(1, Ordering::SeqCst);
        }

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

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

    let count = Arc::new(AtomicU8::new(0));

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

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

    let base = ctx.compute_opaque(&BaseK).await?;

    base.projection(&k)?;

    assert_eq!(count.load(Ordering::SeqCst), 1);
    base.projection(&k)?;
    assert_eq!(count.load(Ordering::SeqCst), 1);

    // call base again but technically same key
    let base = ctx.compute_opaque(&BaseK).await?;

    base.projection(&k)?;
    assert_eq!(count.load(Ordering::SeqCst), 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.
 */

mod concurrent;
mod smoke;


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

use allocative::Allocative;
use async_trait::async_trait;
use dice::DetectCycles;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceData;
use dice::DiceProjectionComputations;
use dice::Key;
use dice::ProjectionKey;
use dice::UserComputationData;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use parking_lot::Mutex;

/// We have three keys in this test:
/// * key for a global "configuration"
/// * key for a configuration "property"
/// * key for a "file" which reads a "property" from a "configuration"
///
/// This enum describes types of these keys.
#[derive(PartialEq, Debug)]
enum Computation {
    File,
    Config,
    ConfigProperty,
}

/// Produce side effects during computation of each key.
/// Of course, users should not do that, but we are testing internals.
struct RecordedComputations {
    computations: Vec<Computation>,
}

/// This is what "configuration" key reads from the outside world.
struct GlobalConfig {
    config: HashMap<String, String>,
}

/// "Evaluate" a file.
#[derive(Debug, derive_more::Display, Clone, Hash, PartialEq, Eq, Allocative)]
#[display(fmt = "{}", name)]
struct FileKey {
    name: String,
}

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

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        // Read "config".
        let config = ctx
            .compute_opaque(&ConfigKey)
            .await
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;
        // But use only one "property" of the "config",
        // which is the result of file evaluation.
        // We are testing that file evaluation is not invalidated
        // if unrelated configurations changed.
        let value = config
            .projection(&ConfigPropertyKey {
                key: "x".to_owned(),
            })
            .map_err(|e| Arc::new(anyhow::anyhow!(e)))?;
        // Record we executed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::File);
        Ok(Arc::new(format!("<{}>", value)))
    }

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

/// Global "configuration".
#[derive(
    Debug,
    derive_more::Display,
    Clone,
    Dupe,
    Hash,
    PartialEq,
    Eq,
    Allocative
)]
#[display(fmt = "{:?}", self)]
struct ConfigKey;

#[async_trait]
impl Key for ConfigKey {
    type Value = Arc<HashMap<String, String>>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Arc<HashMap<String, String>> {
        // Record we performed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::Config);
        // And produce a value fetched from the outside world.
        Arc::new(
            ctx.per_transaction_data()
                .data
                .get::<GlobalConfig>()
                .unwrap()
                .config
                .clone(),
        )
    }

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

/// One "property" of the "configuration".
#[derive(Debug, derive_more::Display, Clone, Hash, PartialEq, Eq, Allocative)]
#[display(fmt = "{}", key)]
struct ConfigPropertyKey {
    key: String,
}

impl ProjectionKey for ConfigPropertyKey {
    /// We read a property from the config.
    type DeriveFromKey = ConfigKey;
    /// And produce a string.
    type Value = Arc<String>;

    fn compute(
        &self,
        derive_from: &Arc<HashMap<String, String>>,
        ctx: &DiceProjectionComputations,
    ) -> Arc<String> {
        // Record we performed this computation.
        ctx.global_data()
            .get::<Arc<Mutex<RecordedComputations>>>()
            .unwrap()
            .lock()
            .computations
            .push(Computation::ConfigProperty);
        // Fetch the config property.
        let value = derive_from
            .get(&self.key)
            .map_or_else(|| "NO".to_owned(), |x| x.to_owned());
        Arc::new(value)
    }

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

#[tokio::test]
async fn smoke() -> anyhow::Result<()> {
    let tracker = Arc::new(Mutex::new(RecordedComputations {
        computations: Vec::new(),
    }));

    let mut dice = Dice::modern();
    dice.set(tracker.dupe());
    let dice = dice.build(DetectCycles::Enabled);

    // Part 1: full evaluation. We request a file,
    // and dice evaluates: config -> config property -> file.

    let mut data = DiceData::new();
    data.set(GlobalConfig {
        config: HashMap::from_iter([("x".to_owned(), "X".to_owned())]),
    });
    let ctx = dice.updater_with_data(UserComputationData {
        data,
        ..Default::default()
    });

    let ctx = ctx.commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [
            Computation::Config,
            Computation::ConfigProperty,
            Computation::File
        ]
        .as_slice(),
        tracker.lock().computations.as_slice()
    );

    let mut ctx = ctx.into_updater();
    ctx.changed([ConfigKey])?;
    ctx.commit().await;
    tracker.lock().computations.clear();

    // Part 2: we update the config with the identical config.
    // Dice performs only "config" computation,
    // and the rest remains cached.

    let mut data = UserComputationData::new();
    data.data.set(GlobalConfig {
        config: HashMap::from_iter([("x".to_owned(), "X".to_owned())]),
    });
    let ctx = dice.updater_with_data(data).commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [Computation::Config].as_slice(),
        tracker.lock().computations.as_slice()
    );

    let mut ctx = ctx.into_updater();
    ctx.changed([ConfigKey])?;
    ctx.commit().await;
    tracker.lock().computations.clear();

    // Part 3: we update the config with a different config,
    // which however preserves the config property we are interested in.
    // So dice performs "config" and "config property" computations,
    // but since "config property" result is unchanged, "file" is not reevaluated.

    let mut data = UserComputationData::new();
    data.data.set(GlobalConfig {
        config: HashMap::from_iter([
            ("x".to_owned(), "X".to_owned()),
            ("y".to_owned(), "Y".to_owned()),
        ]),
    });
    let ctx = dice.updater_with_data(data).commit().await;

    let file = ctx
        .compute(&FileKey {
            name: "file.fl".to_owned(),
        })
        .await?
        .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?;
    assert_eq!("<X>", &*file);

    assert_eq!(
        [Computation::Config, Computation::ConfigProperty].as_slice(),
        tracker.lock().computations.as_slice()
    );

    Ok(())
}


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

use std::collections::HashMap;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use crossbeam::queue::SegQueue;
use derivative::Derivative;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dice::Key;
use dupe::Dupe;
use futures::future;
use futures::FutureExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;
use serde::Deserialize;
use serde::Serialize;

#[derive(Clone, Copy, PartialEq, Eq, Hash, Display, Debug, Allocative)]
#[derive(Serialize, Deserialize)]
#[display(fmt = "key{}", _0)]
#[serde(transparent)]
pub struct Var(pub usize);

#[derive(Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize, Allocative)]
pub enum Unit {
    Variable(Var),
    Literal(bool),
}

async fn resolve_units(
    ctx: &DiceComputations,
    units: &[Unit],
    state: Arc<FuzzState>,
) -> anyhow::Result<Vec<bool>> {
    let futs = units.map(|unit| match unit {
        Unit::Variable(var) => ctx.eval(state.dupe(), *var),
        Unit::Literal(lit) => async move { Ok(*lit) }.boxed(),
    });
    future::join_all(futs).await.into_iter().collect()
}

#[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize, Allocative)]
pub enum Expr {
    Unit(Unit),
    Cond {
        test: Unit,
        then: Unit,
        otherwise: Unit,
    },
    Xor(Vec<Unit>),
}

async fn lookup_unit(ctx: &DiceComputations, var: Var) -> anyhow::Result<Arc<Expr>> {
    Ok(ctx.compute(&LookupVar(var)).await?)
}

#[derive(Clone, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "Lookup({})", _0)]
struct LookupVar(Var);
impl InjectedKey for LookupVar {
    type Value = Arc<Expr>;

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

pub trait FuzzEquations {
    fn set_equation(&mut self, var: Var, expr: Expr) -> anyhow::Result<()>;
    fn set_equations(&mut self, expr: impl IntoIterator<Item = (Var, Expr)>) -> anyhow::Result<()>;
}

impl FuzzEquations for DiceTransactionUpdater {
    fn set_equation(&mut self, var: Var, expr: Expr) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(LookupVar(var), Arc::new(expr))])?)
    }
    fn set_equations(
        &mut self,
        exprs: impl IntoIterator<Item = (Var, Expr)>,
    ) -> anyhow::Result<()> {
        Ok(self.changed_to(
            exprs
                .into_iter()
                .map(|(var, expr)| (LookupVar(var), Arc::new(expr)))
                .collect::<Vec<_>>(),
        )?)
    }
}

#[async_trait]
pub trait FuzzMath {
    async fn eval(&self, state: Arc<FuzzState>, var: Var) -> anyhow::Result<bool>;
}

#[async_trait]
impl FuzzMath for DiceComputations {
    async fn eval(&self, state: Arc<FuzzState>, var: Var) -> anyhow::Result<bool> {
        Ok(*self
            .compute(&state.eval_var(var))
            .await?
            .map_err(|e| anyhow::anyhow!(format!("{:#}", e)))?
            .as_ref())
    }
}

#[derive(Clone, Copy, PartialEq, Debug, Serialize, Deserialize)]
pub enum ComputationStep {
    ReturnStable,
    ReturnTransient,
}

#[derive(Allocative)]
pub struct FuzzState {
    #[allocative(skip)]
    pub steps: HashMap<Var, SegQueue<ComputationStep>>,
}

impl FuzzState {
    pub(crate) fn new() -> Self {
        Self {
            steps: HashMap::new(),
        }
    }

    pub(crate) fn eval_var(self: &Arc<Self>, key: Var) -> EvalVar {
        EvalVar {
            key,
            state: self.dupe(),
        }
    }

    pub(crate) fn next_step_for_var(self: &Arc<Self>, var: Var) -> ComputationStep {
        match self.steps.get(&var).and_then(|q| q.pop()) {
            Some(step) => step,
            _ => ComputationStep::ReturnStable,
        }
    }
}

#[derive(Derivative, Clone, Display, Allocative)]
#[derivative(Hash, Debug)]
#[display(fmt = "Eval({})", key)]
pub struct EvalVar {
    key: Var,
    #[derivative(Debug = "ignore", Hash = "ignore")]
    state: Arc<FuzzState>,
}

impl PartialEq for EvalVar {
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key && Arc::ptr_eq(&self.state, &other.state)
    }
}

impl Eq for EvalVar {}

#[derive(PartialEq, Eq, Clone, Dupe, Allocative)]
pub enum MaybeTransient<T> {
    Stable(T),
    Transient(T),
}

impl<T> AsRef<T> for MaybeTransient<T> {
    fn as_ref(&self) -> &T {
        match self {
            MaybeTransient::Stable(v) => v,
            MaybeTransient::Transient(v) => v,
        }
    }
}

impl<T> MaybeTransient<T> {
    fn is_transient(&self) -> bool {
        matches!(self, MaybeTransient::Transient(_))
    }
}

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

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellations: &CancellationContext,
    ) -> Self::Value {
        let step = self.state.next_step_for_var(self.key);
        let ret = match &*lookup_unit(ctx, self.key).await.map_err(Arc::new)? {
            Expr::Unit(unit) => resolve_units(ctx, &[unit.clone()], self.state.dupe())
                .await
                .map_err(Arc::new)?[0],
            Expr::Cond {
                test,
                then,
                otherwise,
            } => {
                if resolve_units(ctx, &[test.clone()], self.state.dupe())
                    .await
                    .map_err(Arc::new)?[0]
                {
                    resolve_units(ctx, &[then.clone()], self.state.dupe())
                        .await
                        .map_err(Arc::new)?[0]
                } else {
                    resolve_units(ctx, &[otherwise.clone()], self.state.dupe())
                        .await
                        .map_err(Arc::new)?[0]
                }
            }
            Expr::Xor(vars) => resolve_units(ctx, vars, self.state.dupe())
                .await
                .map_err(Arc::new)?
                .into_iter()
                .reduce(|x, y| x ^ y)
                .unwrap_or(false),
        };
        Ok(match step {
            ComputationStep::ReturnStable => MaybeTransient::Stable(ret),
            ComputationStep::ReturnTransient => MaybeTransient::Transient(ret),
        })
    }

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

    fn validity(x: &Self::Value) -> bool {
        x.as_ref().map_or(false, |x| !x.is_transient())
    }
}

#[cfg(test)]
mod tests {
    use dice::DetectCycles;
    use dice::Dice;

    use super::*;

    #[tokio::test]
    pub async fn test_smoke() -> anyhow::Result<()> {
        let empty_state = Arc::new(FuzzState::new());
        let dice = Dice::builder().build(DetectCycles::Disabled);
        let ctx = {
            let mut ctx = dice.updater();
            // let x1 = true
            ctx.set_equation(Var(1), Expr::Unit(Unit::Literal(true)))?;
            // let x2 = x1
            ctx.set_equation(Var(2), Expr::Xor(vec![Unit::Variable(Var(1))]))?;
            // let x3 = x1 ^ x1 = false
            ctx.set_equation(
                Var(3),
                Expr::Xor(vec![Unit::Variable(Var(1)), Unit::Variable(Var(1))]),
            )?;
            // let x4 = if x1 then x2 else x3
            ctx.set_equation(
                Var(4),
                Expr::Cond {
                    test: Unit::Variable(Var(1)),
                    then: Unit::Variable(Var(2)),
                    otherwise: Unit::Variable(Var(3)),
                },
            )?;
            ctx.commit().await
        };
        assert!(ctx.eval(empty_state.dupe(), Var(1)).await?);
        assert!(ctx.eval(empty_state.dupe(), Var(2)).await?);
        assert!(!ctx.eval(empty_state.dupe(), Var(3)).await?);
        assert!(ctx.eval(empty_state.dupe(), Var(4)).await?);
        Ok(())
    }
}