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

//! DAG and Id operations (mostly traits)

use std::sync::Arc;

use futures::StreamExt;
use futures::TryStreamExt;

use crate::clone::CloneData;
use crate::default_impl;
use crate::errors::NotFoundError;
use crate::id::Group;
use crate::id::Id;
use crate::id::VertexName;
pub use crate::iddag::IdDagAlgorithm;
use crate::namedag::MemNameDag;
use crate::nameset::id_lazy::IdLazySet;
use crate::nameset::id_static::IdStaticSet;
use crate::nameset::NameSet;
use crate::IdSet;
use crate::Result;
use crate::VerLink;
use crate::VertexListWithOptions;

/// DAG related read-only algorithms.
#[async_trait::async_trait]
pub trait DagAlgorithm: Send + Sync {
    /// Sort a `NameSet` topologically in descending order.
    ///
    /// The returned set should have `dag` and `id_map` hints set to associated
    /// with this dag or its previous compatible version. For example, if a
    /// `set` is sorted on another dag but not in this dag, it should be resorted
    /// using this dag.  If a `set` is empty and not associated to the current
    /// `dag` in its hints, the return value should be a different empty `set`
    /// that has the `dag` and `id_map` hints set to this dag.
    async fn sort(&self, set: &NameSet) -> Result<NameSet>;

    /// Re-create the graph so it looks better when rendered.
    async fn beautify(&self, main_branch: Option<NameSet>) -> Result<MemNameDag> {
        default_impl::beautify(self, main_branch).await
    }

    /// Extract a sub graph containing only specified vertexes.
    async fn subdag(&self, set: NameSet) -> Result<MemNameDag> {
        default_impl::subdag(self, set).await
    }

    /// Get ordered parent vertexes.
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>>;

    /// Returns a set that covers all vertexes tracked by this DAG.
    async fn all(&self) -> Result<NameSet>;

    /// Returns a set that covers all vertexes in the master group.
    async fn master_group(&self) -> Result<NameSet>;

    /// Calculates all ancestors reachable from any name from the given set.
    async fn ancestors(&self, set: NameSet) -> Result<NameSet>;

    /// Calculates parents of the given set.
    ///
    /// Note: Parent order is not preserved. Use [`NameDag::parent_names`]
    /// to preserve order.
    async fn parents(&self, set: NameSet) -> Result<NameSet> {
        default_impl::parents(self, set).await
    }

    /// Calculates the n-th first ancestor.
    async fn first_ancestor_nth(&self, name: VertexName, n: u64) -> Result<Option<VertexName>> {
        default_impl::first_ancestor_nth(self, name, n).await
    }

    /// Calculates ancestors but only follows the first parent.
    async fn first_ancestors(&self, set: NameSet) -> Result<NameSet> {
        default_impl::first_ancestors(self, set).await
    }

    /// Calculates heads of the given set.
    async fn heads(&self, set: NameSet) -> Result<NameSet> {
        default_impl::heads(self, set).await
    }

    /// Calculates children of the given set.
    async fn children(&self, set: NameSet) -> Result<NameSet>;

    /// Calculates roots of the given set.
    async fn roots(&self, set: NameSet) -> Result<NameSet> {
        default_impl::roots(self, set).await
    }

    /// Calculates merges of the selected set (vertexes with >=2 parents).
    async fn merges(&self, set: NameSet) -> Result<NameSet> {
        default_impl::merges(self, set).await
    }

    /// Calculates one "greatest common ancestor" of the given set.
    ///
    /// If there are no common ancestors, return None.
    /// If there are multiple greatest common ancestors, pick one arbitrarily.
    /// Use `gca_all` to get all of them.
    async fn gca_one(&self, set: NameSet) -> Result<Option<VertexName>> {
        default_impl::gca_one(self, set).await
    }

    /// Calculates all "greatest common ancestor"s of the given set.
    /// `gca_one` is faster if an arbitrary answer is ok.
    async fn gca_all(&self, set: NameSet) -> Result<NameSet> {
        default_impl::gca_all(self, set).await
    }

    /// Calculates all common ancestors of the given set.
    async fn common_ancestors(&self, set: NameSet) -> Result<NameSet> {
        default_impl::common_ancestors(self, set).await
    }

    /// Tests if `ancestor` is an ancestor of `descendant`.
    async fn is_ancestor(&self, ancestor: VertexName, descendant: VertexName) -> Result<bool> {
        default_impl::is_ancestor(self, ancestor, descendant).await
    }

    /// Calculates "heads" of the ancestors of the given set. That is,
    /// Find Y, which is the smallest subset of set X, where `ancestors(Y)` is
    /// `ancestors(X)`.
    ///
    /// This is faster than calculating `heads(ancestors(set))` in certain
    /// implementations like segmented changelog.
    ///
    /// This is different from `heads`. In case set contains X and Y, and Y is
    /// an ancestor of X, but not the immediate ancestor, `heads` will include
    /// Y while this function won't.
    async fn heads_ancestors(&self, set: NameSet) -> Result<NameSet> {
        default_impl::heads_ancestors(self, set).await
    }

    /// Calculates the "dag range" - vertexes reachable from both sides.
    async fn range(&self, roots: NameSet, heads: NameSet) -> Result<NameSet>;

    /// Calculates `ancestors(reachable) - ancestors(unreachable)`.
    async fn only(&self, reachable: NameSet, unreachable: NameSet) -> Result<NameSet> {
        default_impl::only(self, reachable, unreachable).await
    }

    /// Calculates `ancestors(reachable) - ancestors(unreachable)`, and
    /// `ancestors(unreachable)`.
    /// This might be faster in some implementations than calculating `only` and
    /// `ancestors` separately.
    async fn only_both(
        &self,
        reachable: NameSet,
        unreachable: NameSet,
    ) -> Result<(NameSet, NameSet)> {
        default_impl::only_both(self, reachable, unreachable).await
    }

    /// Calculates the descendants of the given set.
    async fn descendants(&self, set: NameSet) -> Result<NameSet>;

    /// Calculates `roots` that are reachable from `heads` without going
    /// through other `roots`. For example, given the following graph:
    ///
    /// ```plain,ignore
    ///   F
    ///   |\
    ///   C E
    ///   | |
    ///   B D
    ///   |/
    ///   A
    /// ```
    ///
    /// `reachable_roots(roots=[A, B, C], heads=[F])` returns `[A, C]`.
    /// `B` is not included because it cannot be reached without going
    /// through another root `C` from `F`. `A` is included because it
    /// can be reached via `F -> E -> D -> A` that does not go through
    /// other roots.
    ///
    /// The can be calculated as
    /// `roots & (heads | parents(only(heads, roots & ancestors(heads))))`.
    /// Actual implementation might have faster paths.
    ///
    /// The `roots & ancestors(heads)` portion filters out bogus roots for
    /// compatibility, if the callsite does not provide bogus roots, it
    /// could be simplified to just `roots`.
    async fn reachable_roots(&self, roots: NameSet, heads: NameSet) -> Result<NameSet> {
        default_impl::reachable_roots(self, roots, heads).await
    }

