/*
 * 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::sync::atomic;
use std::sync::atomic::AtomicU64;

use super::IdMapWrite;
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::Result;
use crate::VerLink;

/// Bi-directional mapping between an integer id and a name (`[u8]`).
///
/// Private. Stored in memory.
pub struct MemIdMap {
    core: CoreMemIdMap,
    map_id: String,
    map_version: VerLink,
}

/// Subset of the `MemIdMap` interface that does not have "map_version".
/// or "version" concept.
#[derive(Default, Clone)]
pub(crate) struct CoreMemIdMap {
    id2name: BTreeMap<Id, VertexName>,
    name2id: BTreeMap<VertexName, Id>,
}

impl MemIdMap {
    /// Create an empty [`MemIdMap`].
    pub fn new() -> Self {
        Self {
            core: Default::default(),
            map_id: format!("mem:{}", next_id()),
            map_version: VerLink::new(),
        }
    }
}

impl Clone for MemIdMap {
    fn clone(&self) -> Self {
        Self {
            core: self.core.clone(),
            map_id: self.map_id.clone(),
            map_version: self.map_version.clone(),
        }
    }
}

impl CoreMemIdMap {
    pub fn lookup_vertex_id(&self, name: &VertexName) -> Option<Id> {
        self.name2id.get(&name).copied()
    }

    pub fn lookup_vertex_name(&self, id: Id) -> Option<VertexName> {
        self.id2name.get(&id).cloned()
    }

    pub fn lookup_vertexes_by_hex_prefix(
        &self,
        hex_prefix: &[u8],
        limit: usize,
    ) -> Result<Vec<VertexName>> {
        let start = VertexName::from_hex(hex_prefix)?;
        let mut result = Vec::new();
        for (vertex, _) in self.name2id.range(start..) {
            if !vertex.to_hex().as_bytes().starts_with(hex_prefix) {
                break;
            }
            result.push(vertex.clone());
            if result.len() >= limit {
                break;
            }
        }
        Ok(result)
    }

    pub fn has_vertex_name(&self, name: &VertexName) -> bool {
        self.name2id.contains_key(name)
    }

    pub fn has_vertex_id(&self, id: Id) -> bool {
        self.id2name.contains_key(&id)
    }

    pub fn insert_vertex_id_name(&mut self, id: Id, vertex_name: VertexName) {
        self.name2id.insert(vertex_name.clone(), id);
        self.id2name.insert(id, vertex_name);
    }

    pub fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
        let to_remove: Vec<(Id, VertexName)> = self
            .id2name
            .range(low..=high)
            .map(|(i, n)| (*i, n.clone()))
            .collect();
        for (id, name) in &to_remove {
            self.id2name.remove(id);
            self.name2id.remove(name);
        }
        Ok(to_remove.into_iter().map(|(_, v)| v).collect())
    }
}

#[async_trait::async_trait]
impl IdConvert for MemIdMap {
    async fn vertex_id(&self, name: VertexName) -> Result<Id> {
        self.core
            .lookup_vertex_id(&name)
            .ok_or_else(|| name.not_found_error())
    }
    async fn vertex_id_with_max_group(
        &self,
        name: &VertexName,
        max_group: Group,
    ) -> Result<Option<Id>> {
        let optional_id = self.core.name2id.get(name).and_then(|id| {
            if id.group() <= max_group {
                Some(*id)
            } else {
                None
            }
        });
        Ok(optional_id)
    }
    async fn vertex_name(&self, id: Id) -> Result<VertexName> {
        self.core
            .lookup_vertex_name(id)
            .ok_or_else(|| id.not_found_error())
    }
    async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool> {
        Ok(self.core.has_vertex_name(name))
    }

    async fn contains_vertex_id_locally(&self, ids: &[Id]) -> Result<Vec<bool>> {
        Ok(ids
            .iter()
            .copied()
            .map(|id| self.core.has_vertex_id(id))
            .collect())
    }
    async fn contains_vertex_name_locally(&self, names: &[VertexName]) -> Result<Vec<bool>> {
        Ok(names
            .iter()
            .map(|name| self.core.has_vertex_name(name))
            .collect())
    }

    fn map_id(&self) -> &str {
        &self.map_id
    }

    fn map_version(&self) -> &VerLink {
        &self.map_version
    }
}

// TODO: Reconsider re-assign master cases. Currently they are ignored.
#[async_trait::async_trait]
impl IdMapWrite for MemIdMap {
    async fn insert(&mut self, id: Id, name: &[u8]) -> Result<()> {
        let vertex_name = VertexName::copy_from(name);
        self.core.insert_vertex_id_name(id, vertex_name);
        self.map_version.bump();
        Ok(())
    }
    async fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
        self.map_version = VerLink::new();
        self.core.remove_range(low, high)
    }
}

impl Persist for MemIdMap {
    type Lock = ();

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

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

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

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

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


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

use crate::namedag::AbstractNameDag;
use crate::IdDag;
use crate::IdDagStore;

/// State to build a new `AbstractNameDag`.
pub struct NameDagBuilder<M, D, P, S> {
    map: M,
    dag: D,
    path: P,
    state: S,
    id: String,
}

impl<M, D> NameDagBuilder<M, D, (), ()> {
    /// Create the builder with specified `IdMap` and `IdDag`.
    ///
    /// The callsite must ensure the `IdMap` and `IdMap` are
    /// in sync.
    pub fn new_with_idmap_dag(map: M, dag: D) -> Self {
        Self {
            map,
            dag,
            path: (),
            id: String::new(),
            state: (),
        }
    }
}

impl<M, D, P, S> NameDagBuilder<M, D, P, S> {
    /// Set the `path`, used to re-open the `NameDag`.
    pub fn with_path<P2>(self, path: P2) -> NameDagBuilder<M, D, P2, S> {
        NameDagBuilder {
            map: self.map,
            dag: self.dag,
            path,
            id: self.id,
            state: self.state,
        }
    }

    /// Set the `state`, additional state maintained used by the callsite.
    pub fn with_state<S2>(self, state: S2) -> NameDagBuilder<M, D, P, S2> {
        NameDagBuilder {
            map: self.map,
            dag: self.dag,
            path: self.path,
            id: self.id,
            state,
        }
    }

    /// Set the `id`, used for debugging.
    pub fn with_id(self, id: String) -> NameDagBuilder<M, D, P, S> {
        NameDagBuilder {
            map: self.map,
            dag: self.dag,
            path: self.path,
            id,
            state: self.state,
        }
    }
}

impl<IS, M, P, S> NameDagBuilder<M, IdDag<IS>, P, S>
where
    M: Send + Sync,
    IS: IdDagStore,
    P: Send + Sync,
    S: Send + Sync,
{
    /// Build the `AbstractNameDag`.
    pub fn build(self) -> crate::Result<AbstractNameDag<IdDag<IS>, M, P, S>> {
        let persisted_id_set = self.dag.all()?;
        let overlay_map_id_set = self.dag.master_group()?;
        let dag = AbstractNameDag {
            dag: self.dag,
            map: self.map,
            path: self.path,
            state: self.state,
            id: self.id,

            snapshot: Default::default(),
            pending_heads: Default::default(),
            persisted_id_set,
            overlay_map: Default::default(),
            overlay_map_id_set,
            overlay_map_paths: Default::default(),
            remote_protocol: Arc::new(()),
            missing_vertexes_confirmed_by_remote: Default::default(),
        };
        Ok(dag)
    }
}

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

    use super::*;
    use crate::iddag::IdDag;
    use crate::idmap::MemIdMap;
    use crate::ops::DagAddHeads;
    use crate::VertexListWithOptions;

    #[tokio::test]
    async fn test_builder_absent_path_state_can_use_add_heads() {
        let dag = IdDag::new_in_process();
        let map = MemIdMap::new();
        let builder = NameDagBuilder::new_with_idmap_dag(map, dag);
        let mut dag = builder.build().unwrap();
        dag.add_heads(&HashMap::new(), &VertexListWithOptions::default())
            .await
            .unwrap();
    }
}


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

use indexedlog::multi;
use indexedlog::DefaultOpenOptions;
use indexedlog::OpenWithRepair;

use super::AbstractNameDag;
use super::NameDagBuilder;
use crate::errors::bug;
use crate::iddag::IdDag;
use crate::iddagstore::IndexedLogStore;
use crate::idmap::IdMap;
use crate::ops::IntVersion;
use crate::ops::Open;
use crate::ops::Persist;
use crate::ops::TryClone;
use crate::Result;

/// A DAG that uses VertexName instead of ids as vertexes.
///
/// A high-level wrapper structure. Combination of [`IdMap`] and [`Dag`].
/// Maintains consistency of dag and map internally.
pub type NameDag =
    AbstractNameDag<IdDag<IndexedLogStore>, IdMap, IndexedLogNameDagPath, NameDagState>;

pub struct NameDagState {
    /// `MultiLog` controls on-disk metadata.
    /// `None` for read-only `NameDag`,
    mlog: Option<multi::MultiLog>,
}

/// Address to on-disk NameDag based on indexedlog.
#[derive(Debug, Clone)]
pub struct IndexedLogNameDagPath(pub PathBuf);

impl Open for IndexedLogNameDagPath {
    type OpenTarget = NameDag;

    fn open(&self) -> Result<Self::OpenTarget> {
        crate::failpoint!("dag-namedag-open");
        let path = &self.0;
        let opts = NameDag::default_open_options();
        tracing::debug!(target: "dag::open",  "open at {:?}", path.display());
        let mut mlog = opts.open_with_repair(path)?;
        let mut logs = mlog.detach_logs();
        let dag_log = logs.pop().unwrap();
        let map_log = logs.pop().unwrap();
        let map = IdMap::open_from_log(map_log)?;
        let dag = IdDag::open_from_store(IndexedLogStore::open_from_clean_log(dag_log)?)?;
        let state = NameDagState { mlog: Some(mlog) };
        let id = format!("ilog:{}", self.0.display());
        let dag = NameDagBuilder::new_with_idmap_dag(map, dag)
            .with_path(self.clone())
            .with_state(state)
            .with_id(id)
            .build()?;
        Ok(dag)
    }
}

impl DefaultOpenOptions<multi::OpenOptions> for NameDag {
    fn default_open_options() -> multi::OpenOptions {
        multi::OpenOptions::from_name_opts(vec![
            ("idmap2", IdMap::log_open_options()),
            ("iddag", IndexedLogStore::log_open_options()),
        ])
    }
}

impl NameDag {
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let path = IndexedLogNameDagPath(path);
        path.open()
    }
}

impl Persist for NameDagState {
    type Lock = indexedlog::multi::LockGuard;

    fn lock(&mut self) -> Result<Self::Lock> {
        if self.mlog.is_none() {
            return bug("MultiLog should be Some for read-write NameDag");
        }
        let mlog = self.mlog.as_mut().unwrap();
        // mlog.lock() reloads its MultiMeta, but not Logs.
        //
        // Usually the use pattern is like:
        //
        //    let locked = self.state.prepare_filesystem_sync()?;  // Get the latest MultiMeta
        //    let mut map = self.map.prepare_filesystem_sync()?;   // Get the latest Log
        //    let mut dag = self.dag.prepare_filesystem_sync()?;   // Get the latest Log.
        //
        // The `NameDagState` does not control the `map` or `dag` Logs so it cannot reload
        // them here, or in `reload()`.
        Ok(mlog.lock()?)
    }

    fn reload(&mut self, _lock: &Self::Lock) -> Result<()> {
        // mlog does reload internally. See `lock()`.
        Ok(())
    }

    fn persist(&mut self, lock: &Self::Lock) -> Result<()> {
        self.mlog.as_mut().unwrap().write_meta(&lock)?;
        Ok(())
    }
}

impl IntVersion for NameDagState {
    fn int_version(&self) -> (u64, u64) {
        match &self.mlog {
            Some(mlog) => mlog.version(),
            None => (0, 0),
        }
    }
}

impl TryClone for NameDagState {
    fn try_clone(&self) -> Result<Self> {
        Ok(Self {
            // mlog cannot be cloned.
            mlog: None,
        })
    }
}


/*
 * 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::sync::atomic;
use std::sync::atomic::AtomicU64;

use super::AbstractNameDag;
use super::NameDagBuilder;
use crate::iddag::IdDag;
use crate::iddagstore::InProcessStore;
use crate::idmap::MemIdMap;
use crate::ops::IntVersion;
use crate::ops::Open;
use crate::ops::Persist;
use crate::Result;

/// In-memory version of [`NameDag`].
///
/// Does not support loading from or saving to the filesystem.
/// The graph has to be built from scratch by `add_heads`.
pub type MemNameDag =
    AbstractNameDag<IdDag<InProcessStore>, MemIdMap, MemNameDagPath, MemNameDagState>;

/// Address to open in-memory Dag.
#[derive(Debug, Clone)]
pub struct MemNameDagPath;

#[derive(Debug, Clone)]
pub struct MemNameDagState {
    version: (u64, u64),
}

impl Default for MemNameDagState {
    fn default() -> Self {
        Self {
            version: (rand::random(), 0),
        }
    }
}

impl Open for MemNameDagPath {
    type OpenTarget = MemNameDag;

    fn open(&self) -> Result<Self::OpenTarget> {
        let dag = IdDag::new_in_process();
        let map = MemIdMap::new();
        let id = format!("mem:{}", next_id());
        let state = MemNameDagState::default();
        let result = NameDagBuilder::new_with_idmap_dag(map, dag)
            .with_path(self.clone())
            .with_id(id)
            .with_state(state)
            .build()?;
        Ok(result)
    }
}

impl MemNameDag {
    pub fn new() -> Self {
        MemNameDagPath.open().unwrap()
    }
}

impl Persist for MemNameDagState {
    type Lock = ();

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

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

    fn persist(&mut self, _lock: &Self::Lock) -> Result<()> {
        self.version.1 += 1;
        Ok(())
    }
}

impl IntVersion for MemNameDagState {
    fn int_version(&self) -> (u64, u64) {
        self.version
    }
}

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


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

use futures::StreamExt;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use super::NameSet;
use crate::fmt::write_debug;
use crate::Result;
use crate::VertexName;

/// Subset of `lhs` that does not overlap with `rhs`.
///
/// The iteration order is defined by `lhs`.
pub struct DifferenceSet {
    lhs: NameSet,
    rhs: NameSet,
    hints: Hints,
}

struct Iter {
    iter: BoxVertexStream,
    rhs: NameSet,
}

impl DifferenceSet {
    pub fn new(lhs: NameSet, rhs: NameSet) -> Self {
        let hints = Hints::new_inherit_idmap_dag(lhs.hints());
        // Inherit flags, min/max Ids from lhs.
        hints.add_flags(
            lhs.hints().flags()
                & (Flags::EMPTY
                    | Flags::ID_DESC
                    | Flags::ID_ASC
                    | Flags::TOPO_DESC
                    | Flags::FILTER),
        );
        if let Some(id) = lhs.hints().min_id() {
            hints.set_min_id(id);
        }
        if let Some(id) = lhs.hints().max_id() {
            hints.set_max_id(id);
        }
        Self { lhs, rhs, hints }
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for DifferenceSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let iter = Iter {
            iter: self.lhs.iter().await?,
            rhs: self.rhs.clone(),
        };
        Ok(iter.into_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let iter = Iter {
            iter: self.lhs.iter_rev().await?,
            rhs: self.rhs.clone(),
        };
        Ok(iter.into_stream())
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        Ok(self.lhs.contains(name).await? && !self.rhs.contains(name).await?)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        let lhs_contains = self.lhs.contains_fast(name).await?;
        if lhs_contains == Some(false) {
            return Ok(Some(false));
        }
        let rhs_contains = self.rhs.contains_fast(name).await?;
        let result = match (lhs_contains, rhs_contains) {
            (Some(true), Some(false)) => Some(true),
            (_, Some(true)) | (Some(false), _) => Some(false),
            (Some(true), None) | (None, _) => None,
        };
        Ok(result)
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

impl fmt::Debug for DifferenceSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<diff")?;
        write_debug(f, &self.lhs)?;
        write_debug(f, &self.rhs)?;
        write!(f, ">")
    }
}

impl Iter {
    async fn next(&mut self) -> Option<Result<VertexName>> {
        loop {
            let result = self.iter.as_mut().next().await;
            if let Some(Ok(ref name)) = result {
                match self.rhs.contains(&name).await {
                    Err(err) => break Some(Err(err)),
                    Ok(true) => continue,
                    _ => {}
                }
            }
            break result;
        }
    }

    fn into_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |mut state| async move {
            let result = state.next().await;
            result.map(|r| (r, state))
        }))
    }
}

#[cfg(test)]
mod tests {
    use nonblocking::non_blocking as nb;

    use super::super::tests::*;
    use super::*;

    fn difference(a: &[u8], b: &[u8]) -> DifferenceSet {
        let a = NameSet::from_query(VecQuery::from_bytes(a));
        let b = NameSet::from_query(VecQuery::from_bytes(b));
        DifferenceSet::new(a, b)
    }

    #[test]
    fn test_difference_basic() -> Result<()> {
        let set = difference(b"\x11\x33\x55\x22\x44", b"\x44\x33\x66");
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["11", "55", "22"]);
        assert_eq!(shorten_iter(ni(set.iter_rev())), ["22", "55", "11"]);
        assert!(!nb(set.is_empty())??);
        assert_eq!(nb(set.count())??, 3);
        assert_eq!(shorten_name(nb(set.first())??.unwrap()), "11");
        assert_eq!(shorten_name(nb(set.last())??.unwrap()), "22");
        for &b in b"\x11\x22\x55".iter() {
            assert!(nb(set.contains(&to_name(b)))??);
        }
        for &b in b"\x33\x44\x66".iter() {
            assert!(!nb(set.contains(&to_name(b)))??);
        }
        Ok(())
    }

    quickcheck::quickcheck! {
        fn test_difference_quickcheck(a: Vec<u8>, b: Vec<u8>) -> bool {
            let set = difference(&a, &b);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap().unwrap();
            assert!(count <= a.len());

            assert!(b.iter().all(|&b| nb(set.contains(&to_name(b))).unwrap().ok() == Some(false)));

            true
        }
    }
}


/*
 * 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::AtomicU32;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::Ordering::Release;
use std::sync::Arc;

use bitflags::bitflags;

use crate::ops::DagAlgorithm;
use crate::ops::IdConvert;
use crate::Id;
use crate::VerLink;

bitflags! {
    pub struct Flags: u32 {
        /// Full. A full Set & other set X with compatible Dag results in X.
        const FULL = 0x1;

        /// Empty (also implies ID_DESC | ID_ASC | TOPO_DESC).
        const EMPTY = 0x2;

        /// Sorted by Id. Descending (head -> root).
        const ID_DESC = 0x4;

        /// Sorted by Id. Ascending (root -> head).
        const ID_ASC = 0x8;

        /// Sorted topologically. Descending (head -> root).
        const TOPO_DESC = 0x10;

        /// Minimal Id is known.
        const HAS_MIN_ID = 0x20;

        /// Maximum Id is known.
        const HAS_MAX_ID = 0x40;

        /// A "filter" set. It provides "contains" but iteration is inefficient.
        const FILTER = 0x80;

        /// The set contains ancestors. If X in set, any ancestor of X is also in set.
        const ANCESTORS = 0x100;
    }
}

/// Optimation hints.
#[derive(Default)]
pub struct Hints {
    // Atomic is used for interior mutability.
    flags: AtomicU32,
    min_id: AtomicU64,
    max_id: AtomicU64,
    id_map: IdMapSnapshot,
    dag: DagSnapshot,
}

impl Hints {
    pub fn new_with_idmap_dag(
        id_map: impl Into<IdMapSnapshot>,
        dag: impl Into<DagSnapshot>,
    ) -> Self {
        Self {
            id_map: id_map.into(),
            dag: dag.into(),
            ..Default::default()
        }
    }

    pub fn new_inherit_idmap_dag(hints: &Self) -> Self {
        Self::new_with_idmap_dag(hints, hints)
    }

    /// Attempt to inherit properties (IdMap and Dag snapshots) from a list of
    /// hints. The returned hints have IdMap and Dag set to be compatible with
    /// all other hints in the list (or set to be None if that's not possible).
    pub fn union(hints_list: &[&Self]) -> Self {
        let default = Self::default();
        // Find the id_map that is compatible with all other id_maps.
        debug_assert!(default.id_map().is_none());
        let id_map = hints_list
            .iter()
            .fold(Some(&default), |opt_a, b| {
                opt_a.and_then(
                    |a| match a.id_map_version().partial_cmp(&b.id_map_version()) {
                        None => None, // Incompatible sets
                        Some(cmp::Ordering::Equal) | Some(cmp::Ordering::Greater) => Some(a),
                        Some(cmp::Ordering::Less) => Some(b),
                    },
                )
            })
            .and_then(|a| a.id_map());
        // Find the dag that is compatible with all other dags.
        debug_assert!(default.dag().is_none());
        let dag = hints_list
            .iter()
            .fold(Some(&default), |opt_a, b| {
                opt_a.and_then(|a| match a.dag_version().partial_cmp(&b.dag_version()) {
                    None => None,
                    Some(cmp::Ordering::Equal) | Some(cmp::Ordering::Greater) => Some(a),
                    Some(cmp::Ordering::Less) => Some(b),
                })
            })
            .and_then(|a| a.dag());
        Self {
            id_map: IdMapSnapshot(id_map),
            dag: DagSnapshot(dag),
            ..Self::default()
        }
    }

    pub fn flags(&self) -> Flags {
        let flags = self.flags.load(Relaxed);
        Flags::from_bits_truncate(flags)
    }

    pub fn contains(&self, flags: Flags) -> bool {
        self.flags().contains(flags)
    }

    pub fn min_id(&self) -> Option<Id> {
        if self.contains(Flags::HAS_MIN_ID) {
            Some(Id(self.min_id.load(Acquire)))
        } else {
            None
        }
    }

    pub fn max_id(&self) -> Option<Id> {
        if self.contains(Flags::HAS_MAX_ID) {
            Some(Id(self.max_id.load(Acquire)))
        } else {
            None
        }
    }

    pub fn update_flags_with(&self, func: impl Fn(Flags) -> Flags) -> &Self {
        let mut flags = func(self.flags());
        // Automatically add "derived" flags.
        if flags.contains(Flags::EMPTY) {
            flags.insert(Flags::ID_ASC | Flags::ID_DESC | Flags::TOPO_DESC | Flags::ANCESTORS);
        }
        if flags.contains(Flags::FULL) {
            flags.insert(Flags::ANCESTORS);
        }
        self.flags.store(flags.bits(), Relaxed);
        self
    }

    pub fn add_flags(&self, flags: Flags) -> &Self {
        self.update_flags_with(|f| f | flags);
        self
    }

    pub fn remove_flags(&self, flags: Flags) -> &Self {
        self.update_flags_with(|f| f - flags);
        self
    }

    pub fn set_min_id(&self, min_id: Id) -> &Self {
        self.min_id.store(min_id.0, Release);
        self.add_flags(Flags::HAS_MIN_ID);
        self
    }

    pub fn set_max_id(&self, max_id: Id) -> &Self {
        self.max_id.store(max_id.0, Release);
        self.add_flags(Flags::HAS_MAX_ID);
        self
    }

    pub fn inherit_flags_min_max_id(&self, other: &Hints) -> &Self {
        self.update_flags_with(|_| other.flags());
        if let Some(id) = other.min_id() {
            self.set_min_id(id);
        }
        if let Some(id) = other.max_id() {
            self.set_max_id(id);
        }
        self
    }

    pub fn dag(&self) -> Option<Arc<dyn DagAlgorithm + Send + Sync>> {
        self.dag.0.clone()
    }

    pub fn id_map(&self) -> Option<Arc<dyn IdConvert + Send + Sync>> {
        self.id_map.0.clone()
    }

    /// The `VerLink` of the Dag. `None` if there is no Dag associated.
    pub fn dag_version(&self) -> Option<&VerLink> {
        self.dag.0.as_ref().map(|d| d.dag_version())
    }

    /// The `VerLink` of the IdMap. `None` if there is no IdMap associated.
    pub fn id_map_version(&self) -> Option<&VerLink> {
        self.id_map.0.as_ref().map(|d| d.map_version())
    }
}

#[derive(Clone, Default)]
pub struct DagSnapshot(Option<Arc<dyn DagAlgorithm + Send + Sync>>);

impl From<&Hints> for DagSnapshot {
    fn from(hints: &Hints) -> Self {
        hints.dag.clone()
    }
}

impl From<Arc<dyn DagAlgorithm + Send + Sync>> for DagSnapshot {
    fn from(dag: Arc<dyn DagAlgorithm + Send + Sync>) -> Self {
        DagSnapshot(Some(dag))
    }
}

#[derive(Clone, Default)]
pub struct IdMapSnapshot(Option<Arc<dyn IdConvert + Send + Sync>>);

impl From<&Hints> for IdMapSnapshot {
    fn from(hints: &Hints) -> Self {
        hints.id_map.clone()
    }
}

impl From<Arc<dyn IdConvert + Send + Sync>> for IdMapSnapshot {
    fn from(dag: Arc<dyn IdConvert + Send + Sync>) -> Self {
        IdMapSnapshot(Some(dag))
    }
}

#[derive(Clone, Default)]
pub struct DagVersion<'a>(Option<&'a VerLink>);

impl<'a> From<&'a Hints> for DagVersion<'a> {
    fn from(hints: &'a Hints) -> Self {
        Self(match &hints.dag.0 {
            Some(d) => Some(d.dag_version()),
            None => None,
        })
    }
}

impl<'a> From<&'a VerLink> for DagVersion<'a> {
    fn from(version: &'a VerLink) -> Self {
        Self(Some(version))
    }
}

#[derive(Clone, Default)]
pub struct IdMapVersion<'a>(Option<&'a VerLink>);

impl<'a> From<&'a Hints> for IdMapVersion<'a> {
    fn from(hints: &'a Hints) -> Self {
        Self(match &hints.id_map.0 {
            Some(m) => Some(m.map_version()),
            None => None,
        })
    }
}

impl<'a> From<&'a VerLink> for IdMapVersion<'a> {
    fn from(version: &'a VerLink) -> Self {
        Self(Some(version))
    }
}

impl Clone for Hints {
    fn clone(&self) -> Self {
        Self {
            flags: AtomicU32::new(self.flags.load(Acquire)),
            min_id: AtomicU64::new(self.min_id.load(Acquire)),
            max_id: AtomicU64::new(self.max_id.load(Acquire)),
            id_map: self.id_map.clone(),
            dag: self.dag.clone(),
        }
    }
}

impl fmt::Debug for Hints {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Hints({:?}",
            self.flags() - (Flags::HAS_MIN_ID | Flags::HAS_MAX_ID)
        )?;
        match (self.min_id(), self.max_id()) {
            (Some(min), Some(max)) => write!(f, ", {}..={}", min.0, max.0)?,
            (Some(min), None) => write!(f, ", {}..", min.0)?,
            (None, Some(max)) => write!(f, ", ..={}", max.0)?,
            (None, None) => {}
        }
        write!(f, ")")?;
        Ok(())
    }
}

#[cfg(test)]
#[test]
fn test_incompatilbe_union() {
    use crate::tests::dummy_dag::DummyDag;
    let dag1 = DummyDag::new();
    let dag2 = DummyDag::new();

    let mut hints1 = Hints::default();
    hints1.dag = DagSnapshot(Some(dag1.dag_snapshot().unwrap()));

    let mut hints2 = Hints::default();
    hints2.dag = DagSnapshot(Some(dag2.dag_snapshot().unwrap()));

    assert_eq!(
        Hints::union(&[&hints1, &hints1]).dag_version(),
        hints1.dag_version()
    );

    assert_eq!(Hints::union(&[&hints1, &hints2]).dag_version(), None);
    assert_eq!(Hints::union(&[&hints2, &hints1]).dag_version(), None);
    assert_eq!(
        Hints::union(&[&hints2, &hints1, &hints2]).dag_version(),
        None
    );
}


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

use indexmap::IndexSet;
use nonblocking::non_blocking_result;

use super::hints::Flags;
use super::id_static::IdStaticSet;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use crate::ops::DagAlgorithm;
use crate::ops::IdConvert;
use crate::protocol::disable_remote_protocol;
use crate::Group;
use crate::Id;
use crate::IdSet;
use crate::Result;
use crate::VertexName;

/// A set backed by a lazy iterator of Ids.
pub struct IdLazySet {
    // Mutex: iter() does not take &mut self.
    // Arc: iter() result does not have a lifetime on this struct.
    inner: Arc<Mutex<Inner>>,
    pub map: Arc<dyn IdConvert + Send + Sync>,
    pub(crate) dag: Arc<dyn DagAlgorithm + Send + Sync>,
    hints: Hints,
}

struct Inner {
    iter: Box<dyn Iterator<Item = Result<Id>> + Send + Sync>,
    visited: IndexSet<Id>,
    state: State,
}

impl Inner {
    fn load_more(&mut self, n: usize, mut out: Option<&mut Vec<Id>>) -> Result<()> {
        if matches!(self.state, State::Complete | State::Error) {
            return Ok(());
        }
        for _ in 0..n {
            match self.iter.next() {
                Some(Ok(id)) => {
                    if let Some(ref mut out) = out {
                        out.push(id);
                    }
                    self.visited.insert(id);
                }
                None => {
                    self.state = State::Complete;
                    break;
                }
                Some(Err(err)) => {
                    self.state = State::Error;
                    return Err(err);
                }
            }
        }
        Ok(())
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum State {
    Incomplete,
    Complete,
    Error,
}

pub struct Iter {
    inner: Arc<Mutex<Inner>>,
    index: usize,
    map: Arc<dyn IdConvert + Send + Sync>,
}

impl Iter {
    fn into_box_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |this| this.next()))
    }

    async fn next(mut self) -> Option<(Result<VertexName>, Self)> {
        loop {
            let state = {
                let inner = self.inner.lock().unwrap();
                inner.state
            };
            match state {
                State::Error => break None,
                State::Complete if self.inner.lock().unwrap().visited.len() <= self.index => {
                    break None;
                }
                State::Complete | State::Incomplete => {
                    let opt_id = {
                        let inner = self.inner.lock().unwrap();
                        inner.visited.get_index(self.index).cloned()
                    };
                    match opt_id {
                        Some(id) => {
                            self.index += 1;
                            match self.map.vertex_name(id).await {
                                Err(err) => {
                                    self.inner.lock().unwrap().state = State::Error;
                                    return Some((Err(err), self));
                                }
                                Ok(vertex) => {
                                    break Some((Ok(vertex), self));
                                }
                            }
                        }
                        None => {
                            // Data not available. Load more.
                            let more = {
                                let mut inner = self.inner.lock().unwrap();
                                inner.load_more(1, None)
                            };
                            if let Err(err) = more {
                                return Some((Err(err), self));
                            }
                        }
                    }
                }
            }
        }
    }
}

struct DebugId {
    id: Id,
    name: Option<VertexName>,
}

impl fmt::Debug for DebugId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(name) = &self.name {
            fmt::Debug::fmt(&name, f)?;
            write!(f, "+{:?}", self.id)?;
        } else {
            write!(f, "{:?}", self.id)?;
        }
        Ok(())
    }
}

impl fmt::Debug for IdLazySet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("<lazy ")?;
        let inner = self.inner.lock().unwrap();
        let limit = f.width().unwrap_or(3);
        f.debug_list()
            .entries(inner.visited.iter().take(limit).map(|&id| DebugId {
                id,
                name: disable_remote_protocol(|| {
                    non_blocking_result(self.map.vertex_name(id)).ok()
                }),
            }))
            .finish()?;
        let remaining = inner.visited.len().max(limit) - limit;
        match (remaining, inner.state) {
            (0, State::Incomplete) => f.write_str(" + ? more")?,
            (n, State::Incomplete) => write!(f, "+ {} + ? more", n)?,
            (0, _) => {}
            (n, _) => write!(f, " + {} more", n)?,
        }
        f.write_str(">")?;
        Ok(())
    }
}

impl IdLazySet {
    pub fn from_iter_idmap_dag<I>(
        names: I,
        map: Arc<dyn IdConvert + Send + Sync>,
        dag: Arc<dyn DagAlgorithm + Send + Sync>,
    ) -> Self
    where
        I: IntoIterator<Item = Result<Id>> + 'static,
        <I as IntoIterator>::IntoIter: Send + Sync,
    {
        let iter = names.into_iter();
        let inner = Inner {
            iter: Box::new(iter),
            visited: IndexSet::new(),
            state: State::Incomplete,
        };
        let hints = Hints::new_with_idmap_dag(map.clone(), dag.clone());
        Self {
            inner: Arc::new(Mutex::new(inner)),
            map,
            dag,
            hints,
        }
    }

    /// Convert to an IdStaticSet.
    pub fn to_static(&self) -> Result<IdStaticSet> {
        let inner = self.load_all()?;
        let mut spans = IdSet::empty();
        for &id in inner.visited.iter() {
            spans.push(id);
        }
        Ok(IdStaticSet::from_spans_idmap_dag(
            spans,
            self.map.clone(),
            self.dag.clone(),
        ))
    }

    fn load_all(&self) -> Result<MutexGuard<Inner>> {
        let mut inner = self.inner.lock().unwrap();
        inner.load_more(usize::max_value(), None)?;
        Ok(inner)
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for IdLazySet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let inner = self.inner.clone();
        let map = self.map.clone();
        let iter = Iter {
            inner,
            index: 0,
            map,
        };
        Ok(iter.into_box_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let inner = self.load_all()?;
        struct State {
            map: Arc<dyn IdConvert + Send + Sync>,
            iter: Box<dyn Iterator<Item = Id> + Send>,
        }
        let state = State {
            map: self.map.clone(),
            iter: Box::new(inner.visited.clone().into_iter().rev()),
        };
        async fn next(mut state: State) -> Option<(Result<VertexName>, State)> {
            match state.iter.next() {
                None => None,
                Some(id) => {
                    let result = state.map.vertex_name(id).await;
                    Some((result, state))
                }
            }
        }

        let stream = futures::stream::unfold(state, next);
        Ok(Box::pin(stream))
    }

    async fn count(&self) -> Result<usize> {
        let inner = self.load_all()?;
        Ok(inner.visited.len())
    }

    async fn last(&self) -> Result<Option<VertexName>> {
        let opt_id = {
            let inner = self.load_all()?;
            inner.visited.iter().rev().nth(0).cloned()
        };
        match opt_id {
            Some(id) => Ok(Some(self.map.vertex_name(id).await?)),
            None => Ok(None),
        }
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        let id = match self
            .map
            .vertex_id_with_max_group(name, Group::NON_MASTER)
            .await?
        {
            None => {
                return Ok(false);
            }
            Some(id) => id,
        };
        let mut inner = self.inner.lock().unwrap();
        if inner.visited.contains(&id) {
            return Ok(true);
        } else {
            let mut loaded = Vec::new();
            loop {
                // Fast paths.
                if let Some(&last_id) = inner.visited.iter().rev().next() {
                    let hints = self.hints();
                    if hints.contains(Flags::ID_DESC) {
                        if last_id < id {
                            return Ok(false);
                        }
                    } else if hints.contains(Flags::ID_ASC) {
                        if last_id > id {
                            return Ok(false);
                        }
                    }
                }
                loaded.clear();
                inner.load_more(1, Some(&mut loaded))?;
                debug_assert!(loaded.len() <= 1);
                if loaded.is_empty() {
                    break;
                }
                if loaded.first() == Some(&id) {
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        let id = match self
            .map
            .vertex_id_with_max_group(name, Group::NON_MASTER)
            .await?
        {
            None => {
                return Ok(Some(false));
            }
            Some(id) => id,
        };
        let inner = self.inner.lock().unwrap();
        if inner.visited.contains(&id) {
            return Ok(Some(true));
        } else if inner.state != State::Incomplete {
            return Ok(Some(false));
        }
        Ok(None)
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }

    fn id_convert(&self) -> Option<&dyn IdConvert> {
        Some(self.map.as_ref() as &dyn IdConvert)
    }
}

#[cfg(test)]
pub(crate) mod test_utils {
    use std::sync::atomic::AtomicU64;
    use std::sync::atomic::Ordering::AcqRel;

    use super::*;
    use crate::ops::PrefixLookup;
    use crate::tests::dummy_dag::DummyDag;
    use crate::VerLink;

    static STR_ID_MAP_ID: AtomicU64 = AtomicU64::new(0);

    pub(crate) struct StrIdMap {
        id: String,
        version: VerLink,
    }

    impl StrIdMap {
        pub(crate) fn new() -> Self {
            Self {
                id: format!("str:{}", STR_ID_MAP_ID.fetch_add(1, AcqRel)),
                version: VerLink::new(),
            }
        }
    }

    #[async_trait::async_trait]
    impl PrefixLookup for StrIdMap {
        async fn vertexes_by_hex_prefix(&self, _: &[u8], _: usize) -> Result<Vec<VertexName>> {
            // Dummy implementation.
            Ok(Vec::new())
        }
    }
    #[async_trait::async_trait]
    impl IdConvert for StrIdMap {
        async fn vertex_id(&self, name: VertexName) -> Result<Id> {
            let slice: [u8; 8] = name.as_ref().try_into().unwrap();
            let id = u64::from_le(unsafe { std::mem::transmute(slice) });
            Ok(Id(id))
        }
        async fn vertex_id_with_max_group(
            &self,
            name: &VertexName,
            _max_group: Group,
        ) -> Result<Option<Id>> {
            if name.as_ref().len() == 8 {
                let id = self.vertex_id(name.clone()).await?;
                Ok(Some(id))
            } else {
                Ok(None)
            }
        }
        async fn vertex_name(&self, id: Id) -> Result<VertexName> {
            let buf: [u8; 8] = unsafe { std::mem::transmute(id.0.to_le()) };
            Ok(VertexName::copy_from(&buf))
        }
        async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool> {
            Ok(name.as_ref().len() == 8)
        }
        fn map_id(&self) -> &str {
            &self.id
        }
        fn map_version(&self) -> &VerLink {
            &self.version
        }
        async fn contains_vertex_id_locally(&self, ids: &[Id]) -> Result<Vec<bool>> {
            Ok(ids.iter().map(|_| true).collect())
        }
        async fn contains_vertex_name_locally(&self, names: &[VertexName]) -> Result<Vec<bool>> {
            Ok(names.iter().map(|name| name.as_ref().len() == 8).collect())
        }
    }

    pub fn lazy_set(a: &[u64]) -> IdLazySet {
        let ids: Vec<Id> = a.iter().map(|i| Id(*i as _)).collect();
        IdLazySet::from_iter_idmap_dag(
            ids.into_iter().map(Ok),
            Arc::new(StrIdMap::new()),
            Arc::new(DummyDag::new()),
        )
    }

    pub fn lazy_set_inherit(a: &[u64], set: &IdLazySet) -> IdLazySet {
        let ids: Vec<Id> = a.iter().map(|i| Id(*i as _)).collect();
        IdLazySet::from_iter_idmap_dag(ids.into_iter().map(Ok), set.map.clone(), set.dag.clone())
    }
}

#[cfg(all(test, feature = "indexedlog-backend"))]
#[allow(clippy::redundant_clone)]
pub(crate) mod tests {
    use std::collections::HashSet;

    use nonblocking::non_blocking_result as r;

    use super::super::tests::*;
    use super::super::NameSet;
    use super::test_utils::*;
    use super::*;

    #[test]
    fn test_id_lazy_basic() -> Result<()> {
        let set = lazy_set(&[0x11, 0x33, 0x22, 0x77, 0x55]);
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["11", "33", "22", "77", "55"]);
        assert_eq!(
            shorten_iter(ni(set.iter_rev())),
            ["55", "77", "22", "33", "11"]
        );
        assert!(!nb(set.is_empty())?);
        assert_eq!(nb(set.count())?, 5);
        assert_eq!(shorten_name(nb(set.first())?.unwrap()), "11");
        assert_eq!(shorten_name(nb(set.last())?.unwrap()), "55");
        Ok(())
    }

    #[test]
    fn test_hints_fast_paths() -> Result<()> {
        let set = lazy_set(&[0x20, 0x50, 0x30, 0x70]);

        // Incorrect hints, but useful for testing.
        set.hints().add_flags(Flags::ID_ASC);

        let v = |i: u64| -> VertexName { r(StrIdMap::new().vertex_name(Id(i))).unwrap() };
        assert!(nb(set.contains(&v(0x20)))?);
        assert!(nb(set.contains(&v(0x50)))?);
        assert!(!nb(set.contains(&v(0x30)))?);

        set.hints().add_flags(Flags::ID_DESC);
        assert!(nb(set.contains(&v(0x30)))?);
        assert!(!nb(set.contains(&v(0x70)))?);

        Ok(())
    }

    #[test]
    fn test_debug() {
        let set = lazy_set(&[0]);
        assert_eq!(format!("{:?}", set), "<lazy [] + ? more>");
        nb(set.count()).unwrap();
        assert_eq!(format!("{:?}", set), "<lazy [0000000000000000+0]>");

        let set = lazy_set(&[1, 3, 2]);
        assert_eq!(format!("{:?}", &set), "<lazy [] + ? more>");
        let mut iter = ni(set.iter()).unwrap();
        iter.next();
        assert_eq!(
            format!("{:?}", &set),
            "<lazy [0100000000000000+1] + ? more>"
        );
        iter.next();
        assert_eq!(
            format!("{:?}", &set),
            "<lazy [0100000000000000+1, 0300000000000000+3] + ? more>"
        );
        iter.next();
        assert_eq!(format!("{:2.2?}", &set), "<lazy [01+1, 03+3]+ 1 + ? more>");
        iter.next();
        assert_eq!(format!("{:1.3?}", &set), "<lazy [010+1] + 2 more>");
    }

    #[test]
    fn test_flatten() {
        let set1 = lazy_set(&[3, 2, 4]);
        let set2 = lazy_set_inherit(&[3, 7, 6], &set1);
        let set1 = NameSet::from_query(set1);
        let set2 = NameSet::from_query(set2);

        // Show flatten by names, and flatten by ids.
        // The first should be <static ...>, the second should be <spans ...>.
        let show = |set: NameSet| {
            [
                format!("{:5.2?}", r(set.flatten_names()).unwrap()),
                format!("{:5.2?}", r(set.flatten()).unwrap()),
            ]
        };

        assert_eq!(
            show(set1.clone() | set2.clone()),
            [
                "<static [03, 02, 04, 07, 06]>",
                "<spans [06:07+6:7, 02:04+2:4]>"
            ]
        );
        assert_eq!(
            show(set1.clone() & set2.clone()),
            ["<static [03]>", "<spans [03+3]>"]
        );
        assert_eq!(
            show(set1.clone() - set2.clone()),
            ["<static [02, 04]>", "<spans [04+4, 02+2]>"]
        );
    }

    quickcheck::quickcheck! {
        fn test_id_lazy_quickcheck(a: Vec<u64>) -> bool {
            let set = lazy_set(&a);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap();
            assert!(count <= a.len());

            let set2: HashSet<_> = a.iter().cloned().collect();
            assert_eq!(count, set2.len());

            true
        }
    }
}


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

use nonblocking::non_blocking_result;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use crate::ops::DagAlgorithm;
use crate::ops::IdConvert;
use crate::protocol::disable_remote_protocol;
use crate::Group;
use crate::IdSet;
use crate::IdSetIter;
use crate::IdSpan;
use crate::Result;
use crate::VertexName;

/// A set backed by [`IdSet`] + [`IdMap`].
/// Efficient for DAG calculation.
pub struct IdStaticSet {
    pub(crate) spans: IdSet,
    pub(crate) map: Arc<dyn IdConvert + Send + Sync>,
    pub(crate) dag: Arc<dyn DagAlgorithm + Send + Sync>,
    hints: Hints,
}

struct Iter {
    iter: IdSetIter<IdSet>,
    map: Arc<dyn IdConvert + Send + Sync>,
    reversed: bool,
    buf: Vec<Result<VertexName>>,
}

impl Iter {
    fn into_box_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |this| this.next()))
    }

    async fn next(mut self) -> Option<(Result<VertexName>, Self)> {
        if let Some(name) = self.buf.pop() {
            return Some((name, self));
        }
        let map = &self.map;
        let opt_id = if self.reversed {
            self.iter.next_back()
        } else {
            self.iter.next()
        };
        match opt_id {
            None => None,
            Some(id) => {
                let contains = map
                    .contains_vertex_id_locally(&[id])
                    .await
                    .unwrap_or_default();
                if contains == &[true] {
                    Some((map.vertex_name(id).await, self))
                } else {
                    // On demand prefetch in batch.
                    let batch_size = 131072;
                    let mut ids = Vec::with_capacity(batch_size);
                    ids.push(id);
                    for _ in ids.len()..batch_size {
                        if let Some(id) = if self.reversed {
                            self.iter.next_back()
                        } else {
                            self.iter.next()
                        } {
                            ids.push(id);
                        } else {
                            break;
                        }
                    }
                    ids.reverse();
                    self.buf = match self.map.vertex_name_batch(&ids).await {
                        Err(e) => return Some((Err(e), self)),
                        Ok(names) => names,
                    };
                    if self.buf.len() != ids.len() {
                        let result =
                            crate::errors::bug("vertex_name_batch does not return enough items");
                        return Some((result, self));
                    }
                    let name = self.buf.pop().expect("buf is not empty");
                    Some((name, self))
                }
            }
        }
    }
}

struct DebugSpan {
    span: IdSpan,
    low_name: Option<VertexName>,
    high_name: Option<VertexName>,
}

impl fmt::Debug for DebugSpan {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match (
            self.span.low == self.span.high,
            &self.low_name,
            &self.high_name,
        ) {
            (true, Some(name), _) => {
                fmt::Debug::fmt(&name, f)?;
                write!(f, "+{:?}", self.span.low)?;
            }
            (true, None, _) => {
                write!(f, "{:?}", self.span.low)?;
            }
            (false, Some(low), Some(high)) => {
                fmt::Debug::fmt(&low, f)?;
                write!(f, ":")?;
                fmt::Debug::fmt(&high, f)?;
                write!(f, "+{:?}:{:?}", self.span.low, self.span.high)?;
            }
            (false, _, _) => {
                write!(f, "{:?}:{:?}", self.span.low, self.span.high)?;
            }
        }
        Ok(())
    }
}

impl fmt::Debug for IdStaticSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<spans ")?;
        let spans = self.spans.as_spans();
        let limit = f.width().unwrap_or(3);
        f.debug_list()
            .entries(spans.iter().take(limit).map(|span| DebugSpan {
                span: *span,
                low_name: disable_remote_protocol(|| {
                    non_blocking_result(self.map.vertex_name(span.low)).ok()
                }),
                high_name: disable_remote_protocol(|| {
                    non_blocking_result(self.map.vertex_name(span.high)).ok()
                }),
            }))
            .finish()?;
        match spans.len().max(limit) - limit {
            0 => {}
            1 => write!(f, " + 1 span")?,
            n => write!(f, " + {} spans", n)?,
        }
        write!(f, ">")?;
        Ok(())
    }
}

impl IdStaticSet {
    pub(crate) fn from_spans_idmap_dag(
        spans: IdSet,
        map: Arc<dyn IdConvert + Send + Sync>,
        dag: Arc<dyn DagAlgorithm + Send + Sync>,
    ) -> Self {
        let hints = Hints::new_with_idmap_dag(map.clone(), dag.clone());
        hints.add_flags(Flags::ID_DESC | Flags::TOPO_DESC);
        if spans.is_empty() {
            hints.add_flags(Flags::EMPTY);
        } else {
            hints.set_min_id(spans.min().unwrap());
            hints.set_max_id(spans.max().unwrap());
        }
        Self {
            spans,
            map,
            hints,
            dag,
        }
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for IdStaticSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let iter = Iter {
            iter: self.spans.clone().into_iter(),
            map: self.map.clone(),
            reversed: false,
            buf: Default::default(),
        };
        Ok(iter.into_box_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let iter = Iter {
            iter: self.spans.clone().into_iter(),
            map: self.map.clone(),
            reversed: true,
            buf: Default::default(),
        };
        Ok(iter.into_box_stream())
    }

    async fn count(&self) -> Result<usize> {
        Ok(self.spans.count() as usize)
    }

    async fn first(&self) -> Result<Option<VertexName>> {
        debug_assert_eq!(self.spans.max(), self.spans.iter_desc().nth(0));
        match self.spans.max() {
            Some(id) => {
                let map = &self.map;
                let name = map.vertex_name(id).await?;
                Ok(Some(name))
            }
            None => Ok(None),
        }
    }

    async fn last(&self) -> Result<Option<VertexName>> {
        debug_assert_eq!(self.spans.min(), self.spans.iter_desc().rev().nth(0));
        match self.spans.min() {
            Some(id) => {
                let map = &self.map;
                let name = map.vertex_name(id).await?;
                Ok(Some(name))
            }
            None => Ok(None),
        }
    }

    async fn is_empty(&self) -> Result<bool> {
        Ok(self.spans.is_empty())
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        let result = match self
            .map
            .vertex_id_with_max_group(name, Group::NON_MASTER)
            .await?
        {
            Some(id) => self.spans.contains(id),
            None => false,
        };
        Ok(result)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        self.contains(name).await.map(Some)
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }

    fn id_convert(&self) -> Option<&dyn IdConvert> {
        Some(self.map.as_ref() as &dyn IdConvert)
    }
}

#[cfg(test)]
#[allow(clippy::redundant_clone)]
pub(crate) mod tests {
    use std::ops::Deref;

    use nonblocking::non_blocking_result as r;

    use super::super::tests::*;
    use super::super::NameSet;
    use super::*;
    use crate::tests::build_segments;
    use crate::DagAlgorithm;
    use crate::NameDag;

    /// Test with a predefined DAG.
    pub(crate) fn with_dag<R, F: Fn(&NameDag) -> R>(func: F) -> R {
        let built = build_segments(
            r#"
            A--B--C--D
                \--E--F--G"#,
            "D G",
            2,
        );
        //  0--1--2--3
        //      \--4--5--6
        func(&built.name_dag)
    }

    #[test]
    fn test_dag_invariants() -> Result<()> {
        with_dag(|dag| {
            let bef = r(dag.range("B".into(), "F".into()))?;
            check_invariants(bef.deref())?;

            Ok(())
        })
    }

    #[test]
    fn test_dag_fast_paths() -> Result<()> {
        with_dag(|dag| {
            let abcd = r(dag.ancestors("D".into()))?;
            let abefg = r(dag.ancestors("G".into()))?;

            let ab = abcd.intersection(&abefg);
            check_invariants(ab.deref())?;

            assert!(nb(abcd.contains(&vec![b'A'].into()))?);
            assert!(!nb(abcd.contains(&vec![b'E'].into()))?);

            // should not be "<and <...> <...>>"
            assert_eq!(format!("{:?}", &ab), "<spans [A:B+0:1]>");

            let abcdefg = abcd.union(&abefg);
            check_invariants(abcd.deref())?;
            // should not be "<or <...> <...>>"
            assert_eq!(format!("{:?}", &abcdefg), "<spans [A:G+0:6]>");

            let cd = abcd.difference(&abefg);
            check_invariants(cd.deref())?;
            // should not be "<difference <...> <...>>"
            assert_eq!(format!("{:?}", &cd), "<spans [C:D+2:3]>");

            Ok(())
        })
    }

    #[test]
    fn test_dag_no_fast_paths() -> Result<()> {
        let f = |s: NameSet| -> String { format!("{:?}", s) };
        with_dag(|dag1| -> Result<()> {
            with_dag(|dag2| -> Result<()> {
                let abcd = r(dag1.ancestors("D".into()))?;
                let abefg = r(dag2.ancestors("G".into()))?;

                // Since abcd and abefg are from 2 "separate" Dags, fast paths should not
                // be used for intersection, union, and difference.

                let ab = abcd.intersection(&abefg);
                check_invariants(ab.deref())?;
                // should not be "<spans ...>"
                assert_eq!(
                    format!("{:?}", &ab),
                    "<and <spans [A:D+0:3]> <spans [E:G+4:6, A:B+0:1]>>"
                );

                let abcdefg = abcd.union(&abefg);
                check_invariants(abcd.deref())?;
                // should not be "<spans ...>"
                assert_eq!(
                    format!("{:?}", &abcdefg),
                    "<or <spans [A:D+0:3]> <spans [E:G+4:6, A:B+0:1]>>"
                );

                let cd = abcd.difference(&abefg);
                check_invariants(cd.deref())?;
                // should not be "<spans ...>"
                assert_eq!(
                    format!("{:?}", &cd),
                    "<diff <spans [A:D+0:3]> <spans [E:G+4:6, A:B+0:1]>>"
                );

                // Should not use FULL hint fast paths for "&, |, -" operations, because
                // dag1 and dag2 are not considered compatible.
                let a1 = || r(dag1.all()).unwrap();
                let a2 = || r(dag2.all()).unwrap();
                assert_eq!(f(a1() & a2()), "<and <spans [A:G+0:6]> <spans [A:G+0:6]>>");
                assert_eq!(f(a1() | a2()), "<or <spans [A:G+0:6]> <spans [A:G+0:6]>>");
                assert_eq!(f(a1() - a2()), "<diff <spans [A:G+0:6]> <spans [A:G+0:6]>>");

                // No fast path for manually constructed StaticSet either, because
                // the StaticSets do not have DAG associated to test compatibility.
                // However, "all & z" is changed to "z & all" for performance.
                let z = || NameSet::from("Z");
                assert_eq!(f(z() & a2()), "<and <static [Z]> <spans [A:G+0:6]>>");
                assert_eq!(f(z() | a2()), "<or <static [Z]> <spans [A:G+0:6]>>");
                assert_eq!(f(z() - a2()), "<diff <static [Z]> <spans [A:G+0:6]>>");
                assert_eq!(f(a1() & z()), "<and <static [Z]> <spans [A:G+0:6]>>");
                assert_eq!(f(a1() | z()), "<or <spans [A:G+0:6]> <static [Z]>>");
                assert_eq!(f(a1() - z()), "<diff <spans [A:G+0:6]> <static [Z]>>");

                // EMPTY fast paths can still be used.
                let e = || NameSet::empty();
                assert_eq!(f(e() & a1()), "<empty>");
                assert_eq!(f(e() | a1()), "<spans [A:G+0:6]>");
                assert_eq!(f(e() - a1()), "<empty>");
                assert_eq!(f(a1() & e()), "<empty>");
                assert_eq!(f(a1() | e()), "<spans [A:G+0:6]>");
                assert_eq!(f(a1() - e()), "<spans [A:G+0:6]>");

                Ok(())
            })
        })
    }

    #[test]
    fn test_dag_all() -> Result<()> {
        with_dag(|dag| {
            let all = r(dag.all())?;
            assert_eq!(format!("{:?}", &all), "<spans [A:G+0:6]>");

            let ac: NameSet = "A C".into();
            let ac = r(dag.sort(&ac))?;

            let intersection = all.intersection(&ac);
            // should not be "<and ...>"
            assert_eq!(format!("{:?}", &intersection), "<spans [C+2, A+0]>");
            Ok(())
        })
    }

    #[test]
    fn test_sort() -> Result<()> {
        with_dag(|dag| -> Result<()> {
            let set = "G C A E".into();
            let sorted = r(dag.sort(&set))?;
            assert_eq!(format!("{:?}", &sorted), "<spans [G+6, E+4, C+2] + 1 span>");
            Ok(())
        })
    }

    #[test]
    fn test_dag_hints_ancestors() -> Result<()> {
        with_dag(|dag| -> Result<()> {
            let abc = r(dag.ancestors("B C".into()))?;
            let abe = r(dag.common_ancestors("E".into()))?;
            let f: NameSet = "F".into();
            let all = r(dag.all())?;

            assert!(has_ancestors_flag(abc.clone()));
            assert!(has_ancestors_flag(abe.clone()));
            assert!(has_ancestors_flag(all.clone()));
            assert!(has_ancestors_flag(r(dag.roots(abc.clone()))?));
            assert!(has_ancestors_flag(r(dag.parents(all.clone()))?));

            assert!(!has_ancestors_flag(f.clone()));
            assert!(!has_ancestors_flag(r(dag.roots(f.clone()))?));
            assert!(!has_ancestors_flag(r(dag.parents(f.clone()))?));

            Ok(())
        })
    }

    #[test]
    fn test_dag_hints_ancestors_inheritance() -> Result<()> {
        with_dag(|dag1| -> Result<()> {
            with_dag(|dag2| -> Result<()> {
                let abc = r(dag1.ancestors("B C".into()))?;

                // The ANCESTORS flag is kept by 'sort', 'parents', 'roots' on
                // the same dag.
                assert!(has_ancestors_flag(r(dag1.sort(&abc))?));
                assert!(has_ancestors_flag(r(dag1.parents(abc.clone()))?));
                assert!(has_ancestors_flag(r(dag1.roots(abc.clone()))?));

                // The ANCESTORS flag is removed on a different dag, since the
                // different dag does not assume same graph / ancestry
                // relationship.
                assert!(!has_ancestors_flag(r(dag2.sort(&abc))?));
                assert!(!has_ancestors_flag(r(dag2.parents(abc.clone()))?));
                assert!(!has_ancestors_flag(r(dag2.roots(abc.clone()))?));

                Ok(())
            })
        })
    }

    #[test]
    fn test_dag_hints_ancestors_fast_paths() -> Result<()> {
        with_dag(|dag| -> Result<()> {
            let bfg: NameSet = "B F G".into();

            // Set the ANCESTORS flag. It's incorrect but make it easier to test fast paths.
            bfg.hints().add_flags(Flags::ANCESTORS);

            // Fast paths are not used if the set is not "bound" to the dag.
            assert_eq!(
                format!("{:?}", r(dag.ancestors(bfg.clone()))?),
                "<static [B, F, G]>"
            );
            assert_eq!(format!("{:?}", r(dag.heads(bfg.clone()))?), "<spans [G+6]>");

            // Binding to the Dag enables fast paths.
            let bfg = r(dag.sort(&bfg))?;
            bfg.hints().add_flags(Flags::ANCESTORS);
            assert_eq!(
                format!("{:?}", r(dag.ancestors(bfg.clone()))?),
                "<spans [F:G+5:6, B+1]>"
            );

            // 'heads' has a fast path that uses 'heads_ancestors' to do the calculation.
            // (in this case the result is incorrect because the hints are wrong).
            assert_eq!(format!("{:?}", r(dag.heads(bfg.clone()))?), "<spans [G+6]>");

            // 'ancestors' has a fast path that returns set as-is.
            // (in this case the result is incorrect because the hints are wrong).
            assert_eq!(
                format!("{:?}", r(dag.ancestors(bfg.clone()))?),
                "<spans [F:G+5:6, B+1]>"
            );

            Ok(())
        })
    }

    fn has_ancestors_flag(set: NameSet) -> bool {
        set.hints().contains(Flags::ANCESTORS)
    }
}


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

use futures::StreamExt;

use super::hints::Flags;
use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use super::NameSet;
use crate::fmt::write_debug;
use crate::Id;
use crate::Result;
use crate::VertexName;

/// Intersection of 2 sets.
///
/// The iteration order is defined by the first set.
pub struct IntersectionSet {
    lhs: NameSet,
    rhs: NameSet,
    hints: Hints,
}

struct Iter {
    iter: BoxVertexStream,
    rhs: NameSet,
    ended: bool,

    /// Optional fast path for stop.
    stop_condition: Option<StopCondition>,
}

impl Iter {
    async fn next(&mut self) -> Option<Result<VertexName>> {
        if self.ended {
            return None;
        }
        loop {
            let result = self.iter.as_mut().next().await;
            if let Some(Ok(ref name)) = result {
                match self.rhs.contains(&name).await {
                    Err(err) => break Some(Err(err)),
                    Ok(false) => {
                        // Check if we can stop iteration early using hints.
                        if let Some(ref cond) = self.stop_condition {
                            if let Some(id_convert) = self.rhs.id_convert() {
                                if let Ok(Some(id)) = id_convert.vertex_id_optional(&name).await {
                                    if cond.should_stop_with_id(id) {
                                        self.ended = true;
                                        return None;
                                    }
                                }
                            }
                        }
                        continue;
                    }
                    Ok(true) => {}
                }
            }
            break result;
        }
    }

    fn into_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |mut state| async move {
            let result = state.next().await;
            result.map(|r| (r, state))
        }))
    }
}

struct StopCondition {
    order: Ordering,
    id: Id,
}

impl StopCondition {
    fn should_stop_with_id(&self, id: Id) -> bool {
        id.cmp(&self.id) == self.order
    }
}

impl IntersectionSet {
    pub fn new(lhs: NameSet, rhs: NameSet) -> Self {
        // More efficient if `lhs` is smaller. Swap `lhs` and `rhs` if `lhs` is `FULL`.
        let (lhs, rhs) = if lhs.hints().contains(Flags::FULL)
            && !rhs.hints().contains(Flags::FULL)
            && !rhs.hints().contains(Flags::FILTER)
            && lhs.hints().dag_version() >= rhs.hints().dag_version()
        {
            (rhs, lhs)
        } else {
            (lhs, rhs)
        };

        let hints = Hints::new_inherit_idmap_dag(lhs.hints());
        hints.add_flags(
            lhs.hints().flags()
                & (Flags::EMPTY
                    | Flags::ID_DESC
                    | Flags::ID_ASC
                    | Flags::TOPO_DESC
                    | Flags::FILTER),
        );
        // Only keep the ANCESTORS flag if lhs and rhs use a compatible Dag.
        if lhs.hints().dag_version() >= rhs.hints().dag_version() {
            hints.add_flags(lhs.hints().flags() & rhs.hints().flags() & Flags::ANCESTORS);
        }
        let (rhs_min_id, rhs_max_id) = if hints.id_map_version() >= rhs.hints().id_map_version() {
            // rhs ids are all known by lhs.
            (rhs.hints().min_id(), rhs.hints().max_id())
        } else {
            (None, None)
        };
        match (lhs.hints().min_id(), rhs_min_id) {
            (Some(id), None) | (None, Some(id)) => {
                hints.set_min_id(id);
            }
            (Some(id1), Some(id2)) => {
                hints.set_min_id(id1.max(id2));
            }
            (None, None) => {}
        }
        match (lhs.hints().max_id(), rhs_max_id) {
            (Some(id), None) | (None, Some(id)) => {
                hints.set_max_id(id);
            }
            (Some(id1), Some(id2)) => {
                hints.set_max_id(id1.min(id2));
            }
            (None, None) => {}
        }
        Self { lhs, rhs, hints }
    }

    fn is_rhs_id_map_comapatible(&self) -> bool {
        let lhs_version = self.lhs.hints().id_map_version();
        let rhs_version = self.rhs.hints().id_map_version();
        lhs_version == rhs_version || (lhs_version > rhs_version && rhs_version > None)
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for IntersectionSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let stop_condition = if !self.is_rhs_id_map_comapatible() {
            None
        } else if self.lhs.hints().contains(Flags::ID_ASC) {
            if let Some(id) = self.rhs.hints().max_id() {
                Some(StopCondition {
                    id,
                    order: Ordering::Greater,
                })
            } else {
                None
            }
        } else if self.lhs.hints().contains(Flags::ID_DESC) {
            if let Some(id) = self.rhs.hints().min_id() {
                Some(StopCondition {
                    id,
                    order: Ordering::Less,
                })
            } else {
                None
            }
        } else {
            None
        };

        let iter = Iter {
            iter: self.lhs.iter().await?,
            rhs: self.rhs.clone(),
            ended: false,
            stop_condition,
        };
        Ok(iter.into_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let stop_condition = if !self.is_rhs_id_map_comapatible() {
            None
        } else if self.lhs.hints().contains(Flags::ID_DESC) {
            if let Some(id) = self.rhs.hints().max_id() {
                Some(StopCondition {
                    id,
                    order: Ordering::Greater,
                })
            } else {
                None
            }
        } else if self.lhs.hints().contains(Flags::ID_ASC) {
            if let Some(id) = self.rhs.hints().min_id() {
                Some(StopCondition {
                    id,
                    order: Ordering::Less,
                })
            } else {
                None
            }
        } else {
            None
        };

        let iter = Iter {
            iter: self.lhs.iter_rev().await?,
            rhs: self.rhs.clone(),
            ended: false,
            stop_condition,
        };
        Ok(iter.into_stream())
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        Ok(self.lhs.contains(name).await? && self.rhs.contains(name).await?)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        for set in &[&self.lhs, &self.rhs] {
            let contains = set.contains_fast(name).await?;
            match contains {
                Some(false) | None => return Ok(contains),
                Some(true) => {}
            }
        }
        Ok(Some(true))
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

impl fmt::Debug for IntersectionSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<and")?;
        write_debug(f, &self.lhs)?;
        write_debug(f, &self.rhs)?;
        write!(f, ">")
    }
}

#[cfg(test)]
#[allow(clippy::redundant_clone)]
mod tests {
    use std::collections::HashSet;

    use super::super::id_lazy::test_utils::lazy_set;
    use super::super::id_lazy::test_utils::lazy_set_inherit;
    use super::super::tests::*;
    use super::*;
    use crate::Id;

    fn intersection(a: &[u8], b: &[u8]) -> IntersectionSet {
        let a = NameSet::from_query(VecQuery::from_bytes(a));
        let b = NameSet::from_query(VecQuery::from_bytes(b));
        IntersectionSet::new(a, b)
    }

    #[test]
    fn test_intersection_basic() -> Result<()> {
        let set = intersection(b"\x11\x33\x55\x22\x44", b"\x44\x33\x66");
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["33", "44"]);
        assert_eq!(shorten_iter(ni(set.iter_rev())), ["44", "33"]);
        assert!(!nb(set.is_empty())?);
        assert_eq!(nb(set.count())?, 2);
        assert_eq!(shorten_name(nb(set.first())?.unwrap()), "33");
        assert_eq!(shorten_name(nb(set.last())?.unwrap()), "44");
        for &b in b"\x11\x22\x55\x66".iter() {
            assert!(!nb(set.contains(&to_name(b)))?);
        }
        Ok(())
    }

    #[test]
    fn test_intersection_min_max_id_fast_path() {
        // The min_ids are intentionally wrong to test the fast paths.
        let a = lazy_set(&[0x70, 0x60, 0x50, 0x40, 0x30, 0x20]);
        let b = lazy_set_inherit(&[0x70, 0x65, 0x50, 0x40, 0x35, 0x20], &a);
        let a = NameSet::from_query(a);
        let b = NameSet::from_query(b);
        a.hints().add_flags(Flags::ID_DESC);
        b.hints().set_min_id(Id(0x40));
        b.hints().set_max_id(Id(0x50));

        let set = IntersectionSet::new(a, b.clone());
        // No "20" - filtered out by min id fast path.
        assert_eq!(shorten_iter(ni(set.iter())), ["70", "50", "40"]);
        // No "70" - filtered out by max id fast path.
        assert_eq!(shorten_iter(ni(set.iter_rev())), ["20", "40", "50"]);

        // Test the reversed sort order.
        let a = lazy_set(&[0x20, 0x30, 0x40, 0x50, 0x60, 0x70]);
        let b = lazy_set_inherit(&[0x70, 0x65, 0x50, 0x40, 0x35, 0x20], &a);
        let a = NameSet::from_query(a);
        let b = NameSet::from_query(b);
        a.hints().add_flags(Flags::ID_ASC);
        b.hints().set_min_id(Id(0x40));
        b.hints().set_max_id(Id(0x50));
        let set = IntersectionSet::new(a, b.clone());
        // No "70".
        assert_eq!(shorten_iter(ni(set.iter())), ["20", "40", "50"]);
        // No "20".
        assert_eq!(shorten_iter(ni(set.iter_rev())), ["70", "50", "40"]);

        // If two sets have incompatible IdMap, fast paths are not used.
        let a = NameSet::from_query(lazy_set(&[0x20, 0x30, 0x40, 0x50, 0x60, 0x70]));
        a.hints().add_flags(Flags::ID_ASC);
        let set = IntersectionSet::new(a, b.clone());
        // Should contain "70" and "20".
        assert_eq!(shorten_iter(ni(set.iter())), ["20", "40", "50", "70"]);
        assert_eq!(shorten_iter(ni(set.iter_rev())), ["70", "50", "40", "20"]);
    }

    quickcheck::quickcheck! {
        fn test_intersection_quickcheck(a: Vec<u8>, b: Vec<u8>) -> bool {
            let set = intersection(&a, &b);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap();
            assert!(count <= a.len(), "len({:?}) = {} should <= len({:?})" , &set, count, &a);
            assert!(count <= b.len(), "len({:?}) = {} should <= len({:?})" , &set, count, &b);

            let contains_a: HashSet<u8> = a.into_iter().filter(|&b| nb(set.contains(&to_name(b))).ok() == Some(true)).collect();
            let contains_b: HashSet<u8> = b.into_iter().filter(|&b| nb(set.contains(&to_name(b))).ok() == Some(true)).collect();
            assert_eq!(contains_a, contains_b);

            true
        }
    }
}


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

use futures::lock::Mutex;
use futures::lock::MutexGuard;
use futures::StreamExt;
use indexmap::IndexSet;

use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use crate::Result;
use crate::VertexName;

/// A set backed by a lazy iterator of names.
pub struct LazySet {
    inner: Arc<Mutex<Inner>>,
    hints: Hints,
}

struct Inner {
    iter: BoxVertexStream,
    visited: IndexSet<VertexName>,
    state: State,
}

impl Inner {
    async fn load_more(&mut self, n: usize, mut out: Option<&mut Vec<VertexName>>) -> Result<()> {
        if matches!(self.state, State::Complete | State::Error) {
            return Ok(());
        }
        for _ in 0..n {
            match self.iter.next().await {
                Some(Ok(name)) => {
                    if let Some(ref mut out) = out {
                        out.push(name.clone());
                    }
                    self.visited.insert(name);
                }
                None => {
                    self.state = State::Complete;
                    break;
                }
                Some(Err(err)) => {
                    self.state = State::Error;
                    return Err(err);
                }
            }
        }
        Ok(())
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum State {
    Incomplete,
    Complete,
    Error,
}

pub struct Iter {
    inner: Arc<Mutex<Inner>>,
    index: usize,
}

impl Iter {
    async fn next(&mut self) -> Option<Result<VertexName>> {
        loop {
            let mut inner = self.inner.lock().await;
            match inner.state {
                State::Error => break None,
                State::Complete if inner.visited.len() <= self.index => break None,
                State::Complete | State::Incomplete => {
                    let value = inner.visited.get_index(self.index).cloned();
                    match value {
                        Some(value) => {
                            self.index += 1;
                            break Some(Ok(value));
                        }
                        None => {
                            // Data not available. Load more.
                            if let Err(err) = inner.load_more(1, None).await {
                                return Some(Err(err));
                            }
                            continue;
                        }
                    }
                }
            }
        }
    }

    fn into_stream(self) -> BoxVertexStream {
        Box::pin(futures::stream::unfold(self, |mut state| async move {
            let result = state.next().await;
            result.map(|r| (r, state))
        }))
    }
}

impl fmt::Debug for LazySet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("<lazy ")?;
        if let Some(inner) = self.inner.try_lock() {
            let limit = f.width().unwrap_or(3);
            f.debug_list()
                .entries(inner.visited.iter().take(limit))
                .finish()?;
            let remaining = inner.visited.len().max(limit) - limit;
            match (remaining, inner.state) {
                (0, State::Incomplete) => f.write_str(" + ? more")?,
                (n, State::Incomplete) => write!(f, "+ {} + ? more", n)?,
                (0, _) => {}
                (n, _) => write!(f, " + {} more", n)?,
            }
        } else {
            f.write_str(" ?")?;
        }
        f.write_str(">")?;
        Ok(())
    }
}

impl LazySet {
    pub fn from_iter<I>(names: I, hints: Hints) -> Self
    where
        I: IntoIterator<Item = Result<VertexName>> + 'static,
        <I as IntoIterator>::IntoIter: Send + Sync,
    {
        let stream = futures::stream::iter(names);
        Self::from_stream(Box::pin(stream), hints)
    }

    pub fn from_stream(names: BoxVertexStream, hints: Hints) -> Self {
        let inner = Inner {
            iter: names,
            visited: IndexSet::new(),
            state: State::Incomplete,
        };
        Self {
            inner: Arc::new(Mutex::new(inner)),
            hints,
        }
    }

    async fn load_all(&self) -> Result<MutexGuard<'_, Inner>> {
        let mut inner = self.inner.lock().await;
        inner.load_more(usize::max_value(), None).await?;
        Ok(inner)
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for LazySet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        let inner = self.inner.clone();
        let iter = Iter { inner, index: 0 };
        Ok(iter.into_stream())
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        let inner = self.load_all().await?;
        let iter = inner.visited.clone().into_iter().rev().map(Ok);
        let stream = futures::stream::iter(iter);
        Ok(Box::pin(stream))
    }

    async fn count(&self) -> Result<usize> {
        let inner = self.load_all().await?;
        Ok(inner.visited.len())
    }

    async fn last(&self) -> Result<Option<VertexName>> {
        let inner = self.load_all().await?;
        Ok(inner.visited.iter().rev().nth(0).cloned())
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        let mut inner = self.inner.lock().await;
        if inner.visited.contains(name) {
            return Ok(true);
        } else {
            let mut loaded = Vec::new();
            loop {
                loaded.clear();
                inner.load_more(1, Some(&mut loaded)).await?;
                debug_assert!(loaded.len() <= 1);
                if loaded.is_empty() {
                    break;
                }
                if loaded.first() == Some(name) {
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        let inner = self.inner.lock().await;
        if inner.visited.contains(name) {
            return Ok(Some(true));
        } else if inner.state != State::Incomplete {
            return Ok(Some(false));
        }
        Ok(None)
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

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

    use super::super::tests::*;
    use super::*;

    fn lazy_set(a: &[u8]) -> LazySet {
        LazySet::from_iter(
            a.to_vec().into_iter().map(|b| Ok(to_name(b))),
            Hints::default(),
        )
    }

    #[test]
    fn test_lazy_basic() -> Result<()> {
        let set = lazy_set(b"\x11\x33\x22\x77\x22\x55\x11");
        check_invariants(&set)?;
        assert_eq!(shorten_iter(ni(set.iter())), ["11", "33", "22", "77", "55"]);
        assert_eq!(
            shorten_iter(ni(set.iter_rev())),
            ["55", "77", "22", "33", "11"]
        );
        assert!(!nb(set.is_empty())?);
        assert_eq!(nb(set.count())?, 5);
        assert_eq!(shorten_name(nb(set.first())?.unwrap()), "11");
        assert_eq!(shorten_name(nb(set.last())?.unwrap()), "55");
        Ok(())
    }

    #[test]
    fn test_debug() {
        let set = lazy_set(b"");
        assert_eq!(format!("{:?}", &set), "<lazy [] + ? more>");
        nb(set.count()).unwrap();
        assert_eq!(format!("{:?}", &set), "<lazy []>");

        let set = lazy_set(b"\x11\x33\x22");
        assert_eq!(format!("{:?}", &set), "<lazy [] + ? more>");
        let mut iter = ni(set.iter()).unwrap();
        iter.next();
        assert_eq!(format!("{:?}", &set), "<lazy [1111] + ? more>");
        iter.next();
        assert_eq!(format!("{:?}", &set), "<lazy [1111, 3333] + ? more>");
        iter.next();
        assert_eq!(format!("{:2.2?}", &set), "<lazy [11, 33]+ 1 + ? more>");
        iter.next();
        assert_eq!(format!("{:1.3?}", &set), "<lazy [111] + 2 more>");
    }

    quickcheck::quickcheck! {
        fn test_lazy_quickcheck(a: Vec<u8>) -> bool {
            let set = lazy_set(&a);
            check_invariants(&set).unwrap();

            let count = nb(set.count()).unwrap();
            assert!(count <= a.len());

            let set2: HashSet<_> = a.iter().cloned().collect();
            assert_eq!(count, set2.len());

            true
        }
    }
}


/*
 * 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 super::super::NameDag;
use super::IdStaticSet;
use super::NameSet;
use crate::IdSet;

/// A legacy token that enables conversion between IdSet (id-based)
/// and NameSet (hash-based). It should not be used for new Rust code.
#[derive(Copy, Clone)]
pub struct LegacyCodeNeedIdAccess;

// This is ideally not provided.  However revision numbers in revset still have
// large use-cases in Python and for now we provide this way to convert IdStaticSet
// to IdSet using "revision" numbers.
impl<'a> From<(LegacyCodeNeedIdAccess, &'a IdStaticSet)> for IdSet {
    fn from(value: (LegacyCodeNeedIdAccess, &'a IdStaticSet)) -> IdSet {
        let set = value.1;
        set.spans.clone()
    }
}

impl<'a> From<(LegacyCodeNeedIdAccess, IdSet, &'a NameDag)> for NameSet {
    fn from(value: (LegacyCodeNeedIdAccess, IdSet, &'a NameDag)) -> NameSet {
        NameSet::from_spans_dag(value.1, value.2).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use nonblocking::non_blocking_result as r;

    use super::super::id_static::tests::with_dag;
    use super::*;
    use crate::DagAlgorithm;
    use crate::Result;

    #[test]
    fn test_legacy_convert() -> Result<()> {
        use LegacyCodeNeedIdAccess as L;
        with_dag(|dag| -> Result<()> {
            let set1 = r(dag.ancestors("G".into()))?;
            let spans: IdSet = (
                L,
                set1.as_any().downcast_ref::<IdStaticSet>().unwrap().clone(),
            )
                .into();
            let set2: NameSet = (L, spans.clone(), dag).into();
            assert_eq!(format!("{:?}", &set1), "<spans [E:G+4:6, A:B+0:1]>");
            assert_eq!(format!("{:?}", &set2), "<spans [E:G+4:6, A:B+0:1]>");
            assert_eq!(format!("{:?}", &spans), "0 1 4 5 6");
            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.
 */

