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

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

use futures::future::BoxFuture;
use futures::TryStreamExt;

use crate::errors::programming;
use crate::Result;
use crate::Set;
use crate::Vertex;

/// Pre-process a parent function that might have cycles.
/// Return a new parent function that won't have cycles.
///
/// This function is not fast. Only use it on small graphs.
pub fn break_parent_func_cycle<F>(parent_func: F) -> impl Fn(Vertex) -> Result<Vec<Vertex>>
where
    F: Fn(Vertex) -> Result<Vec<Vertex>>,
{
    #[derive(Default)]
    struct State {
        /// Previously calculated parents.
        known: HashMap<Vertex, Vec<Vertex>>,
    }
    impl State {
        fn is_ancestor(&self, ancestor: &Vertex, descentant: &Vertex) -> bool {
            if ancestor == descentant {
                return true;
            }
            let mut to_visit = vec![descentant];
            let mut visited = HashSet::new();
            while let Some(v) = to_visit.pop() {
                if !visited.insert(v) {
                    continue;
                }
                if let Some(parents) = self.known.get(&v) {
                    for p in parents {
                        if p == ancestor {
                            return true;
                        }
                        to_visit.push(p);
                    }
                }
            }
            false
        }
    }
    let state: Mutex<State> = Default::default();

    move |v: Vertex| -> Result<Vec<Vertex>> {
        let mut state = state.lock().unwrap();
        if let Some(parents) = state.known.get(&v) {
            return Ok(parents.clone());
        }
        let parents = parent_func(v.clone())?;
        let mut result = Vec::with_capacity(parents.len());
        for p in parents {
            if !state.is_ancestor(&v, &p) {
                // Not a cycle.
                result.push(p);
            }
        }
        state.known.insert(v, result.clone());
        Ok(result)
    }
}

/// Given a `set` (sub-graph) and a filter function that selects "known"
/// subset of its input, apply filter to `set`.
///
/// The filter funtion must have following properties:
/// - filter(xs) + filter(ys) = filter(xs + ys)
/// - If its input contains both X and Y and X is an ancestor of Y in the
///   sub-graph, and its output contains Y, then its output must also
///   contain Y's ancestor X.
///   In other words, if vertex X is considered known, then ancestors
///   of X are also known.
///
/// This function has a similar signature with `filter`, but it utilizes
/// the above properties to test (much) less vertexes for a large input
/// set.
///
/// The idea of the algorithm comes from Mercurial's `setdiscovery.py`,
/// introduced by [1]. `setdiscovery.py` is used to figure out what
/// commits are needed to be pushed or pulled.
///
/// [1]: https://www.mercurial-scm.org/repo/hg/rev/cb98fed52495
pub async fn filter_known<'a>(
    set: Set,
    filter_known_func: &(
         dyn (Fn(&[Vertex]) -> BoxFuture<'a, Result<Vec<Vertex>>>) + Send + Sync + 'a
     ),
) -> Result<Set> {
    // Figure out unassigned (missing) vertexes that do need to be inserted.
    //
    // remaining:  subset not categorized.
    // known:      subset categorized as "known"
    // unknown:    subset categorized as "unknown"
    //
    // See [1] for the algorithm, basically:
    // - Take a subset (sample) of "remaining".
    // - Check the subset (sample). Divide it into (new_known, new_unknown).
    // - known   |= ancestors(new_known)
    // - unknown |= descendants(new_unknown)
    // - remaining -= known | unknown
    // - Repeat until "remaining" becomes empty.
    let mut remaining = set;
    let subdag = match remaining.dag() {
        Some(dag) => dag,
        None => return programming("filter_known requires set to associate to a Dag"),
    };
    let mut known = Set::empty();

    for i in 1usize.. {
        let remaining_old_len = remaining.count().await?;
        if remaining_old_len == 0 {
            break;
        }

        // Sample: heads, roots, and the "middle point" from "remaining".
        let sample = if i <= 2 {
            // But for the first few queries, let's just check the roots.
            // This could reduce remote lookups, when we only need to
            // query the roots to rule out all `remaining` vertexes.
            subdag.roots(remaining.clone()).await?
        } else {
            subdag
                .roots(remaining.clone())
                .await?
                .union(&subdag.heads(remaining.clone()).await?)
                .union(&remaining.skip((remaining_old_len as u64) / 2).take(1))
        };
        let sample: Vec<Vertex> = sample.iter().await?.try_collect().await?;
        let new_known = filter_known_func(&sample).await?;
        let new_unknown: Vec<Vertex> = {
            let filtered_set: HashSet<Vertex> = new_known.iter().cloned().collect();
            sample
                .iter()
                .filter(|v| !filtered_set.contains(v))
                .cloned()
                .collect()
        };

        let new_known = Set::from_static_names(new_known);
        let new_unknown = Set::from_static_names(new_unknown);

        let new_known = subdag.ancestors(new_known).await?;
        let new_unknown = subdag.descendants(new_unknown).await?;

        remaining = remaining.difference(&new_known.union(&new_unknown));
        let remaining_new_len = remaining.count().await?;

        let known_old_len = known.count().await?;
        known = known.union(&new_known);
        let known_new_len = known.count().await?;

        tracing::trace!(
            target: "dag::utils::filter_known",
            "#{} remaining {} => {}, known: {} => {}",
            i,
            remaining_old_len,
            remaining_new_len,
            known_old_len,
            known_new_len
        );
    }

    Ok(known)
}