    /// Vertexes buffered in memory, not yet written to disk.
    async fn dirty(&self) -> Result<NameSet>;

    /// Returns true if the vertex names might need to be resolved remotely.
    fn is_vertex_lazy(&self) -> bool;

    /// Get a snapshot of the current graph that can operate separately.
    ///
    /// This makes it easier to fight with borrowck.
    fn dag_snapshot(&self) -> Result<Arc<dyn DagAlgorithm + Send + Sync>>;

    /// Get a snapshot of the `IdDag` that can operate separately.
    ///
    /// This is for advanced use-cases. For example, if callsite wants to
    /// do some graph calculation without network, and control how to
    /// batch the vertex name lookups precisely.
    fn id_dag_snapshot(&self) -> Result<Arc<dyn IdDagAlgorithm + Send + Sync>> {
        Err(crate::errors::BackendError::Generic(format!(
            "id_dag_snapshot() is not supported for {}",
            std::any::type_name::<Self>()
        ))
        .into())
    }

    /// Identity of the dag.
    fn dag_id(&self) -> &str;

    /// Version of the dag. Useful to figure out compatibility between two dags.
    fn dag_version(&self) -> &VerLink;
}

#[async_trait::async_trait]
pub trait Parents: Send + Sync {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>>;

    /// A hint of a sub-graph for inserting `heads`.
    ///
    /// This is used to reduce remote fetches in a lazy graph.
    /// The roots will be checked first, if a root is unknown locally then
    /// all its descendants will be considered unknown locally.
    ///
    /// The returned graph is only used to optimize network fetches in
    /// `assign_head`. It is not used to be actually inserted to the graph. So
    /// returning an empty or "incorrect" graph does not hurt correctness. But
    /// might hurt performance.
    async fn hint_subdag_for_insertion(&self, _heads: &[VertexName]) -> Result<MemNameDag>;
}

#[async_trait::async_trait]
impl Parents for Arc<dyn DagAlgorithm + Send + Sync> {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        DagAlgorithm::parent_names(self, name).await
    }

    async fn hint_subdag_for_insertion(&self, heads: &[VertexName]) -> Result<MemNameDag> {
        let scope = self.dirty().await?;
        default_impl::hint_subdag_for_insertion(self, &scope, heads).await
    }
}

#[async_trait::async_trait]
impl Parents for &(dyn DagAlgorithm + Send + Sync) {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        DagAlgorithm::parent_names(*self, name).await
    }

    async fn hint_subdag_for_insertion(&self, heads: &[VertexName]) -> Result<MemNameDag> {
        let scope = self.dirty().await?;
        default_impl::hint_subdag_for_insertion(self, &scope, heads).await
    }
}

#[async_trait::async_trait]
impl<'a> Parents for Box<dyn Fn(VertexName) -> Result<Vec<VertexName>> + Send + Sync + 'a> {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        (self)(name)
    }

    async fn hint_subdag_for_insertion(&self, _heads: &[VertexName]) -> Result<MemNameDag> {
        // No clear way to detect the "dirty" scope.
        Ok(MemNameDag::new())
    }
}

#[async_trait::async_trait]
impl Parents for std::collections::HashMap<VertexName, Vec<VertexName>> {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        match self.get(&name) {
            Some(v) => Ok(v.clone()),
            None => name.not_found(),
        }
    }

    async fn hint_subdag_for_insertion(&self, heads: &[VertexName]) -> Result<MemNameDag> {
        let mut keys: Vec<VertexName> = self.keys().cloned().collect();
        keys.sort_unstable();
        let scope = NameSet::from_static_names(keys);
        default_impl::hint_subdag_for_insertion(self, &scope, heads).await
    }
}

/// Add vertexes recursively to the DAG.
#[async_trait::async_trait]
pub trait DagAddHeads {
    /// Add vertexes and their ancestors to the DAG. This does not persistent
    /// changes to disk.
    async fn add_heads(
        &mut self,
        parents: &dyn Parents,
        heads: &VertexListWithOptions,
    ) -> Result<bool>;
}

/// Remove vertexes and their descendants from the DAG.
#[async_trait::async_trait]
pub trait DagStrip {
    /// Remove the given `set` and their descendants.
    ///
    /// This will reload the DAG from its source (ex. filesystem) and writes
    /// changes back with a lock so there are no other processes adding
    /// new descendants of the stripped set.
    ///
    /// After strip, the `self` graph might contain new vertexes because of
    /// the reload.
    async fn strip(&mut self, set: &NameSet) -> Result<()>;
}

/// Import a generated `CloneData` object into an empty DAG.
#[async_trait::async_trait]
pub trait DagImportCloneData {
    /// Updates the DAG using a `CloneData` object.
    async fn import_clone_data(&mut self, clone_data: CloneData<VertexName>) -> Result<()>;
}

/// Import a generated incremental `CloneData` object into an existing DAG.
/// Ids in the passed CloneData might not match ids in existing DAG.
#[async_trait::async_trait]
pub trait DagImportPullData {
    /// Updates the DAG using a `CloneData` object.
    ///
    /// Only import the given `heads`.
    async fn import_pull_data(
        &mut self,
        clone_data: CloneData<VertexName>,
        heads: &VertexListWithOptions,
    ) -> Result<()>;
}

#[async_trait::async_trait]
pub trait DagExportCloneData {
    /// Export `CloneData` for vertexes in the master group.
    async fn export_clone_data(&self) -> Result<CloneData<VertexName>>;
}

#[async_trait::async_trait]
pub trait DagExportPullData {
    /// Export `CloneData` for vertexes in the given set.
    /// The set is typcially calculated by `only(heads, common)`.
    async fn export_pull_data(&self, set: &NameSet) -> Result<CloneData<VertexName>>;
}

/// Persistent the DAG on disk.
#[async_trait::async_trait]
pub trait DagPersistent {
    /// Write in-memory DAG to disk. This might also pick up changes to
    /// the DAG by other processes.
    async fn flush(&mut self, master_heads: &VertexListWithOptions) -> Result<()>;

    /// Write in-memory IdMap that caches Id <-> Vertex translation from
    /// remote service to disk.
    async fn flush_cached_idmap(&self) -> Result<()>;

    /// A faster path for add_heads, followed by flush.
    async fn add_heads_and_flush(
        &mut self,
        parent_names_func: &dyn Parents,
        heads: &VertexListWithOptions,
    ) -> Result<()>;

    /// Import from another (potentially large) DAG. Write to disk immediately.
    async fn import_and_flush(
        &mut self,
        dag: &dyn DagAlgorithm,
        master_heads: NameSet,
    ) -> Result<()> {
        let heads = dag.heads(dag.all().await?).await?;
        let non_master_heads = heads - master_heads.clone();
        let master_heads: Vec<VertexName> =
            master_heads.iter().await?.try_collect::<Vec<_>>().await?;
        let non_master_heads: Vec<VertexName> = non_master_heads
            .iter()
            .await?
            .try_collect::<Vec<_>>()
            .await?;
        let heads = VertexListWithOptions::from(master_heads)
            .with_highest_group(Group::MASTER)
            .chain(non_master_heads);
        self.add_heads_and_flush(&dag.dag_snapshot()?, &heads).await
    }
}