use std::any::Any;
use std::fmt;
use std::sync::RwLock;

use futures::future::BoxFuture;

use super::AsyncNameSetQuery;
use super::BoxVertexStream;
use super::Hints;
use super::NameSet;
use crate::Result;
use crate::VertexName;

/// A set that is lazily evaluated to another set on access, with an
/// optional fast path for "contains".
///
/// This can be useful for various cases. Especially when "contains" and "iter"
/// have separate fast paths. For example, `merge()` (merge commits) obviously
/// has a cheap "contains" fast path by checking if a given commit has
/// multiple parents. However, if we want to iterating through the set,
/// segmented changelog has a fast path by iterating flat segments, instead
/// of testing commits one by one using the "contains" check.
///
/// `MetaSet` is different from `LazySet`: `MetaSet`'s `evaluate` function can
/// return static or lazy or meta sets. `LazySet` does not support a "contains"
/// fast path (yet).
///
/// `MetaSet` is different from a pure filtering set (ex. only has "contains"
/// fast path), as `MetaSet` supports fast path for iteration.
pub struct MetaSet {
    evaluate: Box<dyn Fn() -> BoxFuture<'static, Result<NameSet>> + Send + Sync>,
    evaluated: RwLock<Option<NameSet>>,

    /// Optional "contains" fast path.
    contains: Option<
        Box<
            dyn for<'a> Fn(&'a MetaSet, &'a VertexName) -> BoxFuture<'a, Result<bool>>
                + Send
                + Sync,
        >,
    >,