/// Produce an order of "nodes" in a graph for cleaner graph output.
/// For example, turn:
///
/// ```plain,ignore
/// 0-1---3---5---7---9---11--12
///             / \
///    2---4---6   8---10
/// # [12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
/// ```
///
/// into:
///
/// ```plain,ignore
/// 0-1-3-5-------7------9-11-12
///              / \
///         2-4-6   8-10
/// # [12, 11, 9, 10, 8, 7, 6, 4, 2, 5, 3, 1, 0]
/// ```
///
/// This function takes a slice of integers as graph "nodes". The callsite
/// should maintain mapping between the integers and real graph "nodes"
/// (vertexes, ids, etc.).
///
/// The algorithm works by DFSing from roots to heads following the
/// parent->child edges. It was originally implemented in `smartlog.py`.
///
/// By default, longer branches get output first (for a fork), or last (for a
/// merge). This makes a typical vertical graph rendering algorithm more
/// likely put longer branches in the first column, result in less indentation
/// overall. This can be overridden by `priorities`. `priorities` and their
/// ancestors are more likely outputted in the first column. For a typical
/// vertical graph rendering algorithm, `priorities` is usually set to the
/// head of the main branch.
pub fn beautify_graph(parents: &[Vec<usize>], priorities: &[usize]) -> Vec<usize> {
    let n = parents.len();
    let mut children_list: Vec<Vec<usize>> = (0..n).map(|_| Vec::new()).collect();
    let mut weight: Vec<isize> = vec![1; n];
    let mut roots: Vec<usize> = Vec::new();

    // Consider user-provided additional_weight.
    for &i in priorities {
        if i < n {
            weight[i] = n as isize;
        }
    }

    // Populate children_list.
    for (i, ps) in parents.iter().enumerate().rev() {
        for &p in ps {
            // i has parent p
            if p < n {
                children_list[p].push(i);
                weight[p] = weight[p].saturating_add(weight[i]);
            }
        }
        if ps.is_empty() {
            roots.push(i);
        }
    }

    // Sort children by weight.
    for children in children_list.iter_mut() {
        children.sort_unstable_by_key(|&c| (weight[c], c));
    }
    roots.sort_unstable_by_key(|&c| (-weight[c], c));
    drop(weight);

    // DFS from roots to heads.
    let mut remaining: Vec<usize> = parents.iter().map(|ps| ps.len()).collect();
    let mut output: Vec<usize> = Vec::with_capacity(n);
    let mut outputted: Vec<bool> = vec![false; n];
    let mut to_visit: VecDeque<usize> = roots.into();
    while let Some(id) = to_visit.pop_front() {
        // Already outputted?
        if outputted[id] {
            continue;
        }

        // Need to wait for visiting other parents first?
        if remaining[id] > 0 {
            // Visit it later.
            to_visit.push_back(id);
            continue;
        }

        // Output this id.
        output.push(id);
        outputted[id] = true;

        // Visit children next.
        for &c in children_list[id].iter().rev() {
            remaining[c] -= 1;
            to_visit.push_front(c);
        }
    }

    output.reverse();
    output
}

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

    #[test]
    fn test_break_parent_func_cycle() -> Result<()> {
        let parent_func = |n: Vertex| -> Result<Vec<Vertex>> { Ok(vec![n, v("1"), v("2")]) };
        let parent_func_no_cycle = break_parent_func_cycle(parent_func);
        assert_eq!(parent_func_no_cycle(v("A"))?, vec![v("1"), v("2")]);
        assert_eq!(parent_func_no_cycle(v("1"))?, vec![v("2")]);
        assert_eq!(parent_func_no_cycle(v("2"))?, vec![]);
        Ok(())
    }

    #[test]
    fn test_break_parent_func_cycle_linear() -> Result<()> {
        let parent_func = |n: Vertex| -> Result<Vec<Vertex>> {
            let list = "0123456789".chars().map(|c| v(c)).collect::<Vec<_>>();
            let parents = match list.iter().position(|x| x == &n) {
                Some(i) if i > 0 => vec![list[i - 1].clone()],
                _ => vec![],
            };
            Ok(parents)
        };
        let parent_func_no_cycle = break_parent_func_cycle(parent_func);
        assert_eq!(parent_func_no_cycle(v("2"))?, vec![v("1")]);
        assert_eq!(parent_func_no_cycle(v("9"))?, vec![v("8")]);
        assert_eq!(parent_func_no_cycle(v("8"))?, vec![v("7")]);
        assert_eq!(parent_func_no_cycle(v("1"))?, vec![v("0")]);
        assert_eq!(parent_func_no_cycle(v("5"))?, vec![v("4")]);
        assert_eq!(parent_func_no_cycle(v("6"))?, vec![v("5")]);
        assert_eq!(parent_func_no_cycle(v("4"))?, vec![v("3")]);
        assert_eq!(parent_func_no_cycle(v("0"))?, vec![]);
        assert_eq!(parent_func_no_cycle(v("3"))?, vec![v("2")]);
        assert_eq!(parent_func_no_cycle(v("7"))?, vec![v("6")]);
        Ok(())
    }

    /// Quickly create a Vertex.
    fn v(name: impl ToString) -> Vertex {
        Vertex::copy_from(name.to_string().as_bytes())
    }

    #[test]
    fn test_beautify_graph() {
        let t = |text: &str| -> Vec<usize> {
            let split: Vec<&str> = text.split("#").chain(std::iter::once("")).take(2).collect();
            let parents = parents_from_drawdag(split[0]);
            let priorities: Vec<usize> = split[1]
                .split_whitespace()
                .map(|s| s.parse::<usize>().unwrap())
                .collect();
            beautify_graph(&parents, &priorities)
        };
        assert_eq!(
            t(r#"
              0-2-4-5
                   /
                1-3"#),
            [5, 3, 1, 4, 2, 0]
        );
        assert_eq!(
            t(r#"
              0-2-4-5
               \   /
                1-3"#),
            [5, 4, 2, 3, 1, 0]
        );
        assert_eq!(
            t(r#"
              0-2-4-5
               \
                1-3"#),
            [5, 4, 2, 3, 1, 0]
        );
        assert_eq!(
            t(r#"
                 0-1-3-5-7-9-11-12
                        / \
                   2-4-6   8-10"#),
            [12, 11, 9, 10, 8, 7, 6, 4, 2, 5, 3, 1, 0]
        );
        assert_eq!(
            t(r#"
                   0-3-5-7-9-12
                        / \
                 1-2-4-6   8-10-11"#),
            [11, 10, 8, 12, 9, 7, 5, 3, 0, 6, 4, 2, 1]
        );

        // Preserve (reversed) order for same-length branches.
        // [0, 1, 2], 2 is before 1.
        assert_eq!(
            t(r#"
                  0-1
                   \
                    2"#),
            [2, 1, 0]
        );
        // [0, 1, 2], 1 is before 0.
        assert_eq!(
            t(r#"
                  0-2
                   /
                  1"#),
            [2, 1, 0]
        );

        // With manual 'priorities'
        assert_eq!(
            t(r#"
              0-2-4-5
                   /
                1-3   # 4"#),
            [5, 3, 1, 4, 2, 0]
        );
        assert_eq!(
            t(r#"
              0-2-4-5
                   /
                1-3   # 3"#),
            [5, 4, 2, 0, 3, 1],
        );

        assert_eq!(
            t(r#"
              0-2-4-5
               \
                1-3   # 3"#),
            [3, 1, 5, 4, 2, 0]
        );
        assert_eq!(
            t(r#"
              0-2-4-5
               \
                1-3   # 5"#),
            [5, 4, 2, 3, 1, 0]
        );
    }

    /// Convert drawdag ASCII to an integer parents array.
    fn parents_from_drawdag(ascii: &str) -> Vec<Vec<usize>> {
        let parents_map = drawdag::parse(ascii);
        let mut parents_vec = Vec::new();
        for i in 0usize.. {
            let s = i.to_string();
            let parents = match parents_map.get(&s) {
                Some(parents) => parents,
                None => break,
            };
            let parents: Vec<usize> = parents
                .into_iter()
                .map(|p| p.parse::<usize>().unwrap())
                .collect();
            parents_vec.push(parents);
        }
        parents_vec
    }
}


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

use std::cmp;
use std::fmt;
use std::sync::atomic;
use std::sync::atomic::AtomicU32;
use std::sync::Arc;

/// A linked list tracking a logic "version" with compatibility rules:
/// - Append-only changes bump the version, the new version is backwards
///   compatible.
/// - Non-append-only changes create a new version that is incompatible with
///   all other versions (in the current process).
/// - Clones (cheaply) preserve the version.
///
/// Supported operations:
/// - `new() -> x`: Create a new version that is not comparable (compatible)
///   with other versions.
/// - `clone(x) -> y`: Clone `x` to `y`. `x == y`. `x` and `y` are compatible.
/// - `bump(x) -> y`: Bump `x` to `y`. `y > x`. `y` is backwards-compatible with
///   `x`. Note: `y` is not comparable (compatible) with other `bump(x)`.
/// - `x > y`: `true` if `x` is backwards compatible with `y`.
///
/// The linked list can be shared in other linked lists. So they form a tree
/// effectively. Comparing to a DAG, there is no "merge" operation.
/// Compatibility questions become reachability questions.
#[derive(Clone)]
pub struct VerLink {
    inner: Arc<Inner>,
}

#[derive(Clone)]
struct Inner {
    parent: Option<VerLink>,

    /// The base number. Two `VerLink`s with different base numbers are incompatible.
    /// Used as an optimization to exit `compatible` early.
    base: u32,

    /// The "generation number", distance to root (the parentless `VerLink`).
    /// If `x.compatible(y)` returns `x`, then `x.gen` must be >= `y.gen`.
    /// Used as an optimization to exit `compatible` early.
    gen: u32,
}

impl VerLink {
    /// Creates a new `VerLink` that is incompatible with all other `VerLink`s
    /// in the process.
    pub fn new() -> Self {
        let inner = Inner {
            parent: None,
            base: next_id(),
            gen: 0,
        };
        Self {
            inner: Arc::new(inner),
        }
    }

    /// Bumps the `VerLink` for backward-compatible (ex. append-only) changes.
    pub fn bump(&mut self) {
        match Arc::get_mut(&mut self.inner) {
            // This is an optimization to avoid increasing the length of the
            // linked list (which slows down "compatible" calculation) if
            // possible.
            Some(_inner) => {
                // Increasing gen is not necessary for correctness.
                // _inner.gen += 1;
            }
            None => {
                let next_inner = Inner {
                    parent: Some(self.clone()),
                    base: self.inner.base,
                    gen: self.inner.gen + 1,
                };
                let next = Self {
                    inner: Arc::new(next_inner),
                };
                *self = next;
            }
        }
    }
}

impl PartialOrd for VerLink {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        if self.inner.base != other.inner.base {
            // Fast path: self and other have different root nodes and are not
            // reachable from each other.
            return None;
        }
        if self.inner.gen < other.inner.gen {
            other.partial_cmp(self).map(|o| o.reverse())
        } else {
            debug_assert!(self.inner.gen >= other.inner.gen);
            if Arc::ptr_eq(&self.inner, &other.inner) {
                return Some(cmp::Ordering::Equal);
            }
            let mut cur = self.inner.parent.as_ref();
            while let Some(this) = cur {
                if this.inner.gen < other.inner.gen {
                    // Fast path: not possible to reach other from here.
                    return None;
                }
                if Arc::ptr_eq(&this.inner, &other.inner) {
                    return Some(cmp::Ordering::Greater);
                }
                cur = this.inner.parent.as_ref();
            }
            None
        }
    }
}

impl PartialEq for VerLink {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.inner, &other.inner)
    }
}

impl fmt::Debug for VerLink {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut cur = Some(self);
        while let Some(this) = cur {
            write!(f, "{:p}", this.inner.as_ref())?;
            cur = this.inner.parent.as_ref();
            f.write_str("->")?;
            if cur.is_none() {
                write!(f, "{}", this.inner.base)?;
            }
        }
        Ok(())
    }
}

fn next_id() -> u32 {
    static ID: AtomicU32 = AtomicU32::new(0);
    ID.fetch_add(1, atomic::Ordering::AcqRel)
}

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

    #[test]
    fn test_individual_news_are_incompatible() {
        let a = VerLink::new();
        let b = VerLink::new();
        assert!(compatible(&a, &b).is_none());
    }

    #[test]
    fn test_clone_compatible() {
        let a = VerLink::new();
        let b = a.clone();
        assert_eq!(&a, &b);
        assert_eq!(compatible(&a, &b), Some(&b));
    }

    #[test]
    fn test_bump_is_different_and_backwards_compatible() {
        let a = VerLink::new();
        let mut b = a.clone();
        b.bump();
        assert_ne!(&b, &a);
        assert_eq!(compatible(&a, &b), Some(&b));

        b.bump();
        b.bump();
        assert_eq!(compatible(&a, &b), Some(&b));
    }

    #[test]
    fn test_clone_bump_twice() {
        let a = VerLink::new();
        let mut b = a.clone();
        b.bump();
        let mut c = b.clone();
        c.bump();
        assert_eq!(compatible(&a, &c), Some(&c));
        assert_eq!(compatible(&b, &c), Some(&c));
    }

    #[test]
    fn test_bump_independently_become_incompatible() {
        let mut a = VerLink::new();
        let mut b = a.clone();
        b.bump();
        a.bump();
        assert_eq!(compatible(&a, &b), None);
    }

    #[test]
    fn test_bump_avoid_increase_len_if_possible() {
        let a = VerLink::new();
        let mut b = a.clone();
        assert_eq!(a.chain_len(), 1);
        assert_eq!(b.chain_len(), 1);
        b.bump(); // Increases chain_len by 1.
        assert_eq!(b.chain_len(), 2);
        b.bump(); // Does not change chain len.
        b.bump();
        assert_eq!(b.chain_len(), 2);
    }

    /// Find the more compatible version.
    #[allow(clippy::neg_cmp_op_on_partial_ord)]
    fn compatible<'a>(a: &'a VerLink, b: &'a VerLink) -> Option<&'a VerLink> {
        if a == b {
            assert!(!(a != b));
            assert!(!(a < b));
            assert!(!(b < a));
            assert!(!(a > b));
            assert!(!(b > a));
            Some(a)
        } else if a < b {
            assert!(!(a == b));
            assert!(a != b);
            assert!(!(b < a));
            assert!(!(a > b));
            assert!(b > a);
            Some(b)
        } else if a > b {
            assert!(!(a == b));
            assert!(a != b);
            assert!(!(a < b));
            assert!(b < a);
            assert!(!(b > a));
            Some(a)
        } else {
            assert!(!(a == b));
            assert!(a != b);
            assert!(!(a < b));
            assert!(!(a > b));
            assert!(!(b < a));
            assert!(!(b > a));
            None
        }
    }

    impl VerLink {
        /// Length of the linked list.
        fn chain_len(&self) -> usize {
            let mut len = 0;
            let mut cur = Some(self);
            while let Some(this) = cur {
                len += 1;
                cur = this.inner.parent.as_ref();
            }
            len
        }
    }
}


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

use std::collections::HashSet;

use serde::Deserialize;

use crate::Group;
use crate::VertexName;

/// A list of [`VertexName`]s (usually heads) with options attached to each vertex.
#[derive(Default, Debug, Clone, Deserialize)]
#[serde(transparent)]
pub struct VertexListWithOptions {
    list: Vec<(VertexName, VertexOptions)>,
}

/// Options attached to a vertex. Usually the vertex is a head. The head and its
/// ancestors are going to be inserted to the graph. The options controls some
/// details about the insertion.
#[derive(Debug, Clone, Deserialize)]
#[non_exhaustive]
pub struct VertexOptions {
    /// How many ids to reserve for this vertex. Suppose this vertex has id `n`,
    /// then `n+1..=n+reserve_size` can only be used when inserting this vertex
    /// and its ancestors in the same batch.
    ///
    /// Note: if any id `j` in the `n+1..=n+reserve_size` range were already
    /// taken, then the reserve range becomes `n+1..j` instead. This avoids
    /// fragmentation.
    #[serde(default = "Default::default")]
    pub reserve_size: u32,

    /// The highest [`Group`] for this vertex. If set to `NON_MASTER` then
    /// this vertex could end up in `MASTER` or `NON_MASTER`. If set to
    /// `MASTER` then this vertex will end up in `MASTER` group.
    #[serde(default = "default_highest_group")]
    pub highest_group: Group,
}

const fn default_highest_group() -> Group {
    Group::NON_MASTER
}

impl Default for VertexOptions {
    fn default() -> Self {
        Self {
            reserve_size: 0,
            highest_group: default_highest_group(),
        }
    }
}

impl<'a> From<&'a [VertexName]> for VertexListWithOptions {
    fn from(list: &'a [VertexName]) -> Self {
        // Just use default options.
        Self {
            list: list
                .iter()
                .map(|v| (v.clone(), VertexOptions::default()))
                .collect(),
        }
    }
}

impl From<Vec<VertexName>> for VertexListWithOptions {
    fn from(list: Vec<VertexName>) -> Self {
        // Just use default options.
        Self {
            list: list
                .into_iter()
                .map(|v| (v, VertexOptions::default()))
                .collect(),
        }
    }
}

impl From<Vec<(VertexName, VertexOptions)>> for VertexListWithOptions {
    fn from(list: Vec<(VertexName, VertexOptions)>) -> Self {
        Self { list }
    }
}

impl VertexListWithOptions {
    /// Get the vertexes and their options.
    pub fn vertex_options(&self) -> Vec<(VertexName, VertexOptions)> {
        self.list.clone()
    }

    /// Get the vertexes.
    pub fn vertexes(&self) -> Vec<VertexName> {
        self.list.iter().map(|i| i.0.clone()).collect()
    }

    /// Get the vertexes, filter by the `highest_group` option.
    pub fn vertexes_by_group(&self, group: Group) -> Vec<VertexName> {
        self.list
            .iter()
            .filter_map(|(v, o)| {
                if o.highest_group == group {
                    Some(v.clone())
                } else {
                    None
                }
            })
            .collect()
    }

    /// Test if this list is empty.
    pub fn is_empty(&self) -> bool {
        self.list.is_empty()
    }

    /// Add a new item to the list.
    pub fn push(&mut self, head_opts: (VertexName, VertexOptions)) {
        self.list.push(head_opts);
    }

    /// Set the `highest_group` option for all vertexes.
    pub fn with_highest_group(mut self, group: Group) -> Self {
        for (_v, opts) in self.list.iter_mut() {
            opts.highest_group = group;
        }
        self
    }

    /// Chain another list. Vertexes that are already in this list are skipped.
    pub fn chain(mut self, other: impl Into<Self>) -> Self {
        let other = other.into();
        let existed: HashSet<_> = self.vertexes().into_iter().collect();
        for (v, o) in other.list {
            if !existed.contains(&v) {
                self.list.push((v, o))
            }
        }
        self
    }
}


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

use std::any::Any;
use std::fs;
use std::fs::File;
use std::io;
use std::io::Cursor;
use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Release;
use std::sync::Mutex;

use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use fs2::FileExt;
use indexedlog::log;
use indexedlog::log::Fold;
use minibytes::Bytes;

use super::IdDagStore;
use crate::errors::bug;
use crate::id::Group;
use crate::id::Id;
use crate::ops::Persist;
use crate::segment::describe_segment_bytes;
use crate::segment::hex;
use crate::segment::Segment;
use crate::segment::SegmentFlags;
use crate::spanset::Span;
use crate::IdSet;
use crate::Level;
use crate::Result;

pub struct IndexedLogStore {
    log: log::Log,
    path: PathBuf,
    cached_max_level: AtomicU8,
}

/// Fold (accumulator) that tracks IdSet covered in groups.
/// The state is stored as part in `log`.
#[derive(Debug, Default)]
struct CoveredIdSetFold {
    inner: Mutex<CoveredIdSetInner>,
}

#[derive(Debug, Default, Clone)]
struct CoveredIdSetInner {
    /// Covered id set, including pending removals.
    /// Use `get_id_set_by_group` to get an up-to-date view without removals.
    id_set_by_group: [IdSet; Group::COUNT],

    /// Pending removals. This avoids O(N*M) when there are N spans in id_set
    /// and M segments are being removed.
    ///
    /// Changing this field from `CoveredIdSetFold` must take `&mut self` to
    /// avoid race conditions.
    id_set_pending_remove_by_group: [IdSet; Group::COUNT],
}

const MAX_LEVEL_UNKNOWN: u8 = 0;
const LEVEL_BYTES: usize = std::mem::size_of::<Level>();

impl Fold for CoveredIdSetFold {
    fn load(&mut self, bytes: &[u8]) -> io::Result<()> {
        let id_sets = mincode::deserialize(bytes)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        let mut locked = self.inner.lock().unwrap();
        locked.id_set_by_group = id_sets;
        locked.id_set_pending_remove_by_group = Default::default();
        Ok(())
    }

    fn dump(&self) -> io::Result<Vec<u8>> {
        let mut inner = self.inner.lock().unwrap();
        inner.apply_removals();
        mincode::serialize(&inner.id_set_by_group)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }

    #[allow(clippy::for_loops_over_fallibles)]
    fn accumulate(&mut self, data: &[u8]) -> indexedlog::Result<()> {
        let mut inner = self.inner.lock().unwrap();
        if data.starts_with(IndexedLogStore::MAGIC_REMOVE_SEGMENT) {
            let data_start = IndexedLogStore::MAGIC_REMOVE_SEGMENT.len() + LEVEL_BYTES;
            if let Some(seg_data) = data.get(data_start..) {
                let seg = Segment(Bytes::copy_from_slice(seg_data));
                let span = match seg.span() {
                    Ok(span) => span,
                    Err(e) => return Err(("cannot decode segment", e).into()),
                };
                if let Some(set) = inner
                    .id_set_pending_remove_by_group
                    .get_mut(span.low.group().0)
                {
                    set.push(span);
                } else {
                    let msg = format!(
                        "unsupported group {} in segment {:?}",
                        span.low.group().0,
                        seg
                    );
                    return Err(msg.as_str().into());
                }
            }
            return Ok(());
        }
        inner.apply_removals();

        // See log_open_options for how other index functions read the entry.
        if data == IndexedLogStore::MAGIC_CLEAR_NON_MASTER {
            inner.id_set_by_group[Group::NON_MASTER.0] = IdSet::empty();
            return Ok(());
        }
        let data = if data.starts_with(IndexedLogStore::MAGIC_REWRITE_LAST_FLAT) {
            // See MAGIC_REWRITE_LAST_FLAT for format.
            let data_start = IndexedLogStore::MAGIC_REWRITE_LAST_FLAT.len() + Segment::OFFSET_DELTA
                - Segment::OFFSET_LEVEL;
            &data[data_start..]
        } else {
            data
        };
        let seg = Segment(Bytes::copy_from_slice(data));
        let span = match seg.span() {
            Ok(s) => s,
            Err(e) => return Err(("cannot parse segment in CoveredIdSetFold", e).into()),
        };
        if let Some(set) = inner.id_set_by_group.get_mut(span.low.group().0) {
            set.push(span);
        }
        Ok(())
    }

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

    fn clone_boxed(&self) -> Box<dyn Fold> {
        let mut inner = self.inner.lock().unwrap();
        inner.apply_removals();
        let cloned_inner = inner.clone();
        let cloned = CoveredIdSetFold {
            inner: Mutex::new(cloned_inner),
        };
        Box::new(cloned)
    }
}

impl CoveredIdSetInner {
    /// Apply pending removals.
    fn apply_removals(&mut self) {
        for group in Group::ALL {
            let group = group.0;
            if !self.id_set_pending_remove_by_group[group].is_empty() {
                let new_set = self.id_set_by_group[group]
                    .difference(&self.id_set_pending_remove_by_group[group]);
                self.id_set_by_group[group] = new_set;
                self.id_set_pending_remove_by_group[group] = IdSet::empty();
            }
        }
    }
}

// Required functionality
impl IdDagStore for IndexedLogStore {
    fn max_level(&self) -> Result<Level> {
        let max_level = self.cached_max_level.load(Acquire);
        if max_level != MAX_LEVEL_UNKNOWN {
            return Ok(max_level);
        }
        let max_level = match self
            .log
            .lookup_range(Self::INDEX_LEVEL_HEAD, ..)?
            .rev()
            .nth(0)
        {
            None => 0,
            Some(key) => key?.0.get(0).cloned().unwrap_or(0),
        };
        self.cached_max_level.store(max_level, Release);
        Ok(max_level)
    }

    fn find_segment_by_head_and_level(&self, head: Id, level: u8) -> Result<Option<Segment>> {
        let key = Self::serialize_head_level_lookup_key(head, level);
        match self.log.lookup(Self::INDEX_LEVEL_HEAD, &key)?.nth(0) {
            None => Ok(None),
            Some(bytes) => Ok(Some(self.segment_from_slice(bytes?))),
        }
    }

    fn find_flat_segment_including_id(&self, id: Id) -> Result<Option<Segment>> {
        let level = 0;
        let low = Self::serialize_head_level_lookup_key(id, level);
        let high = [level + 1];
        let iter = self
            .log
            .lookup_range(Self::INDEX_LEVEL_HEAD, &low[..]..&high[..])?;
        for entry in iter {
            let (_, entries) = entry?;
            for entry in entries {
                let entry = entry?;
                let seg = self.segment_from_slice(entry);
                if seg.span()?.low > id {
                    return Ok(None);
                }
                // low <= rev
                debug_assert!(seg.high()? >= id); // by range query
                return Ok(Some(seg));
            }
        }
        Ok(None)
    }

    fn insert_segment(&mut self, segment: Segment) -> Result<()> {
        let level = segment.level()?;
        self.cached_max_level.fetch_max(level, AcqRel);
        // When inserting a new flat segment, consider merging it with the last
        // flat segment on disk.
        //
        // Turn:
        //
        //   [last segment] [(new) segment]
        //
        // Into:
        //
        //   [------------]
        //    (removed)
        //   [(new, merged) segment       ]
        //    (in memory)
        if level == 0 {
            if self.maybe_insert_merged_flat_segment(&segment)? {
                return Ok(());
            }
        }
        self.log.append(&segment.0)?;
        Ok(())
    }

    fn remove_flat_segment_unchecked(&mut self, segment: &Segment) -> Result<()> {
        let max_level = self.max_level()?;
        let mut data =
            Vec::with_capacity(segment.0.len() + Self::MAGIC_REMOVE_SEGMENT.len() + LEVEL_BYTES);
        data.extend_from_slice(Self::MAGIC_REMOVE_SEGMENT);
        data.push(max_level);
        data.extend_from_slice(segment.0.as_ref());
        // The actual remove operation is done by index functions.
        // See log_open_options().
        self.log.append(&data)?;
        Ok(())
    }

    fn all_ids_in_groups(&self, groups: &[Group]) -> Result<IdSet> {
        let fold = self
            .log
            .fold(Self::FOLD_COVERED_ID_SET)?
            .as_any()
            .downcast_ref::<CoveredIdSetFold>()
            .expect("should downcast to CoveredIdSetFold defined by OpenOptions");
        let mut result = IdSet::empty();
        let mut inner = fold.inner.lock().unwrap();
        inner.apply_removals();
        let id_sets = &inner.id_set_by_group;
        for group in groups {
            result = result.union(&id_sets[group.0]);
        }
        Ok(result)
    }

    fn next_segments(&self, id: Id, level: Level) -> Result<Vec<Segment>> {
        let lower_bound = Self::serialize_head_level_lookup_key(id, level);
        let upper_bound = Self::serialize_head_level_lookup_key(id.group().max_id(), level);
        let mut result = Vec::new();
        for entry in self
            .log
            .lookup_range(Self::INDEX_LEVEL_HEAD, &lower_bound[..]..=&upper_bound)?
        {
            let (_, values) = entry?;
            for value in values {
                result.push(self.segment_from_slice(value?));
            }
        }
        Ok(result)
    }

    fn iter_segments_descending<'a>(
        &'a self,
        max_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>> {
        let lower_bound = Self::serialize_head_level_lookup_key(Id::MIN, level);
        let upper_bound = Self::serialize_head_level_lookup_key(max_high_id, level);
        let iter = self
            .log
            .lookup_range(Self::INDEX_LEVEL_HEAD, &lower_bound[..]..=&upper_bound[..])?
            .rev();
        let iter = iter.flat_map(move |entry| match entry {
            Ok((_key, values)) => values
                .into_iter()
                .map(|value| {
                    let value = value?;
                    Ok(self.segment_from_slice(value))
                })
                .collect(),
            Err(err) => vec![Err(err.into())],
        });
        Ok(Box::new(iter))
    }

    fn iter_segments_ascending<'a>(
        &'a self,
        min_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a + Send + Sync>> {
        let lower_bound = Self::serialize_head_level_lookup_key(min_high_id, level);
        let upper_bound = Self::serialize_head_level_lookup_key(Id::MAX, level);
        let iter = self
            .log
            .lookup_range(Self::INDEX_LEVEL_HEAD, &lower_bound[..]..=&upper_bound[..])?;
        let iter = iter.flat_map(move |entry| match entry {
            Ok((_key, values)) => values
                .map(|value| {
                    let value = value?;
                    Ok(self.segment_from_slice(value))
                })
                .collect(),
            Err(err) => vec![Err(err.into())],
        });
        Ok(Box::new(iter))
    }

    fn iter_flat_segments_with_parent_span<'a>(
        &'a self,
        parent_span: Span,
    ) -> Result<Box<dyn Iterator<Item = Result<(Id, Segment)>> + 'a>> {
        let mut result: Vec<(Id, Segment)> = Vec::new();
        for group in Group::ALL {
            let low = index_parent_key(parent_span.low, group.min_id());
            let high = index_parent_key(parent_span.high, group.max_id());
            let range = &low[..]..=&high[..];
            let range_iter = self.log.lookup_range(Self::INDEX_PARENT, range)?;
            for entry in range_iter {
                let (key, _segments) = entry?;
                let parent_id = {
                    let bytes: [u8; 8] = key[1..9].try_into().unwrap();
                    Id(u64::from_be_bytes(bytes))
                };
                let child_id = {
                    let bytes: [u8; 8] = key[9..].try_into().unwrap();
                    Id(u64::from_be_bytes(bytes))
                };
                if let Some(seg) = self.find_flat_segment_including_id(child_id)? {
                    result.push((parent_id, seg));
                }
            }
        }
        Ok(Box::new(result.into_iter().map(Ok)))
    }

    fn iter_flat_segments_with_parent<'a>(
        &'a self,
        parent: Id,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>> {
        let iter = self.iter_flat_segments_with_parent_span(parent.into())?;
        Ok(Box::new(iter.map(|item| item.map(|(_, seg)| seg))))
    }

    /// Mark non-master ids as "removed".
    fn remove_non_master(&mut self) -> Result<()> {
        self.log.append(Self::MAGIC_CLEAR_NON_MASTER)?;
        let non_master_ids = self.all_ids_in_groups(&[Group::NON_MASTER])?;
        if !non_master_ids.is_empty() {
            return bug("remove_non_master did not take effect");
        }
        Ok(())
    }
}

impl Persist for IndexedLogStore {
    type Lock = File;

    fn lock(&mut self) -> Result<File> {
        // Take a filesystem lock. The file name 'lock' is taken by indexedlog
        // running on Windows, so we choose another file name here.
        let lock_file = {
            let mut path = self.path.clone();
            path.push("wlock");
            File::open(&path).or_else(|_| {
                fs::OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&path)
            })?
        };
        lock_file.lock_exclusive()?;
        Ok(lock_file)
    }

    fn reload(&mut self, _lock: &Self::Lock) -> Result<()> {
        self.log.clear_dirty()?;
        self.log.sync()?;
        Ok(())
    }

    fn persist(&mut self, _lock: &Self::Lock) -> Result<()> {
        self.log.sync()?;
        Ok(())
    }
}

impl IndexedLogStore {
    /// Attempt to merge the flat `segment` with the last flat segment to reduce
    /// fragmentation. Insert the merged segment.
    ///
    /// Return true if the merged segment was inserted.
    fn maybe_insert_merged_flat_segment(&mut self, segment: &Segment) -> Result<bool> {
        if let Some(merged) = self.maybe_merged_flat_segment(segment)? {
            let mut bytes = Vec::with_capacity(merged.0.len() + 10);
            let span = segment.span()?;
            bytes.extend_from_slice(IndexedLogStore::MAGIC_REWRITE_LAST_FLAT);
            bytes.extend_from_slice(&Self::serialize_head_level_lookup_key(
                span.low - 1, /* = last_segment.high */
                0,            /* level */
            ));
            bytes.extend_from_slice(&merged.0);
            self.log.append(&bytes)?;

            Ok(true)
        } else {
            Ok(false)
        }
    }

    // Used internally to generate the index key for lookup
    fn serialize_head_level_lookup_key(value: Id, level: u8) -> [u8; Self::KEY_LEVEL_HEAD_LEN] {
        let mut buf = [0u8; Self::KEY_LEVEL_HEAD_LEN];
        {
            let mut cur = Cursor::new(&mut buf[..]);
            cur.write_u8(level).unwrap();
            cur.write_u64::<BigEndian>(value.0).unwrap();
            debug_assert_eq!(cur.position(), Self::KEY_LEVEL_HEAD_LEN as u64);
        }
        buf
    }

    fn segment_from_slice(&self, bytes: &[u8]) -> Segment {
        let bytes = if bytes.starts_with(IndexedLogStore::MAGIC_REWRITE_LAST_FLAT) {
            let start = Self::MAGIC_REWRITE_LAST_FLAT.len() + Self::KEY_LEVEL_HEAD_LEN;
            &bytes[start..]
        } else {
            bytes
        };
        Segment(self.log.slice_to_bytes(bytes))
    }
}

/// Describe bytes of an indexedlog entry.
/// This is only for troubleshooting purpose.
pub fn describe_indexedlog_entry(data: &[u8]) -> String {
    let mut message = String::new();
    if data == IndexedLogStore::MAGIC_CLEAR_NON_MASTER {
        message += &format!("# {}: MAGIC_CLEAR_NON_MASTER\n", hex(data),);
    } else if data.starts_with(IndexedLogStore::MAGIC_REMOVE_SEGMENT) {
        message += &format!(
            "# {}: MAGIC_REMOVE_SEGMENT\n",
            hex(IndexedLogStore::MAGIC_REMOVE_SEGMENT)
        );
        if let Some(max_level) = data.get(IndexedLogStore::MAGIC_REMOVE_SEGMENT.len()) {
            message += &format!("# {}: Max Level = {}\n", hex(&[*max_level]), max_level);
            let end = IndexedLogStore::MAGIC_REMOVE_SEGMENT.len() + LEVEL_BYTES;
            message += &describe_indexedlog_entry(&data[end..]);
        }
    } else if data.starts_with(IndexedLogStore::MAGIC_REWRITE_LAST_FLAT) {
        message += &format!(
            "# {}: MAGIC_REWRITE_LAST_FLAT\n",
            hex(IndexedLogStore::MAGIC_REWRITE_LAST_FLAT)
        );
        let start = IndexedLogStore::MAGIC_REWRITE_LAST_FLAT.len();
        let end = start + Segment::OFFSET_DELTA - Segment::OFFSET_LEVEL;
        let previous_index = &data[start..end];
        let previous_level = previous_index[0];
        let previous_head = (&previous_index[1..]).read_u64::<BigEndian>().unwrap_or(0);

        message += &format!(
            "# {}: Previous index Level = {}, Head = {}\n",
            hex(&data[start..end]),
            previous_level,
            Id(previous_head),
        );

        message += &describe_indexedlog_entry(&data[end..]);
    } else {
        message += &describe_segment_bytes(data);
    }
    message
}

// Implementation details
impl IndexedLogStore {
    const INDEX_LEVEL_HEAD: usize = 0;
    const INDEX_PARENT: usize = 1;
    const FOLD_COVERED_ID_SET: usize = 0;
    const KEY_LEVEL_HEAD_LEN: usize = Segment::OFFSET_DELTA - Segment::OFFSET_LEVEL;

    // "Normal" format is just the plain bytes in `Segment`. See `Segment` for details.
    // Basically, FLAG (1B) + LEVEL (1B) + HIGH (8B) + ...

    /// Magic bytes in `Log` that indicates "remove all non-master segments".
    /// A Segment entry has at least KEY_LEVEL_HEAD_LEN (9) bytes so it does
    /// not conflict with this.
    const MAGIC_CLEAR_NON_MASTER: &'static [u8] = b"CLRNM";

    /// Magic bytes in `Log` that indicates this entry replaces a previous flat
    /// segment.
    ///
    /// Format:
    ///
    /// ```plain,ignore
    /// MAGIC_REWRITE_LAST_FLAT + LEVEL (0u8) + PREVIOUS_HEAD (u64) + SEGMENT
    /// ```
    ///
    /// The `LEVEL + PREVIOUS_HEAD` part is used to remove the segment from the
    /// `(level, head)` index.
    const MAGIC_REWRITE_LAST_FLAT: &'static [u8] = &[0xf0];

    /// Magic prefix that indicates removing a segment and its related indexes.
    ///
    /// Format:
    ///
    /// ```plain,ignore
    /// MAGIC_REMOVE_SEGMENT + MAX_LEVEL (u8) + SEGMENT
    /// ```
    const MAGIC_REMOVE_SEGMENT: &'static [u8] = &[0xf1];

    #[allow(clippy::assertions_on_constants)]
    pub fn log_open_options() -> log::OpenOptions {
        assert!(Self::MAGIC_CLEAR_NON_MASTER.len() < Segment::OFFSET_DELTA);
        assert!(Group::BITS == 8);
        for magic in [Self::MAGIC_REWRITE_LAST_FLAT, Self::MAGIC_REMOVE_SEGMENT] {
            assert_ne!(
                SegmentFlags::all().bits() & magic[Segment::OFFSET_FLAGS],
                magic[Segment::OFFSET_FLAGS],
                "magic prefix should not conflict with possible flags (first byte in Segment)"
            )
        }
        log::OpenOptions::new()
            .create(true)
            .index("level-head", |data| {
                // (level, high)
                if data == Self::MAGIC_CLEAR_NON_MASTER {
                    let max_level = 255;
                    (0..=max_level)
                        .map(|level| {
                            log::IndexOutput::RemovePrefix(Box::new([
                                level,
                                Group::NON_MASTER.0 as u8,
                            ]))
                        })
                        .collect()
                } else if data.starts_with(Self::MAGIC_REMOVE_SEGMENT) {
                    // data: 0xf1 + MAX_LEVEL (u8) + SEGMENT
                    let mut index_output = Vec::new();
                    if let Some(seg_data) =
                        data.get(Self::MAGIC_REMOVE_SEGMENT.len() + LEVEL_BYTES..)
                    {
                        let max_level = data[1];
                        let seg = Segment(Bytes::copy_from_slice(seg_data));
                        // Remove head indexes for all levels.
                        index_output.reserve(max_level as usize + 1);
                        if let Ok(head) = seg.head() {
                            for level in 0..=max_level {
                                let index_key = Self::serialize_head_level_lookup_key(head, level);
                                index_output.push(log::IndexOutput::Remove(index_key.into()));
                            }
                        }
                    }
                    index_output
                } else if data.starts_with(Self::MAGIC_REWRITE_LAST_FLAT) {
                    // See MAGIC_REWRITE_LAST_FLAT for format.
                    let start = Self::MAGIC_REWRITE_LAST_FLAT.len();
                    let end = start + Segment::OFFSET_DELTA - Segment::OFFSET_LEVEL;
                    let previous_index = &data[start..end];
                    vec![
                        log::IndexOutput::Remove(previous_index.to_vec().into_boxed_slice()),
                        log::IndexOutput::Reference(
                            (end + Segment::OFFSET_LEVEL) as u64
                                ..(end + Segment::OFFSET_DELTA) as u64,
                        ),
                    ]
                } else {
                    vec![log::IndexOutput::Reference(
                        Segment::OFFSET_LEVEL as u64..Segment::OFFSET_DELTA as u64,
                    )]
                }
            })
            .index("group-parent-child", |data| {
                //  child-group parent child  -> child for flat segments
                //  ^^^^^^^^^^^ ^^^^^^ ^^^^^^
                //  u8          u64 BE u64 BE
                //
                //  The "child-group" prefix is used for invalidating index when
                //  non-master Ids get re-assigned.
                if data == Self::MAGIC_CLEAR_NON_MASTER {
                    // Invalidate child-group == 1 entries
                    return vec![log::IndexOutput::RemovePrefix(Box::new([
                        Group::NON_MASTER.0 as u8,
                    ]))];
                }

                if data.starts_with(Self::MAGIC_REMOVE_SEGMENT) {
                    // data: 0xf1 + MAX_LEVEL (u8) + SEGMENT
                    let mut index_output = Vec::new();
                    if let Some(data) = data.get(2..) {
                        let seg = Segment(Bytes::copy_from_slice(data));
                        if let (Ok(child), Ok(parents)) = (seg.low(), seg.parents()) {
                            // Remove parent->child indexes.
                            index_output.reserve(parents.len());
                            for parent in parents {
                                let index_key = index_parent_key(parent, child);
                                index_output.push(log::IndexOutput::Remove(index_key.into()));
                            }
                        }
                    }
                    return index_output;
                }

                if data.starts_with(Self::MAGIC_REWRITE_LAST_FLAT) {
                    // NOTE: Segments this index points to will have wrong `high`.
                    //
                    // But we never use the segments from this index. Instead,
                    // we use the index key to figure out (parent, child) and do
                    // an extra lookup of `child` to figure out the segments.
                    return Vec::new();
                }

                let seg = Segment(Bytes::copy_from_slice(data));
                let mut result = Vec::new();
                if seg.level().ok() == Some(0) {
                    // This should never pass since MAGIC_CLEAR_NON_MASTER[0] != 0.
                    // ([0] stores level: u8).
                    assert_ne!(
                        data,
                        Self::MAGIC_CLEAR_NON_MASTER,
                        "bug: MAGIC_CLEAR_NON_MASTER conflicts with data"
                    );
                    if let (Ok(parents), Ok(span)) = (seg.parents(), seg.span()) {
                        assert_eq!(
                            span.low.group(),
                            span.high.group(),
                            "Cross-group segment is unexpected"
                        );
                        let child_id = span.low;
                        for parent_id in parents {
                            let bytes = index_parent_key(parent_id, child_id);
                            result.push(log::IndexOutput::Owned(bytes.into()));
                        }
                    }
                }
                result
            })
            .fold_def("cover", || Box::new(CoveredIdSetFold::default()))
    }

    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let log = Self::log_open_options().open(path.clone())?;
        let iddag = Self {
            log,
            path,
            cached_max_level: AtomicU8::new(MAX_LEVEL_UNKNOWN),
        };
        Ok(iddag)
    }

    pub fn open_from_clean_log(log: log::Log) -> Result<Self> {
        let path = log.path().as_opt_path().unwrap().to_path_buf();
        if log.iter_dirty().next().is_some() {
            return bug("open_from_clean_log got a dirty log");
        }
        let iddag = Self {
            log,
            path,
            cached_max_level: AtomicU8::new(MAX_LEVEL_UNKNOWN),
        };
        Ok(iddag)
    }

    pub fn try_clone(&self) -> Result<IndexedLogStore> {
        let log = self.log.try_clone()?;
        let store = IndexedLogStore {
            log,
            path: self.path.clone(),
            cached_max_level: AtomicU8::new(self.cached_max_level.load(Acquire)),
        };
        Ok(store)
    }

    pub fn try_clone_without_dirty(&self) -> Result<IndexedLogStore> {
        let log = self.log.try_clone_without_dirty()?;
        let store = IndexedLogStore {
            log,
            path: self.path.clone(),
            cached_max_level: AtomicU8::new(MAX_LEVEL_UNKNOWN),
        };
        Ok(store)
    }
}

// Build index key for the INDEX_PARENT (group-parent-child) index.
fn index_parent_key(parent_id: Id, child_id: Id) -> [u8; 17] {
    let group = child_id.group();
    let mut result = [0u8; 1 + 8 + 8];
    debug_assert!(group.0 <= 0xff);
    result[0] = group.0 as u8;
    result[1..9].copy_from_slice(&parent_id.0.to_be_bytes());
    result[9..].copy_from_slice(&child_id.0.to_be_bytes());
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::iddagstore::tests::dump_store_state;

    #[test]
    fn test_merge_persisted_segments() -> Result<()> {
        // Test that a persisted segment can still be mutable and merged.
        //
        // Persisted                    | Not persisted
        // [0..=5] [6..=10, parents=[3]] [11..=20, parents=[10]]
        //         [6..                       =20, parents=[3] ] <- merged
        let tmp = tempfile::tempdir()?;
        let mut iddag = IndexedLogStore::open(tmp.path())?;
        let seg1 = Segment::new(SegmentFlags::HAS_ROOT, 0, Id(0), Id(5), &[]);
        let seg2 = Segment::new(SegmentFlags::empty(), 0, Id(6), Id(10), &[Id(3)]);
        iddag.insert_segment(seg1)?;
        iddag.insert_segment(seg2)?;
        let locked = iddag.lock()?;
        iddag.persist(&locked)?;

        let seg3 = Segment::new(SegmentFlags::ONLY_HEAD, 0, Id(11), Id(20), &[Id(10)]);
        iddag.insert_segment(seg3)?;
        iddag.persist(&locked)?;

        // Reload.
        let iddag2 = IndexedLogStore::open(tmp.path())?;

        // Check the merged segments.
        assert_eq!(
            dbg_iter(iddag2.iter_segments_descending(Id(20), 0)?),
            "[H6-20[3], R0-5[]]"
        );

        // Check parent -> child index.
        // 10 -> 11 parent index wasn't inserted.
        assert_eq!(
            dbg_iter(iddag2.iter_flat_segments_with_parent(Id(10))?),
            "[]"
        );
        // 3 -> 6 parent index only returns the new segment.
        assert_eq!(
            dbg_iter(iddag2.iter_flat_segments_with_parent(Id(3))?),
            "[H6-20[3]]"
        );

        // Check (level, head) -> segment index.
        // Check lookup by "including_id". Should all return the new merged segment.
        assert_eq!(
            dbg(iddag2.find_flat_segment_including_id(Id(7))?),
            "Some(H6-20[3])"
        );
        assert_eq!(
            dbg(iddag2.find_flat_segment_including_id(Id(13))?),
            "Some(H6-20[3])"
        );
        assert_eq!(
            dbg(iddag2.find_flat_segment_including_id(Id(20))?),
            "Some(H6-20[3])"
        );
        // Check lookup by head.
        // By head 20 returns the new merged segment.
        assert_eq!(
            dbg(iddag2.find_segment_by_head_and_level(Id(20), 0)?),
            "Some(H6-20[3])"
        );
        // By head 10 does not return the old segment.
        assert_eq!(
            dbg(iddag2.find_segment_by_head_and_level(Id(10), 0)?),
            "None"
        );

        Ok(())
    }

    #[test]
    fn test_backwards_compatibility() {
        // Test that data written by older versions of this struct can still
        // be understood by the current version.
        let tmp = tempfile::tempdir().unwrap();
        let mut iddag = IndexedLogStore::open(tmp.path()).unwrap();

        // Code used to print "described_entries". They might base on older APIs
        // and no longer compile.
        // #[cfg(any())] makes compiler skip the code block.
        #[cfg(any())]
        {
            const ROOT: SegmentFlags = SegmentFlags::HAS_ROOT;
            const HEAD: SegmentFlags = SegmentFlags::ONLY_HEAD;
            const EMPTY: SegmentFlags = SegmentFlags::empty();
            let segs = vec![
                Segment::new(ROOT | HEAD, 0, Id(0), Id(10), &[]),
                Segment::new(HEAD, 0, Id(11), Id(20), &[Id(10)]), // merge with previous
                Segment::new(ROOT, 0, Id(21), Id(30), &[]),
                Segment::new(EMPTY, 0, nid(0), nid(10), &[Id(20)]),
                Segment::new(EMPTY, 0, nid(11), nid(20), &[nid(10)]),
                Segment::new(EMPTY, 0, nid(5), nid(15), &[Id(10)]),
            ];
            for seg in &segs[..5] {
                iddag.insert_segment(seg.clone()).unwrap();
            }
            iddag.remove_non_master().unwrap();
            iddag.insert_segment(segs[5].clone()).unwrap();
            iddag.remove_flat_segment(&segs[2]).unwrap();
            for item in iddag.log.iter() {
                let data = item.unwrap();
                let s = describe_indexedlog_entry(data);
                eprintln!("{}", s);
            }
        }

        let described_entries = r#"
# 03: Flags = HAS_ROOT | ONLY_HEAD
# 00: Level = 0
# 00 00 00 00 00 00 00 0a: High = 10
# 0a: Delta = 10 (Low = 0)
# 00: Parent count = 0

# f0: MAGIC_REWRITE_LAST_FLAT
# 00 00 00 00 00 00 00 00 0a: Previous index Level = 0, Head = 10
# 03: Flags = HAS_ROOT | ONLY_HEAD
# 00: Level = 0
# 00 00 00 00 00 00 00 14: High = 20
# 14: Delta = 20 (Low = 0)
# 00: Parent count = 0

# 01: Flags = HAS_ROOT
# 00: Level = 0
# 00 00 00 00 00 00 00 1e: High = 30
# 09: Delta = 9 (Low = 21)
# 00: Parent count = 0

# 00: Flags = (empty)
# 00: Level = 0
# 01 00 00 00 00 00 00 0a: High = N10
# 0a: Delta = 10 (Low = N0)
# 01: Parent count = 1
# 14: Parents[0] = 20

# 00: Flags = (empty)
# 00: Level = 0
# 01 00 00 00 00 00 00 14: High = N20
# 09: Delta = 9 (Low = N11)
# 01: Parent count = 1
# 8a 80 80 80 80 80 80 80 01: Parents[0] = N10

# 43 4c 52 4e 4d: MAGIC_CLEAR_NON_MASTER

# 00: Flags = (empty)
# 00: Level = 0
# 01 00 00 00 00 00 00 0f: High = N15
# 0a: Delta = 10 (Low = N5)
# 01: Parent count = 1
# 0a: Parents[0] = 10

# f1: MAGIC_REMOVE_SEGMENT
# 00: Max Level = 0
# 01: Flags = HAS_ROOT
# 00: Level = 0
# 00 00 00 00 00 00 00 1e: High = 30
# 09: Delta = 9 (Low = 21)
# 00: Parent count = 0
"#;

        for described_entry in described_entries.split("\n\n") {
            let data = undescribe(described_entry);
            iddag.log.append(data).unwrap();
        }

        let all = iddag.all_ids_in_groups(&Group::ALL).unwrap();
        assert_eq!(format!("{:?}", &all), "0..=20 N5..=N15");

        let state = dump_store_state(&iddag, &all);
        assert_eq!(state, "\nLv0: RH0-20[], N5-N15[10]\nP->C: 10->N5");
    }

    /// Turn a string generated by "describe_indexedlog_entry" back to bytes.
    fn undescribe(s: &str) -> Vec<u8> {
        let mut data = Vec::new();
        for line in s.lines() {
            // line looks like: "# 00 00: Something".
            if let Some(line) = line.get(2..) {
                if let Some((hex_bytes, _)) = line.split_once(':') {
                    for hex_byte in hex_bytes.split(' ') {
                        let byte = u8::from_str_radix(hex_byte, 16).unwrap();
                        data.push(byte);
                    }
                }
            }
        }
        data
    }

    fn nid(id: u64) -> Id {
        Group::NON_MASTER.min_id() + id
    }

    #[test]
    fn test_describe() -> Result<()> {
        let tmp = tempfile::tempdir()?;
        let mut iddag = IndexedLogStore::open(tmp.path())?;
        let seg1 = Segment::new(SegmentFlags::HAS_ROOT, 0, Id(0), Id(5), &[]);
        let seg2 = Segment::new(SegmentFlags::empty(), 0, Id(6), Id(10), &[Id(5)]);
        iddag.insert_segment(seg1)?;
        iddag.insert_segment(seg2)?;
        let bytes = iddag.log.iter_dirty().nth(1).unwrap()?;
        assert_eq!(
            describe_indexedlog_entry(&bytes),
            r#"# f0: MAGIC_REWRITE_LAST_FLAT
# 00 00 00 00 00 00 00 00 05: Previous index Level = 0, Head = 5
# 01: Flags = HAS_ROOT
# 00: Level = 0
# 00 00 00 00 00 00 00 0a: High = 10
# 0a: Delta = 10 (Low = 0)
# 00: Parent count = 0
"#
        );
        let seg = iddag.find_flat_segment_including_id(Id(10))?.unwrap();
        iddag.remove_flat_segment(&seg)?;
        let bytes = iddag.log.iter_dirty().nth(2).unwrap()?;
        assert_eq!(
            describe_indexedlog_entry(bytes),
            r#"# f1: MAGIC_REMOVE_SEGMENT
# 00: Max Level = 0
# 01: Flags = HAS_ROOT
# 00: Level = 0
# 00 00 00 00 00 00 00 0a: High = 10
# 0a: Delta = 10 (Low = 0)
# 00: Parent count = 0
"#
        );

        Ok(())
    }

    fn dbg_iter<'a, T: std::fmt::Debug>(iter: Box<dyn Iterator<Item = Result<T>> + 'a>) -> String {
        let v = iter.map(|s| s.unwrap()).collect::<Vec<_>>();
        dbg(v)
    }

    fn dbg<T: std::fmt::Debug>(t: T) -> String {
        format!("{:?}", t)
    }
}


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

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::fmt;
use std::iter;
use std::result::Result as StdResult;