/// Import ASCII graph to DAG.
pub trait ImportAscii {
    /// Import vertexes described in an ASCII graph.
    fn import_ascii(&mut self, text: &str) -> Result<()> {
        self.import_ascii_with_heads_and_vertex_fn(text, <Option<&[&str]>>::None, None)
    }

    /// Import vertexes described in an ASCII graph.
    /// `heads` optionally specifies the order of heads to insert.
    /// Useful for testing. Panic if the input is invalid.
    fn import_ascii_with_heads(
        &mut self,
        text: &str,
        heads: Option<&[impl AsRef<str>]>,
    ) -> Result<()> {
        self.import_ascii_with_heads_and_vertex_fn(text, heads, None)
    }

    /// Import vertexes described in an ASCII graph.
    /// `vertex_fn` specifies how to generate Vertex from a str from the ASCII graph.
    /// Useful for testing when we need to generate `HgId` (fixed length) from vertex.
    fn import_ascii_with_vertex_fn(
        &mut self,
        text: &str,
        vertex_fn: fn(&str) -> VertexName,
    ) -> Result<()> {
        self.import_ascii_with_heads_and_vertex_fn(text, <Option<&[&str]>>::None, Some(vertex_fn))
    }

    /// Import vertexes described in an ASCII graph.
    /// `heads` optionally specifies the order of heads to insert.
    /// `vertex_fn` specifies how to generate Vertex from a str from the ASCII graph.
    ///
    /// This method is a helper function of other APIs, choose other APIs if possible.
    fn import_ascii_with_heads_and_vertex_fn(
        &mut self,
        text: &str,
        heads: Option<&[impl AsRef<str>]>,
        vertex_fn: Option<fn(&str) -> VertexName>,
    ) -> Result<()>;
}

/// Lookup vertexes by prefixes.
#[async_trait::async_trait]
pub trait PrefixLookup {
    /// Lookup vertexes by hex prefix.
    async fn vertexes_by_hex_prefix(
        &self,
        hex_prefix: &[u8],
        limit: usize,
    ) -> Result<Vec<VertexName>>;
}

/// Convert between `Vertex` and `Id`.
#[async_trait::async_trait]
pub trait IdConvert: PrefixLookup + Sync {
    async fn vertex_id(&self, name: VertexName) -> Result<Id>;
    async fn vertex_id_with_max_group(
        &self,
        name: &VertexName,
        max_group: Group,
    ) -> Result<Option<Id>>;
    async fn vertex_name(&self, id: Id) -> Result<VertexName>;
    async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool>;

    /// Test if an `id` is present locally. Do not trigger remote fetching.
    async fn contains_vertex_id_locally(&self, id: &[Id]) -> Result<Vec<bool>>;

    /// Test if an `name` is present locally. Do not trigger remote fetching.
    async fn contains_vertex_name_locally(&self, name: &[VertexName]) -> Result<Vec<bool>>;

    async fn vertex_id_optional(&self, name: &VertexName) -> Result<Option<Id>> {
        self.vertex_id_with_max_group(name, Group::NON_MASTER).await
    }

    /// Convert [`Id`]s to [`VertexName`]s in batch.
    async fn vertex_name_batch(&self, ids: &[Id]) -> Result<Vec<Result<VertexName>>> {
        // This is not an efficient implementation in an async context.
        let mut names = Vec::with_capacity(ids.len());
        for &id in ids {
            names.push(self.vertex_name(id).await);
        }
        Ok(names)
    }

    /// Convert [`VertexName`]s to [`Id`]s in batch.
    async fn vertex_id_batch(&self, names: &[VertexName]) -> Result<Vec<Result<Id>>> {
        // This is not an efficient implementation in an async context.
        let mut ids = Vec::with_capacity(names.len());
        for name in names {
            ids.push(self.vertex_id(name.clone()).await);
        }
        Ok(ids)
    }

    /// Identity of the map.
    fn map_id(&self) -> &str;

    /// Version of the map. Useful to figure out compatibility between two maps.
    fn map_version(&self) -> &VerLink;
}

/// Integrity check functions.
#[async_trait::async_trait]
pub trait CheckIntegrity {
    /// Verify that universally known `Id`s (parents of merges) are actually
    /// known locally.
    ///
    /// Returns set of `Id`s that should be universally known but missing.
    /// An empty set means all universally known `Id`s are known locally.
    ///
    /// Check `FirstAncestorConstraint::KnownUniversally` for concepts of
    /// "universally known".
    async fn check_universal_ids(&self) -> Result<Vec<Id>>;

    /// Check segment properties: no cycles, no overlaps, no gaps etc.
    /// This is only about the `Id`s, not about the vertex names.
    ///
    /// Returns human readable messages about problems.
    /// No messages indicates there are no problems detected.
    async fn check_segments(&self) -> Result<Vec<String>>;

    /// Check that the subset of the current graph (ancestors of `heads`)
    /// is isomorphic with the subset in the `other` graph.
    ///
    /// Returns messages about where two graphs are different.
    /// No messages indicates two graphs are likely isomorphic.
    ///
    /// Note: For performance, this function only checks the "shape"
    /// of the graph, without checking the (potentially lazy) vertex
    /// names.
    async fn check_isomorphic_graph(
        &self,
        other: &dyn DagAlgorithm,
        heads: NameSet,
    ) -> Result<Vec<String>>;
}

impl<T> ImportAscii for T
where
    T: DagAddHeads,
{
    fn import_ascii_with_heads_and_vertex_fn(
        &mut self,
        text: &str,
        heads: Option<&[impl AsRef<str>]>,
        vertex_fn: Option<fn(&str) -> VertexName>,
    ) -> Result<()> {
        let vertex_fn = match vertex_fn {
            Some(vertex_fn) => vertex_fn,
            None => |s: &str| VertexName::copy_from(s.as_bytes()),
        };
        let parents = drawdag::parse(&text);
        let heads: Vec<_> = match heads {
            Some(heads) => heads.iter().map(|s| vertex_fn(s.as_ref())).collect(),
            None => {
                let mut heads: Vec<_> = parents.keys().map(|s| vertex_fn(s)).collect();
                heads.sort();
                heads
            }
        };

        let parents: std::collections::HashMap<VertexName, Vec<VertexName>> = parents
            .iter()
            .map(|(k, vs)| (vertex_fn(k), vs.iter().map(|v| vertex_fn(v)).collect()))
            .collect();
        nonblocking::non_blocking_result(self.add_heads(&parents, &heads[..].into()))?;
        Ok(())
    }
}

#[async_trait::async_trait]
pub trait ToIdSet {
    /// Converts [`NameSet`] to [`IdSet`].
    async fn to_id_set(&self, set: &NameSet) -> Result<IdSet>;
}

pub trait ToSet {
    /// Converts [`IdSet`] to [`NameSet`].
    fn to_set(&self, set: &IdSet) -> Result<NameSet>;
}

