/*
 * 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::cmp::Ordering;
use std::hash::Hash;

use allocative::Allocative;
use starlark_map::small_map;
use starlark_map::small_map::SmallMap;
use starlark_map::Equivalent;

/// Wrapper for `SmallMap` which considers map equal if iteration order is equal.
#[derive(Debug, Clone, Allocative)]
pub struct OrderedMap<K, V>(SmallMap<K, V>);

impl<K, V> OrderedMap<K, V> {
    #[inline]
    pub fn new() -> OrderedMap<K, V> {
        OrderedMap(SmallMap::new())
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> OrderedMap<K, V> {
        OrderedMap(SmallMap::with_capacity(capacity))
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    #[inline]
    pub fn iter(&self) -> small_map::Iter<K, V> {
        self.0.iter()
    }

    #[inline]
    pub fn iter_mut(&mut self) -> small_map::IterMut<K, V> {
        self.0.iter_mut()
    }

    #[inline]
    pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> {
        self.0.keys()
    }

    #[inline]
    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> {
        self.0.values()
    }

    #[inline]
    pub fn values_mut(&mut self) -> impl ExactSizeIterator<Item = &mut V> {
        self.0.values_mut()
    }

    #[inline]
    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.0.get(k)
    }

    #[inline]
    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.0.get_mut(k)
    }

    /// Find an entry by an index.
    #[inline]
    pub fn get_index(&self, index: usize) -> Option<(&K, &V)> {
        self.0.get_index(index)
    }

    /// Find an entry index for a given key.
    #[inline]
    pub fn get_index_of<Q>(&self, key: &Q) -> Option<usize>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.0.get_index_of(key)
    }

    #[inline]
    pub fn contains_key<Q>(&self, k: &Q) -> bool
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.0.contains_key(k)
    }

    #[inline]
    pub fn insert(&mut self, k: K, v: V) -> Option<V>
    where
        K: Hash + Eq,
    {
        self.0.insert(k, v)
    }

    #[inline]
    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.0.remove(k)
    }

    #[inline]
    pub fn clear(&mut self) {
        self.0.clear()
    }

    #[inline]
    pub fn entry(&mut self, k: K) -> small_map::Entry<'_, K, V>
    where
        K: Hash + Eq,
    {
        self.0.entry(k)
    }

    #[inline]
    pub fn sort_keys(&mut self)
    where
        K: Ord,
    {
        self.0.sort_keys()
    }
}

impl<K, V> Default for OrderedMap<K, V> {
    #[inline]
    fn default() -> OrderedMap<K, V> {
        OrderedMap(SmallMap::default())
    }
}

impl<K: PartialEq, V: PartialEq> PartialEq for OrderedMap<K, V> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.0.eq_ordered(&other.0)
    }
}

impl<K: Eq, V: Eq> Eq for OrderedMap<K, V> {}

impl<K: Hash, V: Hash> Hash for OrderedMap<K, V> {
    #[inline]
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.hash_ordered(state)
    }
}

impl<K: PartialOrd, V: PartialOrd> PartialOrd for OrderedMap<K, V> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.iter().partial_cmp(other.iter())
    }
}

impl<K: Eq + Ord, V: Eq + Ord> Ord for OrderedMap<K, V> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.iter().cmp(other.iter())
    }
}

impl<K, V> FromIterator<(K, V)> for OrderedMap<K, V>
where
    K: Hash + Eq,
{
    #[inline]
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> OrderedMap<K, V> {
        OrderedMap(SmallMap::from_iter(iter))
    }
}

impl<K, V> IntoIterator for OrderedMap<K, V> {
    type Item = (K, V);
    type IntoIter = small_map::IntoIter<K, V>;

    #[inline]
    fn into_iter(self) -> small_map::IntoIter<K, V> {
        self.0.into_iter()
    }
}

impl<'a, K, V> IntoIterator for &'a OrderedMap<K, V> {
    type Item = (&'a K, &'a V);
    type IntoIter = small_map::Iter<'a, K, V>;

    #[inline]
    fn into_iter(self) -> small_map::Iter<'a, K, V> {
        self.0.iter()
    }
}

impl<'a, K, V> IntoIterator for &'a mut OrderedMap<K, V> {
    type Item = (&'a K, &'a mut V);
    type IntoIter = small_map::IterMut<'a, K, V>;

    #[inline]
    fn into_iter(self) -> small_map::IterMut<'a, K, V> {
        self.0.iter_mut()
    }
}

#[cfg(test)]
mod tests {
    use std::cell::Cell;
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;

    use crate::collections::ordered_map::OrderedMap;

    #[test]
    fn test_keys_are_not_hashed_when_map_is_hashed() {
        struct Tester {
            /// Number of times `hash` was called.
            hash_count: Cell<u32>,
        }

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

        impl Eq for Tester {}

        impl Hash for Tester {
            fn hash<H: std::hash::Hasher>(&self, _state: &mut H) {
                self.hash_count.set(self.hash_count.get() + 1);
            }
        }

        let map = OrderedMap::from_iter([(
            Tester {
                hash_count: Cell::new(0),
            },
            Tester {
                hash_count: Cell::new(0),
            },
        )]);
        assert_eq!(1, map.keys().next().unwrap().hash_count.get());
        assert_eq!(0, map.values().next().unwrap().hash_count.get());

        let mut hasher = DefaultHasher::new();

        map.hash(&mut hasher);
        assert_eq!(1, map.keys().next().unwrap().hash_count.get());
        assert_eq!(1, map.values().next().unwrap().hash_count.get());

        map.hash(&mut hasher);
        assert_eq!(1, map.keys().next().unwrap().hash_count.get());
        assert_eq!(2, map.values().next().unwrap().hash_count.get());

        map.hash(&mut hasher);
        assert_eq!(1, map.keys().next().unwrap().hash_count.get());
        assert_eq!(3, map.values().next().unwrap().hash_count.get());
    }
}


/*
 * 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::cmp::Ordering;
use std::hash::Hash;

use allocative::Allocative;
use starlark_map::small_set;
use starlark_map::small_set::SmallSet;
use starlark_map::Equivalent;
use starlark_map::Hashed;

/// `SmallSet` wrapper, but equality and hash of self depends on iteration order.
#[derive(Debug, Clone, Allocative)]
pub struct OrderedSet<T>(SmallSet<T>);

#[derive(Debug)]
pub struct OccupiedError<'a, T> {
    pub value: T,
    pub occupied: &'a T,
}

impl<T> OrderedSet<T> {
    #[inline]
    pub fn new() -> OrderedSet<T> {
        OrderedSet(SmallSet::new())
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> OrderedSet<T> {
        OrderedSet(SmallSet::with_capacity(capacity))
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    #[inline]
    pub fn get<Q>(&self, value: &Q) -> Option<&T>
    where
        Q: Hash + Equivalent<T> + ?Sized,
        T: Eq,
    {
        self.0.get(value)
    }

    #[inline]
    pub fn contains<Q>(&self, value: &Q) -> bool
    where
        Q: Hash + Equivalent<T> + ?Sized,
        T: Eq,
    {
        self.0.contains(value)
    }

    #[inline]
    pub fn get_index(&self, index: usize) -> Option<&T> {
        self.0.get_index(index)
    }

    #[inline]
    pub fn get_index_of<Q>(&self, value: &Q) -> Option<usize>
    where
        Q: Hash + Equivalent<T> + ?Sized,
        T: Eq,
    {
        self.0.get_index_of(value)
    }

    #[inline]
    pub fn take<Q>(&mut self, value: &Q) -> Option<T>
    where
        Q: Hash + Equivalent<T> + ?Sized,
        T: Eq,
    {
        self.0.take(value)
    }

    #[inline]
    pub fn iter(&self) -> small_set::Iter<T> {
        self.0.iter()
    }

    #[inline]
    pub fn first(&self) -> Option<&T> {
        self.0.first()
    }

    #[inline]
    pub fn last(&self) -> Option<&T> {
        self.0.last()
    }

    #[inline]
    pub fn insert(&mut self, value: T) -> bool
    where
        T: Hash + Eq,
    {
        self.0.insert(value)
    }

    /// Insert an element if element is not present in the set,
    /// otherwise return the element.
    #[inline]
    pub fn try_insert<'a>(&'a mut self, value: T) -> Result<(), OccupiedError<'a, T>>
    where
        T: Hash + Eq,
    {
        let value = Hashed::new(value);

        // SAFETY: Extending the lifetime of the `self` reference here is valid because:
        //  1. The `self` reference already has lifetime `'a`.
        //  2. Within this block, we only use `this`, not `self`.
        //  3. Unless we return from this block, we do not store any references
        //     derived from `this` which might later conflict with `self`.
        // Writing the equivalent safe code does not work,
        // due to a well-known limitation of the borrow-checker:
        // https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=703ece3ad1d40f3c0d6ba1cd8b18194e
        // `-Zpolonius` fixes the issue.
        let this: &'a mut Self = unsafe { &mut *(self as *mut Self) };
        if let Some(occupied) = this.0.get_hashed(value.as_ref()) {
            return Err(OccupiedError {
                value: value.into_key(),
                occupied,
            });
        }

        self.0.insert_hashed_unique_unchecked(value);
        Ok(())
    }

    #[inline]
    pub fn clear(&mut self) {
        self.0.clear()
    }

    #[inline]
    pub fn sort(&mut self)
    where
        T: Ord,
    {
        self.0.sort()
    }

    #[inline]
    pub fn union<'a>(&'a self, other: &'a Self) -> impl Iterator<Item = &'a T>
    where
        T: Eq + Hash,
    {
        self.0.union(&other.0)
    }
}

impl<T> Default for OrderedSet<T> {
    #[inline]
    fn default() -> OrderedSet<T> {
        OrderedSet::new()
    }
}

impl<T: Eq> PartialEq for OrderedSet<T> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.0.eq_ordered(&other.0)
    }
}

impl<T: Eq> Eq for OrderedSet<T> {}

impl<T: Eq + PartialOrd> PartialOrd for OrderedSet<T> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.0.iter().partial_cmp(other.0.iter())
    }
}

impl<T: Eq + Ord> Ord for OrderedSet<T> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.iter().cmp(other.0.iter())
    }
}

impl<T: Hash> Hash for OrderedSet<T> {
    #[inline]
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.hash_ordered(state)
    }
}

impl<T> IntoIterator for OrderedSet<T> {
    type Item = T;
    type IntoIter = small_set::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<'a, T> IntoIterator for &'a OrderedSet<T> {
    type Item = &'a T;
    type IntoIter = small_set::Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<T> From<SmallSet<T>> for OrderedSet<T> {
    #[inline]
    fn from(set: SmallSet<T>) -> OrderedSet<T> {
        OrderedSet(set)
    }
}

impl<T> FromIterator<T> for OrderedSet<T>
where
    T: Eq + Hash,
{
    #[inline]
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        OrderedSet(SmallSet::from_iter(iter))
    }
}

#[cfg(test)]
mod tests {
    use std::cell::Cell;
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;
    use std::sync::Arc;

    use dupe::Dupe;

    use crate::collections::ordered_set::OrderedSet;

    #[test]
    fn test_keys_are_not_hashed_when_map_is_hashed() {
        struct Tester {
            /// Number of times `hash` was called.
            hash_count: Cell<u32>,
        }

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

        impl Eq for Tester {}

        impl Hash for Tester {
            fn hash<H: std::hash::Hasher>(&self, _state: &mut H) {
                self.hash_count.set(self.hash_count.get() + 1);
            }
        }

        let set = OrderedSet::from_iter([Tester {
            hash_count: Cell::new(0),
        }]);
        assert_eq!(1, set.iter().next().unwrap().hash_count.get());

        let mut hasher = DefaultHasher::new();

        set.hash(&mut hasher);
        assert_eq!(1, set.iter().next().unwrap().hash_count.get());

        set.hash(&mut hasher);
        assert_eq!(1, set.iter().next().unwrap().hash_count.get());

        set.hash(&mut hasher);
        assert_eq!(1, set.iter().next().unwrap().hash_count.get());
    }

    #[test]
    fn test_insert_unique() {
        let mut set: OrderedSet<Arc<u32>> = OrderedSet::new();
        let inserted = set.try_insert(Arc::new(1));
        assert!(inserted.is_ok());

        let one = Arc::new(1);
        let inserted = set.try_insert(one.dupe());
        assert!(Arc::ptr_eq(&inserted.unwrap_err().value, &one));
    }
}


/*
 * 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 allocative::Allocative;
use starlark_map::small_map;
use starlark_map::Equivalent;

use crate::collections::ordered_map::OrderedMap;

/// `IndexMap` but with keys sorted.
#[derive(Debug, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Allocative)]
pub struct SortedMap<K, V> {
    map: OrderedMap<K, V>,
}

impl<K: Ord + Hash, V> Default for SortedMap<K, V> {
    #[inline]
    fn default() -> Self {
        SortedMap {
            map: OrderedMap::default(),
        }
    }
}

impl<K, V> SortedMap<K, V>
where
    K: Ord + Hash,
{
    #[inline]
    pub fn new() -> SortedMap<K, V> {
        SortedMap {
            map: OrderedMap::new(),
        }
    }

    #[inline]
    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&K, &V)> {
        self.map.iter()
    }

    #[inline]
    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = (&K, &mut V)> {
        self.map.iter_mut()
    }

    #[inline]
    pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> {
        self.map.keys()
    }

    #[inline]
    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> {
        self.map.values()
    }

    #[inline]
    pub fn values_mut(&mut self) -> impl ExactSizeIterator<Item = &mut V> {
        self.map.values_mut()
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.map.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    #[inline]
    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
    where
        Q: Hash + Equivalent<K>,
    {
        self.map.get(key)
    }

    #[inline]
    pub fn contains_key<Q>(&self, k: &Q) -> bool
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.map.contains_key(k)
    }
}

impl<K: Ord + Hash, V> FromIterator<(K, V)> for SortedMap<K, V> {
    #[inline]
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
        let map = OrderedMap::from_iter(iter);
        SortedMap::from(map)
    }
}

impl<K: Ord + Hash, V> From<OrderedMap<K, V>> for SortedMap<K, V> {
    #[inline]
    fn from(mut map: OrderedMap<K, V>) -> SortedMap<K, V> {
        map.sort_keys();
        SortedMap { map }
    }
}

impl<K: Ord + Hash, V> IntoIterator for SortedMap<K, V> {
    type Item = (K, V);
    type IntoIter = small_map::IntoIter<K, V>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.map.into_iter()
    }
}

impl<'a, K: Ord + Hash, V> IntoIterator for &'a SortedMap<K, V> {
    type Item = (&'a K, &'a V);
    type IntoIter = small_map::Iter<'a, K, V>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.map.iter()
    }
}

impl<'a, K: Ord + Hash, V> IntoIterator for &'a mut SortedMap<K, V> {
    type Item = (&'a K, &'a mut V);
    type IntoIter = small_map::IterMut<'a, K, V>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.map.iter_mut()
    }
}

#[cfg(test)]
mod tests {
    use crate::collections::sorted_map::SortedMap;

    #[test]
    fn test_from_iter() {
        let map = SortedMap::from_iter([(1, 2), (5, 6), (3, 4)]);
        assert_eq!(
            vec![(&1, &2), (&3, &4), (&5, &6)],
            map.iter().collect::<Vec<_>>()
        );
    }
}


/*
 * 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 allocative::Allocative;
use starlark_map::small_set;
use starlark_map::small_set::SmallSet;
use starlark_map::Equivalent;

use crate::collections::ordered_set::OrderedSet;
use crate::collections::sorted_vec::SortedVec;

/// An immutable IndexSet with values guaranteed to be sorted.
#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd, Allocative)]
pub struct SortedSet<T: Eq + Hash> {
    inner: OrderedSet<T>,
}

impl<T> SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    pub fn new() -> SortedSet<T> {
        SortedSet {
            inner: OrderedSet::new(),
        }
    }

    #[inline]
    pub fn new_unchecked(inner: OrderedSet<T>) -> Self {
        Self { inner }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    #[inline]
    pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
        self.inner.iter()
    }

    #[inline]
    pub fn get<Q: ?Sized>(&self, value: &Q) -> Option<&T>
    where
        Q: Hash + Equivalent<T>,
    {
        self.inner.get(value)
    }

    #[inline]
    pub fn contains<Q: ?Sized>(&self, value: &Q) -> bool
    where
        Q: Hash + Equivalent<T>,
    {
        self.inner.contains(value)
    }

    #[inline]
    pub fn get_index(&self, index: usize) -> Option<&T> {
        self.inner.get_index(index)
    }

    #[inline]
    pub fn union<'a>(&'a self, other: &'a Self) -> impl Iterator<Item = &'a T> {
        // TODO(nga): return sorted.
        self.inner.union(&other.inner)
    }
}

impl<T> From<OrderedSet<T>> for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    fn from(mut inner: OrderedSet<T>) -> SortedSet<T> {
        inner.sort();
        SortedSet { inner }
    }
}

impl<T> From<SmallSet<T>> for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    fn from(inner: SmallSet<T>) -> SortedSet<T> {
        SortedSet::from(OrderedSet::from(inner))
    }
}

impl<T> From<SortedVec<T>> for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    fn from(inner: SortedVec<T>) -> SortedSet<T> {
        SortedSet {
            inner: OrderedSet::from_iter(inner),
        }
    }
}

impl<T> FromIterator<T> for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let mut inner = OrderedSet::from_iter(iter);
        inner.sort();
        SortedSet { inner }
    }
}

impl<T> IntoIterator for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    type Item = T;
    type IntoIter = small_set::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.inner.into_iter()
    }
}

impl<'a, T> IntoIterator for &'a SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    type Item = &'a T;
    type IntoIter = small_set::Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.inner.iter()
    }
}

impl<T> Default for SortedSet<T>
where
    T: Eq + Ord + Hash,
{
    #[inline]
    fn default() -> Self {
        SortedSet::new()
    }
}


/*
 * 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::slice;
use std::vec;

use allocative::Allocative;

/// Type which enfoces that its elements are sorted. That's it.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Allocative)]
pub struct SortedVec<T> {
    vec: Vec<T>,
}

impl<T> SortedVec<T> {
    #[inline]
    pub fn new() -> SortedVec<T> {
        SortedVec { vec: Vec::new() }
    }

    /// Construct without checking that the elements are sorted.
    #[inline]
    pub fn new_unchecked(vec: Vec<T>) -> SortedVec<T>
    where
        T: Ord,
    {
        debug_assert!(vec.iter().zip(vec.iter().skip(1)).all(|(a, b)| a <= b));
        SortedVec { vec }
    }

    #[inline]
    pub fn iter(&self) -> slice::Iter<T> {
        self.vec.iter()
    }
}

impl<T: Ord> From<Vec<T>> for SortedVec<T> {
    #[inline]
    fn from(mut vec: Vec<T>) -> Self {
        vec.sort();
        SortedVec { vec }
    }
}

impl<T: Ord> FromIterator<T> for SortedVec<T> {
    #[inline]
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let vec = Vec::from_iter(iter);
        SortedVec::from(vec)
    }
}

impl<T> IntoIterator for SortedVec<T> {
    type Item = T;
    type IntoIter = vec::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.vec.into_iter()
    }
}

#[cfg(test)]
mod tests {
    use crate::collections::sorted_vec::SortedVec;

    /// Test `new_unchecked` panics in debug mode when the elements are not sorted.
    #[cfg(debug_assertions)]
    #[test]
    #[should_panic]
    fn test_new_unchecked() {
        SortedVec::new_unchecked(vec![1, 3, 2]);
    }
}


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

use std::fmt::Debug;
use std::hash::Hash;
use std::hash::Hasher;
use std::mem;

use allocative::Allocative;
use hashbrown::raw::RawTable;
use starlark_map::Equivalent;
use starlark_map::StarlarkHasher;

/// Hash map which does not expose any insertion order-specific behavior
/// (except `Debug`).
#[derive(Clone, Allocative)]
pub struct UnorderedMap<K, V>(RawTable<(K, V)>);

impl<K, V> Default for UnorderedMap<K, V> {
    #[inline]
    fn default() -> UnorderedMap<K, V> {
        UnorderedMap::new()
    }
}

impl<K, V> UnorderedMap<K, V> {
    #[inline]
    pub fn new() -> UnorderedMap<K, V> {
        UnorderedMap(RawTable::new())
    }

    #[inline]
    pub fn with_capacity(n: usize) -> UnorderedMap<K, V> {
        UnorderedMap(RawTable::with_capacity(n))
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    #[inline]
    fn hash<Q: Hash + ?Sized>(k: &Q) -> u64 {
        let mut hasher = StarlarkHasher::new();
        k.hash(&mut hasher);
        hasher.finish()
    }

    #[inline]
    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        let hash = Self::hash(k);
        self.0
            .get(hash, |(next_k, _v)| k.equivalent(next_k))
            .map(|(_, v)| v)
    }

    #[inline]
    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        let hash = Self::hash(k);
        self.0
            .get_mut(hash, |(next_k, _v)| k.equivalent(next_k))
            .map(|(_, v)| v)
    }

    #[inline]
    pub fn contains_key<Q>(&self, k: &Q) -> bool
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        self.get(k).is_some()
    }

    #[inline]
    pub fn insert(&mut self, k: K, v: V) -> Option<V>
    where
        K: Hash + Eq,
    {
        let hash = Self::hash(&k);
        if let Some((_k, existing_value)) =
            self.0.get_mut(hash, |(next_k, _v)| k.equivalent(next_k))
        {
            Some(mem::replace(existing_value, v))
        } else {
            self.0.insert(hash, (k, v), |(k, _v)| Self::hash(k));
            None
        }
    }

    #[inline]
    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where
        Q: Hash + Equivalent<K> + ?Sized,
    {
        let hash = Self::hash(k);
        self.0
            .remove_entry(hash, |(next_k, _v)| k.equivalent(next_k))
            .map(|(_, v)| v)
    }

    #[inline]
    pub fn clear(&mut self) {
        self.0.clear();
    }

    /// This function is private.
    #[inline]
    fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
        unsafe { self.0.iter().map(|e| (&e.as_ref().0, &e.as_ref().1)) }
    }
}

impl<K: Debug, V: Debug> Debug for UnorderedMap<K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

impl<K: Eq + Hash, V: Eq> PartialEq for UnorderedMap<K, V> {
    fn eq(&self, other: &Self) -> bool {
        self.len() == other.len() && self.iter().all(|(k, v)| other.get(k) == Some(v))
    }
}

impl<K: Eq + Hash, V: Eq> Eq for UnorderedMap<K, V> {}

impl<K: Hash, V: Hash> Hash for UnorderedMap<K, V> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.len().hash(state);
        let mut sum: u64 = 0;
        for (k, v) in self.iter() {
            let mut hasher = StarlarkHasher::new();
            (k, v).hash(&mut hasher);
            sum = sum.wrapping_add(hasher.finish());
        }
        sum.hash(state);
    }
}

impl<K: Eq + Hash, V> FromIterator<(K, V)> for UnorderedMap<K, V> {
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> UnorderedMap<K, V> {
        let iter = iter.into_iter();
        let mut map = UnorderedMap::with_capacity(iter.size_hint().0);
        for (k, v) in iter {
            map.insert(k, v);
        }
        map
    }
}

#[cfg(test)]
mod tests {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::Hash;
    use std::hash::Hasher;

    use crate::collections::unordered_map::UnorderedMap;

    #[test]
    fn test_hash() {
        let a = UnorderedMap::from_iter([(1, 2), (3, 4)]);
        let b = UnorderedMap::from_iter([(3, 4), (1, 2)]);

        fn hash<H: Hash>(x: &H) -> u64 {
            let mut hasher = DefaultHasher::new();
            x.hash(&mut hasher);
            hasher.finish()
        }

        assert_eq!(hash(&a), hash(&b));
    }

    #[test]
    fn test_eq() {
        let a = UnorderedMap::from_iter([(1, 2), (3, 4)]);
        let b = UnorderedMap::from_iter([(3, 4), (1, 2)]);
        let c = UnorderedMap::from_iter([(1, 2), (9, 10)]);
        let d = UnorderedMap::from_iter([(1, 2), (3, 4), (5, 6)]);
        assert_eq!(a, b);
        assert_ne!(a, c);
        assert_ne!(a, d);
    }

    #[test]
    fn test_insert_remove() {
        let mut map = UnorderedMap::new();
        assert_eq!(map.insert(1, 2), None);
        assert_eq!(UnorderedMap::from_iter([(1, 2)]), map);
        assert_eq!(map.insert(1, 3), Some(2));
        assert_eq!(UnorderedMap::from_iter([(1, 3)]), map);
        assert_eq!(map.remove(&1), Some(3));
        assert_eq!(UnorderedMap::new(), map);
    }
}


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

use std::hash::Hash;

use allocative::Allocative;
use starlark_map::Equivalent;

use crate::collections::unordered_map::UnorderedMap;

#[derive(Clone, Allocative, Debug)]
pub struct UnorderedSet<T> {
    map: UnorderedMap<T, ()>,
}

impl<T> UnorderedSet<T> {
    #[inline]
    pub fn new() -> UnorderedSet<T> {
        UnorderedSet {
            map: UnorderedMap::new(),
        }
    }

    #[inline]
    pub fn with_capacity(n: usize) -> UnorderedSet<T> {
        UnorderedSet {
            map: UnorderedMap::with_capacity(n),
        }
    }

    #[inline]
    pub fn insert(&mut self, k: T) -> bool
    where
        T: Hash + Eq,
    {
        self.map.insert(k, ()).is_none()
    }

    #[inline]
    pub fn clear(&mut self) {
        self.map.clear();
    }

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub fn contains<Q>(&self, value: &Q) -> bool
    where
        Q: Hash + Equivalent<T> + ?Sized,
    {
        self.map.contains_key(value)
    }
}

impl<T: Eq + Hash> PartialEq for UnorderedSet<T> {
    fn eq(&self, other: &Self) -> bool {
        self.map == other.map
    }
}

impl<T: Eq + Hash> Eq for UnorderedSet<T> {}

impl<T: Eq + Hash> FromIterator<T> for UnorderedSet<T> {
    #[allow(clippy::from_iter_instead_of_collect)]
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> UnorderedSet<T> {
        UnorderedSet {
            map: UnorderedMap::from_iter(iter.into_iter().map(|v| (v, ()))),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::collections::unordered_set::UnorderedSet;

    #[test]
    fn test_insert() {
        let mut set = UnorderedSet::new();
        assert!(set.insert(10));
        assert!(!set.insert(10));
        assert!(set.insert(20));
        assert!(!set.insert(20));
        assert_eq!(set.len(), 2);
    }
}


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

fn main() -> io::Result<()> {
    let proto_files = &["worker.proto"];

    buck2_protoc_dev::configure()
        .setup_protoc()
        .compile(proto_files, &["."])
}


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

tonic::include_proto!("worker");


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

use std::collections::hash_map::DefaultHasher;
use std::env;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hash;
use std::hash::Hasher;
use std::str::FromStr;

use allocative::Allocative;
use anyhow::Context;
use dupe::Dupe;
use uuid::Uuid;

use crate::BUCK_WRAPPER_UUID_ENV_VAR;

#[derive(Debug, thiserror::Error)]
enum TraceIdError {
    #[error("`{}` environment variable is not UTF-8", BUCK_WRAPPER_UUID_ENV_VAR)]
    EnvVarNotUtf8,
}

/// A TraceId is a unique identifier for a trace. Trace IDs are globally unique; their textual form is a v4 UUID.
///
/// TraceIds generally correspond to commands, but they do not have to, e.g. in the case of a Buck daemon producing
/// events even when a command is not running.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Allocative)]
pub struct TraceId(
    #[allocative(skip)] // `Uuid` is inline.
    Uuid,
);

// UUID is 16 bytes.
impl Dupe for TraceId {}

impl serde::ser::Serialize for TraceId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.to_string())
    }
}

impl<'de> serde::de::Deserialize<'de> for TraceId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;

        <&str>::deserialize(deserializer)?
            .parse()
            .map_err(D::Error::custom)
    }
}

impl Display for TraceId {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let mut buf = Uuid::encode_buffer();
        f.write_str(self.0.hyphenated().encode_lower(&mut buf))
    }
}

impl FromStr for TraceId {
    type Err = uuid::Error;
    fn from_str(s: &str) -> Result<TraceId, Self::Err> {
        Ok(TraceId(Uuid::parse_str(s)?))
    }
}

impl TraceId {
    /// Generates a new TraceId, suitable for identifying a particular trace.
    pub fn new() -> TraceId {
        TraceId(Uuid::new_v4())
    }

    pub fn null() -> TraceId {
        TraceId(Uuid::nil())
    }

    /// Fetch `TraceId` from environment variable or generate a new one.
    pub fn from_env_or_new() -> anyhow::Result<TraceId> {
        match env::var(BUCK_WRAPPER_UUID_ENV_VAR) {
            Ok(s) => Ok(TraceId::from_str(&s).with_context(|| {
                format!(
                    "Parsing buck2 invocation id from env variable {}",
                    BUCK_WRAPPER_UUID_ENV_VAR
                )
            })?),
            Err(env::VarError::NotPresent) => Ok(TraceId::new()),
            Err(env::VarError::NotUnicode(_)) => Err(TraceIdError::EnvVarNotUtf8.into()),
        }
    }

    /// Generate short hash to be used as a message key for a Scribe client.
    pub fn hash(&self) -> i64 {
        let mut hasher = DefaultHasher::new();
        Hash::hash(self, &mut hasher);
        hasher.finish() as i64
    }
}

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

    #[test]
    fn roundtrip() {
        let a = TraceId::new();
        let s = serde_json::to_string(&a).unwrap();
        let b = serde_json::from_str::<TraceId>(&s).unwrap();
        assert_eq!(a, b);
    }
}


/*
 * 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::env;
use std::ffi::OsStr;
use std::path::Path;
use std::path::PathBuf;

use dupe::Dupe;
use once_cell::sync::OnceCell;

#[derive(Copy, Clone, Dupe)]
pub enum WhoIsAsking {
    Buck2,
    BuckWrapper,
}

pub(crate) fn is_buck2_exe(path: &Path, who_is_asking: WhoIsAsking) -> bool {
    let Some(file_stem) = path.file_stem() else {
        return false;
    };
    if path.file_stem() == Some(OsStr::new("buck2")) {
        true
    } else {
        match who_is_asking {
            WhoIsAsking::BuckWrapper => {
                // We don't know another name of the buck2 executable in the wrapper.
                false
            }
            WhoIsAsking::Buck2 => {
                static CURRENT_EXE: OnceCell<PathBuf> = OnceCell::new();
                if let Ok(current_exe) = CURRENT_EXE.get_or_try_init(env::current_exe) {
                    if let Some(current_exe_file_stem) = current_exe.file_stem() {
                        if current_exe_file_stem == file_stem {
                            return true;
                        }
                    }
                }
                false
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use std::path::Path;

    use crate::is_buck2::is_buck2_exe;
    use crate::is_buck2::WhoIsAsking;

    #[test]
    fn test_is_buck2_exe() {
        let (fake_buck, other_path) = if cfg!(windows) {
            ("C:\\dir\\buck2.exe", "C:\\dir\\other.exe")
        } else {
            ("/dir/buck2", "/dir/other")
        };

        assert!(is_buck2_exe(Path::new(fake_buck), WhoIsAsking::Buck2));
        assert!(is_buck2_exe(Path::new(fake_buck), WhoIsAsking::BuckWrapper));

        let current_exe = env::current_exe().unwrap();

        assert!(is_buck2_exe(&current_exe, WhoIsAsking::Buck2));
        assert!(!is_buck2_exe(&current_exe, WhoIsAsking::BuckWrapper));

        assert!(!is_buck2_exe(Path::new(other_path), WhoIsAsking::Buck2));
        assert!(!is_buck2_exe(
            Path::new(other_path),
            WhoIsAsking::BuckWrapper
        ));
    }
}


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

//! Cross-platform process killing.

pub fn process_exists(pid: u32) -> anyhow::Result<bool> {
    os_specific::process_exists(pid)
}

/// Send `KILL` or call `TerminateProcess` on the given process.
///
/// Returns a KilledProcessHandle that can be used to observe the termination of the killed process.
pub fn kill(pid: u32) -> anyhow::Result<Box<dyn KilledProcessHandle>> {
    match os_specific::kill(pid)? {
        Some(handle) => Ok(Box::new(handle) as _),
        None => Ok(Box::new(NoProcess) as _),
    }
}

pub trait KilledProcessHandle {
    fn has_exited(&self) -> anyhow::Result<bool>;

    fn status(&self) -> Option<String>;
}

/// Get the status of a given process according to sysinfo.
pub fn get_sysinfo_status(pid: impl TryInto<u32>) -> Option<String> {
    use sysinfo::Pid;
    use sysinfo::PidExt;
    use sysinfo::ProcessExt;
    use sysinfo::ProcessRefreshKind;
    use sysinfo::System;
    use sysinfo::SystemExt;

    let pid = pid.try_into().ok()?;
    let pid = Pid::from_u32(pid);

    let mut system = System::new();
    system.refresh_process_specifics(pid, ProcessRefreshKind::new());

    let proc = system.process(pid)?;
    Some(proc.status().to_string())
}

/// Returned when os_specific::kill reports that nothing was killed because the process wasn't even
/// running.
struct NoProcess;

impl KilledProcessHandle for NoProcess {
    fn has_exited(&self) -> anyhow::Result<bool> {
        Ok(true)
    }

    fn status(&self) -> Option<String> {
        Some("NoProcess".to_owned())
    }
}

#[cfg(unix)]
mod os_specific {
    use anyhow::Context;
    use nix::sys::signal::Signal;

    use crate::kill::get_sysinfo_status;
    use crate::kill::KilledProcessHandle;

    pub(crate) fn process_exists(pid: u32) -> anyhow::Result<bool> {
        let pid = nix::unistd::Pid::from_raw(
            pid.try_into()
                .with_context(|| format!("Integer overflow converting pid {} to pid_t", pid))?,
        );
        match nix::sys::signal::kill(pid, None) {
            Ok(_) => Ok(true),
            Err(nix::errno::Errno::ESRCH) => Ok(false),
            Err(e) => Err(e)
                .with_context(|| format!("unexpected error checking if process {} exists", pid)),
        }
    }

    fn process_exists_impl(pid: nix::unistd::Pid) -> anyhow::Result<bool> {
        match nix::sys::signal::kill(pid, None) {
            Ok(_) => Ok(true),
            Err(nix::errno::Errno::ESRCH) => Ok(false),
            Err(e) => Err(e)
                .with_context(|| format!("unexpected error checking if process {} exists", pid)),
        }
    }

    pub(super) fn kill(pid: u32) -> anyhow::Result<Option<impl KilledProcessHandle>> {
        let pid = nix::unistd::Pid::from_raw(
            pid.try_into()
                .with_context(|| format!("Integer overflow converting pid {} to pid_t", pid))?,
        );

        match nix::sys::signal::kill(pid, Signal::SIGKILL) {
            Ok(()) => Ok(Some(UnixKilledProcessHandle { pid })),
            Err(nix::errno::Errno::ESRCH) => Ok(None),
            Err(e) => Err(e).with_context(|| format!("Failed to kill pid {}", pid)),
        }
    }

    struct UnixKilledProcessHandle {
        pid: nix::unistd::Pid,
    }

    impl KilledProcessHandle for UnixKilledProcessHandle {
        fn has_exited(&self) -> anyhow::Result<bool> {
            Ok(!process_exists_impl(self.pid)?)
        }

        fn status(&self) -> Option<String> {
            get_sysinfo_status(self.pid.as_raw())
        }
    }
}

#[cfg(windows)]
pub mod os_specific {
    use std::time::Duration;

    use crate::kill::get_sysinfo_status;
    use crate::kill::KilledProcessHandle;
    use crate::winapi_process::WinapiProcessHandle;

    pub(crate) fn process_exists(pid: u32) -> anyhow::Result<bool> {
        Ok(WinapiProcessHandle::open_for_info(pid).is_some())
    }

    pub(super) fn kill(pid: u32) -> anyhow::Result<Option<impl KilledProcessHandle>> {
        let handle = match WinapiProcessHandle::open_for_terminate(pid) {
            Some(proc_handle) => proc_handle,
            None => return Ok(None),
        };

        handle.terminate()?;

        Ok(Some(WindowsKilledProcessHandle { handle }))
    }

    /// Windows reuses PIDs more aggressively than UNIX, so there we add an extra guard in the form
    /// of the process creation time.
    struct WindowsKilledProcessHandle {
        handle: WinapiProcessHandle,
    }

    impl KilledProcessHandle for WindowsKilledProcessHandle {
        fn has_exited(&self) -> anyhow::Result<bool> {
            self.handle.has_exited()
        }

        fn status(&self) -> Option<String> {
            // Maybe there is a better way to get this via the handle, but for now this'll do.
            get_sysinfo_status(self.handle.pid())
        }
    }

    /// Returns process creation time with 100 ns precision.
    pub fn process_creation_time(pid: u32) -> Option<Duration> {
        let proc_handle = WinapiProcessHandle::open_for_info(pid)?;
        proc_handle.process_creation_time().ok()
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use std::time::Instant;

    use buck2_util::process::background_command;

    use crate::kill::kill;
    use crate::kill::process_exists;

    #[test]
    fn test_process_exists_kill() {
        let mut command = if !cfg!(windows) {
            let mut command = background_command("sh");
            command.args(["-c", "sleep 10000"]);
            command
        } else {
            let mut command = background_command("powershell");
            command.args(["-c", "Start-Sleep -Seconds 10000"]);
            command
        };
        let mut child = command.spawn().unwrap();
        let pid = child.id().try_into().unwrap();
        for _ in 0..5 {
            assert!(process_exists(pid).unwrap());
        }

        let handle = kill(pid).unwrap();

        child.wait().unwrap();
        // Drop child to ensure the Windows handle is closed.
        drop(child);

        if !cfg!(windows) {
            assert!(handle.has_exited().unwrap());
        } else {
            let start = Instant::now();
            loop {
                if handle.has_exited().unwrap() {
                    break;
                }
                assert!(
                    start.elapsed() < Duration::from_secs(20),
                    "Timed out waiting for process to die"
                );
                std::thread::sleep(Duration::from_millis(100));
            }
        }
    }
}


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

//! Code shared between `buck2_wrapper` and `buck2`.
//!
//! Careful! The wrapper is not released as part of the regular buck version bumps,
//! meaning code changes here are not "atomically" updated.

use std::collections::HashSet;
use std::thread;
use std::time::Duration;
use std::time::Instant;

use is_buck2::WhoIsAsking;
use sysinfo::Pid;
use sysinfo::PidExt;
use sysinfo::ProcessExt;
use sysinfo::System;
use sysinfo::SystemExt;

use crate::is_buck2::is_buck2_exe;

pub mod invocation_id;
pub mod is_buck2;
pub mod kill;
pub mod winapi_handle;
pub(crate) mod winapi_process;

pub const BUCK2_WRAPPER_ENV_VAR: &str = "BUCK2_WRAPPER";
pub const BUCK_WRAPPER_UUID_ENV_VAR: &str = "BUCK_WRAPPER_UUID";

/// Because `sysinfo::Process` is not `Clone`.
struct ProcessInfo {
    pid: u32,
    name: String,
    cmd: Vec<String>,
}

/// Find all buck2 processes in the system.
fn find_buck2_processes(who_is_asking: WhoIsAsking) -> Vec<ProcessInfo> {
    let mut system = System::new();
    system.refresh_processes();

    let mut current_parents = HashSet::new();
    let mut parent = Some(Pid::from_u32(std::process::id()));
    while let Some(pid) = parent {
        // There is a small chance on Windows that the PID of a dead parent
        // was reused by some of its descendants, and this can create a loop.
        if !current_parents.insert(pid) {
            break;
        }
        parent = system.process(pid).and_then(|p| p.parent());
    }

    let mut buck2_processes = Vec::new();
    for (pid, process) in system.processes() {
        if is_buck2_exe(process.exe(), who_is_asking) && !current_parents.contains(pid) {
            buck2_processes.push(ProcessInfo {
                pid: pid.as_u32(),
                name: process.name().to_owned(),
                cmd: process.cmd().to_vec(),
            });
        }
    }

    buck2_processes
}

/// Kills all running Buck2 processes, except this process's hierarchy. Returns whether it
/// succeeded without errors.
pub fn killall(who_is_asking: WhoIsAsking, write: impl Fn(String)) -> bool {
    let buck2_processes = find_buck2_processes(who_is_asking);

    if buck2_processes.is_empty() {
        write("No buck2 processes found".to_owned());
        return true;
    }

    struct Printer<F> {
        write: F,
        /// All processes were killed successfully.
        ok: bool,
    }

    impl<F: Fn(String)> Printer<F> {
        fn fmt_status(&mut self, process: &ProcessInfo, status: &str) -> String {
            format!(
                "{} {} ({}). {}",
                status,
                process.name,
                process.pid,
                shlex::join(process.cmd.iter().map(|s| s.as_str())),
            )
        }

        fn failed_to_kill(&mut self, process: &ProcessInfo, error: anyhow::Error) {
            let mut message = self.fmt_status(process, "Failed to kill");
            for line in format!("{:?}", error).lines() {
                message.push_str("\n  ");
                message.push_str(line);
            }
            (self.write)(message);

            self.ok = false;
        }

        fn killed(&mut self, process: &ProcessInfo) {
            let message = self.fmt_status(process, "Killed");
            (self.write)(message);
        }
    }

    let mut printer = Printer { write, ok: true };

    // Send a kill signal and collect the processes that are still alive.

    let mut processes_still_alive: Vec<(ProcessInfo, _)> = Vec::new();
    for process in buck2_processes {
        match kill::kill(process.pid) {
            Ok(handle) => processes_still_alive.push((process, handle)),
            Err(e) => printer.failed_to_kill(&process, e),
        };
    }

    // Wait for the processes to exit.

    let start = Instant::now();
    while !processes_still_alive.is_empty() {
        let timeout_secs = 10;

        processes_still_alive.retain(|(process, handle)| match handle.has_exited() {
            Err(e) => {
                printer.failed_to_kill(process, e);
                false
            }
            Ok(true) => {
                printer.killed(process);
                false
            }
            Ok(false) => true,
        });

        if start.elapsed() > Duration::from_secs(timeout_secs) {
            for process in processes_still_alive {
                printer.failed_to_kill(
                    &process.0,
                    anyhow::anyhow!("Process still alive after {timeout_secs}s after kill sent"),
                );
            }
            break;
        }

        thread::sleep(Duration::from_millis(100));
    }

    printer.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(windows)]

use winapi::um::handleapi::CloseHandle;
use winapi::um::winnt::HANDLE;

/// Close handle on drop.
pub struct WinapiHandle {
    handle: HANDLE,
}

impl WinapiHandle {
    /// Unsafe because it closes the handle on drop.
    pub unsafe fn new(handle: HANDLE) -> WinapiHandle {
        WinapiHandle { handle }
    }

    pub fn handle(&self) -> HANDLE {
        self.handle
    }
}

impl Drop for WinapiHandle {
    fn drop(&mut self) {
        unsafe {
            if !self.handle.is_null() {
                let res = CloseHandle(self.handle);
                assert!(res != 0, "CloseHandle failed");
            }
        };
    }
}


/*
 * 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(windows)]

use std::io;
use std::time::Duration;

use anyhow::Context;
use winapi::shared::minwindef::FILETIME;
use winapi::um::minwinbase::STILL_ACTIVE;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::processthreadsapi::GetProcessTimes;
use winapi::um::processthreadsapi::OpenProcess;
use winapi::um::processthreadsapi::TerminateProcess;
use winapi::um::winnt::PROCESS_QUERY_INFORMATION;
use winapi::um::winnt::PROCESS_TERMINATE;

use crate::winapi_handle::WinapiHandle;

/// `HANDLE` which points to a process.
pub(crate) struct WinapiProcessHandle {
    handle: WinapiHandle,
    pid: u32,
}

impl WinapiProcessHandle {
    /// Open a process handle to query. `None` if process doesn't exist.
    pub(crate) fn open_for_info(pid: u32) -> Option<WinapiProcessHandle> {
        WinapiProcessHandle::open_impl(pid, PROCESS_QUERY_INFORMATION)
    }

    /// Open a process handle to terminate. `None` if process doesn't exist.
    pub(crate) fn open_for_terminate(pid: u32) -> Option<WinapiProcessHandle> {
        WinapiProcessHandle::open_impl(pid, PROCESS_TERMINATE | PROCESS_QUERY_INFORMATION)
    }

    fn open_impl(pid: u32, desired_access: u32) -> Option<WinapiProcessHandle> {
        let proc_handle = unsafe { OpenProcess(desired_access, 0, pid) };
        if proc_handle.is_null() {
            // If proc_handle is null, process died already, or other error like access denied.
            // TODO(nga): handle error properly.
            return None;
        }
        let handle = unsafe { WinapiHandle::new(proc_handle) };
        Some(WinapiProcessHandle { handle, pid })
    }

    /// Terminate the process. Do not fail if process is already dead.
    pub(crate) fn terminate(&self) -> anyhow::Result<()> {
        unsafe {
            if TerminateProcess(self.handle.handle(), 1) == 0 {
                // Stash the error before calling `exit_code` to avoid overwriting it.
                let os_error = io::Error::last_os_error();

                // From WinAPI doc:
                // After a process has terminated, call to `TerminateProcess` with open handles
                // to the process fails with `ERROR_ACCESS_DENIED` (5) error code.
                if let Ok(true) = self.has_exited() {
                    return Ok(());
                }

                Err(os_error).with_context(|| format!("Failed to kill pid {}", self.pid))
            } else {
                Ok(())
            }
        }
    }

    pub(crate) fn process_creation_time(&self) -> anyhow::Result<Duration> {
        let mut creation_time: FILETIME = unsafe { std::mem::zeroed() };
        let mut exit_time: FILETIME = unsafe { std::mem::zeroed() };
        let mut kernel_time: FILETIME = unsafe { std::mem::zeroed() };
        let mut user_time: FILETIME = unsafe { std::mem::zeroed() };

        let result = unsafe {
            GetProcessTimes(
                self.handle.handle(),
                &mut creation_time,
                &mut exit_time,
                &mut kernel_time,
                &mut user_time,
            )
        };

        if result == 0 {
            return Err(io::Error::last_os_error())
                .with_context(|| format!("Failed to call GetProcessTimes for pid {}", self.pid));
        }

        // `creation_time` stores intervals of 100 ns, so multiply by 100 to obtain
        // proper nanoseconds. The u64 type will overflow around the year 2185.
        let intervals =
            ((creation_time.dwHighDateTime as u64) << 32) | (creation_time.dwLowDateTime as u64);
        Ok(Duration::from_nanos(intervals * 100))
    }

    /// Exit code, or `None` if process is still running.
    fn exit_code(&self) -> anyhow::Result<Option<u32>> {
        let mut exit_code = 0;

        if unsafe { GetExitCodeProcess(self.handle.handle(), &mut exit_code) } != 0 {
            if exit_code == STILL_ACTIVE {
                return Ok(None);
            } else {
                return Ok(Some(exit_code));
            }
        }

        Err(io::Error::last_os_error())
            .with_context(|| format!("Failed to call GetExitCodeProcess for pid {}", self.pid))
    }

    pub(crate) fn has_exited(&self) -> anyhow::Result<bool> {
        Ok(self.exit_code()?.is_some())
    }

    pub(crate) fn pid(&self) -> u32 {
        self.pid
    }
}


/*
 * 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::fs::File;
use std::path::PathBuf;

use clap::Parser;
use dice::introspection::graph::SerializedGraphNodesForKey;

#[derive(Debug, clap::Parser)]
#[clap(name = "read_dump", about = "dice dump reader")]
pub(crate) struct Opt {
    #[clap(name = "DICE_DUMP", help = "The dice dump")]
    file: PathBuf,
    #[clap(long = "out", help = "Copy the output to this path")]
    out: Option<PathBuf>,
}

fn main() -> anyhow::Result<()> {
    let clap = Opt::clap();
    let matches = clap.get_matches_from(std::env::args().collect::<Vec<String>>());
    let opt = Opt::from_clap(&matches);

    let file = File::open(opt.file)?;

    let out: Vec<SerializedGraphNodesForKey> = bincode::deserialize_from(&file)?;

    match opt.out {
        Some(path) => {
            serde_json::to_writer_pretty(File::create(path)?, &out)?;
        }
        None => {
            serde_json::to_writer_pretty(std::io::stdout(), &out)?;
        }
    };

    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::ops::Deref;
use std::ptr;

use allocative::Allocative;
use dupe::Dupe;
use lock_free_hashtable::atomic_value::AtomicValue;

#[derive(Debug, Eq, PartialEq, Hash, Allocative)]
pub(crate) struct Arc<T>(triomphe::Arc<T>);

impl<T> Arc<T> {
    #[inline]
    pub(crate) fn new(value: T) -> Self {
        Arc(triomphe::Arc::new(value))
    }

    #[inline]
    #[cfg(test)]
    pub(crate) fn ptr_eq(this: &Self, other: &Self) -> bool {
        triomphe::Arc::ptr_eq(&this.0, &other.0)
    }
}

impl<T> Clone for Arc<T> {
    #[inline]
    fn clone(&self) -> Self {
        Arc(self.0.clone())
    }
}
impl<T> Dupe for Arc<T> {}

impl<T> Deref for Arc<T> {
    type Target = triomphe::Arc<T>;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> AtomicValue for Arc<T> {
    type Raw = *const T;
    type Ref<'a> = &'a T where Self: 'a;

    #[inline]
    fn null() -> Self::Raw {
        ptr::null()
    }

    #[inline]
    fn is_null(this: Self::Raw) -> bool {
        this.is_null()
    }

    #[inline]
    fn into_raw(this: Self) -> Self::Raw {
        triomphe::Arc::into_raw(this.0)
    }

    #[inline]
    unsafe fn from_raw(raw: Self::Raw) -> Self {
        Arc(triomphe::Arc::from_raw(raw))
    }

    #[inline]
    unsafe fn deref<'a>(raw: Self::Raw) -> Self::Ref<'a> {
        &*raw
    }
}


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

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

use allocative::Allocative;
use futures::FutureExt;

use crate::api::data::DiceData;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::opaque::OpaqueValue;
use crate::api::transaction::DiceTransactionUpdater;
use crate::api::user_data::UserComputationData;
use crate::api::user_data::UserCycleDetectorGuard;
use crate::impls::ctx::PerComputeCtx;
use crate::legacy::ctx::DiceComputationsImplLegacy;
use crate::opaque::OpaqueValueImpl;
use crate::transaction_update::DiceTransactionUpdaterImpl;
use crate::versions::VersionNumber;

#[derive(Allocative)]
pub(crate) enum DiceComputationsImpl {
    Legacy(Arc<DiceComputationsImplLegacy>),
    Modern(PerComputeCtx),
}

impl DiceComputationsImpl {
    /// Gets all the result of of the given computation key.
    /// recorded as dependencies of the current computation for which this
    /// context is for.
    pub(crate) fn compute<'a, K>(
        &'a self,
        key: &K,
    ) -> impl Future<Output = DiceResult<<K as Key>::Value>> + 'a
    where
        K: Key,
    {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate
                .compute_opaque(key)
                .map(|r| r.map(|x| x.into_value()))
                .left_future(),
            DiceComputationsImpl::Modern(delegate) => delegate.compute(key).right_future(),
        }
    }

    /// Compute "opaque" value where the value is only accessible via projections.
    /// Projections allow accessing derived results from the "opaque" value,
    /// where the dependency of reading a projection is the projection value rather
    /// than the entire opaque value.
    pub(crate) fn compute_opaque<'a, K>(
        &'a self,
        key: &K,
    ) -> impl Future<Output = DiceResult<OpaqueValue<'a, K>>> + 'a
    where
        K: Key,
    {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate
                .compute_opaque(key)
                .map(|r| r.map(|x| OpaqueValue::new(OpaqueValueImpl::Legacy(x))))
                .left_future(),
            DiceComputationsImpl::Modern(delegate) => delegate
                .compute_opaque(key)
                .map(|r| r.map(|x| OpaqueValue::new(OpaqueValueImpl::Modern(x))))
                .right_future(),
        }
    }

    /// Data that is static per the entire lifetime of Dice. These data are initialized at the
    /// time that Dice is initialized via the constructor.
    pub(crate) fn global_data(&self) -> &DiceData {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate.global_data(),
            DiceComputationsImpl::Modern(delegate) => delegate.global_data(),
        }
    }

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

    pub(crate) fn cycle_guard<T: UserCycleDetectorGuard>(&self) -> DiceResult<Option<&T>> {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate.cycle_guard(),
            DiceComputationsImpl::Modern(delegate) => delegate.cycle_guard(),
        }
    }

    pub fn store_evaluation_data<T: Send + Sync + 'static>(&self, value: T) -> DiceResult<()> {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate.store_evaluation_data(value),
            DiceComputationsImpl::Modern(delegate) => delegate.store_evaluation_data(value),
        }
    }

    pub(crate) fn get_version(&self) -> VersionNumber {
        match self {
            DiceComputationsImpl::Legacy(delegate) => delegate.get_version(),
            DiceComputationsImpl::Modern(delegate) => delegate.get_version(),
        }
    }

    pub(crate) fn into_updater(self) -> DiceTransactionUpdater {
        DiceTransactionUpdater(match self {
            DiceComputationsImpl::Legacy(delegate) => DiceTransactionUpdaterImpl::Legacy(delegate),
            DiceComputationsImpl::Modern(delegate) => {
                DiceTransactionUpdaterImpl::Modern(delegate.into_updater())
            }
        })
    }
}


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

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

#![feature(async_closure)]
#![feature(entry_insert)]
#![feature(fn_traits)]
#![feature(test)]
#![feature(map_try_insert)]
#![feature(map_entry_replace)]
// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]
// This sometimes flag false positives where proc-macros expand pass by value into pass by refs
#![allow(clippy::trivially_copy_pass_by_ref)]

#[macro_use]
extern crate gazebo;

#[macro_use]
extern crate tracing;

mod api;
pub(crate) mod arc;
mod ctx;
mod impls;
pub mod introspection;
mod legacy;
pub(crate) mod metrics;
mod opaque;
pub(crate) mod result;
mod transaction;
mod transaction_update;
mod versions;

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

use allocative::Allocative;
pub(crate) use fnv::FnvHashMap as HashMap;
pub(crate) use fnv::FnvHashSet as HashSet;
use futures::future::Future;
use futures::FutureExt;
use legacy::dice_futures::future_handle::WeakDiceFutureHandle;
use legacy::incremental::graph::GraphNode;
use legacy::incremental::transaction_ctx::TransactionCtx;
use legacy::key::StoragePropertiesForKey;
use metrics::Metrics;
pub use more_futures::cancellation::future::CancellationHandle; // expose cancellation handle as api
pub use more_futures::cancellation::CancellationContext; // expose cancellation context as api
pub use more_futures::spawn::CancellableJoinHandle; // expose cancellation context as api
pub use more_futures::spawn::FutureAndCancellationHandle;
pub use more_futures::spawn::WeakFutureError; // expose future errors as api
use serde::Serializer;

pub use crate::api::activation_tracker::ActivationData;
pub use crate::api::activation_tracker::ActivationTracker;
pub use crate::api::computations::DiceComputations;
pub use crate::api::cycles::DetectCycles;
pub use crate::api::data::DiceData;
pub use crate::api::dice::Dice;
pub use crate::api::dice::DiceDataBuilder;
pub use crate::api::error::DiceError;
pub use crate::api::error::DiceResult;
pub use crate::api::events::DiceEvent;
pub use crate::api::events::DiceEventListener;
pub use crate::api::injected::InjectedKey;
pub use crate::api::key::Key;
pub use crate::api::opaque::OpaqueValue;
pub use crate::api::projection::DiceProjectionComputations;
pub use crate::api::projection::ProjectionKey;
pub use crate::api::transaction::DiceEquality;
pub use crate::api::transaction::DiceTransaction;
pub use crate::api::transaction::DiceTransactionUpdater;
pub use crate::api::user_data::UserComputationData;
pub use crate::api::user_data::UserCycleDetector;
pub use crate::api::user_data::UserCycleDetectorGuard;
pub use crate::api::which::WhichDice;
pub use crate::api::which::WhichSpawner;
use crate::impls::dice::DiceModern;
use crate::impls::dice::DiceModernDataBuilder;
use crate::introspection::graph::GraphIntrospectable;
use crate::introspection::serialize_dense_graph;
use crate::introspection::serialize_graph;
use crate::legacy::DiceLegacy;
use crate::legacy::DiceLegacyDataBuilder;
use crate::transaction_update::DiceTransactionUpdaterImpl;

#[derive(Allocative, Debug)]
pub(crate) enum DiceImplementation {
    Legacy(Arc<DiceLegacy>),
    Modern(Arc<DiceModern>),
}

impl DiceImplementation {
    pub fn updater(&self) -> DiceTransactionUpdater {
        match self {
            DiceImplementation::Legacy(dice) => dice.updater(),
            DiceImplementation::Modern(dice) => {
                DiceTransactionUpdater(DiceTransactionUpdaterImpl::Modern(dice.updater()))
            }
        }
    }

    pub fn updater_with_data(&self, extra: UserComputationData) -> DiceTransactionUpdater {
        match self {
            DiceImplementation::Legacy(dice) => dice.updater_with_data(extra),
            DiceImplementation::Modern(dice) => DiceTransactionUpdater(
                DiceTransactionUpdaterImpl::Modern(dice.updater_with_data(extra)),
            ),
        }
    }

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

    pub fn serialize_serde<S>(&self, serializer: S) -> Result<(), S::Error>
    where
        S: Serializer,
    {
        serialize_dense_graph(&self.to_introspectable(), serializer)?;
        Ok(())
    }

    fn to_introspectable(&self) -> GraphIntrospectable {
        match self {
            DiceImplementation::Legacy(dice) => dice.to_introspectable(),
            DiceImplementation::Modern(dice) => dice.to_introspectable(),
        }
    }

    pub fn detect_cycles(&self) -> &DetectCycles {
        match self {
            DiceImplementation::Legacy(dice) => dice.detect_cycles(),
            DiceImplementation::Modern(dice) => dice.detect_cycles(),
        }
    }

    pub fn metrics(&self) -> Metrics {
        match self {
            DiceImplementation::Legacy(dice) => dice.metrics(),
            DiceImplementation::Modern(dice) => dice.metrics(),
        }
    }

    /// Wait until all active versions have exited.
    pub fn wait_for_idle(&self) -> impl Future<Output = ()> + 'static {
        match self {
            DiceImplementation::Legacy(dice) => dice.wait_for_idle().left_future(),
            DiceImplementation::Modern(dice) => dice.wait_for_idle().right_future(),
        }
    }

    pub async fn is_idle(&self) -> bool {
        match self {
            DiceImplementation::Legacy(dice) => dice.is_idle(),
            DiceImplementation::Modern(dice) => dice.is_idle().await,
        }
    }
}

pub(crate) enum DiceDataBuilderImpl {
    Legacy(DiceLegacyDataBuilder),
    Modern(DiceModernDataBuilder),
}

impl DiceDataBuilderImpl {
    pub(crate) fn new_legacy() -> Self {
        Self::Legacy(DiceLegacyDataBuilder::new())
    }

    pub(crate) fn new_modern() -> Self {
        Self::Modern(DiceModernDataBuilder::new())
    }

    pub fn set<K: Send + Sync + 'static>(&mut self, val: K) {
        match self {
            DiceDataBuilderImpl::Legacy(d) => d.set(val),
            DiceDataBuilderImpl::Modern(d) => d.set(val),
        }
    }

    pub fn build(self, detect_cycles: DetectCycles, which_spawner: WhichSpawner) -> Arc<Dice> {
        Dice::new(match self {
            DiceDataBuilderImpl::Legacy(d) => {
                DiceImplementation::Legacy(d.build(detect_cycles, which_spawner))
            }
            DiceDataBuilderImpl::Modern(d) => DiceImplementation::Modern(d.build(detect_cycles)),
        })
    }
}

pub mod testing {
    pub use crate::api::dice::testing::DiceBuilder;
}


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

/// Dice metrics.
#[derive(Debug)]
pub struct Metrics {
    pub key_count: usize,
    /// The number of keys currently active in the per transaction cache
    pub currently_active_key_count: usize,
    pub active_transaction_count: u32,
}


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

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

use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::projection::ProjectionKey;
use crate::impls::opaque::OpaqueValueModern;
use crate::legacy::opaque::OpaqueValueImplLegacy;

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

impl<'a, K> Debug for OpaqueValueImpl<'a, K>
where
    K: Key,
    K::Value: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            OpaqueValueImpl::Legacy(delegate) => delegate.fmt(f),
            OpaqueValueImpl::Modern(delegate) => delegate.fmt(f),
        }
    }
}

impl<'a, K: Key> OpaqueValueImpl<'a, K> {
    pub(crate) fn projection<P>(&self, projection_key: &P) -> DiceResult<P::Value>
    where
        P: ProjectionKey<DeriveFromKey = K>,
    {
        match self {
            OpaqueValueImpl::Legacy(delegate) => delegate.projection(projection_key),
            OpaqueValueImpl::Modern(delegate) => delegate.projection(projection_key),
        }
    }
}


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

//! Cancellable results of DICE

use allocative::Allocative;
use derive_more::Display;
use dupe::Dupe;
use thiserror::Error;

#[allow(unused)] // TODO temporary
pub(crate) type CancellableResult<T> = Result<T, Cancelled>;

#[allow(unused)] // TODO temporary
#[derive(Clone, Dupe, Display, Debug, Error, Allocative)]
#[display(fmt = "{:?}", self)]
pub(crate) struct Cancelled;


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

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

use crate::ctx::DiceComputationsImpl;
use crate::impls::ctx::BaseComputeCtx;
use crate::versions::VersionNumber;
use crate::DiceComputations;
use crate::DiceTransactionUpdater;

#[derive(Allocative)]
pub(crate) enum DiceTransactionImpl {
    Legacy(DiceComputations),
    Modern(BaseComputeCtx),
}

impl Clone for DiceTransactionImpl {
    fn clone(&self) -> Self {
        match self {
            DiceTransactionImpl::Legacy(ctx) => match &ctx.0 {
                // since 'DiceComputations' should not be clone or dupe, we manually implement
                // clone and dupe for the 'DiceTransaction'
                DiceComputationsImpl::Legacy(ctx) => DiceTransactionImpl::Legacy(DiceComputations(
                    DiceComputationsImpl::Legacy(ctx.dupe()),
                )),
                DiceComputationsImpl::Modern(_) => {
                    unreachable!("wrong dice")
                }
            },
            DiceTransactionImpl::Modern(ctx) => DiceTransactionImpl::Modern(ctx.clone()),
        }
    }
}

impl Dupe for DiceTransactionImpl {}

impl DiceTransactionImpl {
    pub(crate) fn get_version(&self) -> VersionNumber {
        match self {
            DiceTransactionImpl::Legacy(ctx) => ctx.0.get_version(),
            DiceTransactionImpl::Modern(ctx) => ctx.get_version(),
        }
    }

    pub(crate) fn into_updater(self) -> DiceTransactionUpdater {
        match self {
            DiceTransactionImpl::Legacy(delegate) => delegate.0.into_updater(),
            DiceTransactionImpl::Modern(delegate) => delegate.into_updater(),
        }
    }

    pub(crate) fn as_computations(&self) -> &DiceComputations {
        match self {
            DiceTransactionImpl::Legacy(ctx) => ctx,
            DiceTransactionImpl::Modern(ctx) => ctx.as_computations(),
        }
    }

    pub(crate) fn as_computations_mut(&mut self) -> &mut DiceComputations {
        match self {
            DiceTransactionImpl::Legacy(ctx) => ctx,
            DiceTransactionImpl::Modern(ctx) => ctx.as_computations_mut(),
        }
    }
}


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

use std::future::Future;
use std::sync::Arc;
use std::thread;

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

use crate::api::computations::DiceComputations;
use crate::api::error::DiceResult;
use crate::api::key::Key;
use crate::api::user_data::UserComputationData;
use crate::ctx::DiceComputationsImpl;
use crate::impls::transaction::TransactionUpdater;
use crate::legacy::ctx::DiceComputationsImplLegacy;
use crate::transaction::DiceTransactionImpl;
use crate::DiceTransaction;

/// The struct for which we build transactions. This is where changes are recorded, and committed
/// to DICE, which returns the Transaction where we spawn computations.
#[derive(Allocative)]
pub(crate) enum DiceTransactionUpdaterImpl {
    Legacy(Arc<DiceComputationsImplLegacy>),
    #[allow(unused)]
    Modern(TransactionUpdater),
}

impl DiceTransactionUpdaterImpl {
    pub(crate) fn existing_state(&self) -> impl Future<Output = DiceTransaction> + '_ {
        match self {
            DiceTransactionUpdaterImpl::Legacy(ctx) => {
                futures::future::ready(DiceTransaction(DiceTransactionImpl::Legacy(
                    DiceComputations(DiceComputationsImpl::Legacy(ctx.dupe())),
                )))
                .left_future()
            }
            DiceTransactionUpdaterImpl::Modern(delegate) => delegate
                .existing_state()
                .map(|d| DiceTransaction(DiceTransactionImpl::Modern(d)))
                .right_future(),
        }
    }

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

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

    /// Commit the changes registered via 'changed' and 'changed_to' to the current newest version.
    pub(crate) fn commit(self) -> impl Future<Output = DiceTransaction> {
        match self {
            DiceTransactionUpdaterImpl::Legacy(ctx) => {
                futures::future::ready(DiceTransaction(DiceTransactionImpl::Legacy(
                    DiceComputations(DiceComputationsImpl::Legacy(ctx.commit())),
                )))
                .left_future()
            }
            DiceTransactionUpdaterImpl::Modern(delegate) => delegate
                .commit()
                .map(|x| DiceTransaction(DiceTransactionImpl::Modern(x)))
                .right_future(),
        }
    }

    /// Commit the changes registered via 'changed' and 'changed_to' to the current newest version,
    /// replacing the user data with the given set
    pub(crate) fn commit_with_data(
        self,
        extra: UserComputationData,
    ) -> impl Future<Output = DiceTransaction> {
        match self {
            DiceTransactionUpdaterImpl::Legacy(ctx) => {
                futures::future::ready(DiceTransaction(DiceTransactionImpl::Legacy(
                    DiceComputations(DiceComputationsImpl::Legacy(ctx.commit_with_data(extra))),
                )))
                .left_future()
            }
            DiceTransactionUpdaterImpl::Modern(delegate) => delegate
                .commit_with_data(extra)
                .map(|x| DiceTransaction(DiceTransactionImpl::Modern(x)))
                .right_future(),
        }
    }

    /// Clears the entire DICE state. The dropping of values from memory happens asynchronously.
    pub fn unstable_take(self) -> Self {
        match self {
            DiceTransactionUpdaterImpl::Legacy(ctx) => {
                let map = ctx.unstable_take();
                // Destructors can be slow, so we do this in a separate thread.
                thread::spawn(|| drop(map));

                DiceTransactionUpdaterImpl::Legacy(ctx)
            }
            DiceTransactionUpdaterImpl::Modern(delegate) => {
                delegate.unstable_take();

                DiceTransactionUpdaterImpl::Modern(delegate)
            }
        }
    }
}