use serde::de::Error;
use serde::de::SeqAccess;
use serde::de::Visitor;
use serde::ser::SerializeSeq;
use serde::Deserialize;
use serde::Deserializer;
use serde::Serialize;
use serde::Serializer;

use super::IdDagStore;
use crate::errors::bug;
use crate::id::Group;
use crate::id::Id;
use crate::ops::Persist;
use crate::segment::Segment;
use crate::spanset::Span;
use crate::IdSet;
use crate::Level;
use crate::Result;

#[derive(Clone)]
pub struct InProcessStore {
    master_segments: Vec<Segment>,
    non_master_segments: Vec<Segment>,
    // level -> head -> serialized Segment
    level_head_index: Vec<BTreeMap<Id, StoreId>>,
    // (child-group, parent) -> child (in a flat segment).
    parent_index: BTreeMap<(Group, Id), BTreeSet<Id>>,
    // IdSet covered by flat segments in specified groups.
    id_set_by_group: [IdSet; Group::COUNT],
    // Removed StoreIds.
    // Written by `remove_flat_segment_unchecked`. Read by `serialize`.
    removed_store_ids: BTreeSet<StoreId>,
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(Serialize, Deserialize)]
enum StoreId {
    Master(usize),
    NonMaster(usize),
}

impl IdDagStore for InProcessStore {
    fn max_level(&self) -> Result<Level> {
        Ok((self.level_head_index.len().max(1) - 1) as Level)
    }