pub trait IdMapSnapshot {
    /// Get a snapshot of IdMap.
    fn id_map_snapshot(&self) -> Result<Arc<dyn IdConvert + Send + Sync>>;
}

/// Describes how to persist state to disk.
pub trait Persist {
    /// Return type of `lock()`.
    type Lock: Send + Sync;

    /// Obtain an exclusive lock for writing.
    /// This should prevent other writers.
    fn lock(&mut self) -> Result<Self::Lock>;

    /// Reload from the source of truth. Drop pending changes.
    ///
    /// This requires a lock and is usually called before `persist()`.
    fn reload(&mut self, _lock: &Self::Lock) -> Result<()>;

    /// Write pending changes to the source of truth.
    ///
    /// This requires a lock.
    fn persist(&mut self, _lock: &Self::Lock) -> Result<()>;
}

/// Address that can be used to open things.
///
/// The address type decides the return type of `open`.
pub trait Open: Clone {
    type OpenTarget;

    fn open(&self) -> Result<Self::OpenTarget>;
}

/// Has an integer tuple version that can be used to test if the data was
/// changed. If the first number changes, it means incompatible changes.
/// If only the second number increases, it means append-only changes.
pub trait IntVersion {
    fn int_version(&self) -> (u64, u64);
}

/// Fallible clone.
pub trait TryClone {
    fn try_clone(&self) -> Result<Self>
    where
        Self: Sized;
}

impl<T: Clone> TryClone for T {
    fn try_clone(&self) -> Result<Self> {
        Ok(self.clone())
    }
}

#[async_trait::async_trait]
impl<T: IdConvert + IdMapSnapshot> ToIdSet for T {
    /// Converts [`NameSet`] to [`IdSet`].
    async fn to_id_set(&self, set: &NameSet) -> Result<IdSet> {
        let version = set.hints().id_map_version();

        // Fast path: extract IdSet from IdStaticSet.
        if let Some(set) = set.as_any().downcast_ref::<IdStaticSet>() {
            if None < version && version <= Some(self.map_version()) {
                return Ok(set.spans.clone());
            }
        }

        // Convert IdLazySet to IdStaticSet. Bypass hash lookups.
        if let Some(set) = set.as_any().downcast_ref::<IdLazySet>() {
            if None < version && version <= Some(self.map_version()) {
                let set: IdStaticSet = set.to_static()?;
                return Ok(set.spans);
            }
        }

        // Slow path: iterate through the set and convert it to a non-lazy
        // IdSet. Does not bypass hash lookups.
        let mut spans = IdSet::empty();
        let mut iter = set.iter().await?.chunks(1 << 17);
        while let Some(names) = iter.next().await {
            let names = names.into_iter().collect::<Result<Vec<_>>>()?;
            let ids = self.vertex_id_batch(&names).await?;
            for id in ids {
                spans.push(id?);
            }
        }
        Ok(spans)
    }
}

impl IdMapSnapshot for Arc<dyn IdConvert + Send + Sync> {
    fn id_map_snapshot(&self) -> Result<Arc<dyn IdConvert + Send + Sync>> {
        Ok(self.clone())
    }
}

impl<T: IdMapSnapshot + DagAlgorithm> ToSet for T {
    /// Converts [`IdSet`] to [`NameSet`].
    fn to_set(&self, set: &IdSet) -> Result<NameSet> {
        NameSet::from_spans_dag(set.clone(), 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.
 */

//! # protocol
//!
//! Abstractions used for communication between `(sparse_idmap1, segments1)`
//! (usually, a client) and `(complete_idmap2, segments2)` (usually, a server).
//!
//! When the sparse idmap gets asked to convert unknown id or name, it goes
//! through the following flow to find the answer:
//!
//! - Id -> Name: Id -> RequestLocationToName -> ResponseIdNamePair -> Name
//! - Name -> Id: Name -> RequestNameToLocation -> ResponseIdNamePair -> Id

use std::cell::RefCell;
use std::fmt;
use std::thread_local;

use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;

use crate::id::VertexName;
use crate::iddag::FirstAncestorConstraint;
use crate::iddag::IdDag;
use crate::iddagstore::IdDagStore;
use crate::ops::IdConvert;
use crate::Group;
use crate::Id;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::IdMap;
use crate::IdSet;
use crate::Result;

// Request and Response structures -------------------------------------------

/// Request for locating names (commit hashes) in a IdDag.
/// Useful for converting names to ids.
#[derive(Debug, Clone)]
pub struct RequestNameToLocation {
    pub names: Vec<VertexName>,
    pub heads: Vec<VertexName>,
}

/// Request for converting locations to names (commit hashes).
/// Useful for converting ids to names.
#[derive(Debug, Clone)]
pub struct RequestLocationToName {
    pub paths: Vec<AncestorPath>,
}

/// Response for converting names to ids or converting names to ids.
#[derive(Debug, Clone)]
pub struct ResponseIdNamePair {
    // For converting Id -> Name, the client provides AncestorPath, the server provides
    // Vec<Box<[u8]>>.
    //
    // For converting Name -> Id, the client provides Box<[u8]>, the server provides
    // AncestorPath.
    pub path_names: Vec<(AncestorPath, Vec<VertexName>)>,
}

/// The `n`-th first ancestor of `x`. `x~n` in hg revset syntax.
/// Usually, `x` is commonly known by the client and the server.
///
/// This can be seen as a kind of "location".
#[derive(Clone)]
pub struct AncestorPath {
    pub x: VertexName,

    pub n: u64,

    // Starting from x~n, get a chain of commits following p1.
    pub batch_size: u64,
}

impl fmt::Display for AncestorPath {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}~{}", self.x, self.n)
    }
}

impl fmt::Debug for AncestorPath {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)?;
        if self.batch_size != 1 {
            write!(f, "(+{})", self.batch_size)?;
        }
        Ok(())
    }
}

// Async Remote Protocols ----------------------------------------------------

/// Abstraction of network protocols.
#[async_trait::async_trait]
pub trait RemoteIdConvertProtocol: Send + Sync + 'static {
    /// Ask the server to convert names to "x~n" relative paths.
    ///
    /// If a "name" cannot be resolved using "x~n" form in "::heads", aka. the
    /// "heads" are known to the server, and the server can calculate "::heads",
    /// and knows all names (commit hashes) in "::heads". And the server
    /// confirms "name" is outside "::heads" (either because "name" is unknown
    /// to the server's IdMap, or because "name" is known in the server's IdMap,
    /// but the matching Id is outside "::heads"), this method should skip it in
    /// the resulting list (instead of returning an error).
    async fn resolve_names_to_relative_paths(
        &self,
        heads: Vec<VertexName>,
        names: Vec<VertexName>,
    ) -> Result<Vec<(AncestorPath, Vec<VertexName>)>>;

    /// Ask the server to convert "x~n" relative paths back to commit hashes.
    ///
    /// Unlike resolve_names_to_relative_paths, failures are not expected.
    /// They usually indicate rare events like master moving backwards.
    async fn resolve_relative_paths_to_names(
        &self,
        paths: Vec<AncestorPath>,
    ) -> Result<Vec<(AncestorPath, Vec<VertexName>)>>;

    /// Return `true` if the protocol is local and queries do not need to
    /// optimize for batching or latency.
    fn is_local(&self) -> bool {
        false
    }
}

