/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::VecDeque;
use std::fmt::Display;
use std::hash::Hash;
use std::time::Duration;
use std::time::Instant;

use dashmap::DashMap;

struct WeightSlidingWindow {
    start: Instant,
    window: Duration,
    // Always monotonically increasing times
    entries: VecDeque<(Instant, f64)>,
    // Sum of the weights of entries
    total_weight: f64,
}

impl WeightSlidingWindow {
    fn start_at(start: Instant, window: Duration) -> Self {
        Self {
            start,
            window,
            entries: VecDeque::new(),
            total_weight: 0.,
        }
    }

    fn prune(&mut self) -> Instant {
        let now = Instant::now();
        while let Some((last, weight)) = self.entries.front() {
            if *last
                < now.checked_sub(self.window).unwrap_or_else(|| {
                    panic!(
                        "Duration now ({:?}) is less than window ({:?})",
                        now, self.window
                    )
                })
            {
                self.total_weight -= weight;
                self.entries.pop_front();
            } else {
                break;
            }
        }
        now
    }

    /// Average weight over the window, if any
    pub fn avg(&mut self) -> Option<f64> {
        self.prune();
        if self.entries.is_empty() {
            None
        } else {
            Some(self.total_weight / (self.entries.len() as f64))
        }
    }

    /// How much weight is consumed per second
    /// If weight is number of operations, then this means speed.
    /// If weight is seconds doing something, then this is (approximately)
    /// saturation (less accurate with longer weights or shorter windows).
    pub fn wps(&mut self) -> f64 {
        let now = self.prune();
        self.total_weight / ((now - self.start).min(self.window).as_secs_f64())
    }

    pub fn add_entry(&mut self, weight: f64) -> Instant {
        let now = self.prune();
        self.total_weight += weight;
        self.entries.push_back((now, weight));
        now
    }
}

pub struct AvgTimeSlidingWindows<Key> {
    start: Instant,
    window: Duration,
    sliding_windows: DashMap<Key, WeightSlidingWindow>,
}

impl<Key: Eq + Hash> AvgTimeSlidingWindows<Key> {
    pub fn start(window: Duration) -> Self {
        Self {
            start: Instant::now(),
            window,
            sliding_windows: DashMap::new(),
        }
    }

    pub fn add_entry(&self, key: Key, duration: Duration) -> Instant {
        self.sliding_windows
            .entry(key)
            .or_insert_with(|| WeightSlidingWindow::start_at(self.start, self.window))
            .add_entry(duration.as_secs_f64())
    }
}

impl<Key: Eq + Hash + Display> Display for AvgTimeSlidingWindows<Key> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut first = true;
        for mut r in self.sliding_windows.iter_mut() {
            if !first {
                write!(f, ", ")?;
            }
            first = false;
            write!(f, "{} avg ", r.key())?;
            if let Some(avg) = r.value_mut().avg() {
                write!(f, "{:.2?}", Duration::from_secs_f64(avg))?;
            } else {
                write!(f, "<NO ENTRIES>")?;
            }
        }
        Ok(())
    }
}

pub struct BasicSpeedTracker {
    s_10m: WeightSlidingWindow,
    s_1h: WeightSlidingWindow,
    s_1d: WeightSlidingWindow,
}

impl BasicSpeedTracker {
    pub fn start() -> Self {
        let now = Instant::now();
        Self {
            s_10m: WeightSlidingWindow::start_at(now, Duration::from_secs(60 * 10)),
            s_1h: WeightSlidingWindow::start_at(now, Duration::from_secs(60 * 60)),
            s_1d: WeightSlidingWindow::start_at(now, Duration::from_secs(60 * 60 * 24)),
        }
    }

    #[allow(dead_code)]
    pub fn add_entry(&mut self) {
        self.add_entries(1)
    }

    pub fn add_entries(&mut self, count: usize) {
        let count = count as f64;
        let Self { s_10m, s_1h, s_1d } = self;
        s_10m.add_entry(count);
        s_1h.add_entry(count);
        s_1d.add_entry(count);
    }

    pub fn human_readable(&mut self) -> String {
        let Self { s_10m, s_1h, s_1d } = self;

        format!(
            "Sliding window speeds: {:.2}/s (10m) {:.2}/s (1h) {:.2}/s (1d)",
            s_10m.wps(),
            s_1h.wps(),
            s_1d.wps()
        )
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::time::Duration;
use std::time::Instant;

use anyhow::Error;

#[derive(thiserror::Error, Debug)]
pub enum TimeseriesError {
    #[error("Conversion is out of bounds")]
    ConversionOutOfBounds(#[source] Error),
}

pub trait TimeseriesInstant<D>:
    std::ops::Sub<Self, Output = D>
    + std::ops::Add<D, Output = Self>
    + std::cmp::PartialOrd
    + std::cmp::Ord
    + Copy
    + Sized
{
}

pub trait TimeseriesDelta: Copy + Sized {
    fn div(&self, other: Self) -> Result<usize, Error>;
}

pub trait TimeseriesAccumulator: Default {
    type Value;

    fn insert(&mut self, value: Self::Value);
}

#[derive(Debug)]
pub struct Timeseries<A, D, I> {
    /// buckets represent [start, start + duration) intervals.
    buckets: Vec<A>,
    /// The start of this timeseries
    start_instant: I,
    /// The index in buckets corresponding to start_instant.
    start_idx: usize,
    /// How many buckets are valid, starting from the start bucket. If this is zero, then no
    /// buckets are valid at all.
    valid_count: usize,
    /// The duration of time that each bucket represents.
    interval: D,
}

impl<A, D, I> Timeseries<A, D, I>
where
    A: TimeseriesAccumulator,
    I: TimeseriesInstant<D>,
    D: TimeseriesDelta,
{
    pub fn new(start_instant: I, interval: D, buckets: usize) -> Self {
        let buckets = (0..buckets).map(|_| A::default()).collect();

        Self {
            start_instant,
            start_idx: 0,
            valid_count: 0,
            interval,
            buckets,
        }
    }

    /// Insert a new value at ts.
    pub fn insert(&mut self, ts: I, value: A::Value) -> Result<(), TimeseriesError> {
        if let Some(ref mut bucket) = self.bucket_for_ts(ts)? {
            bucket.insert(value);
        }

        Ok(())
    }

    /// Extend the time range to include ts, flagging all buckets before that interval as valid,
    /// such that they will be returned by [`Self::iter()`].
    pub fn update(&mut self, ts: I) -> Result<(), TimeseriesError> {
        self.bucket_for_ts(ts)?;

        Ok(())
    }

    pub fn iter(&self) -> impl Iterator<Item = &A> {
        self.buckets[self.start_idx..self.buckets.len()]
            .iter()
            .chain(self.buckets[0..self.start_idx].iter())
            .take(self.valid_count)
    }

    fn bucket_for_ts(&mut self, ts: I) -> Result<Option<&mut A>, TimeseriesError> {
        if ts < self.start_instant {
            return Ok(None);
        }

        let offset = (ts - self.start_instant)
            .div(self.interval)
            .map_err(TimeseriesError::ConversionOutOfBounds)?;

        let buckets_to_add = offset.saturating_sub(self.buckets.len() - 1);

        let offset = if buckets_to_add >= self.buckets.len() {
            self.reset_buckets(ts);
            0
        } else {
            self.add_buckets(buckets_to_add);
            offset - buckets_to_add
        };

        self.valid_count = std::cmp::max(offset + 1, self.valid_count);

        let pos = (self.start_idx + offset) % self.buckets.len();

        Ok(Some(&mut self.buckets[pos]))
    }

    fn reset_buckets(&mut self, ts: I) {
        for b in self.buckets.iter_mut() {
            *b = A::default();
        }

        self.start_instant = ts;
        self.start_idx = 0;
        self.valid_count = 0;
    }

    fn add_buckets(&mut self, n: usize) {
        let mut idx = self.start_idx;

        for _ in 0..n {
            let next_idx = (idx + 1) % self.buckets.len();
            self.buckets[idx] = A::default();
            self.start_idx = next_idx;
            self.valid_count = self.valid_count.saturating_sub(1);
            self.start_instant = self.start_instant + self.interval;
            idx = next_idx;
        }
    }
}

impl TimeseriesInstant<Duration> for Instant {}

impl TimeseriesDelta for Duration {
    fn div(&self, other: Self) -> Result<usize, Error> {
        let res = (self.as_micros() / other.as_micros()).try_into()?;
        Ok(res)
    }
}

impl TimeseriesInstant<usize> for usize {}

impl TimeseriesDelta for usize {
    fn div(&self, other: Self) -> Result<usize, Error> {
        Ok(self / other)
    }
}

impl TimeseriesInstant<u64> for u64 {}

impl TimeseriesDelta for u64 {
    fn div(&self, other: Self) -> Result<usize, Error> {
        (self / other).try_into().map_err(Error::from)
    }
}

impl<T> TimeseriesAccumulator for Vec<T> {
    type Value = T;

    fn insert(&mut self, value: Self::Value) {
        self.push(value);
    }
}

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

    #[test]
    fn test_iter() -> Result<(), Error> {
        let mut ts = Timeseries::<Vec<u64>, _, _>::new(0usize, 1usize, 3);

        ts.insert(0, 0)?;
        ts.insert(1, 1)?;

        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![0], vec![1]]);

        ts.insert(3, 3)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![1], vec![], vec![3]]);

        Ok(())
    }

    #[test]
    fn test_basic() -> Result<(), Error> {
        let mut ts = Timeseries::<Vec<u64>, _, _>::new(0usize, 2usize, 2);
        assert_eq!(ts.start_instant, 0);
        assert_eq!(ts.start_idx, 0);
        assert_eq!(ts.buckets, vec![vec![], vec![]]);

        ts.insert(0, 1)?;
        assert_eq!(ts.start_instant, 0);
        assert_eq!(ts.start_idx, 0);
        assert_eq!(ts.buckets, vec![vec![1], vec![]]);

        ts.insert(1, 2)?;
        assert_eq!(ts.start_instant, 0);
        assert_eq!(ts.start_idx, 0);
        assert_eq!(ts.buckets, vec![vec![1, 2], vec![]]);

        ts.insert(2, 3)?;
        assert_eq!(ts.start_instant, 0);
        assert_eq!(ts.start_idx, 0);
        assert_eq!(ts.buckets, vec![vec![1, 2], vec![3]]);

        ts.insert(3, 4)?;
        assert_eq!(ts.start_instant, 0);
        assert_eq!(ts.start_idx, 0);
        assert_eq!(ts.buckets, vec![vec![1, 2], vec![3, 4]]);

        ts.insert(4, 5)?;
        assert_eq!(ts.start_instant, 2);
        assert_eq!(ts.start_idx, 1);
        assert_eq!(ts.buckets, vec![vec![5], vec![3, 4]]);

        Ok(())
    }

    #[test]
    fn test_add_buckets() -> Result<(), Error> {
        let mut ts = Timeseries::<Vec<u64>, _, _>::new(0usize, 2usize, 3);

        ts.insert(0, 1)?;
        ts.insert(2, 2)?;
        ts.insert(4, 3)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![1], vec![2], vec![3]]);

        ts.insert(8, 10)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![3], vec![], vec![10]]);

        Ok(())
    }

    #[test]
    fn test_reset_buckets() -> Result<(), Error> {
        let mut ts = Timeseries::<Vec<u64>, _, _>::new(0usize, 2usize, 3);

        ts.insert(0, 1)?;
        ts.insert(2, 2)?;
        ts.insert(4, 3)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![1], vec![2], vec![3]]);

        ts.insert(20, 20)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![20]]);

        ts.insert(22, 22)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![20], vec![22]]);

        ts.insert(30, 30)?;
        let d = ts.iter().cloned().collect::<Vec<_>>();
        assert_eq!(d, vec![vec![30]]);

        Ok(())
    }

    #[test]
    fn test_sequence() -> Result<(), Error> {
        #[derive(Default, Debug)]
        struct Last(usize);

        impl TimeseriesAccumulator for Last {
            type Value = usize;

            fn insert(&mut self, value: Self::Value) {
                self.0 = value;
            }
        }

        let mut ts = Timeseries::<Last, _, _>::new(0usize, 1usize, 15);

        for i in 0..30 {
            ts.insert(i, i)?;

            // Check that the timeseries always has x, x+1, x+2, etc.

            let mut last = None;

            for v in ts.iter() {
                let v = v.0;

                // Not initialized yet.
                if v == 0 {
                    continue;
                }

                if let Some(last) = last {
                    assert_eq!(v, last + 1);
                }
                last = Some(v);
            }

            assert_eq!(ts.iter().count(), std::cmp::min(i + 1, 15));

            // Check that it ends with i (which we just inserted). This will only be true stating
            // from the 14th element since we have 15 buckets so the last bucket is [14-15).

            if i >= 14 {
                assert_eq!(ts.iter().last().map(|v| v.0), Some(i));
            }
        }

        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::borrow::Borrow;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::fmt::Debug;
use std::hash::Hash;

pub trait GenericMap<K, V> {
    fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Ord + Hash;
    fn keys(&self) -> Box<dyn Iterator<Item = &K> + '_>;
}

impl<K, V> GenericMap<K, V> for &HashMap<K, V>
where
    K: Hash + Eq,
{
    fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Ord + Hash,
    {
        <HashMap<K, V>>::get(self, k)
    }

    fn keys(&self) -> Box<dyn Iterator<Item = &K> + '_> {
        Box::new(<HashMap<K, V>>::keys(self))
    }
}