    fn find_segment_by_head_and_level(&self, head: Id, level: Level) -> Result<Option<Segment>> {
        let answer = self
            .get_head_index(level)
            .and_then(|head_index| head_index.get(&head))
            .map(|store_id| self.get_segment(store_id));
        Ok(answer)
    }

    fn find_flat_segment_including_id(&self, id: Id) -> Result<Option<Segment>> {
        let level = 0;
        let answer = self
            .get_head_index(level)
            .and_then(|head_index| head_index.range(id..).next())
            .map(|(_, store_id)| self.get_segment(store_id));
        if let Some(ref seg) = &answer {
            if seg.span()?.low > id {
                return Ok(None);
            }
        }
        Ok(answer)
    }

    fn insert_segment(&mut self, segment: Segment) -> Result<()> {
        let span = segment.span()?;
        let high = span.high;
        let level = segment.level()?;
        let parents = segment.parents()?;
        let group = high.group();

        // Can we merge the segment with the nearby flat segment?
        if let Some(merged) = self.maybe_merged_flat_segment(&segment)? {
            let (&last_high, &last_store_id) = match self
                .get_head_index(0)
                // This does a duplicated lookup that was done in maybe_merged_flat_segment.
                // But the overhead is probably okay, and the code is not hard to read.
                // The motivation was to share logic for stores (ex. `maybe_merged_flat_segment`).
                .and_then(|index| index.range(..span.low).rev().next())
            {
                Some((high, store_id)) => {
                    if cfg!(debug_assertions) {
                        let seg = self.get_segment(store_id);
                        debug_assert_eq!(seg.span()?.low, merged.span()?.low);
                        debug_assert_eq!(seg.parents()?, merged.parents()?);
                    }
                    (high, store_id)
                }
                None => return bug("last segment should exist if segments are mergeable"),
            };

            // Sanity check. This should pass since `last_store_id` should not be
            // obtained from indexes if it's already removed.
            if self.removed_store_ids.contains(&last_store_id) {
                return bug("insert_segment: reused removed store_id for segment merging");
            }

            // Store the merged segment.
            self.set_segment(&last_store_id, merged);

            // Update the "head" index.
            let index = self.get_head_index_mut(level);
            index.remove(&last_high);
            index.insert(high, last_store_id);

            // No need to update "parents" index.

            // Update "covered" IdSet.
            self.id_set_by_group[group.0].push(span);

            return Ok(());
        }

        let store_id = match high.group() {
            Group::MASTER => {
                self.master_segments.push(segment);
                StoreId::Master(self.master_segments.len() - 1)
            }
            _ => {
                self.non_master_segments.push(segment);
                StoreId::NonMaster(self.non_master_segments.len() - 1)
            }
        };
        if level == 0 {
            for parent in parents {
                let children = self
                    .parent_index
                    .entry((group, parent))
                    .or_insert_with(BTreeSet::new);
                children.insert(span.low);
            }
            // Update "covered" IdSet.
            self.id_set_by_group[group.0].push(span);
        }
        self.get_head_index_mut(level).insert(high, store_id);

        // Sanity check. This should pass because `store_id` is auto incremental.
        if self.removed_store_ids.contains(&store_id) {
            return bug("insert_segment: reused removed store_id");
        }

        Ok(())
    }