#[async_trait::async_trait]
impl RemoteIdConvertProtocol for () {
    async fn resolve_names_to_relative_paths(
        &self,
        _heads: Vec<VertexName>,
        _names: Vec<VertexName>,
    ) -> Result<Vec<(AncestorPath, Vec<VertexName>)>> {
        Ok(Default::default())
    }

    async fn resolve_relative_paths_to_names(
        &self,
        paths: Vec<AncestorPath>,
    ) -> Result<Vec<(AncestorPath, Vec<VertexName>)>> {
        let msg = format!(
            "Asked to resolve {:?} in graph but remote protocol is not configured",
            paths
        );
        crate::errors::programming(msg)
    }

    fn is_local(&self) -> bool {
        true
    }
}

// Traits --------------------------------------------------------------------

/// Similar to `From::from(I) -> O`, but with `self` as context.
///
/// Example use-cases:
/// - Convert a query to a request (client-side).
/// - Convert a request to a response (server-side).
/// - Handle a response from the server (client-side).
#[async_trait::async_trait]
pub(crate) trait Process<I, O> {
    async fn process(self, input: I) -> Result<O>;
}

// Basic implementation ------------------------------------------------------

// Name -> Id, step 1: Name -> RequestNameToLocation
// Works on an incomplete IdMap, client-side.
#[async_trait::async_trait]
impl<M: IdConvert, DagStore: IdDagStore> Process<Vec<VertexName>, RequestNameToLocation>
    for (&M, &IdDag<DagStore>)
{
    async fn process(self, names: Vec<VertexName>) -> Result<RequestNameToLocation> {
        let map = &self.0;
        let dag = &self.1;
        // Only provides heads in the master group, since it's expected that the
        // non-master group is already locally known.
        let heads = stream::iter(dag.heads_ancestors(dag.master_group()?)?.into_iter()).boxed();
        let heads = heads
            .then(|id| map.vertex_name(id))
            .try_collect::<Vec<VertexName>>()
            .await
            .map_err(|e| {
                let msg = format!(
                    concat!(
                        "Cannot resolve heads in master group to vertex name. ",
                        "The vertex name is required for remote vertex resolution. ",
                        "This probably indicates the Dag update logic does not ensure the ",
                        "vertex name of heads exist as it should. ",
                        "(Error: {})",
                    ),
                    e
                );
                crate::Error::Programming(msg)
            })?;
        Ok(RequestNameToLocation { names, heads })
    }
}

// Id -> Name, step 1: Id -> RequestLocationToName
// Works on an incomplete IdMap, client-side.
#[async_trait::async_trait]
impl<M: IdConvert, DagStore: IdDagStore> Process<IdSet, RequestLocationToName>
    for (&M, &IdDag<DagStore>)
{
    async fn process(self, ids: IdSet) -> Result<RequestLocationToName> {
        let map = &self.0;
        let dag = &self.1;
        let heads = dag.heads_ancestors(dag.master_group()?)?;

        let mut id_path: Vec<(Id, u64, u64)> = Vec::with_capacity(ids.as_spans().len());
        let mut last_id_opt = None;
        for id in ids.into_iter() {
            if let Some(last_id) = last_id_opt {
                if dag.try_first_ancestor_nth(last_id, 1)? == Some(id) {
                    // Reuse the last path.
                    if let Some(last) = id_path.last_mut() {
                        last.2 += 1;
                        last_id_opt = Some(id);
                        continue;
                    }
                }
            }
            let (x, n) = dag
                .to_first_ancestor_nth(
                    id,
                    FirstAncestorConstraint::KnownUniversally {
                        heads: heads.clone(),
                    },
                )?
                .ok_or_else(|| {
                    if id.group() == Group::MASTER {
                        let msg = format!(
                            concat!(
                                "Cannot convert {} to x~n form using heads {:?}. ",
                                "This is unexpected. It indicates some serious bugs in graph ",
                                "calculation or the graph is corrupted (ex. has cycles).",
                            ),
                            id, &heads,
                        );
                        crate::Error::Bug(msg)
                    } else {
                        let msg = format!(
                            concat!(
                                "Cannot convert {} to x~n form. This is unexpected for non-master ",
                                "vertexes since they are expected to be non-lazy.",
                            ),
                            id
                        );
                        crate::Error::Programming(msg)
                    }
                })?;
            id_path.push((x, n, 1));
            last_id_opt = Some(id);
        }

        let paths = stream::iter(id_path)
            .then(|(x, n, batch_size)| async move {
                let x = map.vertex_name(x).await.map_err(|e| {
                    let msg = format!(
                        concat!(
                            "Cannot resolve {} in to vertex name (Error: {}). ",
                            "The vertex name is required for remote vertex resolution. ",
                            "This probably indicates the Dag clone or update logic does ",
                            "not maintain \"universally known\" vertexes as it should.",
                        ),
                        x, e,
                    );
                    crate::Error::Programming(msg)
                })?;
                Ok::<_, crate::Error>(AncestorPath { x, n, batch_size })
            })
            .try_collect::<Vec<_>>()
            .await?;

        Ok(RequestLocationToName { paths })
    }
}

// Name -> Id, step 2: RequestNameToLocation -> ResponseIdNamePair
// Works on a complete IdMap, server-side.
#[async_trait::async_trait]
impl<M: IdConvert, DagStore: IdDagStore> Process<RequestNameToLocation, ResponseIdNamePair>
    for (&M, &IdDag<DagStore>)
{
    async fn process(self, request: RequestNameToLocation) -> Result<ResponseIdNamePair> {
        let map = &self.0;
        let dag = &self.1;

        let heads: IdSet = {
            let heads = stream::iter(request.heads.into_iter());
            let heads = heads
                .then(|s| map.vertex_id(s))
                .try_collect::<Vec<Id>>()
                .await?;
            IdSet::from_spans(heads)
        };
        let resolvable = dag.ancestors(heads.clone())?;

        let id_names: Vec<(Id, VertexName)> = {
            let ids_result = map.vertex_id_batch(&request.names).await?;
            let mut id_names = Vec::with_capacity(ids_result.len());
            for (name, id_result) in request.names.into_iter().zip(ids_result) {
                match id_result {
                    // If one of the names cannot be resolved to id, just skip it.
                    Err(crate::Error::VertexNotFound(n)) => {
                        tracing::trace!(
                            "RequestNameToLocation -> ResponseIdNamePair: skip unknown name {:?}",
                            &n
                        );
                        continue;
                    }
                    Err(e) => {
                        return Err(e);
                    }
                    Ok(id) => {
                        if resolvable.contains(id) {
                            id_names.push((id, name))
                        }
                    }
                }
            }
            id_names
        };

        let path_names: Vec<(AncestorPath, Vec<VertexName>)> = {
            let x_n_names: Vec<(Id, u64, VertexName)> = id_names
                .into_iter()
                .filter_map(|(id, name)| {
                    match dag.to_first_ancestor_nth(
                        id,
                        FirstAncestorConstraint::KnownUniversally {
                            heads: heads.clone(),
                        },
                    ) {
                        Err(e) => Some(Err(e)),
                        // Skip ids that cannot be translated.
                        Ok(None) => None,
                        Ok(Some((x, n))) => Some(Ok((x, n, name))),
                    }
                })
                .collect::<Result<Vec<_>>>()?;

            // Convert x from Id to VertexName.
            stream::iter(x_n_names)
                .then(|(x, n, name)| async move {
                    let x = map.vertex_name(x).await?;
                    Ok::<_, crate::Error>((
                        AncestorPath {
                            x,
                            n,
                            batch_size: 1,
                        },
                        vec![name],
                    ))
                })
                .try_collect()
                .await?
        };

        Ok(ResponseIdNamePair { path_names })
    }
}