impl<K, V> GenericMap<K, V> for &BTreeMap<K, V>
where
    K: Ord,
{
    fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Ord + Hash,
    {
        <BTreeMap<K, V>>::get(self, k)
    }

    fn keys(&self) -> Box<dyn Iterator<Item = &K> + '_> {
        Box::new(<BTreeMap<K, V>>::keys(self))
    }
}

/// Sort nodes of DAG topologically. Implemented as depth-first search with tail-call
/// eliminated. Complexity: `O(N)` from number of nodes.
/// It returns None if graph has a cycle.
/// Nodes with no outgoing edges will be *first* in the resulting vector i.e. ancestors go first
///
/// The function accepts either BTreeMap or HashMap. BTreeMap has advantage of providing
/// stable output.
pub fn sort_topological<T, M>(dag: M) -> Option<Vec<T>>
where
    T: Clone + Eq + Hash + Ord + 'static + Debug,
    M: GenericMap<T, Vec<T>>, // either &BTreeMap or &HashMap
{
    sort_topological_starting_with_heads(dag, &[])
}

/// This method works like  sort topological but also takes a hint about which
/// heads should be traversed first with DFS.
///
/// There are many valid topological orders for non-linear DAGs. When there
/// are merges the starting point for the traversal matters a lot. If we
/// start from one of the heads the merge branches would be continous parts
/// of the sorted list.  If we start from the leaves we might end up
/// interleaving merges which matters for some consumer of this sorted order.
pub fn sort_topological_starting_with_heads<T, M>(dag: M, heads: &[T]) -> Option<Vec<T>>
where
    T: Clone + Eq + Hash + Ord + 'static + Debug,
    M: GenericMap<T, Vec<T>>, // either &BTreeMap or &HashMap
{
    /// Current state of the node in the DAG
    #[derive(Debug)]
    enum Mark {
        /// DFS is currently visiting the sub-DAG, reachable from this node
        /// *and* it entered this sub-DAG from this node. (the node is present
        /// deeper in the stack with `Action::Mark`). If there is an
        /// `InProgress` node at the end of the edge we are traversing,
        /// this means that the graph has a cycle.
        InProgress,
        /// The node has been visited before, and we have already visited
        /// the entire sub-DAG, reachable from this node
        Visited,
    }

    /// Action to be applied to the node, once we pop it from the stack
    #[derive(Debug)]
    enum Action<T> {
        /// Visit the node and every node, reachable from it, which has
        /// not been visited yet. Mark the node as `Mark::InProgress`
        Visit(T),
        /// Mark the node `Mark::Visited`, as we have just finished
        /// processing the entire sub-DAG reachable from it. This node
        /// no longer needs visiting
        Mark(T),
    }

    let mut marks = HashMap::new();
    let mut stack = Vec::new();
    let mut output = Vec::new();
    for node in heads.iter().chain(dag.keys()) {
        stack.push(Action::Visit(node));
        while let Some(action) = stack.pop() {
            match action {
                Action::Visit(node) => {
                    if let Some(mark) = marks.get(node) {
                        match mark {
                            Mark::InProgress => return None, // cycle
                            Mark::Visited => continue,
                        }
                    }
                    marks.insert(node, Mark::InProgress);
                    stack.push(Action::Mark(node));
                    if let Some(children) = dag.get(node) {
                        for child in children {
                            stack.push(Action::Visit(child));
                        }
                    }
                }
                Action::Mark(node) => {
                    marks.insert(node, Mark::Visited);
                    output.push(node.clone());
                }
            }
        }
    }

    Some(output)
}

// Wrapper that allows traversing commits in toposorted order. The biggest difference
// from sort_topological function above is that it allows visiting multiple commits
// in parallel, however it requires the caller to mark which commits were already visited.
// E.g. for a graph like
// A
// |\
// B C
//
// `sort_topological` returns a single order [B, C, A]. It's not clear that B and C are independent
// of each other and can be processed in parallel.
//
// On the other hand TopoSortedDagTraversal::drain() method first returns [B, C],
// and once B and C are marked as visited then commit A is returned. This allows processing
// commits B and C in parallel.
pub struct TopoSortedDagTraversal<T> {
    child_to_parents: HashMap<T, BTreeSet<T>>,
    parent_to_children: HashMap<T, BTreeSet<T>>,
    q: VecDeque<T>,
}

impl<T> TopoSortedDagTraversal<T>
where
    T: Copy + Clone + Eq + Hash + Ord,
{
    pub fn new(child_to_parents: HashMap<T, Vec<T>>) -> Self {
        let child_to_parents = child_to_parents
            .into_iter()
            .map(|(v, parents)| {
                let parents = parents.into_iter().collect::<BTreeSet<_>>();
                (v, parents)
            })
            .collect::<HashMap<_, _>>();

        // Find revert mapping - from parent to child
        let mut parent_to_children: HashMap<_, BTreeSet<_>> = HashMap::new();
        for (child, parents) in &child_to_parents {
            for p in parents {
                parent_to_children.entry(*p).or_default().insert(*child);
            }
        }

        // If we have no parents, then we can just add it to the queue
        let mut q = VecDeque::new();
        for (child, parents) in &child_to_parents {
            if parents.is_empty() {
                q.push_back(*child);
            }
            // An entry from `parents` does not have to be a key in child_to_parents.
            // e.g. child_to_parents - {1 => {2}}. `2` is not a key in child_to_parents,
            // but we still want to return it, and that's why below we are
            // adding it to the queue.
            for p in parents.iter() {
                if !child_to_parents.contains_key(p) {
                    q.push_back(*p);
                }
            }
        }

        Self {
            child_to_parents,
            parent_to_children,
            q,
        }
    }

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

    pub fn drain(&mut self, max_to_drain: usize) -> impl Iterator<Item = T> + '_ {
        self.q.drain(..std::cmp::min(self.q.len(), max_to_drain))
    }

    pub fn visited(&mut self, visited: T) {
        let children = match self.parent_to_children.get_mut(&visited) {
            Some(children) => children,
            None => {
                return;
            }
        };

        for child in children.iter() {
            if let Some(parents) = self.child_to_parents.get_mut(child) {
                parents.remove(&visited);
                if parents.is_empty() {
                    self.q.push_back(*child);
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use maplit::btreemap;
    use maplit::hashmap;
    use maplit::hashset;

    use super::*;

    #[test]
    fn sort_topological_test() {
        let res = sort_topological(&hashmap! {1 => vec![2]});
        assert_eq!(Some(vec![2, 1]), res);

        let res = sort_topological(&hashmap! {1 => vec![1]});
        assert_eq!(None, res);

        let res = sort_topological(&hashmap! {1 => vec![2], 2 => vec![3]});
        assert_eq!(Some(vec![3, 2, 1]), res);

        let res = sort_topological(&hashmap! {1 => vec![2, 3], 2 => vec![3]});
        assert_eq!(Some(vec![3, 2, 1]), res);

        let res = sort_topological(&hashmap! {1 => vec![2, 3], 2 => vec![4], 3 => vec![4]});
        assert!(Some(vec![4, 3, 2, 1]) == res || Some(vec![4, 2, 3, 1]) == res);

        let res = sort_topological_starting_with_heads(
            &hashmap! {1 => vec![2, 4], 2 => vec![3], 4 => vec![5]},
            &[1],
        );
        assert!(Some(vec![5, 4, 3, 2, 1]) == res || Some(vec![3, 2, 5, 4, 1]) == res);

        let res = sort_topological_starting_with_heads(
            &hashmap! {1 => vec![2, 4], 2 => vec![3], 4 => vec![5]},
            &[4],
        );
        assert!(Some(vec![5, 4, 3, 2, 1]) == res);

        let res = sort_topological_starting_with_heads(
            &btreemap! {1 => vec![2, 4], 2 => vec![3], 4 => vec![5]},
            &[2],
        );
        assert!(Some(vec![3, 2, 5, 4, 1]) == res);

        let res = sort_topological_starting_with_heads(
            &hashmap! {1 => vec![2, 4], 2 => vec![3], 4 => vec![5]},
            &[5, 3, 4],
        );
        assert!(Some(vec![5, 3, 4, 2, 1]) == res);
    }

    #[test]
    fn topo_sorted_traversal() {
        let mut dag = TopoSortedDagTraversal::new(hashmap! {1 => vec![]});
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![1]);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![]);

        let mut dag = TopoSortedDagTraversal::new(hashmap! {1 => vec![2], 2 => vec![3]});
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![3]);
        dag.visited(3);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![2]);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![]);
        dag.visited(2);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![1]);

        //   1
        //  / \
        // 2   3
        //  \ /
        //   4
        let mut dag = TopoSortedDagTraversal::new(
            hashmap! {1 => vec![2, 3], 2 => vec![4], 3 => vec![4], 4 => vec![]},
        );
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![4]);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![]);
        dag.visited(4);
        assert_eq!(dag.drain(10).collect::<HashSet<_>>(), hashset![2, 3]);
        dag.visited(2);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![]);
        dag.visited(3);
        assert_eq!(dag.drain(10).collect::<Vec<_>>(), vec![1]);

        //   1
        //  / \
        // 2   3
        // |   |
        // 4   5
        let mut dag = TopoSortedDagTraversal::new(
            hashmap! {1 => vec![2, 3], 2 => vec![4], 3 => vec![5], 4 => vec![], 5 => vec![]},
        );
        assert_eq!(dag.drain(2).collect::<HashSet<_>>(), hashset![4, 5]);
        dag.visited(4);
        dag.visited(5);
        assert_eq!(dag.drain(2).collect::<HashSet<_>>(), hashset![2, 3]);
        dag.visited(2);
        dag.visited(3);
        assert_eq!(dag.drain(2).collect::<HashSet<_>>(), hashset![1]);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! TypeMap is a heterogeneous collection (i.e can store object of an aribitrary type).