    hints: Hints,
}

impl fmt::Debug for MetaSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("<meta ")?;
        if let Some(set) = self.evaluated() {
            set.fmt(f)?;
        } else {
            f.write_str("?")?;
        }
        f.write_str(">")?;
        Ok(())
    }
}

impl MetaSet {
    /// Constructs `MetaSet` from an `evaluate` function that returns a
    /// `NameSet`. The `evaluate` function is not called immediately.
    pub fn from_evaluate_hints(
        evaluate: Box<dyn Fn() -> BoxFuture<'static, Result<NameSet>> + Send + Sync + 'static>,
        hints: Hints,
    ) -> Self {
        Self {
            evaluate,
            evaluated: Default::default(),
            contains: None,
            hints,
        }
    }

    /// Provides a fast path for `contains`. Be careful to make sure "contains"
    /// matches "evaluate".
    pub fn with_contains(
        mut self,
        contains: Box<
            dyn for<'a> Fn(&'a MetaSet, &'a VertexName) -> BoxFuture<'a, Result<bool>>
                + Send
                + Sync,
        >,
    ) -> Self {
        self.contains = Some(contains);
        self
    }

    /// Evaluate the set. Returns a new set.
    pub async fn evaluate(&self) -> Result<NameSet> {
        if let Some(s) = &*self.evaluated.read().unwrap() {
            return Ok(s.clone());
        }
        let s = (self.evaluate)().await?;
        *self.evaluated.write().unwrap() = Some(s.clone());
        Ok(s)
    }