// Id -> Name, step 2: RequestLocationToName -> ResponseIdNamePair
// Works on a complete IdMap, server-side.
#[async_trait::async_trait]
impl<M: IdConvert, DagStore: IdDagStore> Process<RequestLocationToName, ResponseIdNamePair>
    for (&M, &IdDag<DagStore>)
{
    async fn process(self, request: RequestLocationToName) -> Result<ResponseIdNamePair> {
        let map = &self.0;
        let dag = &self.1;

        let path_names: Vec<(AncestorPath, Vec<VertexName>)> =
            stream::iter(request.paths.into_iter())
                .then(|path| async move {
                    let id = map.vertex_id(path.x.clone()).await?;
                    let mut id = dag.first_ancestor_nth(id, path.n)?;
                    let mut ids = Vec::with_capacity(path.batch_size as _);
                    for i in 0..path.batch_size {
                        if i > 0 {
                            id = dag.first_ancestor_nth(id, 1)?;
                        }
                        ids.push(id);
                    }
                    let fallible_names = map.vertex_name_batch(&ids).await?;
                    let mut names = Vec::with_capacity(fallible_names.len());
                    for name in fallible_names {
                        names.push(name?);
                    }
                    debug_assert_eq!(path.batch_size, names.len() as u64);
                    Ok::<_, crate::Error>((path, names))
                })
                .try_collect()
                .await?;
        Ok(ResponseIdNamePair { path_names })
    }
}

// Name -> Id or Id -> Name, step 3: Apply RequestNameToLocation to a local IdMap.
// Works on an incomplete IdMap, client-side.
#[cfg(any(test, feature = "indexedlog-backend"))]
#[async_trait::async_trait]
impl<'a, DagStore: IdDagStore> Process<ResponseIdNamePair, ()>
    for (&'a mut IdMap, &'a IdDag<DagStore>)
{
    async fn process(mut self, res: ResponseIdNamePair) -> Result<()> {
        use crate::errors::NotFoundError;

        let map = &mut self.0;
        let dag = &self.1;
        for (path, names) in res.path_names.iter() {
            let x: Id = map
                .find_id_by_name(path.x.as_ref())?
                .ok_or_else(|| path.x.not_found_error())?;
            let mut id = dag.first_ancestor_nth(x, path.n)?;
            tracing::trace!("insert path {:?} names {:?} (x = {})", &path, &names, id);
            for (i, name) in names.iter().enumerate() {
                if i > 0 {
                    id = dag.first_ancestor_nth(id, 1)?;
                }
                tracing::trace!(" insert {:?} = {:?}", id, &name);
                map.insert(id, name.as_ref())?;
            }
        }
        Ok(())
    }
}

// Disable remote protocol temporarily ---------------------------------------
// This can be useful for Debug::fmt to disable remote fetching which might
// panic (ex. calling tokio without tokio runtime) when executing futures
// via nonblocking.

thread_local! {
    static NON_BLOCKING_DEPTH: RefCell<usize> = RefCell::new(0);
}

/// Check if the current future is running inside a "non-blocking" block.
pub(crate) fn disable_remote_protocol<F, R>(f: F) -> R
where
    F: FnOnce() -> R,
{
    NON_BLOCKING_DEPTH.with(|v| *v.borrow_mut() += 1);
    let result = f();
    NON_BLOCKING_DEPTH.with(|v| *v.borrow_mut() -= 1);
    result
}

pub(crate) fn is_remote_protocol_disabled() -> bool {
    NON_BLOCKING_DEPTH.with(|v| *v.borrow() != 0)
}


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

#[cfg(any(test, feature = "indexedlog-backend"))]
use std::cmp::Ordering;
#[cfg(any(test, feature = "indexedlog-backend"))]
use std::io::Write;

use anyhow::Result;
use nonblocking::non_blocking_result;
// Re-export
pub use renderdag::Ancestor;
pub use renderdag::GraphRow;
pub use renderdag::GraphRowRenderer;
pub use renderdag::Renderer;
pub use renderdag::*;

use crate::nameset::SyncNameSetQuery;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::ops::IdConvert;
use crate::DagAlgorithm;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::Group;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::IdSpan;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::Level;
#[cfg(any(test, feature = "indexedlog-backend"))]
use crate::NameDag;
use crate::Set;
use crate::VertexName;

/// Render a NameDag or MemNameDag into a String.
pub fn render_namedag(
    dag: &(impl DagAlgorithm + ?Sized),
    get_message: impl Fn(&VertexName) -> Option<String>,
) -> Result<String> {
    let mut renderer = GraphRowRenderer::new().output().build_box_drawing();

    let mut out = String::new();
    let next_rows = dag_to_renderer_next_rows(dag, None)?;
    for (node, parents) in next_rows {
        let mut name = format!("{:?}", &node);
        let message = get_message(&node).unwrap_or_default();
        let row = if name.len() == 1 {
            renderer.next_row(node, parents, name, message)
        } else {
            if !message.is_empty() {
                name += &format!(" {}", message);
            }
            renderer.next_row(node, parents, String::from("o"), name)
        };
        out.push_str(&row);
    }

    let output = format!(
        "\n{}",
        out.trim_end()
            .lines()
            .map(|l| format!("            {}", l))
            .collect::<Vec<_>>()
            .join("\n")
    );
    Ok(output)
}

/// Render a NameDag or MemNameDag into structured `GraphRow`s.
/// The `GraphRow` can serialize to other formats.
///
/// If `subset` is provided, only render a subset of the graph.
pub fn render_namedag_structured(
    dag: &dyn DagAlgorithm,
    subset: Option<Set>,
) -> Result<Vec<GraphRow<VertexName>>> {
    let mut renderer = GraphRowRenderer::new();
    let next_rows = dag_to_renderer_next_rows(dag, subset)?;
    let mut out = Vec::with_capacity(next_rows.len());
    for (node, parents) in next_rows {
        let name = String::from_utf8_lossy(node.as_ref()).into_owned();
        let message = Default::default();
        let row = renderer.next_row(node, parents, name, message);
        out.push(row);
    }
    Ok(out)
}