    fn remove_flat_segment_unchecked(&mut self, segment: &Segment) -> Result<()> {
        let span = segment.span()?;
        for level in 0..=self.max_level()? {
            // Remove from "level_head_index".
            let index = match self.level_head_index.get_mut(level as usize) {
                Some(index) => index,
                None => continue,
            };
            let store_id = match index.remove(&span.high) {
                Some(id) => id,
                None => continue,
            };
            // Remove from "parent_index".
            if level == 0 {
                let parents = segment.parents()?;
                let child = span.low;
                let child_group = child.group();
                for parent in parents {
                    let index_key = (child_group, parent);
                    let index = match self.parent_index.get_mut(&index_key) {
                        Some(index) => index,
                        None => continue,
                    };
                    index.remove(&child);
                }
            }
            // Mark as removed.
            self.removed_store_ids.insert(store_id);
        }
        // Update "id_set_by_group".
        //
        // PERF: This could be O(N^M) (N: len(id_set.spans), M:
        // len(segments) to remove) when removing many segments.
        // Practically N might be small so it might not be an issue.
        // In case this becomes an issue, we should make `id_set_by_group`
        // update lazy (i.e. track what id_set to remove, and update
        // `id_set_by_group` only when `id_set_by_group` needs to
        // be accessed.
        let id_set = &mut self.id_set_by_group[span.low.group().0 as usize];
        *id_set = id_set.difference(&span.into());
        // Not updating self.non_master_segments and master_segments
        // to keep existing StoreIds valid. Removed entries will be
        // skipped during serialization.
        Ok(())
    }