    /// Returns the evaluated set if it was evaluated.
    /// Returns None if the set was not evaluated.
    pub fn evaluated(&self) -> Option<NameSet> {
        self.evaluated.read().unwrap().clone()
    }
}

#[async_trait::async_trait]
impl AsyncNameSetQuery for MetaSet {
    async fn iter(&self) -> Result<BoxVertexStream> {
        self.evaluate().await?.iter().await
    }

    async fn iter_rev(&self) -> Result<BoxVertexStream> {
        self.evaluate().await?.iter_rev().await
    }

    async fn count(&self) -> Result<usize> {
        self.evaluate().await?.count().await
    }

    async fn last(&self) -> Result<Option<VertexName>> {
        self.evaluate().await?.last().await
    }

    async fn contains(&self, name: &VertexName) -> Result<bool> {
        match self.evaluated() {
            Some(set) => set.contains(name).await,
            None => match &self.contains {
                Some(f) => f(self, name).await,
                None => self.evaluate().await?.contains(name).await,
            },
        }
    }

    async fn contains_fast(&self, name: &VertexName) -> Result<Option<bool>> {
        match &self.contains {
            Some(f) => Ok(Some(f(self, name).await?)),
            None => Ok(None),
        }
    }

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

    fn hints(&self) -> &Hints {
        &self.hints
    }
}