/// Produce inputs (node, parents) for graph_row.
///
/// If `subset` is provided, only render a subset of the graph in
/// the `subset` order.
fn dag_to_renderer_next_rows(
    dag: &(impl DagAlgorithm + ?Sized),
    subset: Option<Set>,
) -> Result<Vec<(VertexName, Vec<Ancestor<VertexName>>)>> {
    let subset = match subset {
        None => non_blocking_result(dag.all())?,
        Some(set) => set,
    };
    let iter: Vec<_> = subset.iter()?.collect::<crate::Result<_>>()?;
    let subdag = non_blocking_result(dag.subdag(subset))?;

    let mut out = Vec::with_capacity(iter.len());
    for node in iter {
        let direct_parents = non_blocking_result(dag.parent_names(node.clone()))?;
        let subdag_parents = non_blocking_result(subdag.parent_names(node.clone()))?;
        let mut parents: Vec<Ancestor<VertexName>> = subdag_parents
            .iter()
            .cloned()
            .map(|p| {
                if direct_parents.contains(&p) {
                    Ancestor::Parent(p)
                } else {
                    Ancestor::Ancestor(p)
                }
            })
            .collect();
        if parents.len() < direct_parents.len() {
            let subdag_ancestors =
                non_blocking_result(dag.ancestors(Set::from_static_names(subdag_parents.clone())))?;
            for p in &direct_parents {
                if subdag_parents.contains(p) {
                    continue;
                }
                // Is the direct parent connect to any of the subdag parents?
                let direct_reachable = non_blocking_result(dag.ancestors(p.into()))?;
                if direct_reachable
                    .intersection(&subdag_ancestors)
                    .is_empty()?
                {
                    parents.push(Ancestor::Anonymous)
                }
            }
        }
        out.push((node, parents));
    }

    Ok(out)
}

#[cfg(any(test, feature = "indexedlog-backend"))]
pub fn render_segment_dag(
    mut out: impl Write,
    dag: &NameDag,
    level: Level,
    group: Group,
) -> Result<()> {
    let mut renderer = renderdag::GraphRowRenderer::new()
        .output()
        .build_box_drawing();
    let segs = dag.dag.next_segments(group.min_id(), level)?;

    for seg in segs.iter().rev() {
        let mut parents = vec![];
        for parent_id in seg.parents()? {
            // For each parent Id, look for the containing segment.
            let parent_span: IdSpan = parent_id.into();
            let parent_idx = segs.binary_search_by(|s| {
                let span = s.span().unwrap();
                if span.contains(parent_id) {
                    Ordering::Equal
                } else {
                    span.cmp(&parent_span)
                }
            });

            if let Ok(parent_idx) = parent_idx {
                parents.push(Ancestor::Parent(&segs[parent_idx]));
            } else {
                // Probably a non-master segment with master parent.
                parents.push(Ancestor::Anonymous);
            }
        }

        let span = seg.span()?;
        let get_hex = |id| -> String {
            non_blocking_result(dag.vertex_name(id))
                .map(|s| format!("{:.12?}", s))
                .unwrap_or_default()
        };
        let name = format!(
            "{}({})-{}({})",
            get_hex(span.low),
            span.low,
            get_hex(span.high),
            span.high,
        );
        let row = renderer.next_row(seg, parents, String::from("o"), name);
        write!(out, "{}", row)?;
    }

    Ok(())
}


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

//! # segment
//!
//! Segmented DAG. See [`IdDag`] for the main structure.

use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::io::Cursor;

use bitflags::bitflags;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
pub use dag_types::segment::FlatSegment;
pub use dag_types::segment::PreparedFlatSegments;
use minibytes::Bytes;
use serde::Deserialize;
use serde::Serialize;
use vlqencoding::VLQDecode;
use vlqencoding::VLQDecodeAt;
use vlqencoding::VLQEncode;

use crate::errors::bug;
use crate::errors::programming;
use crate::id::Id;
use crate::IdSpan;
use crate::Level;
use crate::Result;

/// [`Segment`] represents a range of [`Id`]s in an [`IdDag`] graph.
/// It provides methods to access properties of the segments, including the range itself,
/// parents, and level information.
///
/// This is the main struct used by the `dag` crate. It is mostly read from a store.
/// Once read, it's immutable.
#[derive(Clone, Eq, Serialize, Deserialize)]
pub struct Segment(pub(crate) Bytes);

/// In memory representation of a segment. Mostly useful for external use-cases.
///
/// Unlike `Segment`, it is not necessarily read from a store and can be
/// constructed on the fly.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct IdSegment {
    pub low: Id,
    pub high: Id,
    pub parents: Vec<Id>,
    pub has_root: bool,
    pub level: Level,
}

// Serialization format for Segment:
//
// ```plain,ignore
// SEGMENT := FLAG (1B) + LEVEL (1B) + HIGH (8B) + vlq(HIGH-LOW) + vlq(PARENT_COUNT) + vlq(VLQ, PARENTS)
// ```
//
// The reason HIGH is not stored in VLQ is because it's used by range lookup,
// and vlq `[u8]` order does not match integer order.
//
// The reason HIGH-LOW is used instead of LOW is because it is more compact
// for the worse case (i.e. each flat segment has length 1). Each segment has
// only 1 byte overhead.

bitflags! {
    pub struct SegmentFlags: u8 {
        /// This segment has roots (i.e. there is at least one id in
        /// `low..=high`, `parents(id)` is empty).
        ///
        /// This flag must be set correct for correctness.
        const HAS_ROOT = 0b1;

        /// This segment is the only head in `0..=high`.
        /// In other words, `heads(0..=high)` is `[high]`.
        ///
        /// This flag should not be set if the segment is either a high-level
        /// segment, or in a non-master group.
        ///
        /// This flag is an optimization. Not setting it might hurt performance
        /// but not correctness.
        const ONLY_HEAD = 0b10;
    }
}

impl Segment {
    pub(crate) const OFFSET_FLAGS: usize = 0;
    pub(crate) const OFFSET_LEVEL: usize = Self::OFFSET_FLAGS + 1;
    pub(crate) const OFFSET_HIGH: usize = Self::OFFSET_LEVEL + 1;
    pub(crate) const OFFSET_DELTA: usize = Self::OFFSET_HIGH + 8;

    pub(crate) fn flags(&self) -> Result<SegmentFlags> {
        match self.0.get(Self::OFFSET_FLAGS) {
            Some(bits) => Ok(SegmentFlags::from_bits_truncate(*bits)),
            None => bug("cannot read Segment::flags"),
        }
    }

    pub(crate) fn has_root(&self) -> Result<bool> {
        Ok(self.flags()?.contains(SegmentFlags::HAS_ROOT))
    }

    pub(crate) fn only_head(&self) -> Result<bool> {
        Ok(self.flags()?.contains(SegmentFlags::ONLY_HEAD))
    }

    pub(crate) fn high(&self) -> Result<Id> {
        match self.0.get(Self::OFFSET_HIGH..Self::OFFSET_HIGH + 8) {
            Some(slice) => Ok(Id(BigEndian::read_u64(slice))),
            None => bug("cannot read Segment::high"),
        }
    }