    fn remove_non_master(&mut self) -> Result<()> {
        for segment in self.non_master_segments.iter() {
            let level = segment.level()?;
            let head = segment.head()?;
            self.level_head_index
                .get_mut(level as usize)
                .map(|head_index| head_index.remove(&head));
        }
        let group = Group::NON_MASTER;
        for (_key, children) in self
            .parent_index
            .range_mut((group, group.min_id())..=(group, group.max_id()))
        {
            children.clear();
        }
        self.non_master_segments = Vec::new();
        self.id_set_by_group[Group::NON_MASTER.0] = IdSet::empty();
        Ok(())
    }

    fn all_ids_in_groups(&self, groups: &[Group]) -> Result<IdSet> {
        let mut result = IdSet::empty();
        for group in groups {
            result = result.union(&self.id_set_by_group[group.0]);
        }
        Ok(result)
    }

    fn next_segments(&self, id: Id, level: Level) -> Result<Vec<Segment>> {
        match self.get_head_index(level) {
            None => Ok(vec![]),
            Some(head_index) => {
                let segments = head_index
                    .range(id..id.group().max_id())
                    .map(|(_, store_id)| self.get_segment(store_id))
                    .collect();
                Ok(segments)
            }
        }
    }

    fn iter_segments_descending<'a>(
        &'a self,
        max_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>> {
        match self.get_head_index(level) {
            None => Ok(Box::new(iter::empty())),
            Some(head_index) => {
                let iter = head_index
                    .range(Id::MIN..=max_high_id)
                    .rev()
                    .map(move |(_, store_id)| Ok(self.get_segment(store_id)));
                Ok(Box::new(iter))
            }
        }
    }

    fn iter_segments_ascending<'a>(
        &'a self,
        min_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a + Send + Sync>> {
        match self.get_head_index(level) {
            None => Ok(Box::new(iter::empty())),
            Some(head_index) => {
                let iter = head_index
                    .range(min_high_id..=Id::MAX)
                    .map(move |(_, store_id)| Ok(self.get_segment(store_id)));
                Ok(Box::new(iter))
            }
        }
    }