#[cfg(test)]
mod tests {
    use nonblocking::non_blocking_result as r;

    use super::super::tests::*;
    use super::*;

    fn meta_set(v: &[impl ToString]) -> MetaSet {
        let v: Vec<_> = v.iter().map(|s| s.to_string()).collect();
        let f = move || -> BoxFuture<_> {
            let s = NameSet::from_static_names(v.clone().into_iter().map(Into::into));
            Box::pin(async move { Ok(s) })
        };
        MetaSet::from_evaluate_hints(Box::new(f), Hints::default())
    }

    #[test]
    fn test_meta_basic() -> Result<()> {
        let set = meta_set(&["1", "3", "2", "7", "5"]);

        assert!(set.evaluated().is_none());
        assert!(nb(set.contains(&"2".into()))?);
        // The set is evaluated after a "contains" check without a "contains" fast path.
        assert!(set.evaluated().is_some());

        check_invariants(&set)?;
        Ok(())
    }

    #[test]
    fn test_meta_contains() -> Result<()> {
        let set = meta_set(&["1", "3", "2", "7", "5"]).with_contains(Box::new(|_, v| {
            let r = Ok(v.as_ref().len() == 1 && b"12357".contains(&v.as_ref()[0]));
            Box::pin(async move { r })
        }));

        assert!(nb(set.contains(&"2".into()))?);
        assert!(!nb(set.contains(&"6".into()))?);
        // The set is not evaluated - contains fast path takes care of the checks.
        assert!(set.evaluated().is_none());

        check_invariants(&set)?;
        Ok(())
    }

    #[test]
    fn test_debug() {
        let set = meta_set(&["1", "3", "2", "7", "5"]);
        assert_eq!(format!("{:?}", &set), "<meta ?>");
        r(set.evaluate()).unwrap();
        assert_eq!(
            format!("{:5?}", &set),
            "<meta <static [31, 33, 32, 37, 35]>>"
        );
    }

    quickcheck::quickcheck! {
        fn test_meta_quickcheck(v: Vec<String>) -> bool {
            let set = meta_set(&v);
            check_invariants(&set).unwrap();
            true
        }
    }
}