//!
//! `TypeMap` can store single instance of an object per type. Any object of `Arc<T>` type
//! can be stored inside `TypeMap` as long as `T` implements `std::any::Any` trait.
//!
//! On a high level `TypeMap` is basically a mapping from `TypeId` to `Arc<dyn Any>`.
//! When object is inserted, its type is erased by converting `Arc<T>` to `Arc<dyn Any>`.
//! When we want to retrieve an object by its type it is done again by `TypeId` and `Arc<dyn Any>`
//! is converted back to `Arc<T>`.
//!
//! Special care is taken to make sure this collection works well with `Send` and `Sync` objects,
//! as well as `Arc<?Sized>` types (i.e `Arc<dyn Trait>` types).

use std::any::Any;
use std::any::TypeId;
use std::collections::HashMap;
use std::sync::Arc;

struct Handle<T: ?Sized + Send + Sync + 'static>(Arc<T>);

/// Heterogeneous collection of objects.
#[derive(Clone, Default)]
pub struct TypeMap {
    mapping: HashMap<TypeId, Arc<dyn Any + Send + Sync>>,
}

impl TypeMap {
    /// Create an empty collection.
    pub fn new() -> Self {
        Default::default()
    }

    /// Adds new object to the collection.
    ///
    /// Returns old value associated with this type, if any.
    pub fn insert<T: ?Sized + Send + Sync + 'static>(&mut self, value: Arc<T>) -> Option<Arc<T>> {
        self.mapping
            .insert(TypeId::of::<Handle<T>>(), Arc::new(Handle(value)))
            .and_then(|value| {
                let handle = value.downcast_ref::<Handle<T>>()?;
                Some(handle.0.clone())
            })
    }

    /// Get an object by its type.
    pub fn get<T: ?Sized + Send + Sync + 'static>(&self) -> Option<&Arc<T>> {
        self.mapping
            .get(&TypeId::of::<Handle<T>>())
            .and_then(|value| {
                let handle = value.downcast_ref::<Handle<T>>()?;
                Some(&handle.0)
            })
    }

    /// Number of elements stored in the collection.
    pub fn len(&self) -> usize {
        self.mapping.len()
    }

    /// Check if the collection is empty.
    pub fn is_empty(&self) -> bool {
        self.mapping.is_empty()
    }
}

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

    #[derive(Debug)]
    struct Type(i32);

    trait Trait: Send + Sync {
        fn method(&self) -> i32;
    }

    impl Trait for Type {
        fn method(&self) -> i32 {
            self.0
        }
    }

    #[test]
    fn basic_test() {
        // make sure it works with unsized types
        let mut m = TypeMap::new();
        assert!(m.insert::<dyn Trait>(Arc::new(Type(42))).is_none());

        // ensure that it is Send + Sync
        let handle = std::thread::spawn(move || m.get::<dyn Trait>().cloned());
        assert_eq!(handle.join().unwrap().unwrap().method(), 42)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BinaryHeap;
use std::collections::HashSet;
use std::hash::Hash;

#[derive(Clone, Debug)]
pub struct UniqueHeap<T>
where
    T: Clone + Ord + Hash + Eq,
{
    sorted_vals: BinaryHeap<T>,
    unique_vals: HashSet<T>,
}

impl<T> UniqueHeap<T>
where
    T: Clone + Ord + Hash + Eq + Clone,
{
    pub fn new() -> Self {
        UniqueHeap {
            sorted_vals: BinaryHeap::new(),
            unique_vals: HashSet::new(),
        }
    }

    pub fn push(&mut self, val: T) {
        if !self.unique_vals.contains(&val) {
            self.unique_vals.insert(val.clone());
            self.sorted_vals.push(val);
        }
    }

    pub fn pop(&mut self) -> Option<T> {
        if let Some(v) = self.sorted_vals.pop() {
            self.unique_vals.remove(&v);
            Some(v)
        } else {
            None
        }
    }

    pub fn peek(&self) -> Option<&T> {
        self.sorted_vals.peek()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::ops::DerefMut;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::Result;
use cached_config::ConfigHandle;
use cached_config::ConfigStore;
use fbinit::FacebookInit;
use futures::try_join;
use metaconfig_types::BlobConfig;
#[cfg(fbcode_build)]
use metaconfig_types::BlobstoreId;
#[cfg(fbcode_build)]
use metaconfig_types::MultiplexedStoreType;
#[cfg(fbcode_build)]
use metaconfig_types::ShardableRemoteDatabaseConfig;
use metaconfig_types::ShardedDatabaseConfig;
use metaconfig_types::StorageConfig;
use replication_lag_config::ReplicationLagBlobstoreConfig;
use replication_lag_config::ReplicationLagTableConfig;
use slog::info;
use slog::Logger;
#[cfg(fbcode_build)]
use sql_ext::facebook::MyAdmin;
use sql_ext::replication::NoReplicaLagMonitor;
use sql_ext::replication::ReplicaLagMonitor;
use sql_ext::replication::WaitForReplicationConfig;
use tokio::sync::Mutex;

#[derive(Default)]
struct State {
    last_sync_queue_lag: Option<(Instant, Duration)>,
    last_xdb_blobstore_lag: Option<(Instant, Duration)>,
}

#[derive(Clone)]
pub struct WaitForReplication {
    config_handle: ConfigHandle<ReplicationLagBlobstoreConfig>,
    sync_queue_monitor: Arc<dyn ReplicaLagMonitor>,
    xdb_blobstore_monitor: Arc<dyn ReplicaLagMonitor>,
    state: Arc<Mutex<State>>,
}

const CONFIGS_PATH: &str = "scm/mononoke/mysql/replication_lag/config";

#[cfg(fbcode_build)]
fn blobstore_monitor(
    my_admin: &MyAdmin,
    blobstores: Vec<(BlobstoreId, MultiplexedStoreType, BlobConfig)>,
) -> Arc<dyn ReplicaLagMonitor> {
    blobstores
        .into_iter()
        .find_map(|(_, _, config)| match config {
            BlobConfig::Mysql {
                remote: ShardableRemoteDatabaseConfig::Unsharded(remote),
            } => Some(
                Arc::new(my_admin.single_shard_lag_monitor(remote.db_address))
                    as Arc<dyn ReplicaLagMonitor>,
            ),
            BlobConfig::Mysql {
                remote: ShardableRemoteDatabaseConfig::Sharded(remote),
            } => Some(Arc::new(my_admin.shardmap_lag_monitor(remote.shard_map))),
            _ => None,
        })
        .unwrap_or_else(|| Arc::new(NoReplicaLagMonitor()))
}

impl WaitForReplication {
    pub fn new(
        fb: FacebookInit,
        config_store: &ConfigStore,
        storage_config: StorageConfig,
        config_name: &'static str,
    ) -> Result<Self> {
        let config_handle =
            config_store.get_config_handle(format!("{}/{}", CONFIGS_PATH, config_name))?;
        let (sync_queue_monitor, xdb_blobstore_monitor) = match storage_config.blobstore {
            BlobConfig::MultiplexedWal {
                blobstores,
                queue_db: ShardedDatabaseConfig::Sharded(remote),
                ..
            } => {
                #[cfg(fbcode_build)]
                {
                    let my_admin = MyAdmin::new(fb)?;
                    let sync_queue = Arc::new(my_admin.shardmap_lag_monitor(remote.shard_map))
                        as Arc<dyn ReplicaLagMonitor>;
                    let xdb_blobstore = blobstore_monitor(&my_admin, blobstores);
                    (sync_queue, xdb_blobstore)
                }
                #[cfg(not(fbcode_build))]
                {
                    let _ = (fb, remote, blobstores);
                    unimplemented!()
                }
            }
            _ => (
                Arc::new(NoReplicaLagMonitor()) as Arc<dyn ReplicaLagMonitor>,
                Arc::new(NoReplicaLagMonitor()) as Arc<dyn ReplicaLagMonitor>,
            ),
        };
        Ok(Self {
            config_handle,
            sync_queue_monitor,
            xdb_blobstore_monitor,
            state: Arc::new(Mutex::new(State::default())),
        })
    }

    pub async fn wait_for_replication(&self, logger: &Logger) -> Result<()> {
        let mut state_lock = self.state.lock().await;
        let State {
            last_sync_queue_lag,
            last_xdb_blobstore_lag,
        } = state_lock.deref_mut();
        try_join!(
            self.wait_for_table(
                logger,
                "sync queue",
                last_sync_queue_lag,
                &self.sync_queue_monitor,
                || self.config_handle.get().sync_queue.clone(),
            ),
            self.wait_for_table(
                logger,
                "XDB blobstore",
                last_xdb_blobstore_lag,
                &self.xdb_blobstore_monitor,
                || self.config_handle.get().xdb_blobstore.clone(),
            ),
        )?;
        Ok(())
    }

    async fn wait_for_table<'a>(
        &'a self,
        logger: &'a Logger,
        name: &'static str,
        last_lag: &'a mut Option<(Instant, Duration)>,
        monitor: &'a Arc<dyn ReplicaLagMonitor>,
        config_getter: impl Fn() -> Option<ReplicationLagTableConfig> + Sync,
    ) -> Result<()> {
        if let Some(raw_config) = config_getter() {
            let max_replication_lag_allowed =
                Duration::from_millis(raw_config.max_replication_lag_allowed_ms.try_into()?);
            let poll_interval = Duration::from_millis(raw_config.poll_interval_ms.try_into()?);
            match last_lag.as_mut() {
                // If queried too recently, just assume it's all ok.
                Some((instant, duration))
                    if instant.elapsed() < poll_interval
                        && *duration < max_replication_lag_allowed =>
                {
                    return Ok(());
                }
                // If impossible to have surpassed replication_lag, don't query
                Some((instant, duration))
                    if *duration + instant.elapsed() < max_replication_lag_allowed =>
                {
                    return Ok(());
                }
                _ => {}
            }
            info!(
                logger,
                "Waiting for replication lag on {} to drop below {:?}",
                name,
                max_replication_lag_allowed
            );
            let new_last_lag = monitor
                .wait_for_replication(&|| {
                    // Get the most up to date config, but default to the previous one
                    // in case of errors.
                    let (max_replication_lag_allowed, poll_interval) =
                        if let Some(raw_config) = config_getter() {
                            (
                                raw_config
                                    .max_replication_lag_allowed_ms
                                    .try_into()
                                    .map_or(max_replication_lag_allowed, Duration::from_millis),
                                raw_config
                                    .poll_interval_ms
                                    .try_into()
                                    .map_or(poll_interval, Duration::from_millis),
                            )
                        } else {
                            (max_replication_lag_allowed, poll_interval)
                        };
                    WaitForReplicationConfig::new(
                        max_replication_lag_allowed,
                        poll_interval,
                        logger,
                    )
                })
                .await?;
            *last_lag = Some((Instant::now(), new_last_lag.delay));
        }
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::pin::Pin;

use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;

#[pin_project]
pub struct YieldStream<S, F>
where
    S: Stream,
    F: Fn(&S::Item) -> usize,
{
    generated: usize,
    yield_every: usize,
    size_fn: F,
    #[pin]
    inner: S,
}

impl<S, F> YieldStream<S, F>
where
    S: Stream,
    F: Fn(&S::Item) -> usize,
{
    pub fn new(inner: S, yield_every: usize, size_fn: F) -> Self {
        Self {
            generated: 0,
            yield_every,
            inner,
            size_fn,
        }
    }
}

impl<S, F> Stream for YieldStream<S, F>
where
    S: Stream,
    F: Fn(&S::Item) -> usize,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        if *this.generated >= *this.yield_every {
            *this.generated %= *this.yield_every;
            ctx.waker().wake_by_ref();
            return Poll::Pending;
        }

        let ret = futures::ready!(this.inner.poll_next_unpin(ctx));
        if let Some(item) = &ret {
            *this.generated += (this.size_fn)(item);
        }

        Poll::Ready(ret)
    }
}

pub trait YieldStreamExt: Stream + Sized {
    /// Modify the stream to yield every time a certain amount of data is
    /// generated by the stream.
    ///
    /// The `size_fn` closure is called on each item and should return an
    /// approximation of the size of the item.
    ///
    /// Whenever the stream cumulatively produces `yield_every` values (that
    /// is, the sum of the values returned from `size_fn` exceeds
    /// `yield_every`), the stream will yield back to the executor.
    fn yield_every<F>(self, yield_every: usize, size_fn: F) -> YieldStream<Self, F>
    where
        F: Fn(&Self::Item) -> usize + Sized,
    {
        YieldStream::new(self, yield_every, size_fn)
    }
}

impl<S: Stream> YieldStreamExt for S {}

#[cfg(test)]
mod test {
    use bytes::Bytes;
    use futures::stream;

    use super::*;

    #[tokio::test]
    async fn test_yield_every() {
        // NOTE: This tests that the yield probably wakes up but assumes it yields.

        let data = &[b"foo".as_ref(), b"bar2".as_ref()];
        let data = stream::iter(
            data.iter()
                .map(|d| Result::<_, ()>::Ok(Bytes::copy_from_slice(d))),
        );
        let mut stream = data.yield_every(1, |data| data.as_ref().map_or(0, |b| b.len()));

        assert_eq!(
            stream.next().await,
            Some(Ok(Bytes::copy_from_slice(b"foo")))
        );

        assert!(stream.generated > stream.yield_every);

        assert_eq!(
            stream.next().await,
            Some(Ok(Bytes::copy_from_slice(b"bar2")))
        );

        assert_eq!(stream.next().await, None,);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! See docs/basename_suffix_skeleton_manifest.md for more information

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

use anyhow::Result;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::Storable;
use cloned::cloned;
use context::CoreContext;
use derived_data_manager::DerivationContext;
use futures::future::FutureExt;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use manifest::derive_manifest_with_io_sender;
use manifest::Entry;
use manifest::LeafInfo;
use manifest::TreeInfo;
use mononoke_types::basename_suffix_skeleton_manifest::BasenameSuffixSkeletonManifest;
use mononoke_types::basename_suffix_skeleton_manifest::BssmDirectory;
use mononoke_types::basename_suffix_skeleton_manifest::BssmEntry;
use mononoke_types::BasenameSuffixSkeletonManifestId;
use mononoke_types::BlobstoreValue;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ContentId;
use mononoke_types::FileType;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use skeleton_manifest::mapping::get_file_changes;

use crate::mapping::RootBasenameSuffixSkeletonManifest;
use crate::path::BssmPath;

/// Calculate a list of changes of the changeset, but putting the basename first and
/// reversing it.
fn get_fixed_up_changes(bcs: &BonsaiChangeset) -> Vec<(MPath, Option<(ContentId, FileType)>)> {
    get_file_changes(bcs)
        .into_iter()
        .map(|(path, content)| (BssmPath::transform(path).into_raw(), content))
        .collect()
}

fn split_by_root_dir<X>(changes: Vec<(MPath, X)>) -> HashMap<MPathElement, Vec<(MPath, X)>> {
    let mut map = HashMap::new();
    for (p, x) in changes {
        let (root_dir, rest) = p.split_first();
        let rest = rest.expect("We always add a sentinel suffix to the path");
        map.entry(root_dir.clone())
            .or_insert_with(Vec::new)
            .push((rest, x));
    }
    map
}

async fn empty_mf(
    ctx: &CoreContext,
    blobstore: &impl Blobstore,
) -> Result<BasenameSuffixSkeletonManifestId> {
    let leaf = BasenameSuffixSkeletonManifest::empty();
    leaf.into_blob().store(ctx, blobstore).await
}

async fn inner_derive(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    parents: Vec<BssmDirectory>,
    changes: Vec<(MPath, Option<(ContentId, FileType)>)>,
) -> Result<Option<BssmDirectory>> {
    // Types to help understand how to use derive_manifest helper
    type Leaf = (ContentId, FileType);
    type LeafId = ();
    type TreeId = BssmDirectory;
    type IntermediateLeafId = LeafId;
    type Ctx = ();
    derive_manifest_with_io_sender(
        ctx.clone(),
        blobstore.clone(),
        parents,
        changes,
        // create_tree
        {
            cloned!(ctx, blobstore);
            move |info: TreeInfo<TreeId, IntermediateLeafId, Ctx>, fut_sender| {
                cloned!(ctx, blobstore);
                async move {
                    let entries =
                        info.subentries
                            .into_iter()
                            .map(|(path_el, (_maybe_ctx, entry_in))| {
                                let entry = match entry_in {
                                    Entry::Leaf(()) => BssmEntry::File,
                                    Entry::Tree(entry) => BssmEntry::Directory(entry),
                                };
                                (path_el, Some(entry))
                            });

                    let (mf, rollup_count) = BasenameSuffixSkeletonManifest::empty()
                        .update(&ctx, &blobstore, entries.collect())
                        .await?;
                    let entry = {
                        let blob = mf.into_blob();
                        let id = *blob.id();
                        fut_sender
                            .unbounded_send(
                                async move { blob.store(&ctx, &blobstore).await.map(|_| ()) }
                                    .boxed(),
                            )
                            .map_err(|err| {
                                anyhow::anyhow!("failed to send manifest future {}", err)
                            })?;
                        BssmDirectory {
                            id,
                            rollup_count: (1 + rollup_count) as u64,
                        }
                    };
                    anyhow::Ok(((), entry))
                }
            }
        },
        // create_leaf
        {
            move |_leaf_info: LeafInfo<IntermediateLeafId, Leaf>, _fut_sender| async move {
                anyhow::Ok(((), ()))
            }
        },
    )
    .await
}

pub(crate) async fn derive_single(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bonsai: BonsaiChangeset,
    parents: Vec<RootBasenameSuffixSkeletonManifest>,
) -> Result<RootBasenameSuffixSkeletonManifest> {
    const CONCURRENCY: usize = 100;
    let parents = parents.into_iter().map(|root| root.0).collect::<Vec<_>>();
    let changes = get_fixed_up_changes(&bonsai);
    let blobstore = derivation_ctx.blobstore();
    // TODO(T123518092): Once T123518092 is done, this optimisation can be removed and
    // we can call `inner_derive` as in the `else` clause.
    let root = if parents.len() <= 1 {
        let (parent, previous_rollup) = match parents.into_iter().next() {
            Some(p) => (Some(p.load(ctx, blobstore).await?), p.rollup_count),
            None => (None, 1),
        };
        let entries_to_update = stream::iter(split_by_root_dir(changes))
            .map(|(root_dir, changes)| async {
                let new_parent = match parent.as_ref() {
                    None => vec![],
                    Some(p) => match p.lookup(ctx, blobstore, &root_dir).await? {
                        Some(BssmEntry::Directory(dir)) => vec![dir],
                        None => vec![],
                        other => anyhow::bail!("Invalid directory {:?}", other),
                    },
                };
                Ok((
                    root_dir,
                    inner_derive(ctx, blobstore, new_parent, changes)
                        .await?
                        .map(BssmEntry::Directory),
                ))
            })
            .buffer_unordered(CONCURRENCY)
            .try_collect()
            .await?;
        let (mf, rollup_diff) = parent
            .unwrap_or_else(BasenameSuffixSkeletonManifest::empty)
            .update(ctx, blobstore, entries_to_update)
            .await?;
        Some(BssmDirectory {
            id: mf.into_blob().store(ctx, blobstore).await?,
            rollup_count: ((previous_rollup as i64) + rollup_diff) as u64,
        })
    } else {
        inner_derive(ctx, blobstore, parents, changes).await?
    };
    Ok(RootBasenameSuffixSkeletonManifest(match root {
        Some(root) => root,
        // Only happens on empty repo
        None => BssmDirectory {
            id: empty_mf(ctx, blobstore).await?,
            rollup_count: 1,
        },
    }))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

mod derive;
mod mapping;
mod ops;
mod path;
#[cfg(test)]
mod tests;

pub use mapping::format_key;
pub use mapping::RootBasenameSuffixSkeletonManifest;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::BlobstoreGetData;
use context::CoreContext;
use derived_data::impl_bonsai_derived_via_manager;
use derived_data_manager::dependencies;
use derived_data_manager::BonsaiDerivable;
use derived_data_manager::DerivableType;
use derived_data_manager::DerivationContext;
use derived_data_service_if::types as thrift;
use mononoke_types::basename_suffix_skeleton_manifest::BssmDirectory;
use mononoke_types::BasenameSuffixSkeletonManifestId;
use mononoke_types::BlobstoreBytes;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::ThriftConvert;

use crate::derive::derive_single;

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct RootBasenameSuffixSkeletonManifest(pub(crate) BssmDirectory);

pub fn format_key(derivation_ctx: &DerivationContext, changeset_id: ChangesetId) -> String {
    let root_prefix = "derived_root_bssm.";
    let key_prefix = derivation_ctx.mapping_key_prefix::<RootBasenameSuffixSkeletonManifest>();
    format!("{}{}{}", root_prefix, key_prefix, changeset_id)
}

impl TryFrom<BlobstoreBytes> for RootBasenameSuffixSkeletonManifest {
    type Error = Error;
    fn try_from(blob_bytes: BlobstoreBytes) -> Result<Self> {
        BssmDirectory::from_bytes(&blob_bytes.into_bytes()).map(RootBasenameSuffixSkeletonManifest)
    }
}

impl TryFrom<BlobstoreGetData> for RootBasenameSuffixSkeletonManifest {
    type Error = Error;
    fn try_from(blob_val: BlobstoreGetData) -> Result<Self> {
        blob_val.into_bytes().try_into()
    }
}

impl From<RootBasenameSuffixSkeletonManifest> for BlobstoreBytes {
    fn from(root_mf_id: RootBasenameSuffixSkeletonManifest) -> Self {
        BlobstoreBytes::from_bytes(root_mf_id.0.into_bytes())
    }
}

impl RootBasenameSuffixSkeletonManifest {
    pub fn into_inner_id(self) -> BasenameSuffixSkeletonManifestId {
        self.0.id
    }
    pub fn inner_id(&self) -> &BasenameSuffixSkeletonManifestId {
        &self.0.id
    }
}

#[async_trait]
impl BonsaiDerivable for RootBasenameSuffixSkeletonManifest {
    const VARIANT: DerivableType = DerivableType::Bssm;

    type Dependencies = dependencies![];

    async fn derive_single(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsai: BonsaiChangeset,
        parents: Vec<Self>,
    ) -> Result<Self> {
        derive_single(ctx, derivation_ctx, bonsai, parents).await
    }

    async fn store_mapping(
        self,
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<()> {
        let key = format_key(derivation_ctx, changeset_id);
        derivation_ctx.blobstore().put(ctx, key, self.into()).await
    }

    async fn fetch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<Option<Self>> {
        let key = format_key(derivation_ctx, changeset_id);
        derivation_ctx
            .blobstore()
            .get(ctx, &key)
            .await?
            .map(TryInto::try_into)
            .transpose()
    }

    fn from_thrift(data: thrift::DerivedData) -> Result<Self> {
        if let thrift::DerivedData::basename_suffix_skeleton_manifest(
            thrift::DerivedDataBasenameSuffixSkeletonManifest::root_basename_suffix_skeleton_manifest(entry),
        ) = data
        {
            BssmDirectory::from_thrift(entry).map(Self)
        } else {
            Err(anyhow!(
                "Can't convert {} from provided thrift::DerivedData",
                Self::NAME.to_string(),
            ))
        }
    }

    fn into_thrift(data: Self) -> Result<thrift::DerivedData> {
        Ok(thrift::DerivedData::basename_suffix_skeleton_manifest(
            thrift::DerivedDataBasenameSuffixSkeletonManifest::root_basename_suffix_skeleton_manifest(
                data.0.into_thrift(),
            ),
        ))
    }
}

impl_bonsai_derived_via_manager!(RootBasenameSuffixSkeletonManifest);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cmp::Ordering;

use anyhow::Context;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::Loadable;
use cloned::cloned;
use context::CoreContext;
use futures::future;
use futures::stream;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use itertools::EitherOrBoth;
use manifest::After;
use manifest::AsyncManifest;
use manifest::ManifestOps;
use manifest::ManifestOrderedOps;
use manifest::PathOrPrefix;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use vec1::vec1;
use vec1::Vec1;

use crate::path::BssmPath;
use crate::RootBasenameSuffixSkeletonManifest;

type MononokePath = Option<MPath>;

fn normal_to_custom(path: MononokePath) -> MononokePath {
    path.map(|p| BssmPath::transform(p).into_raw())
}

enum BasenameOrSuffix {
    Basename(MPathElement),
    Suffix(MPathElement),
}

impl BasenameOrSuffix {
    fn key(&self) -> (&[u8], bool) {
        match self {
            Self::Basename(b) => (b.as_ref(), false),
            Self::Suffix(s) => (s.as_ref(), true),
        }
    }
    fn cmp(a: &Self, b: &Self) -> Ordering {
        a.key().cmp(&b.key())
    }
}

impl RootBasenameSuffixSkeletonManifest {
    /// Finds all files with given basenames in the given directories.
    pub async fn find_files_filter_basenames<'a>(
        &self,
        ctx: &'a CoreContext,
        blobstore: impl Blobstore + Clone + Sync + Send + 'static,
        prefixes: Vec<MononokePath>,
        basenames_and_suffixes: EitherOrBoth<Vec1<String>, Vec1<String>>,
        ordered: Option<Option<MononokePath>>,
    ) -> Result<impl Stream<Item = Result<MononokePath>> + 'a> {
        let (basenames, suffixes) = basenames_and_suffixes
            .map_any(Vec1::into_vec, Vec1::into_vec)
            .or_default();
        let prefixes: Vec1<PathOrPrefix> = Vec1::try_from_vec(prefixes)
            .unwrap_or_else(|_| vec1![None])
            .mapped(PathOrPrefix::Prefix);
        let mut all: Vec<BasenameOrSuffix> = basenames
            .into_iter()
            .map(|b| {
                Ok(BasenameOrSuffix::Basename(MPathElement::new(
                    b.into_bytes(),
                )?))
            })
            .chain(
                suffixes
                    .into_iter()
                    .map(|s| Ok(BasenameOrSuffix::Suffix(MPathElement::new(s.into_bytes())?))),
            )
            .collect::<Result<_>>()?;
        if ordered.is_some() {
            all.sort_unstable_by(BasenameOrSuffix::cmp);
        }
        let node = self.0.clone().load(ctx, &blobstore).await?;
        Ok(stream::iter(all.into_iter().map(anyhow::Ok))
            .and_then(move |bos| {
                cloned!(blobstore, ordered, prefixes, node);
                async move {
                    const FILE_ERR: &str = "Algorithmic error: Node should always be directory.";
                    use BasenameOrSuffix::*;
                    let nodes = match bos {
                        Basename(mut basename) => {
                            basename.reverse();
                            stream::iter(node.lookup(ctx, &blobstore, &basename).await?.map(
                                |entry| anyhow::Ok((basename, entry.into_dir().context(FILE_ERR)?)),
                            ))
                            .left_stream()
                        }
                        Suffix(mut suffix) => {
                            suffix.reverse();
                            // Why use async_stream::try_stream here to "forward" and not just
                            // return the stream? Because we need to take ownership of blobstore
                            // and suffix, otherwise we return a stream that references those values
                            // but does not own them, which causes a compilation error.
                            // I did not find a simpler way to do this.
                            cloned!(blobstore);
                            async_stream::try_stream! {
                                let mut s = node
                                    .list_prefix(ctx, &blobstore, suffix.as_ref())
                                    .await?;

                                while let Some((name, node)) = s.try_next().await? {
                                    let value = (name, node.into_tree().context(FILE_ERR)?);
                                    yield value;
                                }
                            }
                        }
                        .right_stream(),
                    };

                    Ok(nodes
                        .map_ok(move |(name, node)| {
                            let entries = match ordered.clone() {
                                None => node
                                    .find_entries(ctx.clone(), blobstore.clone(), prefixes.clone())
                                    .left_stream(),
                                Some(after) => {
                                    let after: After = after.map(normal_to_custom).into();
                                    // This will still traverse all distinct basenames even if we're
                                    // skipping a bunch. If this needs optimisation, we should skip
                                    // a suffix if we know all names starting with that suffix are
                                    // gonna be skipped anyway.
                                    let after = if after.skip(&name) {
                                        return stream::empty().left_stream();
                                    } else {
                                        after.enter_dir(&name)
                                    };
                                    node.find_entries_ordered(
                                        ctx.clone(),
                                        blobstore.clone(),
                                        prefixes.clone(),
                                        after,
                                    )
                                    .right_stream()
                                }
                            };
                            entries
                                .try_filter_map(|(path, entry)| {
                                    // No need to "untransform" path because we're already searching
                                    // from the reverse basename directory
                                    future::ok(entry.into_leaf().map(|_| path))
                                })
                                .right_stream()
                        })
                        .try_flatten())
                }
            })
            .try_flatten())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use mononoke_types::MPath;

/// Put reversed basename in beginning
pub(crate) struct BssmPath(MPath);

impl BssmPath {
    pub(crate) fn transform(path: MPath) -> Self {
        let mut rev_basename = path.split_dirname().1.clone();
        rev_basename.reverse();
        // Let's add the basename in the end of the path as well
        // This prevents bugs, otherwise files in top-level become files
        // But they should become directories
        // So a repo with files `file` and `dir/file` will become a repo with files
        // `elif/file` and `elif/dir/file`, which otherwise could cause a file-dir conflict.
        // It also conserves ordering of files, for the same basename.
        Self(MPath::from(rev_basename).join(&path))
    }

    #[cfg(test)]
    pub(crate) fn from_bsm_formatted_path(path: MPath) -> Self {
        Self(path)
    }

    pub(crate) fn into_raw(self) -> MPath {
        self.0
    }

    #[cfg(test)]
    pub(crate) fn untransform(self) -> anyhow::Result<MPath> {
        use anyhow::Context;

        let (rev_basename, rest) = self.0.split_first();
        let rest = rest.with_context(|| format!("Invalid format for path {}", self.0))?;
        let mut rev_basename2 = rest.split_dirname().1.clone();
        rev_basename2.reverse();
        if *rev_basename != rev_basename2 {
            anyhow::bail!(
                "Invalid format for path {}, expected reverse basename in root.",
                self.0
            )
        }
        Ok(rest)
    }
}

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

    use super::*;

    quickcheck! {
        fn roundtrip(path: MPath) -> bool {
            let bsm = BssmPath::transform(path.clone());
            bsm.untransform().unwrap() == path
        }
    }

    fn assert_transform(orig: &str, transform: &str) {
        assert_eq!(
            BssmPath::transform(MPath::new(orig).unwrap()).into_raw(),
            MPath::new(transform).unwrap()
        );
    }

    fn assert_untransform(transform: &str, orig: Option<&str>) {
        let untransform =
            BssmPath::from_bsm_formatted_path(MPath::new(transform).unwrap()).untransform();
        if let Some(orig) = orig {
            assert_eq!(untransform.unwrap(), MPath::new(orig).unwrap(),);
        } else {
            assert!(untransform.is_err());
        }
    }

    #[test]
    fn test_transform() {
        assert_transform("a/b/c", "c/a/b/c");
        assert_transform("dir/file", "elif/dir/file");
        assert_transform("file", "elif/file");
        assert_transform("dir/subdir/hello", "olleh/dir/subdir/hello");
        assert_transform("eden/mononoke/main.rs", "sr.niam/eden/mononoke/main.rs");
    }

    #[test]
    fn test_untransform() {
        assert_untransform("c/a/b/c", Some("a/b/c"));
        assert_untransform("elif/dir/file", Some("dir/file"));
        assert_untransform("elif/file", Some("file"));
        assert_untransform("olleh/dir/subdir/hello", Some("dir/subdir/hello"));
        assert_untransform(
            "sr.niam/eden/mononoke/main.rs",
            Some("eden/mononoke/main.rs"),
        );
        assert_untransform("$", None);
        assert_untransform("c/b/a", None);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::str::FromStr;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::RwLock;

use anyhow::Result;
use async_recursion::async_recursion;
use blobstore::Blobstore;
use blobstore::Loadable;
use bookmarks::BookmarkKey;
use bookmarks::BookmarksRef;
use changesets::ChangesetsRef;
use context::CoreContext;
use fbinit::FacebookInit;
use fixtures::TestRepoFixture;
use futures::stream;
use futures::TryStreamExt;
use mononoke_types::basename_suffix_skeleton_manifest::BssmDirectory;
use mononoke_types::basename_suffix_skeleton_manifest::BssmEntry;
use mononoke_types::BasenameSuffixSkeletonManifestId;
use mononoke_types::ChangesetIdPrefix;
use mononoke_types::ChangesetIdsResolvedFromPrefix;
use mononoke_types::MPathElement;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataRef;

use crate::RootBasenameSuffixSkeletonManifest;

#[async_recursion]
async fn validate(
    cache: &RwLock<HashMap<BasenameSuffixSkeletonManifestId, u64>>,
    ctx: &CoreContext,
    blobstore: &impl Blobstore,
    dir: BssmDirectory,
    rev_basename: Option<&'async_recursion MPathElement>,
) -> Result<()> {
    let from_cache = cache.read().unwrap().get(&dir.id).copied();
    let calculated_rollup_count = if let Some(count) = from_cache {
        count
    } else {
        let mf = dir.load(ctx, blobstore).await?;
        let calculated_rollup_count = AtomicU64::new(1);
        let calculated_rollup_count = &calculated_rollup_count;
        mf.into_subentries(ctx, blobstore)
            .try_for_each_concurrent(None, |(mut path, entry)| async move {
                calculated_rollup_count.fetch_add(entry.rollup_count(), Ordering::Relaxed);
                match entry {
                    BssmEntry::File => {
                        path.reverse();
                        assert_eq!(Some(&path), rev_basename)
                    }
                    BssmEntry::Directory(dir) => {
                        validate(cache, ctx, blobstore, dir, rev_basename.or(Some(&path))).await?
                    }
                }
                Ok(())
            })
            .await?;
        let count = calculated_rollup_count.load(Ordering::Relaxed);
        cache.write().unwrap().insert(dir.id, count);
        count
    };
    assert_eq!(dir.rollup_count, calculated_rollup_count);
    Ok(())
}

async fn test_for_fixture<F: TestRepoFixture + Send>(fb: FacebookInit) -> Result<()> {
    let ctx = CoreContext::test_mock(fb);
    let ctx = &ctx;
    let repo = F::getrepo(fb).await;
    let derived_data = repo.repo_derived_data();
    let blobstore = repo.repo_blobstore();
    let master = repo
        .bookmarks()
        .get(ctx.clone(), &BookmarkKey::from_str("master")?)
        .await?
        .unwrap();
    derived_data
        .derive::<RootBasenameSuffixSkeletonManifest>(ctx, master)
        .await?;
    let all_commits = repo
        .changesets()
        .get_many_by_prefix(ctx, ChangesetIdPrefix::from_bytes("").unwrap(), 1000)
        .await?;
    let all_commits = match all_commits {
        ChangesetIdsResolvedFromPrefix::Multiple(all_commits) => all_commits,
        other => anyhow::bail!("Weird number of commits: {:?}", other),
    };
    let cache = RwLock::new(HashMap::new());
    let cache = &cache;
    stream::iter(all_commits.into_iter().map(anyhow::Ok))
        .try_for_each_concurrent(None, |cs_id| async move {
            let mf: RootBasenameSuffixSkeletonManifest = derived_data.derive(ctx, cs_id).await?;
            validate(cache, ctx, blobstore, mf.0, None).await
        })
        .await?;
    Ok(())
}

#[fbinit::test]
async fn basic_test(fb: FacebookInit) {
    futures::try_join!(
        test_for_fixture::<fixtures::Linear>(fb),
        test_for_fixture::<fixtures::BranchEven>(fb),
        test_for_fixture::<fixtures::BranchUneven>(fb),
        test_for_fixture::<fixtures::BranchWide>(fb),
        test_for_fixture::<fixtures::MergeEven>(fb),
        test_for_fixture::<fixtures::ManyFilesDirs>(fb),
        test_for_fixture::<fixtures::MergeUneven>(fb),
        test_for_fixture::<fixtures::UnsharedMergeEven>(fb),
        test_for_fixture::<fixtures::UnsharedMergeUneven>(fb),
        test_for_fixture::<fixtures::ManyDiamonds>(fb),
    )
    .unwrap();
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::sync::Mutex;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use borrowed::borrowed;
use cloned::cloned;
use context::CoreContext;
use derived_data::batch::split_bonsais_in_linear_stacks;
use derived_data::batch::FileConflicts;
use derived_data::batch::SplitOptions;
use derived_data::batch::DEFAULT_STACK_FILE_CHANGES_LIMIT;
use derived_data_manager::DerivationContext;
use futures::stream::FuturesOrdered;
use futures::stream::TryStreamExt;
use lock_ext::LockExt;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use slog::debug;
use unodes::RootUnodeManifestId;

use crate::derive_v2::derive_blame_v2;
use crate::RootBlameV2;

pub async fn derive_blame_v2_in_batch(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bonsais: Vec<BonsaiChangeset>,
) -> Result<HashMap<ChangesetId, RootBlameV2>, Error> {
    let batch_len = bonsais.len();
    // We must split on any change as blame data must use the parent file, and
    // must split on copy info as this also affects blame.
    let linear_stacks = split_bonsais_in_linear_stacks(
        &bonsais,
        SplitOptions {
            file_conflicts: FileConflicts::AnyChange,
            copy_info: true,
            file_changes_limit: DEFAULT_STACK_FILE_CHANGES_LIMIT,
        },
    )?;
    let bonsais = Mutex::new(
        bonsais
            .into_iter()
            .map(|bcs| (bcs.get_changeset_id(), bcs))
            .collect::<HashMap<_, _>>(),
    );
    borrowed!(bonsais);

    let mut res = HashMap::new();
    for linear_stack in linear_stacks {
        let stack_len = linear_stack.stack_items.len();
        if let Some(item) = linear_stack.stack_items.first() {
            debug!(
                ctx.logger(),
                "derive blame batch at {} (stack of {} from batch of {})",
                item.cs_id.to_hex(),
                stack_len,
                batch_len,
            );
        }

        let new_blames = linear_stack
            .stack_items
            .into_iter()
            .enumerate()
            .map(|(index, item)| {
                // Clone owning copied to pass into the spawned future.
                cloned!(ctx, derivation_ctx);
                async move {
                    let csid = item.cs_id;
                    let bonsai = bonsais
                        .with(|bonsais| bonsais.remove(&csid))
                        .ok_or_else(|| anyhow!("changeset {} should be in bonsai batch", csid))?;
                    let derivation_fut = async move {
                        let root_manifest = derivation_ctx
                            .derive_dependency::<RootUnodeManifestId>(&ctx, csid)
                            .await?;
                        derive_blame_v2(&ctx, &derivation_ctx, bonsai, root_manifest)
                            .await
                            .with_context(|| {
                                format!(
                                    concat!(
                                        "failed to derive blame_v2 for {}, ",
                                        "index {} in stack of {} from batch of {}"
                                    ),
                                    csid, index, stack_len, batch_len
                                )
                            })?;
                        Ok::<_, Error>(root_manifest)
                    };
                    let derivation_handle = tokio::spawn(derivation_fut);
                    let root_manifest = derivation_handle.await??;
                    let derived = RootBlameV2 {
                        csid,
                        root_manifest,
                    };
                    Ok::<_, Error>((csid, derived))
                }
            })
            .collect::<FuturesOrdered<_>>()
            .try_collect::<Vec<_>>()
            .await?;

        res.extend(new_blames);
    }

    Ok(res)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

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

use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::Loadable;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use derived_data_manager::DerivationContext;
use futures::future;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use manifest::find_intersection_of_diffs;
use mononoke_types::blame_v2::store_blame;
use mononoke_types::blame_v2::BlameParent;
use mononoke_types::blame_v2::BlameV2;
use mononoke_types::blame_v2::BlameV2Id;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::FileUnodeId;
use mononoke_types::MPath;
use unodes::find_unode_rename_sources;
use unodes::RootUnodeManifestId;
use unodes::UnodeRenameSource;

use crate::fetch::fetch_content_for_blame_with_limit;
use crate::fetch::FetchOutcome;
use crate::DEFAULT_BLAME_FILESIZE_LIMIT;

pub(crate) async fn derive_blame_v2(
    ctx: &CoreContext,
    derivation_ctx: &DerivationContext,
    bonsai: BonsaiChangeset,
    root_manifest: RootUnodeManifestId,
) -> Result<(), Error> {
    let csid = bonsai.get_changeset_id();
    let parent_manifests = bonsai.parents().map(|csid| async move {
        let parent_root_mf_id = derivation_ctx
            .derive_dependency::<RootUnodeManifestId>(ctx, csid)
            .await?;
        Ok::<_, Error>(parent_root_mf_id.manifest_unode_id().clone())
    });

    let (parent_manifests, renames) = future::try_join(
        future::try_join_all(parent_manifests).err_into(),
        find_unode_rename_sources(ctx, derivation_ctx, &bonsai),
    )
    .await?;

    let filesize_limit = derivation_ctx
        .config()
        .blame_filesize_limit
        .unwrap_or(DEFAULT_BLAME_FILESIZE_LIMIT);
    let renames = Arc::new(renames);
    let blobstore = derivation_ctx.blobstore();
    find_intersection_of_diffs(
        ctx.clone(),
        blobstore.clone(),
        root_manifest.manifest_unode_id().clone(),
        parent_manifests,
    )
    .map_ok(|(path, entry)| Some((path?, entry.into_leaf()?)))
    .try_filter_map(future::ok)
    .map(move |path_and_file_unode| {
        cloned!(ctx, blobstore, renames);
        async move {
            let (path, file_unode) = path_and_file_unode?;
            tokio::spawn(async move {
                create_blame_v2(
                    &ctx,
                    &blobstore,
                    renames,
                    csid,
                    path,
                    file_unode,
                    filesize_limit,
                )
                .await
            })
            .await?
        }
    })
    .buffered(256)
    .try_for_each(|_| future::ok(()))
    .await?;

    Ok(())
}

async fn create_blame_v2(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    renames: Arc<HashMap<MPath, UnodeRenameSource>>,
    csid: ChangesetId,
    path: MPath,
    file_unode_id: FileUnodeId,
    filesize_limit: u64,
) -> Result<BlameV2Id, Error> {
    let file_unode = file_unode_id.load(ctx, blobstore).await?;

    let mut blame_parents = Vec::new();
    if let Some(source) = renames.get(&path) {
        // If the file was copied from another path, then we ignore its
        // contents in the parents, even if it existed there, and just use the
        // copy-from source as a parent.  This matches the Mercurial blame
        // implementation.
        blame_parents.push(fetch_blame_parent(
            ctx,
            blobstore,
            source.parent_index,
            source.from_path.clone(),
            source.unode_id,
            filesize_limit,
        ));
    } else {
        for (parent_index, unode_id) in file_unode.parents().iter().enumerate() {
            blame_parents.push(fetch_blame_parent(
                ctx,
                blobstore,
                parent_index,
                path.clone(),
                *unode_id,
                filesize_limit,
            ));
        }
    }

    let (content, blame_parents) = future::try_join(
        fetch_content_for_blame_with_limit(ctx, blobstore, file_unode_id, filesize_limit),
        future::try_join_all(blame_parents),
    )
    .await?;

    let blame = match content {
        FetchOutcome::Rejected(rejected) => BlameV2::rejected(rejected),
        FetchOutcome::Fetched(content) => BlameV2::new(csid, path, content, blame_parents)?,
    };

    store_blame(ctx, &blobstore, file_unode_id, blame).await
}

async fn fetch_blame_parent(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    parent_index: usize,
    path: MPath,
    unode_id: FileUnodeId,
    filesize_limit: u64,
) -> Result<BlameParent<Bytes>, Error> {
    let (content, blame) = future::try_join(
        fetch_content_for_blame_with_limit(ctx, blobstore, unode_id, filesize_limit),
        BlameV2Id::from(unode_id).load(ctx, blobstore).err_into(),
    )
    .await?;

    Ok(BlameParent::new(
        parent_index,
        path,
        content.into_bytes().ok(),
        blame,
    ))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use blobstore::Blobstore;
use blobstore::Loadable;
use bytes::Bytes;
use context::CoreContext;
use filestore::FetchKey;
use futures::TryStreamExt;
use mononoke_types::blame_v2::BlameRejected;
use mononoke_types::FileUnodeId;
use repo_blobstore::RepoBlobstoreArc;
use repo_derived_data::RepoDerivedDataRef;

use crate::DEFAULT_BLAME_FILESIZE_LIMIT;

pub enum FetchOutcome {
    Fetched(Bytes),
    Rejected(BlameRejected),
}

impl FetchOutcome {
    pub fn into_bytes(self) -> Result<Bytes, BlameRejected> {
        match self {
            FetchOutcome::Fetched(bytes) => Ok(bytes),
            FetchOutcome::Rejected(rejected) => Err(rejected),
        }
    }
}

/// Fetch the content of a file ready for blame.  If the file content is
/// too large or binary data is detected then the fetch may be rejected.
pub async fn fetch_content_for_blame(
    ctx: &CoreContext,
    repo: &(impl RepoBlobstoreArc + RepoDerivedDataRef),
    file_unode_id: FileUnodeId,
) -> Result<FetchOutcome> {
    let filesize_limit = repo
        .repo_derived_data()
        .manager()
        .config()
        .blame_filesize_limit
        .unwrap_or(DEFAULT_BLAME_FILESIZE_LIMIT);
    let blobstore = repo.repo_blobstore_arc() as Arc<dyn Blobstore>;
    fetch_content_for_blame_with_limit(ctx, &blobstore, file_unode_id, filesize_limit).await
}

/// Fetch the content of a file ready for blame.  If the file content is
/// too large or binary data is detected then the fetch may be rejected.
pub async fn fetch_content_for_blame_with_limit(
    ctx: &CoreContext,
    blobstore: &Arc<dyn Blobstore>,
    file_unode_id: FileUnodeId,
    filesize_limit: u64,
) -> Result<FetchOutcome> {
    let file_unode = file_unode_id.load(ctx, blobstore).await?;
    let content_id = *file_unode.content_id();
    let (mut stream, size) =
        filestore::fetch_with_size(blobstore, ctx.clone(), &FetchKey::Canonical(content_id))
            .await?
            .ok_or_else(|| anyhow!("Missing content: {}", content_id))?;
    if size > filesize_limit {
        return Ok(FetchOutcome::Rejected(BlameRejected::TooBig));
    }
    let mut buffer = Vec::with_capacity(size as usize);
    while let Some(bytes) = stream.try_next().await? {
        if bytes.contains(&0u8) {
            return Ok(FetchOutcome::Rejected(BlameRejected::Binary));
        }
        buffer.extend(bytes);
    }
    Ok(FetchOutcome::Fetched(Bytes::from(buffer)))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![type_length_limit = "1441792"]

mod batch_v2;
mod derive_v2;
mod fetch;
mod mapping_v2;

#[cfg(test)]
mod tests;

use anyhow::Error;
use blobstore::Loadable;
use blobstore::LoadableError;
use context::CoreContext;
use derived_data::BonsaiDerived;
use derived_data::DeriveError;
pub use fetch::fetch_content_for_blame;
pub use fetch::FetchOutcome;
use manifest::ManifestOps;
pub use mapping_v2::format_key;
pub use mapping_v2::RootBlameV2;
use metaconfig_types::BlameVersion;
use mononoke_types::blame_v2::BlameRejected;
use mononoke_types::blame_v2::BlameV2;
use mononoke_types::blame_v2::BlameV2Id;
use mononoke_types::ChangesetId;
use mononoke_types::FileUnodeId;
use mononoke_types::MPath;
use repo_blobstore::RepoBlobstoreArc;
use repo_derived_data::RepoDerivedDataRef;
use thiserror::Error;

pub const DEFAULT_BLAME_FILESIZE_LIMIT: u64 = 10 * 1024 * 1024;

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub struct BlameDeriveOptions {
    filesize_limit: u64,
    blame_version: BlameVersion,
}

impl Default for BlameDeriveOptions {
    fn default() -> Self {
        BlameDeriveOptions {
            filesize_limit: DEFAULT_BLAME_FILESIZE_LIMIT,
            blame_version: BlameVersion::default(),
        }
    }
}

#[derive(Debug, Error)]
pub enum BlameError {
    #[error("No such path: {0}")]
    NoSuchPath(MPath),
    #[error("Blame is not available for directories: {0}")]
    IsDirectory(MPath),
    #[error(transparent)]
    Rejected(#[from] BlameRejected),
    #[error(transparent)]
    DeriveError(#[from] DeriveError),
    #[error(transparent)]
    LoadableError(#[from] LoadableError),
    #[error(transparent)]
    Error(#[from] Error),
}

/// Fetch the blame for a file.  Blame will be derived if necessary.
pub async fn fetch_blame_v2(
    ctx: &CoreContext,
    repo: impl RepoBlobstoreArc + RepoDerivedDataRef + Sync + Send + Copy,
    csid: ChangesetId,
    path: MPath,
) -> Result<(BlameV2, FileUnodeId), BlameError> {
    let root_unode = RootBlameV2::derive(ctx, &repo, csid).await?.root_manifest();
    let blobstore = repo.repo_blobstore();
    let file_unode_id = root_unode
        .manifest_unode_id()
        .clone()
        .find_entry(ctx.clone(), blobstore.clone(), Some(path.clone()))
        .await?
        .ok_or_else(|| BlameError::NoSuchPath(path.clone()))?
        .into_leaf()
        .ok_or(BlameError::IsDirectory(path))?;
    let blame = BlameV2Id::from(file_unode_id).load(ctx, &blobstore).await?;
    Ok((blame, file_unode_id))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use context::CoreContext;
use derived_data::impl_bonsai_derived_via_manager;
use derived_data_manager::dependencies;
use derived_data_manager::BonsaiDerivable;
use derived_data_manager::DerivableType;
use derived_data_manager::DerivationContext;
use derived_data_service_if::types as thrift;
use metaconfig_types::BlameVersion;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::ManifestUnodeId;
use unodes::RootUnodeManifestId;

use crate::batch_v2::derive_blame_v2_in_batch;
use crate::derive_v2::derive_blame_v2;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RootBlameV2 {
    pub(crate) csid: ChangesetId,
    pub(crate) root_manifest: RootUnodeManifestId,
}

impl RootBlameV2 {
    pub fn root_manifest(&self) -> RootUnodeManifestId {
        self.root_manifest
    }

    pub fn changeset_id(&self) -> ChangesetId {
        self.csid
    }
}

pub fn format_key(derivation_ctx: &DerivationContext, changeset_id: ChangesetId) -> String {
    let root_prefix = "derived_root_blame_v2.";
    let key_prefix = derivation_ctx.mapping_key_prefix::<RootBlameV2>();
    format!("{}{}{}", root_prefix, key_prefix, changeset_id)
}

#[async_trait]
impl BonsaiDerivable for RootBlameV2 {
    const VARIANT: DerivableType = DerivableType::BlameV2;

    type Dependencies = dependencies![RootUnodeManifestId];

    async fn derive_single(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsai: BonsaiChangeset,
        _parents: Vec<Self>,
    ) -> Result<Self, Error> {
        let csid = bonsai.get_changeset_id();
        let root_manifest = derivation_ctx
            .derive_dependency::<RootUnodeManifestId>(ctx, csid)
            .await?;
        if derivation_ctx.config().blame_version != BlameVersion::V2 {
            return Err(anyhow!(
                "programming error: incorrect blame version (expected V2)"
            ));
        }
        derive_blame_v2(ctx, derivation_ctx, bonsai, root_manifest).await?;
        Ok(RootBlameV2 {
            csid,
            root_manifest,
        })
    }

    async fn derive_batch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        bonsais: Vec<BonsaiChangeset>,
        _gap_size: Option<usize>,
    ) -> Result<HashMap<ChangesetId, Self>, Error> {
        derive_blame_v2_in_batch(ctx, derivation_ctx, bonsais).await
    }

    async fn store_mapping(
        self,
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<()> {
        let key = format_key(derivation_ctx, changeset_id);
        derivation_ctx
            .blobstore()
            .put(ctx, key, self.root_manifest.into())
            .await
    }

    async fn fetch(
        ctx: &CoreContext,
        derivation_ctx: &DerivationContext,
        changeset_id: ChangesetId,
    ) -> Result<Option<Self>> {
        let key = format_key(derivation_ctx, changeset_id);
        match derivation_ctx.blobstore().get(ctx, &key).await? {
            Some(value) => Ok(Some(RootBlameV2 {
                csid: changeset_id,
                root_manifest: value.try_into()?,
            })),
            None => Ok(None),
        }
    }

    fn from_thrift(data: thrift::DerivedData) -> Result<Self> {
        if let thrift::DerivedData::blame(thrift::DerivedDataBlame::root_blame_v2(blame)) = data {
            Ok(Self {
                csid: ChangesetId::from_thrift(blame.changeset_id)?,
                root_manifest: match blame.unode {
                    thrift::DerivedDataUnode::root_unode_manifest_id(id) => {
                        ManifestUnodeId::from_thrift(id).map(RootUnodeManifestId)
                    }
                    thrift::DerivedDataUnode::UnknownField(x) => Err(anyhow!(
                        "Can't convert {} from provided thrift::DerivedData, unknown field: {}",
                        Self::NAME.to_string(),
                        x,
                    )),
                }?,
            })
        } else {
            Err(anyhow!(
                "Can't convert {} from provided thrift::DerivedData",
                Self::NAME.to_string(),
            ))
        }
    }

    fn into_thrift(data: Self) -> Result<thrift::DerivedData> {
        Ok(thrift::DerivedData::blame(
            thrift::DerivedDataBlame::root_blame_v2(thrift::DerivedDataRootBlameV2 {
                changeset_id: data.csid.into_thrift(),
                unode: thrift::DerivedDataUnode::root_unode_manifest_id(
                    data.root_manifest().manifest_unode_id().into_thrift(),
                ),
            }),
        ))
    }
}

impl_bonsai_derived_via_manager!(RootBlameV2);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;

use anyhow::anyhow;
use anyhow::Error;
use bonsai_hg_mapping::BonsaiHgMapping;
use bookmarks::Bookmarks;
use borrowed::borrowed;
use changeset_fetcher::ChangesetFetcher;
use changesets::Changesets;
use context::CoreContext;
use fbinit::FacebookInit;
use filestore::FilestoreConfig;
use maplit::btreemap;
use maplit::hashmap;
use metaconfig_types::BlameVersion;
use mononoke_types::blame_v2::BlameRejected;
use mononoke_types::blame_v2::BlameV2;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use repo_blobstore::RepoBlobstore;
use repo_derived_data::RepoDerivedData;
use test_repo_factory::TestRepoFactory;
use tests_utils::create_commit;
use tests_utils::store_files;
use tests_utils::store_rename;
use tests_utils::CreateCommitContext;

use crate::fetch_blame_v2;

#[facet::container]
struct TestRepo {
    #[facet]
    bonsai_hg_mapping: dyn BonsaiHgMapping,
    #[facet]
    bookmarks: dyn Bookmarks,
    #[facet]
    repo_blobstore: RepoBlobstore,
    #[facet]
    repo_derived_data: RepoDerivedData,
    #[facet]
    filestore_config: FilestoreConfig,
    #[facet]
    changeset_fetcher: dyn ChangesetFetcher,
    #[facet]
    changesets: dyn Changesets,
}

// File with multiple changes and a merge
const F0: &[&str] = &[
    // c0
    r#"|
1 0
1 1
"#,
    // c1
    r#"|
2 0
1 0
2 1
"#,
    // c2
    r#"|
2 0
1 0
3 0
3 1
2 1
3 2
"#,
    // c3
    r#"|
1 0
1 1
3 2
4 0
"#,
    // c4
    r#"|
2 0
1 0
3 0
3 1
2 1
3 2
4 0
"#,
];

const F0_AT_C4: &str = r#"c0: |
c1: 2 0
c0: 1 0
c2: 3 0
c2: 3 1
c1: 2 1
c2: 3 2
c3: 4 0
"#;

// file with multiple change only in one parent and a merge
const F1: &[&str] = &[
    // c0
    r#"|
1 0
1 1
"#,
    // c3
    r#"|
1 0
4 0
1 1
"#,
];

const F1_AT_C4: &str = r#"c0: |
c0: 1 0
c3: 4 0
c0: 1 1
"#;

// renamed file
const F2: &[&str] = &[
    // c0 as _f2
    r#"|
1 0
1 1
"#,
    // c1 as _f2 => f2
    r#"|
1 0
2 0
1 1
"#,
    // c3 as new f2
    r#"|
1 0
4 0
1 1
"#,
    // c4 as f2
    r#"|
5 0
1 0
2 0
4 0
1 1
"#,
];

const F2_AT_C4: &str = r#"c0: |
c4: 5 0
c0: 1 0
c1: 2 0
c3: 4 0
c0: 1 1
"#;

#[fbinit::test]
async fn test_blame_v2(fb: FacebookInit) -> Result<(), Error> {
    test_blame_version(fb, BlameVersion::V2).await
}

async fn test_blame_version(fb: FacebookInit, version: BlameVersion) -> Result<(), Error> {
    // Commits structure
    //
    //   0
    //  / \
    // 1   3
    // |   |
    // 2   |
    //  \ /
    //   4
    //
    let ctx = CoreContext::test_mock(fb);
    let repo: TestRepo = TestRepoFactory::new(fb)?
        .with_config_override(|config| {
            config
                .derived_data_config
                .get_active_config()
                .expect("No enabled derived data types config")
                .blame_version = version
        })
        .build()
        .await?;
    borrowed!(ctx, repo);

    let c0 = create_commit(
        ctx.clone(),
        repo,
        vec![],
        store_files(
            ctx,
            btreemap! {
                "f0" => Some(F0[0]),
                "f1" => Some(F1[0]),
                "_f2" => Some(F2[0]),
            },
            repo,
        )
        .await,
    )
    .await;

    let mut c1_changes = store_files(ctx, btreemap! {"f0" => Some(F0[1])}, repo).await;
    let (f2_path, f2_change) = store_rename(ctx, (MPath::new("_f2")?, c0), "f2", F2[1], repo).await;
    c1_changes.insert(f2_path, f2_change);
    let c1 = create_commit(ctx.clone(), repo, vec![c0], c1_changes).await;

    let c2 = create_commit(
        ctx.clone(),
        repo,
        vec![c1],
        store_files(ctx, btreemap! {"f0" => Some(F0[2])}, repo).await,
    )
    .await;

    let c3 = create_commit(
        ctx.clone(),
        repo,
        vec![c0],
        store_files(
            ctx,
            btreemap! {
                "f0" => Some(F0[3]),
                "f1" => Some(F1[1]),
                "f2" => Some(F2[2]),
            },
            repo,
        )
        .await,
    )
    .await;

    let c4 = create_commit(
        ctx.clone(),
        repo,
        vec![c2, c3],
        store_files(
            ctx,
            btreemap! {
                "f0" => Some(F0[4]),
                "f1" => Some(F1[1]), // did not change after c3
                "f2" => Some(F2[3]),
            },
            repo,
        )
        .await,
    )
    .await;

    let names = hashmap! {
        c0 => "c0",
        c1 => "c1",
        c2 => "c2",
        c3 => "c3",
        c4 => "c4",
    };

    let (blame, _) = fetch_blame_v2(ctx, repo, c4, MPath::new("f0")?).await?;
    assert_eq!(annotate(F0[4], blame, &names)?, F0_AT_C4);

    let (blame, _) = fetch_blame_v2(ctx, repo, c4, MPath::new("f1")?).await?;
    assert_eq!(annotate(F1[1], blame, &names)?, F1_AT_C4);

    let (blame, _) = fetch_blame_v2(ctx, repo, c4, MPath::new("f2")?).await?;
    assert_eq!(annotate(F2[3], blame, &names)?, F2_AT_C4);

    Ok(())
}

#[fbinit::test]
async fn test_blame_size_rejected_v2(fb: FacebookInit) -> Result<(), Error> {
    test_blame_size_rejected_version(fb, BlameVersion::V2).await
}

async fn test_blame_size_rejected_version(
    fb: FacebookInit,
    version: BlameVersion,
) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo: TestRepo = test_repo_factory::build_empty(fb).await.unwrap();
    borrowed!(ctx, repo);
    let file1 = "file1";
    let content = "content";
    let c1 = CreateCommitContext::new_root(ctx, &repo)
        .add_file(file1, content)
        .commit()
        .await?;

    // Default file size is 10MiB, so blame should be computed
    // without problems.
    let (blame, _) = fetch_blame_v2(ctx, repo, c1, MPath::new(file1)?).await?;
    let _ = blame.ranges()?;

    let repo: TestRepo = TestRepoFactory::new(fb)?
        .with_config_override(|config| {
            config
                .derived_data_config
                .get_active_config()
                .expect("No enabled derived data types config")
                .blame_version = version;
            config
                .derived_data_config
                .get_active_config()
                .expect("No enabled derived data types config")
                .blame_filesize_limit = Some(4);
        })
        .build()
        .await?;

    let file2 = "file2";
    let c2 = CreateCommitContext::new_root(ctx, &repo)
        .add_file(file2, content)
        .commit()
        .await?;

    // This repo has a decreased limit, so derivation should fail now
    let (blame, _) = fetch_blame_v2(ctx, &repo, c2, MPath::new(file2)?).await?;

    match blame.ranges() {
        Err(BlameRejected::TooBig) => {}
        _ => {
            return Err(anyhow!("unexpected result"));
        }
    }

    Ok(())
}

#[fbinit::test]
async fn test_blame_copy_source(fb: FacebookInit) -> Result<(), Error> {
    let ctx = CoreContext::test_mock(fb);
    let repo: TestRepo = TestRepoFactory::new(fb)?
        .with_config_override(|config| {
            config
                .derived_data_config
                .get_active_config()
                .expect("No enabled derived data types config")
                .blame_version = BlameVersion::V2
        })
        .build()
        .await?;
    borrowed!(ctx, repo);

    let c1 = CreateCommitContext::new_root(ctx, &repo)
        .add_file("file1", "one\ntwo\nthree\n")
        .add_file("file2", "zero\none\ntwo\nfour\n")
        .commit()
        .await?;

    let data = "none\none\ntwo\nthree\n";
    let c2 = CreateCommitContext::new(ctx, &repo, vec![c1])
        .add_file_with_copy_info("file1", data, (c1, "file2"))
        .commit()
        .await?;

    let (blame, _) = fetch_blame_v2(ctx, repo, c2, MPath::new("file1")?).await?;
    let lines = blame
        .lines()?
        .map(|line| (line.changeset_id, line.path.to_string(), line.origin_offset))
        .collect::<Vec<_>>();

    // The "one" and "two" lines are blamed to the copy source, and not the
    // parent.  The "three" line blames to the commit that performed the copy,
    // and not the parent.
    assert_eq!(
        lines,
        vec![
            (&c2, "file1".to_string(), 0),
            (&c1, "file2".to_string(), 1),
            (&c1, "file2".to_string(), 2),
            (&c2, "file1".to_string(), 3),
        ]
    );
    Ok(())
}

fn annotate(
    content: &str,
    blame: BlameV2,
    names: &HashMap<ChangesetId, &'static str>,
) -> Result<String, Error> {
    let mut result = String::new();
    let mut ranges = blame.ranges()?;
    let mut range = ranges
        .next()
        .ok_or_else(|| Error::msg("empty blame for non empty content"))?;
    for (index, line) in content.lines().enumerate() {
        if index as u32 >= range.offset + range.length {
            range = ranges
                .next()
                .ok_or_else(|| Error::msg("not enough ranges in a blame"))?;
        }
        let name = names
            .get(&range.csid)
            .ok_or_else(|| Error::msg("unresolved csid"))?;
        result.push_str(name);
        result.push_str(": ");
        result.push_str(line);
        result.push('\n');
    }
    Ok(result)
}