    fn iter_flat_segments_with_parent_span<'a>(
        &'a self,
        parent_span: Span,
    ) -> Result<Box<dyn Iterator<Item = Result<(Id, Segment)>> + 'a>> {
        let mut iter: Box<dyn Iterator<Item = Result<(Id, Segment)>> + 'a> =
            Box::new(std::iter::empty());
        for group in Group::ALL {
            let range = (group, parent_span.low)..=(group, parent_span.high);
            let group_iter =
                self.parent_index
                    .range(range)
                    .flat_map(move |((_group, parent_id), child_ids)| {
                        let parent_id = *parent_id;
                        child_ids.iter().filter_map(move |&id| {
                            let seg = match self.find_flat_segment_including_id(id) {
                                Ok(Some(s)) => s,
                                Err(e) => return Some(Err(e)),
                                Ok(None) => return None,
                            };
                            Some(Ok((parent_id, seg)))
                        })
                    });
            iter = Box::new(iter.chain(group_iter));
        }
        Ok(Box::new(iter))
    }

    fn iter_flat_segments_with_parent<'a>(
        &'a self,
        parent: Id,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>> {
        let get_iter = |group: Group| -> Result<Box<dyn Iterator<Item = Result<_>> + 'a>> {
            match self.parent_index.get(&(group, parent)) {
                None => Ok(Box::new(iter::empty())),
                Some(children) => {
                    let iter = children.iter().filter_map(move |&id| {
                        let seg = match self.find_flat_segment_including_id(id) {
                            Ok(Some(s)) => s,
                            Err(e) => return Some(Err(e)),
                            Ok(None) => return None,
                        };
                        Some(Ok(seg))
                    });
                    Ok(Box::new(iter))
                }
            }
        };
        let iter = get_iter(Group::MASTER)?.chain(get_iter(Group::NON_MASTER)?);
        Ok(Box::new(iter))
    }
}

impl Persist for InProcessStore {
    type Lock = ();

    fn lock(&mut self) -> Result<()> {
        Ok(())
    }

    fn reload(&mut self, _lock: &Self::Lock) -> Result<()> {
        Ok(())
    }

    fn persist(&mut self, _lock: &Self::Lock) -> Result<()> {
        Ok(())
    }
}

impl InProcessStore {
    fn get_head_index(&self, level: Level) -> Option<&BTreeMap<Id, StoreId>> {
        self.level_head_index.get(level as usize)
    }

    fn get_head_index_mut(&mut self, level: Level) -> &mut BTreeMap<Id, StoreId> {
        if self.level_head_index.len() <= level as usize {
            self.level_head_index
                .resize(level as usize + 1, BTreeMap::new());
        }
        &mut self.level_head_index[level as usize]
    }

    fn get_segment(&self, store_id: &StoreId) -> Segment {
        match store_id {
            &StoreId::Master(offset) => self.master_segments[offset].clone(),
            &StoreId::NonMaster(offset) => self.non_master_segments[offset].clone(),
        }
    }

    fn set_segment(&mut self, store_id: &StoreId, segment: Segment) {
        match store_id {
            &StoreId::Master(offset) => self.master_segments[offset] = segment,
            &StoreId::NonMaster(offset) => self.non_master_segments[offset] = segment,
        }
    }
}

impl InProcessStore {
    pub fn new() -> Self {
        InProcessStore {
            master_segments: Vec::new(),
            non_master_segments: Vec::new(),
            level_head_index: Vec::new(),
            parent_index: BTreeMap::new(),
            id_set_by_group: [IdSet::empty(), IdSet::empty()],
            removed_store_ids: Default::default(),
        }
    }
}

impl Serialize for InProcessStore {
    fn serialize<S>(&self, serializer: S) -> StdResult<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut seq = serializer.serialize_seq(Some(
            self.master_segments.len() + self.non_master_segments.len()
                - self.removed_store_ids.len(),
        ))?;
        for (i, e) in self.master_segments.iter().enumerate() {
            if !self.removed_store_ids.contains(&StoreId::Master(i)) {
                seq.serialize_element(e)?;
            }
        }
        for (i, e) in self.non_master_segments.iter().enumerate() {
            if !self.removed_store_ids.contains(&StoreId::NonMaster(i)) {
                seq.serialize_element(e)?;
            }
        }
        seq.end()
    }
}

impl<'de> Deserialize<'de> for InProcessStore {
    fn deserialize<D>(deserializer: D) -> StdResult<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct InProcessStoreVisitor;
        impl<'de> Visitor<'de> for InProcessStoreVisitor {
            type Value = InProcessStore;
            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a list of segments")
            }
            fn visit_seq<A>(self, mut access: A) -> StdResult<Self::Value, A::Error>
            where
                A: SeqAccess<'de>,
            {
                let mut store = InProcessStore::new();
                while let Some(segment) = access.next_element()? {
                    store.insert_segment(segment).map_err(|e| {
                        A::Error::custom(format!("failed to deserialize IdDagStore: {} ", e))
                    })?;
                }
                Ok(store)
            }
        }

        deserializer.deserialize_seq(InProcessStoreVisitor)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::iddagstore::tests::dump_store_state;
    use crate::iddagstore::tests::test_remove_segment;

    #[test]
    fn test_remove_segment_serialize() {
        // Test remove_flat_segment is still effective after serialize->deserialize.
        let mut store = InProcessStore::new();
        test_remove_segment(&mut store);
        assert!(!store.removed_store_ids.is_empty());

        let all = store.all_ids_in_groups(&Group::ALL).unwrap();
        let old_state = dump_store_state(&store, &all);

        // Check store state after serialize -> deserialize round-trip.
        let data = mincode::serialize(&store).unwrap();
        let store: InProcessStore = mincode::deserialize(&data).unwrap();
        assert!(store.removed_store_ids.is_empty());

        let new_state = dump_store_state(&store, &all);
        assert_eq!(old_state, new_state);
    }
}


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

use std::fmt;
use std::fs;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::path::PathBuf;

use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use fs2::FileExt;
use indexedlog::log;
use vlqencoding::VLQDecode;
use vlqencoding::VLQEncode;

use super::IdMapWrite;
use crate::errors::bug;
use crate::errors::programming;
use crate::errors::NotFoundError;
use crate::id::Group;
use crate::id::Id;
use crate::id::VertexName;
use crate::ops::IdConvert;
use crate::ops::Persist;
use crate::ops::PrefixLookup;
use crate::ops::TryClone;
use crate::Result;
use crate::VerLink;

/// Bi-directional mapping between an integer id and a name (`[u8]`).
///
/// Backed by the filesystem.
pub struct IdMap {
    pub(crate) log: log::Log,
    path: PathBuf,
    map_id: String,
    map_version: VerLink,
}

impl IdMap {
    // Format:
    //
    // - Insertion:
    //   id (8 bytes, BE) + group (1 byte) + name (n bytes)
    // - Deletion:
    //   u64::MAX (8 bytes, BE) + n (VLQ) + [id (VLQ) + len(name) (VLQ) + name ] * n
    // - Clear non-master (only id->name mappings, being deprecated):
    //   CLRNM

    const INDEX_ID_TO_NAME: usize = 0;
    const INDEX_GROUP_NAME_TO_ID: usize = 1;

    /// Magic bytes in `Log` that indicates "remove all non-master id->name
    /// mappings". A valid entry has at least 8 bytes so does not conflict
    /// with this.
    const MAGIC_CLEAR_NON_MASTER: &'static [u8] = b"CLRNM";

    /// Magic prefix for deletion. It's u64::MAX id, which does not conflict
    /// with a valid id because it's > `Id::MAX`.
    const MAGIC_DELETION_PREFIX: &'static [u8] = &u64::MAX.to_be_bytes();

    /// Start offset in an entry for "name".
    const NAME_OFFSET: usize = 8 + Group::BYTES;

    /// Create an [`IdMap`] backed by the given directory.
    ///
    /// By default, only read-only operations are allowed. For writing
    /// access, call [`IdMap::make_writable`] to get a writable instance.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref();
        let log = Self::log_open_options().open(path)?;
        Self::open_from_log(log)
    }
}

impl TryClone for IdMap {
    fn try_clone(&self) -> Result<Self> {
        let result = Self {
            log: self.log.try_clone()?,
            path: self.path.clone(),
            map_id: self.map_id.clone(),
            map_version: self.map_version.clone(),
        };
        Ok(result)
    }
}

impl IdMap {
    pub(crate) fn open_from_log(log: log::Log) -> Result<Self> {
        let path = log.path().as_opt_path().unwrap().to_path_buf();
        let map_id = format!("ilog:{}", path.display());
        Ok(Self {
            log,
            path,
            map_id,
            map_version: VerLink::new(),
        })
    }

    pub(crate) fn log_open_options() -> log::OpenOptions {
        assert!(Self::MAGIC_DELETION_PREFIX > &Id::MAX.0.to_be_bytes()[..]);
        log::OpenOptions::new()
            .create(true)
            .index("id", |data| {
                assert!(Self::MAGIC_CLEAR_NON_MASTER.len() < 8);
                assert!(Group::BITS == 8);
                if data.starts_with(Self::MAGIC_DELETION_PREFIX) {
                    let items =
                        decode_deletion_entry(data).expect("deletion entry should be valid");
                    items
                        .into_iter()
                        .map(|(id, _name)| log::IndexOutput::Remove(id.0.to_be_bytes().into()))
                        .collect()
                } else if data.len() < 8 {
                    if data == Self::MAGIC_CLEAR_NON_MASTER {
                        vec![log::IndexOutput::RemovePrefix(Box::new([
                            Group::NON_MASTER.0 as u8,
                        ]))]
                    } else {
                        panic!("bug: invalid segment {:?}", &data);
                    }
                } else {
                    vec![log::IndexOutput::Reference(0..8)]
                }
            })
            .index("group-name", |data| {
                if data.starts_with(Self::MAGIC_DELETION_PREFIX) {
                    let items =
                        decode_deletion_entry(data).expect("deletion entry should be valid");
                    items
                        .into_iter()
                        .map(|(id, name)| {
                            let mut key = Vec::with_capacity(name.len() + 1);
                            key.extend_from_slice(&id.group().bytes());
                            key.extend_from_slice(name);
                            log::IndexOutput::Remove(key.into())
                        })
                        .collect()
                } else if data.len() >= 8 {
                    vec![log::IndexOutput::Reference(8..(data.len() as u64))]
                } else {
                    if data == Self::MAGIC_CLEAR_NON_MASTER {
                        vec![log::IndexOutput::RemovePrefix(Box::new([
                            Group::NON_MASTER.0 as u8,
                        ]))]
                    } else {
                        panic!("bug: invalid segment {:?}", &data);
                    }
                }
            })
            .flush_filter(Some(|_, _| {
                panic!("programming error: idmap changed by other process")
            }))
    }

    /// Find name by a specified integer id.
    pub fn find_name_by_id(&self, id: Id) -> Result<Option<&[u8]>> {
        let key = id.0.to_be_bytes();
        let key = self.log.lookup(Self::INDEX_ID_TO_NAME, &key)?.nth(0);
        match key {
            Some(Ok(entry)) => {
                if entry.len() < 8 {
                    return bug("index key should have 8 bytes at least");
                }
                Ok(Some(&entry[Self::NAME_OFFSET..]))
            }
            None => Ok(None),
            Some(Err(err)) => Err(err.into()),
        }
    }