    // high - low
    fn delta(&self) -> Result<u64> {
        let (len, _) = self.0.read_vlq_at(Self::OFFSET_DELTA)?;
        Ok(len)
    }

    pub(crate) fn span(&self) -> Result<IdSpan> {
        let high = self.high()?;
        let delta = self.delta()?;
        let low = high - delta;
        Ok((low..=high).into())
    }

    pub(crate) fn head(&self) -> Result<Id> {
        self.high()
    }

    pub(crate) fn low(&self) -> Result<Id> {
        Ok(self.span()?.low)
    }

    pub(crate) fn level(&self) -> Result<Level> {
        match self.0.get(Self::OFFSET_LEVEL) {
            Some(level) => Ok(*level),
            None => bug("cannot read Segment::level"),
        }
    }

    pub(crate) fn parent_count(&self) -> Result<usize> {
        let mut cur = Cursor::new(&self.0);
        cur.set_position(Self::OFFSET_DELTA as u64);
        let _: u64 = cur.read_vlq()?;
        let parent_count: usize = cur.read_vlq()?;
        Ok(parent_count)
    }

    pub(crate) fn parents(&self) -> Result<Vec<Id>> {
        let mut cur = Cursor::new(&self.0);
        cur.set_position(Self::OFFSET_DELTA as u64);
        let _: u64 = cur.read_vlq()?;
        let parent_count: usize = cur.read_vlq()?;
        let mut result = Vec::with_capacity(parent_count);
        for _ in 0..parent_count {
            result.push(Id(cur.read_vlq()?));
        }
        Ok(result)
    }

    /// Duplicate the segment with `high` set to a new value.
    pub(crate) fn with_high(&self, high: Id) -> Result<Self> {
        let span = self.span()?;
        if high.group() != span.high.group() || high < span.low {
            return programming(format!(
                "with_high got invalid input (segment: {:?} new_high: {:?})",
                self, high,
            ));
        }
        let parents = self.parents()?;
        let seg = Self::new(self.flags()?, self.level()?, span.low, high, &parents);
        Ok(seg)
    }

    pub(crate) fn new(
        flags: SegmentFlags,
        level: Level,
        low: Id,
        high: Id,
        parents: &[Id],
    ) -> Self {
        debug_assert!(high >= low);
        debug_assert!(parents.iter().all(|&p| p < low));
        let mut buf = Vec::with_capacity(1 + 8 + (parents.len() + 2) * 4);
        buf.write_u8(flags.bits()).unwrap();
        buf.write_u8(level).unwrap();
        buf.write_u64::<BigEndian>(high.0).unwrap();
        buf.write_vlq(high.0 - low.0).unwrap();
        buf.write_vlq(parents.len()).unwrap();
        for parent in parents {
            buf.write_vlq(parent.0).unwrap();
        }
        Self(buf.into())
    }
}

impl PartialEq for Segment {
    fn eq(&self, other: &Self) -> bool {
        self.0[..] == other.0[..]
    }
}

impl Debug for Segment {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let span = self.span().unwrap();
        if self.has_root().unwrap() {
            write!(f, "R")?;
        }
        if self.only_head().unwrap() {
            write!(f, "H")?;
        }
        let parents = self.parents().unwrap();
        write!(f, "{}-{}{:?}", span.low, span.high, parents,)?;
        if span.low > span.high {
            write!(f, " (Invalid Span!!)")?;
        }
        if parents.iter().any(|&p| p >= span.low) {
            write!(f, " (Invalid Parent!!)")?;
        }
        Ok(())
    }
}

impl Debug for IdSegment {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(
            f,
            "L{} {}..={} {:?}{}",
            self.level,
            self.low,
            self.high,
            &self.parents,
            if self.has_root { "R" } else { "" }
        )
    }
}

/// Describe bytes of a Segment.
/// This is only for troubleshooting purpose.
pub fn describe_segment_bytes(data: &[u8]) -> String {
    let mut message = String::new();
    let mut cur = Cursor::new(data);
    let mut start = 0;
    let mut explain = |cur: &Cursor<_>, m: String| {
        let end = cur.position() as usize;
        message += &format!("# {}: {}\n", hex(&data[start..end]), m);
        start = end;
    };
    if let Ok(flags) = cur.read_u8() {
        let flags = SegmentFlags::from_bits_truncate(flags);
        explain(&cur, format!("Flags = {:?}", flags));
    }
    if let Ok(lv) = cur.read_u8() {
        explain(&cur, format!("Level = {:?}", lv));
    }
    if let Ok(head) = cur.read_u64::<BigEndian>() {
        explain(&cur, format!("High = {}", Id(head)));
        if let Ok(delta) = VLQDecode::<u64>::read_vlq(&mut cur) {
            let low = head - delta;
            explain(&cur, format!("Delta = {} (Low = {})", delta, Id(low)));
        }
    }
    if let Ok(count) = VLQDecode::<usize>::read_vlq(&mut cur) {
        explain(&cur, format!("Parent count = {}", count));
        for i in 0..count {
            if let Ok(p) = VLQDecode::<u64>::read_vlq(&mut cur) {
                explain(&cur, format!("Parents[{}] = {}", i, Id(p)));
            }
        }
    }
    message
}

pub(crate) fn hex(bytes: &[u8]) -> String {
    bytes
        .iter()
        .cloned()
        .map(|b| format!("{:02x}", b))
        .collect::<Vec<String>>()
        .join(" ")
}

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

    use super::*;

    #[test]
    fn test_segment_roundtrip() {
        fn prop(has_root: bool, level: Level, range1: u64, range2: u64, parents: Vec<u64>) -> bool {
            let flags = if has_root {
                SegmentFlags::HAS_ROOT
            } else {
                SegmentFlags::empty()
            };
            let low = u64::min(range1, range2);
            let high = u64::max(range1, range2);
            let parents: Vec<Id> = parents.into_iter().filter(|&p| p < low).map(Id).collect();
            let low = Id(low);
            let high = Id(high);
            let node = Segment::new(flags, level, low, high, &parents);
            node.flags().unwrap() == flags
                && node.level().unwrap() == level
                && node.span().unwrap() == (low..=high).into()
                && node.parents().unwrap() == parents
        }
        quickcheck(prop as fn(bool, Level, u64, u64, Vec<u64>) -> bool);
    }

    #[test]
    fn test_describe() {
        let seg = Segment::new(
            SegmentFlags::ONLY_HEAD,
            3,
            Id(101),
            Id(202),
            &[Id(90), Id(80)],
        );
        assert_eq!(
            describe_segment_bytes(&seg.0),
            r#"# 02: Flags = ONLY_HEAD
# 03: Level = 3
# 00 00 00 00 00 00 00 ca: High = 202
# 65: Delta = 101 (Low = 101)
# 02: Parent count = 2
# 5a: Parents[0] = 90
# 50: Parents[1] = 80
"#
        );
    }

    #[test]
    fn test_invalid_fmt() {
        let bytes = Bytes::from_static(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 0, 1, 10]);
        let segment = Segment(bytes);
        assert_eq!(format!("{:?}", segment), "10-10[10] (Invalid Parent!!)");
    }
}