    /// Find VertexName by a specified integer id.
    pub fn find_vertex_name_by_id(&self, id: Id) -> Result<Option<VertexName>> {
        self.find_name_by_id(id)
            .map(|v| v.map(|n| VertexName(self.log.slice_to_bytes(n))))
    }

    /// Find the integer id matching the given name.
    pub fn find_id_by_name(&self, name: &[u8]) -> Result<Option<Id>> {
        for group in Group::ALL.iter() {
            let mut group_name = Vec::with_capacity(Group::BYTES + name.len());
            group_name.extend_from_slice(&group.bytes());
            group_name.extend_from_slice(name);
            let key = self
                .log
                .lookup(Self::INDEX_GROUP_NAME_TO_ID, group_name)?
                .nth(0);
            match key {
                Some(Ok(mut entry)) => {
                    if entry.len() < 8 {
                        return bug("index key should have 8 bytes at least");
                    }
                    let id = Id(entry.read_u64::<BigEndian>().unwrap());
                    return Ok(Some(id));
                }
                None => {}
                Some(Err(err)) => return Err(err.into()),
            }
        }
        Ok(None)
    }

    /// Similar to `find_name_by_id`, but returns None if group > `max_group`.
    pub fn find_id_by_name_with_max_group(
        &self,
        name: &[u8],
        max_group: Group,
    ) -> Result<Option<Id>> {
        Ok(self.find_id_by_name(name)?.and_then(|id| {
            if id.group() <= max_group {
                Some(id)
            } else {
                None
            }
        }))
    }

    /// Insert a new entry mapping from a name to an id.
    ///
    /// Errors if the new entry conflicts with existing entries.
    pub fn insert(&mut self, id: Id, name: &[u8]) -> Result<()> {
        let existing_name = self.find_name_by_id(id)?;
        if let Some(existing_name) = existing_name {
            if existing_name == name {
                return Ok(());
            } else {
                return bug(format!(
                    "new entry {} = {:?} conflicts with an existing entry {} = {:?}",
                    id, name, id, existing_name
                ));
            }
        }
        let existing_id = self.find_id_by_name(name)?;
        if let Some(existing_id) = existing_id {
            // Allow re-assigning Ids from a higher group to a lower group.
            // For example, when a non-master commit gets merged into the
            // master branch, the id is re-assigned to master. But, the
            // ids in the master group will never be re-assigned to
            // non-master groups.
            if existing_id == id {
                return Ok(());
            } else {
                return bug(format!(
                    "new entry {} = {:?} conflicts with an existing entry {} = {:?}",
                    id, name, existing_id, name
                ));
            }
        }

        let mut data = Vec::with_capacity(8 + Group::BYTES + name.len());
        data.extend_from_slice(&id.0.to_be_bytes());
        data.extend_from_slice(&id.group().bytes());
        data.extend_from_slice(&name);
        self.log.append(data)?;
        self.map_version.bump();
        #[cfg(debug_assertions)]
        {
            let items = self.find_range(id, id).unwrap();
            assert_eq!(items[0], (id, name));
        }
        Ok(())
    }

    /// Find all (id, name) pairs in the `low..=high` range.
    fn find_range(&self, low: Id, high: Id) -> Result<Vec<(Id, &[u8])>> {
        let low = low.0.to_be_bytes();
        let high = high.0.to_be_bytes();
        let range = &low[..]..=&high[..];
        let mut items = Vec::new();
        for entry in self.log.lookup_range(Self::INDEX_ID_TO_NAME, range)? {
            let (key, values) = entry?;
            let key: [u8; 8] = match key.as_ref().try_into() {
                Ok(key) => key,
                Err(_) => {
                    return bug("find_range got non-u64 keys in INDEX_ID_TO_NAME");
                }
            };
            let id = Id(u64::from_be_bytes(key));
            for value in values {
                let value = value?;
                if value.len() < 8 {
                    return bug(format!(
                        "find_range got entry {:?} shorter than expected",
                        &value
                    ));
                }
                let name: &[u8] = &value[9..];
                items.push((id, name));
            }
        }
        Ok(items)
    }

    fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
        // Step 1: Find (id, name) pairs in the range.
        let items = self.find_range(low, high)?;
        let names = items
            .iter()
            .map(|(_, bytes)| VertexName::copy_from(bytes))
            .collect();
        // Step 2: Write a "delete" entry to delete those indexes.
        // The indexedlog index function (defined by log_open_options())
        // will handle it.
        let data = encode_deletion_entry(&items);
        self.log.append(data)?;
        // New map is not an "append-only" version of the previous map.
        // Re-create the VerLink to mark it as incompatible.
        self.map_version = VerLink::new();
        Ok(names)
    }

    /// Lookup names by hex prefix.
    fn find_names_by_hex_prefix(&self, hex_prefix: &[u8], limit: usize) -> Result<Vec<VertexName>> {
        let mut result = Vec::with_capacity(limit);
        for group in Group::ALL.iter().rev() {
            let mut prefix = Vec::with_capacity(Group::BYTES * 2 + hex_prefix.len());
            prefix.extend_from_slice(&group.hex_bytes());
            prefix.extend_from_slice(hex_prefix);
            for entry in self
                .log
                .lookup_prefix_hex(Self::INDEX_GROUP_NAME_TO_ID, prefix)?
            {
                let (k, _v) = entry?;
                let vertex = VertexName(self.log.slice_to_bytes(&k[Group::BYTES..]));
                if !result.contains(&vertex) {
                    result.push(vertex);
                }
                if result.len() >= limit {
                    return Ok(result);
                }
            }
        }
        Ok(result)
    }
}

/// Encode a list of (id, name) pairs as an deletion entry.
/// The deletion entry will be consumed by the index functions defined by
/// `log_open_options()`.
fn encode_deletion_entry(items: &[(Id, &[u8])]) -> Vec<u8> {
    // Rough size for common 20-byte sha1 hashes.
    let len = IdMap::MAGIC_DELETION_PREFIX.len() + 9 + items.len() * 30;
    let mut data = Vec::with_capacity(len);
    data.extend_from_slice(IdMap::MAGIC_DELETION_PREFIX);
    data.write_vlq(items.len()).unwrap();
    for (id, name) in items {
        data.write_vlq(id.0).unwrap();
        data.write_vlq(name.len()).unwrap();
        data.extend_from_slice(name);
    }
    data
}

/// Decode `encode_deletion_entry` result.
/// Used by index functions in `log_open_options()`.
fn decode_deletion_entry(data: &[u8]) -> Result<Vec<(Id, &[u8])>> {
    assert!(data.starts_with(IdMap::MAGIC_DELETION_PREFIX));
    let mut data = &data[IdMap::MAGIC_DELETION_PREFIX.len()..];
    let n = data.read_vlq()?;
    let mut items = Vec::with_capacity(n);
    for _ in 0..n {
        let id: u64 = data.read_vlq()?;
        let id = Id(id);
        let name_len: usize = data.read_vlq()?;
        if name_len > data.len() {
            return bug("decode_deletion_id_names got incomplete input");
        }
        let (name, rest) = data.split_at(name_len);
        data = rest;
        items.push((id, name));
    }
    Ok(items)
}

impl fmt::Debug for IdMap {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "IdMap {{\n")?;
        for data in self.log.iter() {
            if let Ok(mut data) = data {
                let id = data.read_u64::<BigEndian>().unwrap();
                let _group = data.read_u8().unwrap();
                let mut name = Vec::with_capacity(20);
                data.read_to_end(&mut name).unwrap();
                let name = if name.len() >= 20 {
                    VertexName::from(name).to_hex()
                } else {
                    String::from_utf8_lossy(&name).to_string()
                };
                let id = Id(id);
                write!(f, "  {}: {},\n", name, id)?;
            }
        }
        write!(f, "}}\n")?;
        Ok(())
    }
}

#[async_trait::async_trait]
impl IdConvert for IdMap {
    async fn vertex_id(&self, name: VertexName) -> Result<Id> {
        self.find_id_by_name(name.as_ref())?
            .ok_or_else(|| name.not_found_error())
    }
    async fn vertex_id_with_max_group(
        &self,
        name: &VertexName,
        max_group: Group,
    ) -> Result<Option<Id>> {
        self.find_id_by_name_with_max_group(name.as_ref(), max_group)
    }
    async fn vertex_name(&self, id: Id) -> Result<VertexName> {
        self.find_vertex_name_by_id(id)?
            .ok_or_else(|| id.not_found_error())
    }
    async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool> {
        Ok(self.find_id_by_name(name.as_ref())?.is_some())
    }
    async fn contains_vertex_id_locally(&self, ids: &[Id]) -> Result<Vec<bool>> {
        let mut list = Vec::with_capacity(ids.len());
        for &id in ids {
            list.push(self.find_name_by_id(id)?.is_some());
        }
        Ok(list)
    }
    async fn contains_vertex_name_locally(&self, names: &[VertexName]) -> Result<Vec<bool>> {
        let mut list = Vec::with_capacity(names.len());
        for name in names {
            let contains = self.find_id_by_name(name.as_ref())?.is_some();
            tracing::trace!("contains_vertex_name_locally({:?}) = {}", name, contains);
            list.push(contains);
        }
        Ok(list)
    }
    fn map_id(&self) -> &str {
        &self.map_id
    }
    fn map_version(&self) -> &VerLink {
        &self.map_version
    }
}

#[async_trait::async_trait]
impl IdMapWrite for IdMap {
    async fn insert(&mut self, id: Id, name: &[u8]) -> Result<()> {
        IdMap::insert(self, id, name)
    }
    async fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
        IdMap::remove_range(self, low, high)
    }
}

impl Persist for IdMap {
    type Lock = File;

    fn lock(&mut self) -> Result<Self::Lock> {
        if self.log.iter_dirty().next().is_some() {
            return programming("lock() must be called without dirty in-memory entries");
        }
        let lock_file = {
            let mut path = self.path.clone();
            path.push("wlock");
            File::open(&path).or_else(|_| {
                fs::OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(&path)
            })?
        };
        lock_file.lock_exclusive()?;
        Ok(lock_file)
    }

    fn reload(&mut self, _lock: &Self::Lock) -> Result<()> {
        self.log.clear_dirty()?;
        self.log.sync()?;
        Ok(())
    }

    fn persist(&mut self, _lock: &Self::Lock) -> Result<()> {
        self.log.sync()?;
        Ok(())
    }
}

#[async_trait::async_trait]
impl PrefixLookup for IdMap {
    async fn vertexes_by_hex_prefix(
        &self,
        hex_prefix: &[u8],
        limit: usize,
    ) -> Result<Vec<VertexName>> {
        self.find_names_by_hex_prefix(hex_prefix, limit)
    }
}

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

    #[test]
    fn test_encode_decode_deletion_entry() {
        let items: &[(Id, &[u8])] = &[
            (Id(0), b"a"),
            (Id(1), b"bb"),
            (Id(10), b"ccc"),
            (Id(20), b"dd"),
        ];
        let data = encode_deletion_entry(items);
        let decoded = decode_deletion_entry(&data).unwrap();
        assert_eq!(&decoded, items);
    }
}