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

//! # spanset
//!
//! See [`SpanSet`] for the main structure.

use std::cmp::Ordering;
use std::cmp::Ordering::Equal;
use std::cmp::Ordering::Greater;
use std::cmp::Ordering::Less;
use std::collections::BinaryHeap;
use std::collections::VecDeque;
use std::fmt;
use std::fmt::Debug;
use std::iter::Rev;
use std::ops::Bound;
use std::ops::RangeBounds;
use std::ops::RangeInclusive;

use dag_types::FlatSegment;
use serde::Deserialize;
use serde::Serialize;

use crate::bsearch::BinarySearchBy;
use crate::id::Id;

/// Range `low..=high`. `low` must be <= `high`.
#[derive(Copy, Clone, Debug, Eq, Serialize, Deserialize)]
pub struct Span {
    #[serde(with = "flat_id")]
    pub(crate) low: Id,
    #[serde(with = "flat_id")]
    pub(crate) high: Id,
}

/// A set of integer spans.
#[derive(Clone, Serialize, Deserialize, Default)]
pub struct SpanSet {
    /// `spans` are sorted in DESC order.
    spans: VecDeque<Span>,
}

impl PartialOrd for Span {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(match self.high.cmp(&other.high) {
            Less => Less,
            Greater => Greater,
            Equal => self.low.cmp(&other.low),
        })
    }
}

impl PartialEq for Span {
    fn eq(&self, other: &Self) -> bool {
        other.low == self.low && other.high == self.high
    }
}

impl Ord for Span {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.high.cmp(&other.high) {
            Less => Less,
            Greater => Greater,
            Equal => self.low.cmp(&other.low),
        }
    }
}

impl Span {
    pub fn new(low: Id, high: Id) -> Self {
        assert!(low <= high, "low {:?} <= high {:?}", low, high);
        Self { low, high }
    }

    pub fn count(self) -> u64 {
        self.high.0 - self.low.0 + 1
    }

    /// Get the n-th [`Id`] in this [`Span`].
    ///
    /// Similar to [`SpanSet`], ids are sorted in descending order.
    /// The 0-th Id is `high`.
    pub fn nth(self, n: u64) -> Option<Id> {
        if n >= self.count() {
            None
        } else {
            Some(self.high - n)
        }
    }

    pub(crate) fn contains(self, value: Id) -> bool {
        self.low <= value && value <= self.high
    }

    /// Construct a full [`Span`] that contains everything.
    /// Warning: The [`Id`] in this span might be unknown to an actual storage.
    pub fn full() -> Self {
        (Id::MIN..=Id::MAX).into()
    }

    pub(crate) fn try_from_bounds(bounds: impl RangeBounds<Id>) -> Option<Self> {
        use Bound::Excluded;
        use Bound::Included;
        #[cfg(debug_assertions)]
        {
            use Bound::Unbounded;
            match (bounds.start_bound(), bounds.end_bound()) {
                (Excluded(_), _) | (Unbounded, _) | (_, Unbounded) => {
                    panic!("unsupported bound type")
                }
                _ => {}
            }
        }
        match (bounds.start_bound(), bounds.end_bound()) {
            (Included(&low), Included(&high)) if low <= high => Some(Span { low, high }),
            (Included(&low), Excluded(&high_plus_one)) if low < high_plus_one => {
                let high = high_plus_one - 1;
                Some(Span { low, high })
            }
            _ => None,
        }
    }
}

/// Subspan is a trait for an object that
/// (a) Can be mapped into Span
/// (b) Can return 'subspan' for any given non-empty subset of it's span
pub trait Subspan {
    fn span(&self) -> Span;

    /// Provided span should be subset of T::span(), otherwise this method behavior is undefined
    fn subspan(&self, span: Span) -> Self;

    /// Overlaps two objects and returns result of overlap and remainders for left and right object
    /// This method is generally defined for any two types that implement Subspan
    /// Type of an overlap if the same as type of Self
    ///
    /// Returns:
    ///  - overlap: overlap between two objects
    ///  - rem_left: remaining non-overlapping part of left object
    ///  - rem_right: remaining non-overlapping part of right object
    ///
    /// L: [123456]
    /// R:    [456789]
    /// overlap(L, R) = ([456], [123], None)
    fn intersect<R: Subspan>(&self, r: &R) -> (Option<Self>, Option<Self>, Option<R>)
    where
        Self: Sized,
    {
        let left = self.span();
        let right = r.span();
        let span_low = left.low.max(right.low);
        let span_high = left.high.min(right.high);
        let overlap = Span::try_from_bounds(span_low..=span_high);

        let rem_left = Span::try_from_bounds(left.low..(left.high + 1).min(span_low));
        let rem_right = Span::try_from_bounds(right.low..(right.high + 1).min(span_low));

        let overlap = overlap.map(|s| self.subspan(s));
        let rem_left = rem_left.map(|s| self.subspan(s));
        let rem_right = rem_right.map(|s| r.subspan(s));
        (overlap, rem_left, rem_right)
    }
}

/// Calculates the intersection of two ordered iterators of span-like objects.
pub fn intersect_iter<
    L: Subspan,
    R: Subspan,
    LI: Iterator<Item = L>,
    RI: Iterator<Item = R>,
    P: FnMut(L),
>(
    mut lhs: LI,
    mut rhs: RI,
    mut push: P,
) {
    let mut next_left = lhs.next();
    let mut next_right = rhs.next();

    while let (Some(left), Some(right)) = (next_left, next_right) {
        // current:
        //   |------- A --------|
        //         |------- B ------|
        //         |--- span ---|
        // next:
        //   |- A -| (remaining part of A)
        //           (next B)
        // note: (A, B) can be either (left, right) or (right, left)
        let (span, rem_left, rem_right) = left.intersect(&right);

        if let Some(span) = span {
            push(span);
        }

        next_right = rem_right.or_else(|| rhs.next());
        next_left = rem_left.or_else(|| lhs.next());
    }
}

impl Subspan for Span {
    fn span(&self) -> Span {
        *self
    }

    fn subspan(&self, span: Span) -> Self {
        assert!(self.low <= span.low);
        assert!(self.high >= span.high);
        span
    }
}

impl Subspan for FlatSegment {
    fn span(&self) -> Span {
        Span::new(self.low, self.high)
    }

    fn subspan(&self, span: Span) -> Self {
        assert!(self.low <= span.low);
        assert!(self.high >= span.high);
        if span.low == self.low {
            FlatSegment {
                low: span.low,
                high: span.high,
                parents: self.parents.clone(),
            }
        } else {
            FlatSegment {
                low: span.low,
                high: span.high,
                parents: vec![span.low - 1],
            }
        }
    }
}

// This is for users who want shorter code than [`Span::new`].
// Internal logic here should use [`Span::new`], or [`Span::try_from_bounds`],
// or construct [`Span`] directly.
impl From<RangeInclusive<Id>> for Span {
    fn from(range: RangeInclusive<Id>) -> Span {
        Span::new(*range.start(), *range.end())
    }
}

impl From<Id> for Span {
    fn from(id: Id) -> Span {
        Span::new(id, id)
    }
}

impl<T: Into<Span>> From<T> for SpanSet {
    fn from(span: T) -> SpanSet {
        SpanSet::from_sorted_spans(std::iter::once(span.into()))
    }
}

impl From<Span> for RangeInclusive<Id> {
    fn from(span: Span) -> RangeInclusive<Id> {
        span.low..=span.high
    }
}

// This is used by `gca(set)` where `set` usually contains 2 ids. The code
// can then be written as `gca((a, b))`.
impl From<(Id, Id)> for SpanSet {
    fn from(ids: (Id, Id)) -> SpanSet {
        SpanSet::from_spans([ids.0, ids.1].iter().cloned())
    }
}

impl SpanSet {
    /// Construct a [`SpanSet`] containing given spans.
    /// Overlapped or adjacent spans will be merged automatically.
    pub fn from_spans<T: Into<Span>, I: IntoIterator<Item = T>>(spans: I) -> Self {
        let mut heap: BinaryHeap<Span> = spans.into_iter().map(|span| span.into()).collect();
        let mut spans = VecDeque::with_capacity(heap.len().min(64));
        while let Some(span) = heap.pop() {
            push_with_union(&mut spans, span);
        }
        let result = SpanSet { spans };
        // `result` should be valid because the use of `push_with_union`.
        #[cfg(debug_assertions)]
        result.validate();
        result
    }

    /// Construct a [`SpanSet`] containing given spans.
    /// The given spans must be already sorted (i.e. larger ids first), and do
    /// not have overlapped spans.
    /// Adjacent spans will be merged automatically.
    pub fn from_sorted_spans<T: Into<Span>, I: IntoIterator<Item = T>>(span_iter: I) -> Self {
        let mut spans = VecDeque::<Span>::new();
        for span in span_iter {
            let span = span.into();
            push_with_union(&mut spans, span);
        }
        let result = Self { spans };
        #[cfg(debug_assertions)]
        result.validate();
        result
    }

    /// Construct an empty [`SpanSet`].
    pub fn empty() -> Self {
        let spans = VecDeque::new();
        SpanSet { spans }
    }

    /// Construct a full [`SpanSet`] that contains everything.
    /// Warning: The [`Id`] in this set might be unknown to an actual storage.
    pub fn full() -> Self {
        Span::full().into()
    }

    /// Check if this [`SpanSet`] contains nothing.
    pub fn is_empty(&self) -> bool {
        self.spans.is_empty()
    }

    /// Validate the spans are in the expected order and there are no mergable
    /// adjacent spans.
    #[cfg(debug_assertions)]
    fn validate(&self) {
        for (i, span) in self.spans.iter().enumerate() {
            assert!(span.low <= span.high);
            if i > 0 {
                assert!(
                    span.high + 1 < self.spans[i - 1].low,
                    "{:?} is not in DESC order or has mergable adjacent spans (around #{})",
                    &self.spans,
                    i
                );
            }
        }
    }

    /// Count integers covered by this [`SpanSet`].
    pub fn count(&self) -> u64 {
        self.spans.iter().fold(0, |acc, span| acc + span.count())
    }

    /// Tests if a given [`Id`] or [`Span`] is covered by this set.
    pub fn contains(&self, value: impl Into<Span>) -> bool {
        self.span_contains(value).is_some()
    }

    /// Find the [`Span`] that covers the given `value`.
    pub fn span_contains(&self, value: impl Into<Span>) -> Option<&Span> {
        let span = value.into();
        let idx = match self.spans.bsearch_by(|probe| span.low.cmp(&probe.low)) {
            Ok(idx) => idx,
            Err(idx) => idx,
        };
        if let Some(existing_span) = self.spans.get(idx) {
            debug_assert!(existing_span.low <= span.low);
            if existing_span.high >= span.high {
                return Some(existing_span);
            }
        }
        None
    }

    /// Skip the first `n` items.
    pub fn skip(&self, mut n: u64) -> Self {
        #[cfg(test)]
        let expected = n.max(self.count()) - n;
        let mut result = SpanSet::empty();
        for span in self.as_spans() {
            if n == 0 {
                result.push_span(*span);
            } else {
                let count = span.count();
                if count <= n {
                    // This span is skipped entirely.
                    n -= count;
                } else {
                    // This span is skipped partially.
                    debug_assert!(n > 0);
                    debug_assert!(n < count);
                    let high = span.high - n;
                    n = 0;
                    result.push_span((span.low..=high).into());
                }
            }
        }
        #[cfg(test)]
        assert_eq!(result.count(), expected);
        result
    }

    /// Only take the first `n` items.
    pub fn take(&self, mut n: u64) -> Self {
        #[cfg(test)]
        let expected = n.min(self.count());
        let mut result = SpanSet::empty();
        for span in self.as_spans() {
            if n == 0 {
                break;
            } else {
                let count = span.count();
                if count <= n {
                    // This span is taken entirely.
                    n -= count;
                    result.push_span(*span);
                } else {
                    // Part of the span is the last to be taken.
                    debug_assert!(n > 0);
                    debug_assert!(n < count);
                    let low = span.high - (n - 1);
                    n = 0;
                    result.push_span((low..=span.high).into());
                }
            }
        }
        #[cfg(test)]
        assert_eq!(result.count(), expected);
        result
    }

    /// Calculates the union of two sets.
    pub fn union(&self, rhs: &SpanSet) -> SpanSet {
        let mut spans = VecDeque::with_capacity((self.spans.len() + rhs.spans.len()).min(32));
        let mut iter_left = self.spans.iter().cloned();
        let mut iter_right = rhs.spans.iter().cloned();
        let mut next_left = iter_left.next();
        let mut next_right = iter_right.next();
        let mut push = |span: Span| push_with_union(&mut spans, span);

        loop {
            match (next_left, next_right) {
                (Some(left), Some(right)) => {
                    if left.high < right.high {
                        push(right);
                        next_right = iter_right.next();
                    } else {
                        push(left);
                        next_left = iter_left.next();
                    }
                }
                (Some(span), None) => {
                    push(span);
                    next_left = iter_left.next();
                }
                (None, Some(span)) => {
                    push(span);
                    next_right = iter_right.next();
                }
                (None, None) => {
                    let result = SpanSet { spans };
                    #[cfg(debug_assertions)]
                    result.validate();
                    return result;
                }
            }
        }
    }

    /// Calculates the intersection of two sets.
    pub fn intersection(&self, rhs: &SpanSet) -> SpanSet {
        let mut spans = VecDeque::with_capacity(self.spans.len().max(rhs.spans.len()).min(32));
        let push = |span: Span| push_with_union(&mut spans, span);
        intersect_iter(self.spans.iter().cloned(), rhs.spans.iter().cloned(), push);

        let result = SpanSet { spans };
        #[cfg(debug_assertions)]
        result.validate();
        result
    }

    /// Calculates spans that are included only by this set, not `rhs`.
    pub fn difference(&self, rhs: &SpanSet) -> SpanSet {
        let mut spans = VecDeque::with_capacity(self.spans.len().max(rhs.spans.len()).min(32));
        let mut iter_left = self.spans.iter().cloned();
        let mut iter_right = rhs.spans.iter().cloned();
        let mut next_left = iter_left.next();
        let mut next_right = iter_right.next();
        let mut push = |span: Span| push_with_union(&mut spans, span);

        loop {
            match (next_left, next_right) {
                (Some(left), Some(right)) => {
                    if right.low > left.high {
                        next_right = iter_right.next();
                    } else {
                        next_left = if right.high < left.low {
                            push(left);
                            iter_left.next()
                        } else {
                            // |----------------- left ------------------|
                            // |--- span1 ---|--- right ---|--- span2 ---|
                            if let Some(span2) = Span::try_from_bounds(right.high + 1..=left.high) {
                                push(span2);
                            }

                            Span::try_from_bounds(left.low..right.low).or_else(|| iter_left.next())
                        };
                    }
                }
                (Some(left), None) => {
                    push(left);
                    next_left = iter_left.next();
                }
                (None, _) => {
                    let result = SpanSet { spans };
                    #[cfg(debug_assertions)]
                    result.validate();
                    return result;
                }
            }
        }
    }

    /// Iterate `Id`s in descending order.
    pub fn iter_desc(&self) -> SpanSetIter<&SpanSet> {
        let len = self.spans.len();
        let back = (
            len as isize - 1,
            if len == 0 {
                0
            } else {
                let span = self.spans[len - 1];
                span.high.0 - span.low.0
            },
        );
        SpanSetIter {
            span_set: self,
            front: (0, 0),
            back,
        }
    }

    /// Iterate `Id`s in ascending order.
    pub fn iter_asc(&self) -> Rev<SpanSetIter<&Self>> {
        self.iter_desc().rev()
    }

    /// Iterate `Span`s in descending order.
    pub fn iter_span_desc(&self) -> impl Iterator<Item = &Span> {
        self.as_spans().iter()
    }

    /// Iterate `Span`s in ascending order.
    pub fn iter_span_asc(&self) -> impl Iterator<Item = &Span> {
        self.as_spans().iter().rev()
    }

    /// Get the maximum id in this set.
    pub fn max(&self) -> Option<Id> {
        self.spans.get(0).map(|span| span.high)
    }

    /// Get the minimal id in this set.
    pub fn min(&self) -> Option<Id> {
        self.spans
            .get(self.spans.len().max(1) - 1)
            .map(|span| span.low)
    }

    /// Internal use only. Append a span, which must have lower boundaries
    /// than existing spans.
    pub(crate) fn push_span(&mut self, span: Span) {
        push_with_union(&mut self.spans, span);
    }

    /// Internal use only. Append a span, which must have high boundaries
    /// than existing spans. In other words, spans passed to this function
    /// should be in ascending order.
    pub(crate) fn push_span_asc(&mut self, span: Span) {
        if self.spans.is_empty() {
            self.spans.push_back(span);
        } else {
            let last = &mut self.spans[0];
            // | last |
            //     | span |  | span |
            debug_assert!(span.low >= last.low);
            if last.high + 1 >= span.low {
                // Update in-place.
                last.high = span.high.max(last.high);
            } else {
                self.spans.push_front(span);
            }
        }
    }

    /// Internal use only. Append a [`SpanSet`], which must have lower
    /// boundaries than the existing spans.
    ///
    /// This is faster than [`SpanSet::union`]. used when it's known
    /// that the all ids in `set` being added is below the minimal id
    /// in the `self` set.
    pub(crate) fn push_set(&mut self, set: &SpanSet) {
        for span in &set.spans {
            self.push_span(*span);
        }
    }

    /// Get a reference to the spans.
    pub fn as_spans(&self) -> &VecDeque<Span> {
        &self.spans
    }

    /// Make this [`SpanSet`] contain the specified `span`.
    ///
    /// The current implementation works best when spans are pushed in
    /// ascending or descending order.
    pub fn push(&mut self, span: impl Into<Span>) {
        let span = span.into();
        if self.spans.is_empty() {
            self.spans.push_back(span)
        } else {
            let len = self.spans.len();
            {
                // Fast path: pushing to the last span.
                // 30->22 20->12 last H->L
                //               span H------>L union [Case 1]
                //                         H->L new   [Case 2]
                let last = &mut self.spans[len - 1];
                if last.high >= span.high {
                    if last.low <= span.high + 1 {
                        // Union spans in-place [Case 1]
                        last.low = last.low.min(span.low);
                    } else {
                        // New back span [Case 2]
                        self.spans.push_back(span)
                    }
                    return;
                }
            }
            {
                // Fast path: pushing to the last span.
                // first      H->L  20->12 10->12
                // span  H------>L union [Case 3]
                //       H->L      new   [Case 4]
                // Fast path: pushing to the first span.
                let first = &mut self.spans[0];
                if span.low >= first.low {
                    if span.low <= first.high + 1 {
                        // Union [Case 3]
                        first.high = first.high.max(span.high);
                    } else {
                        // New front span [Case 4]
                        self.spans.push_front(span);
                    }
                    return;
                }
            }
            {
                // Fast path: modify a span in-place.
                // higher H1---->L1     cur H2---->L2     lower H3---->L3
                // safe range        L1-2---------------->H3+2
                // Exceeding the safe range would cause spans to overlap and this path cannot
                // handle that.
                let idx = match self
                    .spans
                    .bsearch_by(|probe| (span.high + 1).cmp(&probe.low))
                {
                    Ok(idx) => idx,
                    Err(idx) => idx,
                };
                for idx in [idx] {
                    if let Some(cur) = self.spans.get(idx) {
                        // Not overlap with span?
                        if span.high + 1 < cur.low || cur.high + 1 < span.low {
                            continue;
                        }
                        // Might merge with a higher span? (Not in safe range)
                        if idx > 0 {
                            if let Some(higher) = self.spans.get(idx - 1) {
                                if span.high + 1 >= higher.low {
                                    continue;
                                }
                            }
                        }
                        // Might merge with a lower span? (Not in safe range)
                        if let Some(lower) = self.spans.get(idx + 1) {
                            if lower.high + 1 >= span.low {
                                continue;
                            }
                        }
                        // Passed all checks. Merge the span.
                        let cur = &mut self.spans[idx];
                        cur.high = cur.high.max(span.high);
                        cur.low = cur.low.min(span.low);
                        return;
                    }
                }
            }
            {
                // PERF: There might be a better way to do this by bisecting
                // spans and insert or delete in-place.  For now, this code
                // path remains not optimized since it is rarely used.
                *self = self.union(&SpanSet::from(span))
            }
        }
    }

    /// Intersection with a span. Return the min Id.
    ///
    /// This is not a general purpose API, but useful for internal logic
    /// like DAG descendant calculation.
    pub(crate) fn intersection_span_min(&self, rhs: Span) -> Option<Id> {
        let i = match self.spans.bsearch_by(|probe| rhs.low.cmp(&probe.high)) {
            Ok(idx) => idx,
            Err(idx) => idx.max(1) - 1,
        };
        // Prefer small index so we get the span that might overlap:
        // |----spans[1]----|      |----spans[0]----|
        //                     |----rhs-----|
        //                     (want spans[0], not spans[1])
        if i < self.spans.len() {
            let lhs = self.spans[i];
            if lhs.low <= rhs.high && lhs.high >= rhs.low {
                Some(lhs.low.max(rhs.low))
            } else {
                None
            }
        } else {
            // Happens if the set is empty.
            None
        }
    }
}

/// Push a span to `VecDeque<Span>`. Try to union them in-place.
fn push_with_union(spans: &mut VecDeque<Span>, span: Span) {
    if spans.is_empty() {
        spans.push_back(span);
    } else {
        let len = spans.len();
        let last = &mut spans[len - 1];
        debug_assert!(last.high >= span.high);
        if last.low <= span.high + 1 {
            // Union spans in-place.
            last.low = last.low.min(span.low);
        } else {
            spans.push_back(span)
        }
    }
}

impl Debug for SpanSet {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        // Limit spans to show.
        let limit = f.width().unwrap_or(12);
        let mut ranges: Vec<String> = self
            .spans
            .iter()
            .rev()
            .take(limit)
            .flat_map(|s| {
                if s.low + 2 >= s.high {
                    // "low..=high" form is not shorter.
                    (s.low.to(s.high)).map(|i| format!("{}", i)).collect()
                } else {
                    vec![format!("{}..={}", s.low, s.high)]
                }
            })
            .collect();
        let total = self.spans.len();
        if total == limit + 1 {
            ranges.push("and 1 span".into());
        } else if total > limit {
            ranges.push(format!("and {} spans", total - limit));
        }
        write!(f, "{}", ranges.join(" "))
    }
}

/// Iterator of integers in a [`SpanSet`].
#[derive(Clone)]
pub struct SpanSetIter<T> {
    span_set: T,
    // (index of span_set.spans, index of span_set.spans[i])
    front: (isize, u64),
    back: (isize, u64),
}

impl<T: AsRef<SpanSet>> SpanSetIter<T> {
    fn count_remaining(&self) -> u64 {
        let mut front = self.front;
        let back = self.back;
        let mut count = 0;
        while front <= back {
            let (vec_id, span_id) = front;
            let (back_vec_id, back_span_id) = back;
            if vec_id < back_vec_id {
                let len = self.span_set.as_ref().spans[vec_id as usize].count();
                count += len.max(span_id) - span_id;
                front = (vec_id + 1, 0);
            } else {
                count += back_span_id - span_id + 1;
                front = (vec_id + 1, 0);
            }
        }
        count
    }
}

impl<T: AsRef<SpanSet>> Iterator for SpanSetIter<T> {
    type Item = Id;

    fn next(&mut self) -> Option<Id> {
        if self.front > self.back {
            #[cfg(test)]
            assert_eq!(self.size_hint().0, 0);
            None
        } else {
            #[cfg(test)]
            let old_size = self.size_hint().0;
            let (vec_id, span_id) = self.front;
            let span = &self.span_set.as_ref().spans[vec_id as usize];
            self.front = if span_id == span.high.0 - span.low.0 {
                (vec_id + 1, 0)
            } else {
                (vec_id, span_id + 1)
            };
            #[cfg(test)]
            assert_eq!(self.size_hint().0 + 1, old_size);
            Some(span.high - span_id)
        }
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        #[cfg(test)]
        let expected_size = self.size_hint().0.max(n + 1) - n - 1;
        let mut n = n as u64;
        while self.front <= self.back {
            let (vec_id, span_id) = self.front;
            let span = &self.span_set.as_ref().spans[vec_id as usize];
            let span_remaining = span.count() - span_id;
            if n >= span_remaining {
                n -= span_remaining;
                self.front = (vec_id + 1, 0)
            } else {
                let span_id = span_id + n;
                self.front = (vec_id, span_id);
                let result = if self.front <= self.back {
                    self.front.1 += 1;
                    Some(span.high - span_id)
                } else {
                    None
                };
                #[cfg(test)]
                assert_eq!(self.size_hint().0, expected_size);
                return result;
            };
        }
        None
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let size = self.count_remaining() as _;
        (size, Some(size))
    }

    fn count(self) -> usize {
        self.count_remaining() as _
    }

    fn last(mut self) -> Option<Self::Item> {
        self.next_back()
    }
}

impl<T: AsRef<SpanSet>> DoubleEndedIterator for SpanSetIter<T> {
    fn next_back(&mut self) -> Option<Id> {
        if self.front > self.back {
            #[cfg(test)]
            assert_eq!(self.size_hint().0, 0);
            None
        } else {
            #[cfg(test)]
            let old_size = self.size_hint().0;
            let (vec_id, span_id) = self.back;
            let span = &self.span_set.as_ref().spans[vec_id as usize];
            self.back = if span_id == 0 {
                let span_len = if vec_id > 0 {
                    let span = self.span_set.as_ref().spans[(vec_id - 1) as usize];
                    span.high.0 - span.low.0
                } else {
                    0
                };
                (vec_id - 1, span_len)
            } else {
                (vec_id, span_id - 1)
            };
            #[cfg(test)]
            assert_eq!(self.size_hint().0 + 1, old_size);
            Some(span.high - span_id)
        }
    }

    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
        #[cfg(test)]
        let expected_size = self.size_hint().0.max(n + 1) - n - 1;
        let mut n = n as u64;
        while self.front <= self.back {
            let (vec_id, span_id) = self.back;
            let span = &self.span_set.as_ref().spans[vec_id as usize];
            let span_remaining = span_id + 1;
            if n >= span_remaining {
                n -= span_remaining;
                let span_end = if vec_id > 0 {
                    self.span_set.as_ref().spans[(vec_id - 1) as usize].count() - 1
                } else {
                    0
                };
                self.back = (vec_id - 1, span_end);
            } else {
                let span_id = span_id - n;
                self.back = (vec_id, span_id);
                let result = if self.front <= self.back {
                    if span_id == 0 {
                        let span_end = if vec_id > 0 {
                            self.span_set.as_ref().spans[(vec_id - 1) as usize].count() - 1
                        } else {
                            0
                        };
                        self.back = (vec_id - 1, span_end);
                    } else {
                        self.back.1 -= 1;
                    }
                    Some(span.high - span_id)
                } else {
                    None
                };
                #[cfg(test)]
                assert_eq!(self.size_hint().0, expected_size);
                return result;
            }
        }
        None
    }
}

impl<T: AsRef<SpanSet>> ExactSizeIterator for SpanSetIter<T> {
    fn len(&self) -> usize {
        self.count_remaining() as _
    }
}

impl IntoIterator for SpanSet {
    type Item = Id;
    type IntoIter = SpanSetIter<SpanSet>;

    /// Get an iterator for integers in this [`SpanSet`].
    fn into_iter(self) -> SpanSetIter<SpanSet> {
        let len = self.spans.len();
        let back = (
            len as isize - 1,
            if len == 0 {
                0
            } else {
                let span = self.spans[len - 1];
                span.high.0 - span.low.0
            },
        );
        SpanSetIter {
            span_set: self,
            front: (0, 0),
            back,
        }
    }
}

impl AsRef<SpanSet> for SpanSet {
    fn as_ref(&self) -> &SpanSet {
        self
    }
}

// `#[serde(transparent)]` on the `Id` struct.
// This would be easier if `Id` has `#[serde(transparent)]`.
// But that might be a breaking change now...
mod flat_id {
    use serde::de;
    use serde::de::Visitor;
    use serde::Deserializer;
    use serde::Serializer;

    use super::*;

    pub fn serialize<S: Serializer>(id: &Id, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_u64(id.0)
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Id, D::Error> {
        struct IdVisitor;
        impl<'de> Visitor<'de> for IdVisitor {
            type Value = Id;
            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("u64")
            }
            fn visit_u64<E: de::Error>(self, value: u64) -> Result<Self::Value, E> {
                Ok(Id(value))
            }
        }
        deserializer.deserialize_u64(IdVisitor)
    }
}

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

    impl From<RangeInclusive<u64>> for Span {
        fn from(range: RangeInclusive<u64>) -> Span {
            Span::new(Id(*range.start()), Id(*range.end()))
        }
    }

    impl From<u64> for Span {
        fn from(id: u64) -> Span {
            let id = Id(id);
            Span::new(id, id)
        }
    }

    impl From<(u64, u64)> for SpanSet {
        fn from(ids: (u64, u64)) -> SpanSet {
            SpanSet::from_spans([ids.0, ids.1].iter().cloned().map(Id))
        }
    }

    impl From<Span> for RangeInclusive<u64> {
        fn from(span: Span) -> RangeInclusive<u64> {
            span.low.0..=span.high.0
        }
    }

    #[test]
    fn test_overlapped_spans() {
        let span = SpanSet::from_spans(vec![1..=3, 3..=4]);
        assert_eq!(span.as_spans(), &[Span::from(1..=4)]);
    }

    #[test]
    fn test_valid_spans() {
        SpanSet::empty();
        SpanSet::from_spans(vec![4..=4, 3..=3, 1..=2]);
    }

    #[test]
    fn test_from_sorted_spans_merge() {
        let s = SpanSet::from_sorted_spans(vec![4..=4, 3..=3, 1..=2]);
        assert_eq!(format!("{:?}", s), "1..=4");
    }

    #[test]
    fn test_count() {
        let set = SpanSet::empty();
        assert_eq!(set.count(), 0);

        let set = SpanSet::from_spans(vec![1..=10, 20..=20, 31..=40]);
        assert_eq!(set.count(), 10 + 1 + 10);
    }

    #[test]
    fn test_skip() {
        let set = SpanSet::from_spans(vec![1..=10, 20..=20, 31..=40]);
        let skip = |n| format!("{:?}", set.skip(n));
        assert_eq!(skip(0), "1..=10 20 31..=40");
        assert_eq!(skip(1), "1..=10 20 31..=39");
        assert_eq!(skip(9), "1..=10 20 31");
        assert_eq!(skip(10), "1..=10 20");
        assert_eq!(skip(11), "1..=10");
        assert_eq!(skip(12), "1..=9");
        assert_eq!(skip(20), "1");
        assert_eq!(skip(21), "");
        assert_eq!(skip(22), "");
        assert_eq!(skip(50), "");
    }

    #[test]
    fn test_take() {
        let set = SpanSet::from_spans(vec![1..=10, 20..=20, 31..=40]);
        let take = |n| format!("{:?}", set.take(n));
        assert_eq!(take(0), "");
        assert_eq!(take(1), "40");
        assert_eq!(take(9), "32..=40");
        assert_eq!(take(10), "31..=40");
        assert_eq!(take(11), "20 31..=40");
        assert_eq!(take(12), "10 20 31..=40");
        assert_eq!(take(20), "2..=10 20 31..=40");
        assert_eq!(take(21), "1..=10 20 31..=40");
        assert_eq!(take(22), "1..=10 20 31..=40");
        assert_eq!(take(50), "1..=10 20 31..=40");
    }

    #[test]
    fn test_contains() {
        let set = SpanSet::empty();
        assert!(!set.contains(0));
        assert!(!set.contains(10));

        let set = SpanSet::from_spans(vec![1..=1, 2..=9, 10..=10, 20..=20, 31..=35, 36..=40]);
        assert!(!set.contains(0));
        assert!(set.contains(1));
        assert!(set.contains(5));
        assert!(set.contains(10));
        assert!(!set.contains(11));

        assert!(set.contains(1..=10));
        assert!(set.contains(1..=8));
        assert!(set.contains(3..=10));
        assert!(set.contains(3..=7));
        assert!(!set.contains(1..=11));
        assert!(!set.contains(0..=10));

        assert!(!set.contains(19));
        assert!(!set.contains(19..=20));
        assert!(set.contains(20));
        assert!(!set.contains(20..=21));
        assert!(!set.contains(21));

        assert!(!set.contains(30));
        assert!(set.contains(31));
        assert!(set.contains(32));
        assert!(set.contains(39));
        assert!(set.contains(40));
        assert!(!set.contains(41));

        assert!(set.contains(31..=40));
        assert!(set.contains(32..=40));
        assert!(set.contains(31..=39));
        assert!(set.contains(31..=39));
        assert!(!set.contains(31..=41));
        assert!(!set.contains(30..=40));
        assert!(!set.contains(30..=41));
    }

    fn union(a: Vec<impl Into<Span>>, b: Vec<impl Into<Span>>) -> Vec<RangeInclusive<u64>> {
        let a = SpanSet::from_spans(a);
        let b = SpanSet::from_spans(b);
        let spans1 = a.union(&b).spans;
        let spans2 = b.union(&a).spans;
        assert_eq!(spans1, spans2);
        spans1.into_iter().map(|span| span.into()).collect()
    }

    #[test]
    fn test_union() {
        assert_eq!(union(vec![1..=10], vec![10..=20]), vec![1..=20]);
        assert_eq!(union(vec![1..=30], vec![10..=20]), vec![1..=30]);
        assert_eq!(union(vec![6, 8, 10], vec![5, 7, 9]), vec![5..=10]);
        assert_eq!(
            union(vec![6..=6, 8..=9, 10..=10], vec![5]),
            vec![8..=10, 5..=6]
        );
    }

    fn intersect(a: Vec<impl Into<Span>>, b: Vec<impl Into<Span>>) -> Vec<RangeInclusive<u64>> {
        let a = SpanSet::from_spans(a);
        let b = SpanSet::from_spans(b);
        let spans1 = a.intersection(&b).spans;
        let spans2 = b.intersection(&a).spans;
        assert_eq!(spans1, spans2);
        spans1.into_iter().map(|span| span.into()).collect()
    }

    #[test]
    fn test_intersection() {
        assert_eq!(intersect(vec![1..=10], vec![11..=20]), vec![]);
        assert_eq!(intersect(vec![1..=10], vec![10..=20]), vec![10..=10]);
        assert_eq!(intersect(vec![1..=30], vec![10..=20]), vec![10..=20]);
        assert_eq!(
            intersect(vec![0..=10, 15..=20], vec![0..=30]),
            vec![15..=20, 0..=10]
        );
        assert_eq!(
            intersect(vec![0..=10, 15..=20], vec![5..=19]),
            vec![15..=19, 5..=10]
        );
        assert_eq!(intersect(vec![10, 9, 8, 7], vec![8..=11]), vec![8..=10]);
        assert_eq!(intersect(vec![10, 9, 8, 7], vec![5..=8]), vec![7..=8]);
    }

    fn difference(a: Vec<impl Into<Span>>, b: Vec<impl Into<Span>>) -> Vec<RangeInclusive<u64>> {
        let a = SpanSet::from_spans(a);
        let b = SpanSet::from_spans(b);
        let spans1 = a.difference(&b).spans;
        let spans2 = b.difference(&a).spans;

        // |------------- a -------------------|
        // |--- spans1 ---|--- intersection ---|--- spans2 ---|
        //                |------------------- b -------------|
        let intersected = intersect(
            a.spans.iter().cloned().collect(),
            b.spans.iter().cloned().collect(),
        );
        let unioned = union(
            a.spans.iter().cloned().collect(),
            b.spans.iter().cloned().collect(),
        );
        assert_eq!(
            union(intersected.clone(), spans1.iter().cloned().collect()),
            union(a.spans.iter().cloned().collect(), Vec::<Span>::new())
        );
        assert_eq!(
            union(intersected.clone(), spans2.iter().cloned().collect()),
            union(b.spans.iter().cloned().collect(), Vec::<Span>::new())
        );
        assert_eq!(
            union(
                spans1.iter().cloned().collect(),
                union(intersected.clone(), spans2.iter().cloned().collect())
            ),
            unioned.clone(),
        );

        assert!(
            intersect(
                spans1.iter().cloned().collect(),
                spans2.iter().cloned().collect()
            )
            .is_empty()
        );
        assert!(intersect(spans1.iter().cloned().collect(), intersected.clone()).is_empty());
        assert!(intersect(spans2.iter().cloned().collect(), intersected.clone()).is_empty());

        spans1.into_iter().map(|span| span.into()).collect()
    }

    #[test]
    fn test_difference() {
        assert_eq!(difference(vec![0..=5], Vec::<Span>::new()), vec![0..=5]);
        assert_eq!(difference(Vec::<Span>::new(), vec![0..=5]), vec![]);
        assert_eq!(difference(vec![0..=0], vec![1..=1]), vec![0..=0]);
        assert_eq!(difference(vec![0..=0], vec![0..=1]), vec![]);
        assert_eq!(difference(vec![0..=10], vec![0..=5]), vec![6..=10]);

        assert_eq!(
            difference(vec![0..=10], vec![3..=4, 7..=8]),
            vec![9..=10, 5..=6, 0..=2]
        );
        assert_eq!(
            difference(vec![3..=4, 7..=8, 10..=12], vec![4..=11]),
            vec![12..=12, 3..=3]
        );
    }

    #[test]
    fn test_iter() {
        let set = SpanSet::empty();
        assert!(set.iter_desc().next().is_none());
        assert!(set.iter_desc().rev().next().is_none());
        assert_eq!(set.iter_desc().size_hint(), (0, Some(0)));

        let set = SpanSet::from(0..=1);
        assert_eq!(set.iter_desc().collect::<Vec<Id>>(), vec![1, 0]);
        assert_eq!(set.iter_desc().rev().collect::<Vec<Id>>(), vec![0, 1]);
        assert_eq!(set.iter_desc().size_hint(), (2, Some(2)));
        assert_eq!(set.iter_desc().count(), 2);

        let mut iter = set.iter_desc();
        assert!(iter.next().is_some());
        assert!(iter.next_back().is_some());
        assert!(iter.next_back().is_none());

        let set = SpanSet::from_spans(vec![3..=5, 7..=8]);
        assert_eq!(set.iter_desc().collect::<Vec<Id>>(), vec![8, 7, 5, 4, 3]);
        assert_eq!(
            set.iter_desc().rev().collect::<Vec<Id>>(),
            vec![3, 4, 5, 7, 8]
        );
        assert_eq!(set.iter_desc().size_hint(), (5, Some(5)));
        assert_eq!(set.iter_desc().last(), Some(Id(3)));

        assert_eq!(
            set.clone().into_iter().collect::<Vec<Id>>(),
            vec![8, 7, 5, 4, 3]
        );
        assert_eq!(
            set.clone().into_iter().rev().collect::<Vec<Id>>(),
            vec![3, 4, 5, 7, 8]
        );
        assert_eq!(
            set.clone()
                .into_iter()
                .rev()
                .skip(1)
                .take(2)
                .rev()
                .collect::<Vec<Id>>(),
            vec![5, 4]
        );

        let set = SpanSet::from_spans(vec![3..=5, 7..=8]);
        let mut iter = set.iter_desc();
        assert_eq!(iter.next().unwrap(), 8);
        assert_eq!(iter.next_back().unwrap(), 3);

        let mut iter2 = iter.clone();
        assert_eq!(iter.next().unwrap(), 7);
        assert_eq!(iter.next_back().unwrap(), 4);
        assert_eq!(iter2.next().unwrap(), 7);
        assert_eq!(iter2.next_back().unwrap(), 4);
    }

    #[test]
    fn test_push() {
        let mut set = SpanSet::from(10..=20);
        set.push(5..=15);
        assert_eq!(set.as_spans(), &vec![Span::from(5..=20)]);

        let mut set = SpanSet::from(10..=20);
        set.push(5..=9);
        assert_eq!(set.as_spans(), &vec![Span::from(5..=20)]);

        let mut set = SpanSet::from(10..=20);
        set.push(5..=8);
        assert_eq!(
            set.as_spans(),
            &vec![Span::from(10..=20), Span::from(5..=8)]
        );

        let mut set = SpanSet::from(10..=20);
        set.push(5..=30);
        assert_eq!(set.as_spans(), &vec![Span::from(5..=30)]);

        let mut set = SpanSet::from(10..=20);
        set.push(20..=30);
        assert_eq!(set.as_spans(), &vec![Span::from(10..=30)]);

        let mut set = SpanSet::from(10..=20);
        set.push(10..=20);
        assert_eq!(set.as_spans(), &vec![Span::from(10..=20)]);

        let mut set = SpanSet::from(10..=20);
        set.push(22..=30);
        assert_eq!(
            set.as_spans(),
            &vec![Span::from(22..=30), Span::from(10..=20)]
        );
    }

    #[test]
    fn test_push_brute_force() {
        // Brute force pushing all spans in 1..=45 range to a SpanSet.
        let set = SpanSet::from_spans(vec![5..=10, 15..=16, 18..=20, 23..=23, 26..=30, 35..=40]);
        for low in 1..=45 {
            for high in low..=45 {
                let expected = SpanSet::from_spans(
                    (1..=45)
                        .filter(|&i| (i >= low && i <= high) || set.contains(Id(i)))
                        .map(Id),
                );
                let mut set = set.clone();
                set.push(low..=high);
                assert_eq!(set.as_spans(), expected.as_spans());
            }
        }
    }

    #[test]
    fn test_span_contains_brute_force() {
        let set = SpanSet::from_spans(vec![5..=10, 15..=16, 18..=20, 23..=23, 26..=30, 35..=40]);
        for low in 1..=45 {
            for high in low..=45 {
                let span = Span::from(low..=high);
                let result1 = set.span_contains(span);
                let result2 = set
                    .as_spans()
                    .iter()
                    .find(|s| s.contains(Id(low)) && s.contains(Id(high)));
                assert_eq!(result1, result2);
            }
        }
    }

    #[test]
    fn test_span_iter_nth() {
        let set = SpanSet::from_spans(vec![5..=10, 15..=15, 18..=20, 23..=23, 26..=30, 35..=40]);
        let vec: Vec<Id> = set.iter_desc().collect();
        for n in 0..=(vec.len() + 2) {
            assert_eq!(set.iter_desc().nth(n), vec.get(n).cloned());
        }
    }

    #[test]
    fn test_span_iter_nth_back() {
        let set = SpanSet::from_spans(vec![5..=10, 15..=15, 18..=20, 23..=23, 26..=30, 35..=40]);
        let vec: Vec<Id> = set.iter_asc().collect();
        for n in 0..=(vec.len() + 2) {
            assert_eq!(set.iter_desc().nth_back(n), vec.get(n).cloned());
        }
    }

    #[test]
    fn test_intersection_span_min() {
        let set = SpanSet::from_spans(vec![1..=10, 11..=20, 30..=40]);
        assert_eq!(set.intersection_span_min((15..=45).into()), Some(Id(15)));
        assert_eq!(set.intersection_span_min((20..=32).into()), Some(Id(20)));
        assert_eq!(set.intersection_span_min((21..=29).into()), None);
        assert_eq!(set.intersection_span_min((21..=32).into()), Some(Id(30)));
        assert_eq!(set.intersection_span_min((35..=45).into()), Some(Id(35)));
        assert_eq!(set.intersection_span_min((45..=55).into()), None);
    }

    #[test]
    fn test_debug() {
        let set = SpanSet::from_spans(vec![1..=1, 2..=9, 10..=10, 20..=20, 31..=35, 36..=40]);
        assert_eq!(format!("{:10?}", &set), "1..=10 20 31..=40");
        assert_eq!(format!("{:3?}", &set), "1..=10 20 31..=40");
        assert_eq!(format!("{:2?}", &set), "1..=10 20 and 1 span");
        assert_eq!(format!("{:1?}", &set), "1..=10 and 2 spans");
    }
}


/*
 * 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 nonblocking::non_blocking_result as r;
use tempfile::tempdir;
pub use test_dag::TestDag;

pub use self::drawdag::DrawDag;
use crate::id::Group;
use crate::id::VertexName;
use crate::nameset::SyncNameSetQuery;
use crate::ops::DagAddHeads;
use crate::ops::DagPersistent;
use crate::ops::ImportAscii;
#[cfg(feature = "render")]
use crate::render::render_namedag;
use crate::DagAlgorithm;
use crate::IdMap;
use crate::IdSet;
use crate::NameDag;
use crate::NameSet;
use crate::Result;

mod drawdag;
mod test_dag;

#[cfg(test)]
mod test_integrity;

#[cfg(test)]
mod test_sparse;

#[cfg(test)]
mod test_strip;

#[cfg(test)]
mod test_to_parents;

#[cfg(test)]
mod test_discontinuous;

#[cfg(test)]
mod test_server;

#[cfg(test)]
pub mod dummy_dag;

#[cfg(test)]
pub(crate) use test_dag::ProtocolMonitor;

#[cfg(test)]
use crate::iddag::FirstAncestorConstraint;
use crate::namedag::MemNameDag;
use crate::ops::IdConvert;
#[cfg(test)]
use crate::protocol::Process;
#[cfg(test)]
use crate::protocol::RequestLocationToName;
#[cfg(test)]
use crate::protocol::RequestNameToLocation;
#[cfg(test)]
use crate::render::render_segment_dag;
#[cfg(test)]
use crate::Id;
#[cfg(test)]
use crate::VertexListWithOptions;

// Example from segmented-changelog.pdf
// - DAG1: page 10
// - DAG2: page 11
// - DAG3: page 17
// - DAG4: page 18
// - DAG5: page 19

static ASCII_DAG1: &str = r#"
                C-D-\     /--I--J--\
            A-B------E-F-G-H--------K--L"#;

static ASCII_DAG2: &str = r#"
                      T /---------------N--O---\           T
                     / /                        \           \
               /----E-F-\    /-------L--M--------P--\     S--U---\
            A-B-C-D------G--H--I--J--K---------------Q--R---------V--W
                                   \--N"#;

static ASCII_DAG3: &str = r#"
              B---D---F--\
            A---C---E-----G"#;

static ASCII_DAG4: &str = r#"
             D  C  B
              \  \  \
            A--E--F--G"#;

static ASCII_DAG5: &str = r#"
        B---D---F
         \   \   \
      A---C---E---G"#;

fn test_dag_sort_version<T: DagAlgorithm + IdConvert>(dag: &T) -> Result<()> {
    // Test that sort() returns a set with dag and map hints assigned.
    let sets = [
        nameset(""),
        nameset("A C B"),
        r(from_ascii(MemNameDag::new(), ASCII_DAG3).sort(&nameset("C A B")))?,
    ];
    for set in sets {
        let sorted = r(dag.sort(&set))?;
        let hints = sorted.hints();
        assert_eq!(hints.dag_version(), Some(dag.dag_version()));
        assert_eq!(hints.id_map_version(), Some(dag.map_version()));
    }

    Ok(())
}

fn test_generic_dag1<T: DagAlgorithm + DagAddHeads + IdConvert>(dag: T) -> Result<T> {
    let dag = from_ascii(dag, ASCII_DAG1);
    assert_eq!(expand(r(dag.all())?), "A B C D E F G H I J K L");
    assert_eq!(expand(r(dag.dirty())?), "A B C D E F G H I J K L");
    assert_eq!(
        expand(r(dag.ancestors(nameset("H I")))?),
        "A B C D E F G H I"
    );
    assert_eq!(
        expand(r(dag.sort(&nameset("H E A")))?.skip(1).take(2)),
        "A E"
    );
    assert_eq!(expand(r(dag.first_ancestors(nameset("F")))?), "A B E F");
    assert_eq!(expand(r(dag.parents(nameset("H I E")))?), "B D G");
    assert_eq!(expand(r(dag.children(nameset("G D L")))?), "E H I");
    assert_eq!(expand(r(dag.merges(r(dag.all())?))?), "E K");
    assert_eq!(expand(r(dag.merges(nameset("E F J K")))?), "E K");
    assert_eq!(expand(r(dag.merges(nameset("A B D F H J L")))?), "");
    assert_eq!(expand(r(dag.roots(nameset("A B E F C D I J")))?), "A C I");
    assert_eq!(expand(r(dag.heads(nameset("A B E F C D I J")))?), "F J");
    assert_eq!(expand(r(dag.gca_all(nameset("J K H")))?), "G");

    test_dag_sort_version(&dag)?;

    Ok(dag)
}

fn test_generic_dag_beautify<D: DagAlgorithm + DagAddHeads>(
    new_dag: &impl Fn() -> D,
) -> Result<()> {
    let ascii = r#"
        A C
        | |
        B D
        |/
        E"#;
    let order = ["B", "D", "A", "C"];
    let dag = from_ascii_with_heads(new_dag(), ascii, Some(&order));
    assert_eq!(expand(r(dag.all())?), "A B C D E");

    let dag2 = r(dag.beautify(None))?;
    assert_eq!(expand(r(dag2.all())?), "A B C D E");

    let dag3 = r(dag.beautify(Some(nameset("A B E"))))?;
    assert_eq!(expand(r(dag3.all())?), "A B C D E");

    let dag4 = r(dag.beautify(Some(nameset("C D E"))))?;
    assert_eq!(expand(r(dag4.all())?), "A B C D E");

    let ascii = r#"
        A G
        |/
        B F
        |/
        C E
        |/
        D"#;
    let order = ["C", "E", "G", "F", "A"];
    let dag = from_ascii_with_heads(new_dag(), ascii, Some(&order));
    assert_eq!(expand(r(dag.all())?), "A B C D E F G");

    let dag2 = r(dag.beautify(None))?;
    assert_eq!(expand(r(dag2.all())?), "A B C D E F G");

    let dag3 = r(dag.beautify(Some(r(dag.ancestors(nameset("A")))?)))?;
    assert_eq!(expand(r(dag3.all())?), "A B C D E F G");

    let ascii = r#"
        A---B---C---D---E---F---G
             \
              H---I---J---K
                   \
                    L "#;
    let order = ["D", "J", "L", "K", "G"];
    let dag = from_ascii_with_heads(new_dag(), ascii, Some(&order));
    assert_eq!(expand(r(dag.all())?), "A B C D E F G H I J K L");

    let dag2 = r(dag.beautify(None))?;
    assert_eq!(expand(r(dag2.all())?), "A B C D E F G H I J K L");

    Ok(())
}

fn test_generic_dag_reachable_roots(dag: impl DagAlgorithm + DagAddHeads) -> Result<()> {
    let ascii = r#"
         Z
         |\
         D |
         | F
         C |
         | E
         B |
         |/
         A
         "#;
    let dag = from_ascii_with_heads(dag, ascii, Some(&["Z"][..]));

    // B is not reachable without going through other roots (C).
    // A is reachable through Z -> F -> E -> A.
    assert_eq!(
        expand(r(dag.reachable_roots(nameset("A B C"), nameset("Z")))?),
        "A C"
    );

    // A, E are not reachable without going through other roots (C, F).
    assert_eq!(
        expand(r(dag.reachable_roots(nameset("A C E F"), nameset("Z")))?),
        "C F"
    );

    // roots and heads overlap.
    assert_eq!(
        expand(r(
            dag.reachable_roots(nameset("A B C D E F Z"), nameset("D F"))
        )?),
        "D F"
    );

    // E, F are not reachable.
    assert_eq!(
        expand(r(dag.reachable_roots(nameset("A B E F"), nameset("D")))?),
        "B"
    );

    // "Bogus" root "Z".
    assert_eq!(
        expand(r(dag.reachable_roots(nameset("A Z"), nameset("C")))?),
        "A"
    );

    Ok(())
}

// This test is specific to the implementation in this crate, and requires
// >2 parents support.
fn test_specific_dag_import(dag: impl DagAlgorithm + DagAddHeads) -> Result<()> {
    let ascii = r#"
            J K
           /|\|\
          G H I H
          |/|/|
          E F |
         /|/|\|
        A B C D"#;
    let dag1 = from_ascii_with_heads(dag, ascii, Some(&["J", "K"][..]));

    let dir = tempdir().unwrap();
    let mut dag2 = NameDag::open(&dir.path())?;
    r(dag2.import_and_flush(&dag1, nameset("J")))?;
    #[cfg(feature = "render")]
    assert_eq!(
        render(&dag2),
        r#"
            K
            â”œâ”€â•®
            â”‚ â”‚ J
            â•â”€â”¬â”€â”¤
            â”‚ I â”‚
            â”‚ â”œâ”€â”€â”€â•®
            H â”‚ â”‚ â”‚
            â”œâ”€â”€â”€â”€â”€â•®
            â”‚ â”‚ â”‚ F
            â”‚ â•â”€â”€â”€â”¼â”€â•®
            â”‚ D â”‚ â”‚ â”‚
            â”‚   â”‚ â”‚ â”‚
            â”‚   â”‚ â”‚ C
            â”‚   â”‚ â”‚
            â”‚   G â”‚
            â”œâ”€â”€â”€â•¯ â”‚
            E     â”‚
            â”œâ”€â”€â”€â”€â”€â•®
            â”‚     B
            â”‚
            A"#
    );

    // Check that dag2 is actually flushed to disk.
    let dag3 = NameDag::open(&dir.path())?;
    #[cfg(feature = "render")]
    assert_eq!(
        render(&dag3),
        r#"
            K
            â”œâ”€â•®
            â”‚ â”‚ J
            â•â”€â”¬â”€â”¤
            â”‚ I â”‚
            â”‚ â”œâ”€â”€â”€â•®
            H â”‚ â”‚ â”‚
            â”œâ”€â”€â”€â”€â”€â•®
            â”‚ â”‚ â”‚ F
            â”‚ â•â”€â”€â”€â”¼â”€â•®
            â”‚ D â”‚ â”‚ â”‚
            â”‚   â”‚ â”‚ â”‚
            â”‚   â”‚ â”‚ C
            â”‚   â”‚ â”‚
            â”‚   G â”‚
            â”œâ”€â”€â”€â•¯ â”‚
            E     â”‚
            â”œâ”€â”€â”€â”€â”€â•®
            â”‚     B
            â”‚
            A"#
    );
    Ok(())
}

fn test_generic_dag2<T: DagAlgorithm + DagAddHeads>(dag: T) -> Result<T> {
    let ascii = r#"
            J K
           / \|\
          G H I H
          |/|/|
          E F |
         /|/| |
        A B C D"#;
    let dag = from_ascii_with_heads(dag, ascii, Some(&["J", "K"][..]));

    let v = |name: &str| -> VertexName { VertexName::copy_from(name.as_bytes()) };

    assert_eq!(expand(r(dag.all())?), "A B C D E F G H I J K");
    assert_eq!(expand(r(dag.ancestors(nameset("H I")))?), "A B C D E F H I");
    assert_eq!(expand(r(dag.first_ancestors(nameset("H I")))?), "A D E H I");
    assert_eq!(
        expand(r(dag.first_ancestors(nameset("J G D")))?),
        "A D E G J"
    );
    assert_eq!(expand(r(dag.parents(nameset("H I E")))?), "A B D E F");
    assert_eq!(r(dag.first_ancestor_nth(v("H"), 2))?.unwrap(), v("A"));
    assert!(r(dag.first_ancestor_nth(v("H"), 3))?.is_none());
    assert_eq!(expand(r(dag.heads(nameset("E H F K I D")))?), "K");
    assert_eq!(expand(r(dag.children(nameset("E F I")))?), "G H I J K");
    assert_eq!(expand(r(dag.merges(r(dag.all())?))?), "E F H I J K");
    assert_eq!(expand(r(dag.merges(nameset("E H G D I")))?), "E H I");
    assert_eq!(expand(r(dag.roots(nameset("E G H J I K D")))?), "D E");
    assert_eq!(r(dag.gca_one(nameset("J K")))?, Some(v("I")));
    assert_eq!(expand(r(dag.gca_all(nameset("J K")))?), "E I");
    assert_eq!(expand(r(dag.common_ancestors(nameset("G H")))?), "A B E");
    assert!(r(dag.is_ancestor(v("B"), v("K")))?);
    assert!(!r(dag.is_ancestor(v("K"), v("B")))?);
    assert_eq!(
        expand(r(dag.heads_ancestors(nameset("A E F D G")))?),
        "D F G"
    );
    assert_eq!(expand(r(dag.range(nameset("A"), nameset("K")))?), "A E H K");
    assert_eq!(expand(r(dag.only(nameset("I"), nameset("G")))?), "C D F I");
    let (reachable, unreachable) = r(dag.only_both(nameset("I"), nameset("G")))?;
    assert_eq!(expand(reachable), "C D F I");
    assert_eq!(expand(unreachable), expand(r(dag.ancestors(nameset("G")))?));
    assert_eq!(expand(r(dag.descendants(nameset("F E")))?), "E F G H I J K");

    assert!(r(dag.is_ancestor(v("B"), v("J")))?);
    assert!(r(dag.is_ancestor(v("F"), v("F")))?);
    assert!(!r(dag.is_ancestor(v("K"), v("I")))?);

    Ok(dag)
}

#[test]
fn test_import_ascii_with_vertex_fn() {
    let ascii = r#"
        B C
        |/
        A
    "#;
    let mut dag = MemNameDag::new();
    let vertex_fn = |s: &str| -> VertexName {
        let mut bytes = s.as_bytes().to_vec();
        bytes.push(b'0');
        VertexName::copy_from(&bytes)
    };
    dag.import_ascii_with_vertex_fn(ascii, vertex_fn).unwrap();
    assert_eq!(expand(r(dag.all()).unwrap()), "A0 B0 C0");
}

#[test]
fn test_mem_namedag() {
    let new_dag = MemNameDag::new;
    test_generic_dag(&new_dag);
    test_specific_dag_import(new_dag()).unwrap();
}

#[test]
fn test_namedag() {
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering;

    let dir = tempdir().unwrap();
    let counter = AtomicUsize::new(0);
    let new_dag = move || {
        let count = counter.fetch_add(1, Ordering::AcqRel);
        NameDag::open(dir.path().join(count.to_string())).unwrap()
    };
    test_generic_dag(&new_dag);
    test_specific_dag_import(new_dag()).unwrap();
}

#[test]
fn test_protocols() {
    let mut built = build_segments(ASCII_DAG1, "A C E L", 3);
    assert_eq!(
        built.ascii[3],
        r#"
                1-2-\     /--7--8--\
            0-3------4-5-6-9--------10-11
Lv0: RH0-0[] R1-2[] 3-3[0] H4-8[3, 2] 9-9[6] H10-11[9, 8]
Lv1: R0-0[] R1-8[0]"#
    );

    // Replace "[66]" to "B", "[67]" to "C", etc.
    let replace = |mut s: String| -> String {
        for ch in "ABCDEFGHIJKL".chars() {
            s = s.replace(&format!("[{}]", ch as u8), &format!("{}", ch));
        }
        s
    };

    // [Id] -> RequestLocationToName (useful for getting commit hashes from ids).
    // 3 (D) and 9 (J) are p2 that cannot be resolved.
    let ids: Vec<Id> = b"ABCEFGHI"
        .iter()
        .map(|&b| built.name_dag.map.find_id_by_name(&[b]).unwrap().unwrap())
        .collect();
    let ids = IdSet::from_spans(ids);
    let request1: RequestLocationToName =
        r((&built.name_dag.map, &built.name_dag.dag).process(ids)).unwrap();
    assert_eq!(
        replace(format!("{:?}", &request1)),
        "RequestLocationToName { paths: [L~2, J~1(+5), D~1, B~1] }"
    );

    // [name] -> RequestNameToLocation (useful for getting ids from commit hashes).
    let names = b"ABCEFGHI"
        .iter()
        .map(|&b| VertexName::copy_from(&[b]))
        .collect();
    let request2: RequestNameToLocation =
        r((&built.name_dag.map, &built.name_dag.dag).process(names)).unwrap();
    assert_eq!(
        replace(format!("{:?}", &request2)),
        "RequestNameToLocation { names: [A, B, C, E, F, G, H, I], heads: [L] }"
    );

    // RequestLocationToName -> ResponseIdNamePair
    let response1 = r((&built.name_dag.map, &built.name_dag.dag).process(request1)).unwrap();
    assert_eq!(
        replace(format!("{:?}", &response1)),
        "ResponseIdNamePair { path_names: [(L~2, [H]), (J~1(+5), [I, G, F, E, B]), (D~1, [C]), (B~1, [A])] }"
    );

    // RequestNameToLocation -> ResponseIdNamePair
    // Only B, D, H, J, L are used since they are "universally known".
    let response2 = r((&built.name_dag.map, &built.name_dag.dag).process(request2)).unwrap();
    assert_eq!(
        replace(format!("{:?}", &response2)),
        "ResponseIdNamePair { path_names: [(B~1, [A]), (H~4, [B]), (D~1, [C]), (H~3, [E]), (H~2, [F]), (H~1, [G]), (L~2, [H]), (J~1, [I])] }"
    );

    // Applying responses to IdMap. Should not cause errors.
    r((&mut built.name_dag.map, &built.name_dag.dag).process(response1.clone())).unwrap();
    r((&mut built.name_dag.map, &built.name_dag.dag).process(response2.clone())).unwrap();

    // Try applying response2 to a sparse IdMap.
    // Prepare the sparse IdMap.
    let mut sparse_id_map = IdMap::open(built.dir.path().join("sparse-id")).unwrap();
    r(built
        .name_dag
        .dag
        .write_sparse_idmap(&built.name_dag.map, &mut sparse_id_map))
    .unwrap();
    assert_eq!(
        format!("{:?}", &sparse_id_map),
        "IdMap {\n  D: 2,\n  B: 3,\n  J: 8,\n  H: 9,\n  L: 11,\n}\n"
    );
    // Apply response2.
    r((&mut sparse_id_map, &built.name_dag.dag).process(response2)).unwrap();
    assert_eq!(
        format!("{:?}", &sparse_id_map),
        r#"IdMap {
  D: 2,
  B: 3,
  J: 8,
  H: 9,
  L: 11,
  A: 0,
  C: 1,
  E: 4,
  F: 5,
  G: 6,
  I: 7,
}
"#
    );
}

#[test]
fn test_segment_non_master() {
    let ascii = r#"
a----b----c----d----e----f----g----------h----i
     \                    \             /
      h---i---j---k        l---m---n---o
               \                \
                -----------------p---q"#;
    let built = build_segments(ascii, "i q", 3);
    assert_eq!(
        built.ascii[0],
        r#"
N0---N1---N2---N3---N4---N5---N6---------N11--N12
     \                    \             /
      N11-N12-j---k        N7--N8--N9--N10
               \                \
                -----------------p---q
Lv0: RN0-N6[] N7-N10[N5] N11-N12[N0, N6, N10]"#
    );
    assert_eq!(
        built.ascii[1],
        r#"
N0---N1---N2---N3---N4---N5---N6---------N11--N12
     \                    \             /
      N11-N12-N13-k        N7--N8--N9--N10
               \                \
                -----------------N14-N15
Lv0: RN0-N6[] N7-N10[N5] N11-N12[N0, N6, N10] N13-N13[N12] N14-N15[N13, N8]
Lv1: RN0-N12[]"#
    );
}

#[test]
fn test_segment_examples() {
    assert_eq!(
        build_segments(ASCII_DAG1, "L", 3).ascii[0],
        r#"
                2-3-\     /--8--9--\
            0-1------4-5-6-7--------10-11
Lv0: RH0-1[] R2-3[] H4-7[1, 3] 8-9[6] H10-11[7, 9]
Lv1: R0-7[]"#
    );

    assert_eq!(
        build_segments(ASCII_DAG2, "W", 3).ascii[0],
        r#"
                      19/---------------13-14--\           19
                     / /                        \           \
               /----4-5-\    /-------11-12-------15-\     18-20--\
            0-1-2-3------6--7--8--9--10--------------16-17--------21-22
                                   \--13
Lv0: RH0-3[] 4-5[1] H6-10[3, 5] 11-12[7] 13-14[5, 9] 15-15[12, 14] H16-17[10, 15] R18-18[] 19-19[4] 20-20[18, 19] H21-22[17, 20]
Lv1: R0-10[] 11-15[7, 5, 9] 16-17[10, 15] R18-20[4]
Lv2: R0-17[]"#
    );

    assert_eq!(
        build_segments(ASCII_DAG3, "G", 3).ascii[0],
        r#"
              3---4---5--\
            0---1---2-----6
Lv0: RH0-2[] R3-5[] H6-6[2, 5]"#
    );

    assert_eq!(
        build_segments(ASCII_DAG4, "G", 3).ascii[0],
        r#"
             3  1  0
              \  \  \
            2--4--5--6
Lv0: RH0-0[] R1-1[] R2-2[] R3-3[] 4-4[2, 3] 5-5[1, 4] H6-6[0, 5]
Lv1: R0-0[] R1-1[] R2-4[]"#
    );

    assert_eq!(
        build_segments(ASCII_DAG5, "G", 3).ascii[0],
        r#"
        1---3---5
         \   \   \
      0---2---4---6
Lv0: RH0-0[] R1-1[] H2-2[0, 1] 3-3[1] H4-4[2, 3] 5-5[3] H6-6[4, 5]
Lv1: R0-2[] 3-4[1, 2]"#
    );

    // Examples outside segmented-changelog.pdf

    // For this graph, the numbers should look continuous in one direction.
    let ascii_dag = r#"
            Z---E--M--J--C
                 \  \  \  \
                  O--T--D--L
                   \  \  \  \
                    K--H--P--W
                     \  \  \  \
                      X--R--U--V
                       \  \  \  \
                        A--N--S--Y"#;
    assert_eq!(
        build_segments(ascii_dag, "Y", 3).ascii[0],
        r#"
            0---1--6--11-16
                 \  \  \  \
                  2--7--12-17
                   \  \  \  \
                    3--8--13-18
                     \  \  \  \
                      4--9--14-19
                       \  \  \  \
                        5--10-15-20
Lv0: RH0-5[] 6-6[1] 7-7[6, 2] 8-8[3, 7] 9-9[8, 4] H10-10[5, 9] 11-11[6] 12-12[11, 7] 13-13[12, 8] 14-14[13, 9] H15-15[10, 14] 16-16[11] 17-17[16, 12] 18-18[17, 13] 19-19[14, 18] H20-20[15, 19]
Lv1: R0-5[] 6-8[1, 2, 3] 9-10[8, 4, 5] 11-13[6, 7, 8] 14-15[13, 9, 10] 16-18[11, 12, 13]
Lv2: R0-10[] 11-15[6, 7, 8, 9, 10]"#
    );

    // If a graph looks like this, it's hard to optimize anyway.
    let ascii_dag = r#"
            Z---E--J--C--O--T
                 \     \     \
                  D     L     K
                   \     \     \
                    H--P--W--X--R
                     \     \     \
                      U     V     A
                       \     \     \
                        N--S--Y--B--G"#;
    assert_eq!(
        build_segments(ascii_dag, "G", 3).ascii[0],
        r#"
            0---1--2--3--4--5
                 \     \     \
                  8     7     6
                   \     \     \
                    9--10-11-12-13
                     \     \     \
                      15    18    14
                       \     \     \
                        16-17-19-20-21
Lv0: RH0-6[] 7-7[3] 8-10[1] 11-12[7, 10] H13-14[6, 12] 15-17[9] 18-18[11] 19-20[17, 18] H21-21[14, 20]
Lv1: R0-6[] 7-12[3, 1] 13-14[6, 12] 15-20[9, 11]
Lv2: R0-14[]"#
    );
}

#[test]
fn test_segment_groups() {
    let dag = r#"
A---B---C---D---E---F---G--------H---I
     \               \          /
      h--i--j--k      l--m--n--o
                \            \
                 -------------p---q"#;

    // This test involves many things. Lower-case commits are non-master commits.
    // - D after B: Test incremental build of a master commit with a master parent.
    // - i after D: Test non-master with master parent.
    // - k after i: Test non-master with non-master parent.
    // - q after G: Test non-master with both master and non-master ancestors.
    // - I after q: Test overwriting non-master Ids with master Ids (!).
    let built = build_segments(dag, "B D i k G q I", 3);
    assert_eq!(
        built.ascii.join("\n"),
        r#"
0---1---C---D---E---F---G--------H---I
     \               \          /
      h--i--j--k      l--m--n--o
                \            \
                 -------------p---q
Lv0: RH0-1[]

0---1---2---3---E---F---G--------H---I
     \               \          /
      h--i--j--k      l--m--n--o
                \            \
                 -------------p---q
Lv0: RH0-3[]

0---1---2---3---E---F---G--------H---I
     \               \          /
      N0-N1-j--k      l--m--n--o
                \            \
                 -------------p---q
Lv0: RH0-3[] N0-N1[1]

0---1---2---3---E---F---G--------H---I
     \               \          /
      N0-N1-N2-N3     l--m--n--o
                \            \
                 -------------p---q
Lv0: RH0-3[] N0-N1[1] N2-N3[N1]

0---1---2---3---4---5---6--------H---I
     \               \          /
      N0-N1-N2-N3     l--m--n--o
                \            \
                 -------------p---q
Lv0: RH0-6[] N0-N3[1]

0---1---2---3---4---5---6--------H---I
     \               \          /
      N0-N1-N2-N3     N4-N5-N6-o
                \            \
                 -------------N7--N8
Lv0: RH0-6[] N0-N3[1] N4-N6[5] N7-N8[N3, N6]

0---1---2---3---4---5---6--------11--12
     \               \          /
      N0-N1-N2-N3     7--8--9--10
                \            \
                 -------------N4--N5
Lv0: RH0-6[] 7-10[5] H11-12[6, 10] N0-N3[1] N4-N5[N3, 9]"#
    );

    // Notice that N4 to N6 were re-written in the last step.
    // 'm' only has 1 id: 8 (master). The old id (N5) is now taken by 'q'.
    assert_eq!(
        built.name_dag.map.find_id_by_name(b"m").unwrap().unwrap(),
        Id(8)
    );
    assert_eq!(
        built.name_dag.map.find_name_by_id(Id(8)).unwrap().unwrap(),
        b"m"
    );
    let id = Group::NON_MASTER.min_id() + 5;
    assert_eq!(
        built.name_dag.map.find_name_by_id(id).unwrap().unwrap(),
        b"q"
    );

    // Parent-child indexes work fine.
    assert_eq!(
        format!("{:?}", built.name_dag.dag.children_id(Id(5)).unwrap(),),
        "6 7"
    );
}

#[test]
fn test_namedag_reassign_master() -> crate::Result<()> {
    let dir = tempdir().unwrap();
    let mut dag = NameDag::open(&dir.path())?;
    dag = from_ascii(dag, "A-B-C");

    // The in-memory DAG can answer parent_names questions.
    assert_eq!(format!("{:?}", r(dag.parent_names("A".into()))?), "[]");
    assert_eq!(format!("{:?}", r(dag.parent_names("C".into()))?), "[B]");

    // First flush, A, B, C are non-master.
    r(dag.flush(&Default::default())).unwrap();

    assert_eq!(format!("{:?}", r(dag.vertex_id("A".into()))?), "N0");
    assert_eq!(format!("{:?}", r(dag.vertex_id("C".into()))?), "N2");

    // Second flush, making B master without adding new vertexes.
    let heads =
        VertexListWithOptions::from(vec![VertexName::from("B")]).with_highest_group(Group::MASTER);
    r(dag.flush(&heads)).unwrap();
    assert_eq!(format!("{:?}", r(dag.vertex_id("A".into()))?), "0");
    assert_eq!(format!("{:?}", r(dag.vertex_id("B".into()))?), "1");
    assert_eq!(format!("{:?}", r(dag.vertex_id("C".into()))?), "N0");

    Ok(())
}

#[test]
fn test_namedag_reassign_non_master() {
    let mut t = TestDag::new();

    // A: master; B, Z: non-master.
    t.drawdag("A--B--Z", &["A"]);
    // C, D, E: non-master.
    t.drawdag("B--C--D--E", &[]);
    // Prompt C to master. Triggers non-master reassignment.
    t.drawdag("", &["C"]);

    // Z still exists.
    #[cfg(feature = "render")]
    assert_eq!(
        t.render_graph(),
        r#"
            Z  N2
            â”‚
            â”‚ E  N1
            â”‚ â”‚
            â”‚ D  N0
            â”‚ â”‚
            â”‚ C  2
            â”œâ”€â•¯
            B  1
            â”‚
            A  0"#
    );

    // Z can round-trip in IdMap.
    let z_id = r(t.dag.vertex_id("Z".into())).unwrap();
    let z_vertex = r(t.dag.vertex_name(z_id)).unwrap();
    assert_eq!(format!("{:?}", z_vertex), "Z");
}

#[test]
fn test_segment_ancestors_example1() {
    // DAG from segmented-changelog.pdf
    let ascii_dag = r#"
            2-3-\     /--8--9--\
        0-1------4-5-6-7--------10-11"#;
    let result = build_segments(ascii_dag, "11", 3);
    let dag = result.name_dag.dag;

    for (id, count) in vec![
        (11, 12),
        (10, 11),
        (9, 9),
        (8, 8),
        (7, 8),
        (6, 7),
        (5, 6),
        (4, 5),
        (3, 2),
        (2, 1),
        (1, 2),
        (0, 1),
    ] {
        assert_eq!(dag.ancestors(id.into()).unwrap().count(), count);
    }

    for (a, b, ancestor) in vec![
        (10, 3, 3.into()),
        (11, 0, 0.into()),
        (11, 10, 10.into()),
        (11, 9, 9.into()),
        (3, 0, None),
        (7, 1, 1.into()),
        (9, 2, 2.into()),
        (9, 7, 6.into()),
    ] {
        let ancestor = ancestor.map(Id);
        let a = Id(a);
        let b = Id(b);
        assert_eq!(dag.gca_one((a, b).into()).unwrap(), ancestor);
        assert_eq!(
            dag.gca_all((a, b).into()).unwrap().iter_desc().nth(0),
            ancestor
        );
        assert_eq!(dag.gca_all((a, b).into()).unwrap().iter_desc().nth(1), None);
        assert_eq!(dag.is_ancestor(b, a).unwrap(), ancestor == Some(b));
        assert_eq!(dag.is_ancestor(a, b).unwrap(), ancestor == Some(a));
    }

    for (spans, ancestors) in vec![
        (vec![3..=8], vec![3]),
        (vec![1..=1, 4..=9], vec![1]),
        (vec![1..=4], vec![]),
    ] {
        assert_eq!(
            dag.gca_all(IdSet::from_spans(spans))
                .unwrap()
                .iter_desc()
                .collect::<Vec<Id>>(),
            ancestors.into_iter().map(Id).collect::<Vec<Id>>(),
        );
    }
}

#[test]
fn test_segment_multiple_gcas() {
    let ascii_dag = r#"
        B---C
         \ /
        A---D"#;
    let result = build_segments(ascii_dag, "C D", 3);
    assert_eq!(
        result.ascii[1],
        r#"
        1---2
         \ /
        0---3
Lv0: RH0-0[] R1-1[] H2-2[0, 1] 3-3[0, 1]
Lv1: R0-2[]"#
    );
    let dag = result.name_dag.dag;
    // This is kind of "undefined" whether it's 1 or 0.
    assert_eq!(dag.gca_one((2, 3).into()).unwrap(), Some(Id(1)));
    assert_eq!(
        dag.gca_all((2, 3).into())
            .unwrap()
            .iter_desc()
            .collect::<Vec<_>>(),
        vec![1, 0]
    );
}

#[test]
fn test_parents() {
    let result = build_segments(ASCII_DAG1, "L", 3);
    assert_eq!(
        result.ascii[0],
        r#"
                2-3-\     /--8--9--\
            0-1------4-5-6-7--------10-11
Lv0: RH0-1[] R2-3[] H4-7[1, 3] 8-9[6] H10-11[7, 9]
Lv1: R0-7[]"#
    );

    let dag = result.name_dag.dag;

    let parents = |spans| -> String { format_set(dag.parents(IdSet::from_spans(spans)).unwrap()) };
    let parent_ids = |id| -> String { format!("{:?}", dag.parent_ids(Id(id)).unwrap()) };
    let first_ancestor_nth =
        |id, n| -> String { format!("{:?}", dag.first_ancestor_nth(Id(id), n).unwrap()) };
    let to_first_ancestor_nth = |id| -> String {
        let c = FirstAncestorConstraint::KnownUniversally {
            heads: Id(11).into(),
        };
        let res = dag.to_first_ancestor_nth(Id(id), c);
        match res {
            Ok(s) => format!("{:?}", s),
            Err(e) => e.to_string(),
        }
    };

    assert_eq!(parents(vec![]), "");

    assert_eq!(parents(vec![0..=0]), "");
    assert_eq!(parents(vec![0..=1]), "0");
    assert_eq!(parents(vec![0..=2]), "0");
    assert_eq!(parents(vec![0..=3]), "0 2");
    assert_eq!(parents(vec![0..=4]), "0..=3");
    assert_eq!(parents(vec![0..=5]), "0..=4");
    assert_eq!(parents(vec![0..=6]), "0..=5");
    assert_eq!(parents(vec![0..=7]), "0..=6");
    assert_eq!(parents(vec![0..=8]), "0..=6");
    assert_eq!(parents(vec![0..=9]), "0..=6 8");
    assert_eq!(parents(vec![0..=10]), "0..=9");
    assert_eq!(parents(vec![0..=11]), "0..=10");

    assert_eq!(parents(vec![0..=0, 2..=2]), "");
    assert_eq!(parents(vec![0..=0, 3..=3, 5..=5, 9..=10]), "2 4 7 8 9");
    assert_eq!(parents(vec![1..=1, 4..=4, 6..=6, 8..=11]), "0 1 3 5..=10");

    assert_eq!(parent_ids(0), "[]");
    assert_eq!(parent_ids(1), "[0]");
    assert_eq!(parent_ids(4), "[1, 3]");
    assert_eq!(parent_ids(10), "[7, 9]");
    assert_eq!(parent_ids(11), "[10]");

    assert_eq!(first_ancestor_nth(0, 0), "0");
    assert_eq!(first_ancestor_nth(4, 2), "0");
    assert_eq!(first_ancestor_nth(10, 2), "6");
    assert_eq!(first_ancestor_nth(10, 3), "5");
    assert_eq!(first_ancestor_nth(11, 0), "11");
    assert_eq!(first_ancestor_nth(11, 1), "10");
    assert_eq!(first_ancestor_nth(11, 2), "7");
    assert_eq!(first_ancestor_nth(11, 3), "6");
    assert_eq!(first_ancestor_nth(11, 4), "5");
    assert_eq!(first_ancestor_nth(11, 6), "1");
    assert_eq!(first_ancestor_nth(11, 7), "0");
    assert!(dag.first_ancestor_nth(Id::MIN, 1).is_err());
    assert!(dag.first_ancestor_nth(Id(11), 8).is_err());

    assert_eq!(to_first_ancestor_nth(0), "Some((1, 1))");
    assert_eq!(to_first_ancestor_nth(1), "Some((9, 5))");
    assert_eq!(to_first_ancestor_nth(2), "Some((3, 1))");
    assert_eq!(
        to_first_ancestor_nth(3),
        "ProgrammingError: cannot convert 3 to x~n form (x must be in `H + parents(ancestors(H) & merge())` where H = 11) (trace: in seg R2-3[], 3 has child seg (H4-7[1, 3]), child seg cannot be followed (3 is not p1))"
    );
    assert_eq!(to_first_ancestor_nth(4), "Some((9, 4))");
    assert_eq!(to_first_ancestor_nth(5), "Some((9, 3))");
    assert_eq!(to_first_ancestor_nth(6), "Some((9, 2))");
    assert_eq!(to_first_ancestor_nth(7), "Some((11, 2))");
    assert_eq!(to_first_ancestor_nth(8), "Some((9, 1))");
    assert_eq!(
        to_first_ancestor_nth(9),
        "ProgrammingError: cannot convert 9 to x~n form (x must be in `H + parents(ancestors(H) & merge())` where H = 11) (trace: in seg 8-9[6], 9 has child seg (H10-11[7, 9]), child seg cannot be followed (9 is not p1))"
    );
    assert_eq!(to_first_ancestor_nth(10), "Some((11, 1))");
    assert_eq!(to_first_ancestor_nth(11), "Some((11, 0))");
}

#[test]
fn test_children() {
    let result = build_segments(ASCII_DAG1, "L", 3);
    let dag = result.name_dag.dag;
    let children =
        |spans| -> String { format_set(dag.children(IdSet::from_spans(spans)).unwrap()) };

    // See test_parents above for the ASCII DAG.

    assert_eq!(children(vec![]), "");
    assert_eq!(children(vec![0..=0]), "1");

    assert_eq!(children(vec![0..=1]), "1 4");
    assert_eq!(children(vec![0..=2]), "1 3 4");
    assert_eq!(children(vec![0..=3]), "1 3 4");
    assert_eq!(children(vec![0..=4]), "1 3 4 5");
    assert_eq!(children(vec![0..=5]), "1 3..=6");
    assert_eq!(children(vec![0..=6]), "1 3..=8");
    assert_eq!(children(vec![0..=7]), "1 3..=8 10");
    assert_eq!(children(vec![0..=8]), "1 3..=10");
    assert_eq!(children(vec![0..=9]), "1 3..=10");
    assert_eq!(children(vec![0..=10]), "1 3..=11");
    assert_eq!(children(vec![0..=11]), "1 3..=11");

    assert_eq!(children(vec![1..=10]), "3..=11");
    assert_eq!(children(vec![2..=10]), "3..=11");
    assert_eq!(children(vec![3..=10]), "4..=11");
    assert_eq!(children(vec![4..=10]), "5..=11");
    assert_eq!(children(vec![5..=10]), "6..=11");
    assert_eq!(children(vec![6..=10]), "7..=11");
    assert_eq!(children(vec![7..=10]), "9 10 11");
    assert_eq!(children(vec![8..=10]), "9 10 11");
    assert_eq!(children(vec![9..=10]), "10 11");
    assert_eq!(children(vec![10..=10]), "11");

    assert_eq!(children(vec![0..=0, 2..=2]), "1 3");
    assert_eq!(children(vec![0..=0, 3..=3, 5..=5, 9..=10]), "1 4 6 10 11");
    assert_eq!(children(vec![1..=1, 4..=4, 6..=6, 10..=10]), "4 5 7 8 11");
}

#[test]
fn test_heads() {
    let ascii = r#"
    C G   K L
    | |\  |/
    B E F I J
    | |/  |/
    A D   H"#;

    let result = build_segments(ascii, "C G K L J", 2);
    assert_eq!(
        result.ascii[4],
        r#"
    2 6   9 10
    | |\  |/
    1 5 4 8 11
    | |/  |/
    0 3   7
Lv0: RH0-2[] R3-4[] 5-5[3] 6-6[5, 4] R7-9[] 10-10[8] 11-11[7]
Lv1: R0-2[] R3-4[] 5-6[3, 4] R7-9[] 10-10[8]
Lv2: R0-2[] R3-6[] R7-9[]"#
    );

    let dag = result.name_dag.dag;
    let heads = |spans| -> String { format_set(dag.heads(IdSet::from_spans(spans)).unwrap()) };

    assert_eq!(heads(vec![]), "");
    assert_eq!(heads(vec![0..=11]), "2 6 9 10 11");
    assert_eq!(heads(vec![0..=1, 3..=5, 7..=10]), "1 4 5 9 10");
    assert_eq!(heads(vec![0..=0, 2..=2]), "0 2");
    assert_eq!(heads(vec![1..=2, 4..=6, 7..=7, 11..=11, 9..=9]), "2 6 9 11");
}

#[test]
fn test_roots() {
    let ascii = r#"
    C G   J
    | |\  |\
    B E F I K
    | |/  |\
    A D   H L"#;

    let result = build_segments(ascii, "C G J", 2);
    assert_eq!(
        result.ascii[2],
        r#"
    2 6   11
    | |\  |\
    1 5 4 107
    | |/  |\
    0 3   9 8
Lv0: RH0-2[] R3-4[] 5-5[3] 6-6[5, 4] R7-7[] R8-8[] R9-9[] 10-10[9, 8] 11-11[10, 7]
Lv1: R0-2[] R3-4[] 5-6[3, 4] R7-7[] R8-8[] R9-10[8]
Lv2: R0-2[] R3-6[] R7-7[]"#
    );

    let dag = result.name_dag.dag;
    let roots = |spans| -> String { format_set(dag.roots(IdSet::from_spans(spans)).unwrap()) };

    assert_eq!(roots(vec![]), "");
    assert_eq!(roots(vec![0..=11]), "0 3 7 8 9");
    assert_eq!(roots(vec![1..=2, 4..=6, 8..=10]), "1 4 5 8 9");
    assert_eq!(roots(vec![0..=0, 2..=3, 5..=6, 9..=11]), "0 2 3 9");
    assert_eq!(roots(vec![1..=1, 3..=3, 6..=8, 11..=11]), "1 3 6 7 8");
}

#[test]
fn test_range() {
    let ascii = r#"
            J
           /|\
          G H I
          |/|/
          E F
         /|/|\
        A B C D"#;

    let result = build_segments(ascii, "J", 2);
    assert_eq!(
        result.ascii[0],
        r#"
            9
           /|\
          3 7 8
          |/|/
          2 6
         /|/|\
        0 1 4 5
Lv0: RH0-0[] R1-1[] H2-3[0, 1] R4-4[] R5-5[] 6-6[1, 4, 5] 7-7[2, 6] 8-8[6] H9-9[3, 7, 8]
Lv1: R0-0[] R1-3[0] R4-4[] R5-6[1, 4] 7-7[2, 6]
Lv2: R0-3[] R4-6[1]"#
    );

    let dag = result.name_dag.dag;
    let range = |roots, heads| -> String {
        format_set(
            dag.range(IdSet::from_spans(roots), IdSet::from_spans(heads))
                .unwrap(),
        )
    };

    assert_eq!(range(vec![6], vec![3]), "");
    assert_eq!(range(vec![1], vec![3, 8]), "1 2 3 6 8");
    assert_eq!(range(vec![4], vec![3, 8]), "4 6 8");
    assert_eq!(range(vec![0, 5], vec![7]), "0 2 5 6 7");
    assert_eq!(range(vec![0, 5], vec![3, 8]), "0 2 3 5 6 8");
    assert_eq!(range(vec![0, 1, 4, 5], vec![3, 7, 8]), "0..=8");

    assert_eq!(range(vec![0], vec![0]), "0");
    assert_eq!(range(vec![0], vec![1]), "");
    assert_eq!(range(vec![0], vec![2]), "0 2");
    assert_eq!(range(vec![0], vec![3]), "0 2 3");
    assert_eq!(range(vec![0], vec![4]), "");
    assert_eq!(range(vec![0], vec![5]), "");
    assert_eq!(range(vec![0], vec![6]), "");
    assert_eq!(range(vec![0], vec![7]), "0 2 7");
    assert_eq!(range(vec![0], vec![8]), "");
    assert_eq!(range(vec![0], vec![9]), "0 2 3 7 9");
    assert_eq!(range(vec![1], vec![1]), "1");
    assert_eq!(range(vec![1], vec![2]), "1 2");
    assert_eq!(range(vec![1], vec![3]), "1 2 3");
    assert_eq!(range(vec![1], vec![4]), "");
    assert_eq!(range(vec![1], vec![5]), "");
    assert_eq!(range(vec![1], vec![6]), "1 6");
    assert_eq!(range(vec![1], vec![7]), "1 2 6 7");
    assert_eq!(range(vec![1], vec![8]), "1 6 8");
    assert_eq!(range(vec![1], vec![9]), "1 2 3 6..=9");
    assert_eq!(range(vec![2], vec![2]), "2");
    assert_eq!(range(vec![2], vec![3]), "2 3");
    assert_eq!(range(vec![2], vec![4]), "");
    assert_eq!(range(vec![2], vec![5]), "");
    assert_eq!(range(vec![2], vec![6]), "");
    assert_eq!(range(vec![2], vec![7]), "2 7");
    assert_eq!(range(vec![2], vec![8]), "");
    assert_eq!(range(vec![2], vec![9]), "2 3 7 9");
    assert_eq!(range(vec![3], vec![3]), "3");
    assert_eq!(range(vec![3], vec![4]), "");
    assert_eq!(range(vec![3], vec![5]), "");
    assert_eq!(range(vec![3], vec![6]), "");
    assert_eq!(range(vec![3], vec![7]), "");
    assert_eq!(range(vec![3], vec![8]), "");
    assert_eq!(range(vec![3], vec![9]), "3 9");
    assert_eq!(range(vec![4], vec![4]), "4");
    assert_eq!(range(vec![4], vec![5]), "");
    assert_eq!(range(vec![4], vec![6]), "4 6");
    assert_eq!(range(vec![4], vec![7]), "4 6 7");
    assert_eq!(range(vec![4], vec![8]), "4 6 8");
    assert_eq!(range(vec![4], vec![9]), "4 6..=9");
    assert_eq!(range(vec![5], vec![5]), "5");
    assert_eq!(range(vec![5], vec![6]), "5 6");
    assert_eq!(range(vec![5], vec![7]), "5 6 7");
    assert_eq!(range(vec![5], vec![8]), "5 6 8");
    assert_eq!(range(vec![5], vec![9]), "5..=9");
    assert_eq!(range(vec![6], vec![6]), "6");
    assert_eq!(range(vec![6], vec![7]), "6 7");
    assert_eq!(range(vec![6], vec![8]), "6 8");
    assert_eq!(range(vec![6], vec![9]), "6..=9");
    assert_eq!(range(vec![7], vec![7]), "7");
    assert_eq!(range(vec![7], vec![8]), "");
    assert_eq!(range(vec![7], vec![9]), "7 9");
    assert_eq!(range(vec![8], vec![8]), "8");
    assert_eq!(range(vec![8], vec![9]), "8 9");
    assert_eq!(range(vec![9], vec![9]), "9");

    // Test descendants() and ancestors() against range().
    for bits in 0..(1 << 10) {
        let mut set = IdSet::empty();
        for i in (0..=9).rev() {
            if bits & (1 << i) != 0 {
                set.push_span(i.into());
            }
        }

        let all = dag.all().unwrap();
        assert_eq!(
            dag.range(set.clone(), all.clone()).unwrap().as_spans(),
            dag.descendants(set.clone()).unwrap().as_spans(),
        );

        assert_eq!(
            dag.range(all.clone(), set.clone()).unwrap().as_spans(),
            dag.ancestors(set.clone()).unwrap().as_spans(),
        );
    }
}

#[test]
fn test_render_segment_dag() {
    // For reference in below graphs.
    assert_eq!(
        build_segments(ASCII_DAG2, "W", 3).ascii[0],
        r#"
                      19/---------------13-14--\           19
                     / /                        \           \
               /----4-5-\    /-------11-12-------15-\     18-20--\
            0-1-2-3------6--7--8--9--10--------------16-17--------21-22
                                   \--13
Lv0: RH0-3[] 4-5[1] H6-10[3, 5] 11-12[7] 13-14[5, 9] 15-15[12, 14] H16-17[10, 15] R18-18[] 19-19[4] 20-20[18, 19] H21-22[17, 20]
Lv1: R0-10[] 11-15[7, 5, 9] 16-17[10, 15] R18-20[4]
Lv2: R0-17[]"#
    );

    let built = build_segments(ASCII_DAG2, "W", 3);
    let mut buf = Vec::new();
    buf.push(b'\n');
    render_segment_dag(&mut buf, &built.name_dag, 0, Group::MASTER).unwrap();

    assert_eq!(
        String::from_utf8(buf).unwrap(),
        r#"
o    V(21)-W(22)
â”œâ”€â•®
â”‚ o    U(20)-U(20)
â”‚ â”œâ”€â•®
â”‚ â”‚ o  T(19)-T(19)
â”‚ â”‚ â”‚
â”‚ o â”‚  S(18)-S(18)
â”‚   â”‚
o   â”‚  Q(16)-R(17)
â”œâ”€â•® â”‚
â”‚ o â”‚    P(15)-P(15)
â”‚ â”œâ”€â”€â”€â•®
â”‚ â”‚ â”‚ o  N(13)-O(14)
â•â”€â”€â”€â”¬â”€â•¯
â”‚ o â”‚  L(11)-M(12)
â”œâ”€â•¯ â”‚
o   â”‚  G(6)-K(10)
â”œâ”€â”€â”€â•®
â”‚   o  E(4)-F(5)
â”œâ”€â”€â”€â•¯
o  A(0)-D(3)

"#
    );

    let mut buf = Vec::new();
    buf.push(b'\n');
    render_segment_dag(&mut buf, &built.name_dag, 1, Group::MASTER).unwrap();

    assert_eq!(
        String::from_utf8(buf).unwrap(),
        r#"
o  S(18)-U(20)
â”‚
â”‚ o  Q(16)-R(17)
â•â”€â”¤
â”‚ o  L(11)-P(15)
â•â”€â•¯
o  A(0)-K(10)

"#
    );
}

#[test]
fn test_render_segment_dag_non_master() {
    let mut t = TestDag::new();

    // A: master; B, Z: non-master.
    t.drawdag("A--B--Z", &["A"]);

    let mut buf = Vec::new();
    buf.push(b'\n');
    render_segment_dag(&mut buf, &t.dag, 0, Group::NON_MASTER).unwrap();

    assert_eq!(
        String::from_utf8(buf).unwrap(),
        r#"
o  B(N0)-Z(N1)
â”‚
~
"#
    );
}

#[cfg_attr(test, tokio::test)]
async fn test_subdag() {
    let t = TestDag::draw("A..E");
    let s = t.dag.subdag(nameset("B D E")).await.unwrap();
    #[cfg(feature = "render")]
    assert_eq!(
        render(&s),
        r#"
            E
            â”‚
            D
            â”‚
            B"#
    );

    // Test ordering: preserve the heads order (D before C).
    let t = TestDag::draw("A-X B-X X-C X-D");
    let s1 = t.dag.subdag(nameset("D C B A")).await.unwrap();
    let s2 = t.dag.subdag(nameset("A B C D")).await.unwrap();
    #[cfg(feature = "render")]
    assert_eq!(
        render(&s1),
        r#"
            D
            â”œâ”€â•®
            â”‚ â”‚ C
            â•â”€â”¬â”€â•¯
            â”‚ A
            â”‚
            B"#
    );
    #[cfg(feature = "render")]
    assert_eq!(render(&s1), render(&s2));
}

// Test utilities

fn expand(set: NameSet) -> String {
    let mut names = set
        .iter()
        .unwrap()
        .map(|n| String::from_utf8_lossy(n.unwrap().as_ref()).to_string())
        .collect::<Vec<String>>();
    names.sort();
    names.join(" ")
}

fn nameset(names: &str) -> NameSet {
    let names: Vec<VertexName> = names
        .split_whitespace()
        .map(|n| VertexName::copy_from(n.as_bytes()))
        .collect();
    NameSet::from_static_names(names)
}

fn format_set(set: IdSet) -> String {
    format!("{:?}", set)
}

impl IdMap {
    /// Replace names in an ASCII DAG using the ids assigned.
    fn replace(&self, text: &str) -> String {
        let mut result = text.to_string();
        for &group in Group::ALL.iter() {
            const MAX_ID_IN_ASCII_TEST: u64 = 30;
            for id in group.min_id().to(group.min_id() + MAX_ID_IN_ASCII_TEST) {
                if let Ok(Some(name)) = self.find_name_by_id(id) {
                    let name = String::from_utf8(name.to_vec()).unwrap();
                    let id_str = format!("{:01$}", id, name.len());
                    // Try to replace while maintaining width
                    if name.len() + 2 == id_str.len() {
                        result = result
                            .replace(&format!("{}--", name), &id_str)
                            .replace(&format!("{}  ", name), &id_str);
                    } else if name.len() + 1 == id_str.len() {
                        result = result
                            .replace(&format!("{}-", name), &id_str)
                            .replace(&format!("{} ", name), &id_str);
                    }
                    result = result.replace(&format!("{}", name), &id_str);
                }
            }
        }
        result
    }
}

fn get_parents_func_from_ascii(text: &str) -> impl Fn(VertexName) -> Result<Vec<VertexName>> {
    let parents = ::drawdag::parse(&text);
    move |name: VertexName| -> Result<Vec<VertexName>> {
        Ok(parents[&String::from_utf8(name.as_ref().to_vec()).unwrap()]
            .iter()
            .map(|p| VertexName::copy_from(p.as_bytes()))
            .collect())
    }
}

/// Result of `build_segments`.
pub(crate) struct BuildSegmentResult {
    pub(crate) ascii: Vec<String>,
    pub(crate) name_dag: NameDag,
    pub(crate) dir: tempfile::TempDir,
}

/// Take an ASCII DAG, assign segments from given heads.
/// Return the ASCII DAG and the built NameDag.
pub(crate) fn build_segments(text: &str, heads: &str, segment_size: usize) -> BuildSegmentResult {
    let mut dag = TestDag::new_with_segment_size(segment_size);

    let mut ascii = Vec::new();
    for head in heads.split(' ') {
        // Assign to non-master if the name starts with a lowercase character.
        let master = if head.chars().nth(0).unwrap().is_lowercase() {
            vec![]
        } else {
            vec![head]
        };
        dag.drawdag_with_limited_heads(text, &master[..], Some(&[head]));
        let annotated = dag.annotate_ascii(text);
        let segments = dag.render_segments();
        ascii.push(format!("{}\n{}", annotated, segments));
    }

    BuildSegmentResult {
        ascii,
        name_dag: dag.dag,
        dir: dag.dir,
    }
}

fn from_ascii<D: DagAddHeads>(dag: D, text: &str) -> D {
    from_ascii_with_heads(dag, text, None)
}

fn from_ascii_with_heads<D: DagAddHeads>(mut dag: D, text: &str, heads: Option<&[&str]>) -> D {
    dag.import_ascii_with_heads(text, heads).unwrap();
    dag
}

/// Test a general DAG interface against a few test cases.
pub fn test_generic_dag<D: DagAddHeads + DagAlgorithm + IdConvert + Send + Sync + 'static>(
    new_dag: &impl Fn() -> D,
) {
    test_generic_dag1(new_dag()).unwrap();
    test_generic_dag2(new_dag()).unwrap();
    test_generic_dag_reachable_roots(new_dag()).unwrap();
    test_generic_dag_beautify(new_dag).unwrap();
}

#[cfg(feature = "render")]
fn render(dag: &(impl DagAlgorithm + ?Sized)) -> String {
    render_namedag(dag, |_| None).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.
 */

//! Extensions to types in other crates.

use std::collections::BTreeSet;

use crate::FlatSegment;
use crate::Id;
use crate::PreparedFlatSegments;

pub(crate) trait PreparedFlatSegmentsExt {
    /// Merge with another (newer) `PreparedFlatSegmentsExt`.
    fn merge(&mut self, rhs: Self);

    /// Add graph edges: id -> parent_ids. Used by `assign_head`.
    fn push_edge(&mut self, id: Id, parent_ids: &[Id]) {
        self.push_segment(id, id, parent_ids)
    }

    /// Add segment. Used by `import_pull_data`.
    fn push_segment(&mut self, low: Id, high: Id, parent_ids: &[Id]);
}

impl PreparedFlatSegmentsExt for PreparedFlatSegments {
    fn merge(&mut self, rhs: Self) {
        if rhs.segments.is_empty() {
            return;
        }
        if self.segments.is_empty() {
            *self = rhs;
            return;
        }

        for seg in rhs.segments {
            if !maybe_merge_in_place(&mut self.segments, seg.low, seg.high, &seg.parents) {
                self.segments.insert(seg);
            }
        }

        if cfg!(debug_assertions) {
            ensure_sorted_and_merged(&self.segments);
        }
    }

    fn push_segment(&mut self, low: Id, high: Id, parent_ids: &[Id]) {
        if !maybe_merge_in_place(&mut self.segments, low, high, parent_ids) {
            let new_seg = FlatSegment {
                low,
                high,
                parents: parent_ids.to_vec(),
            };
            self.segments.insert(new_seg);
        }
        if cfg!(debug_assertions) {
            ensure_sorted_and_merged(&self.segments);
        }
    }
}

/// Try to merge a flat segment (low..=high, parents=parents) in place.
/// Return true if it was merged in place.
fn maybe_merge_in_place(
    segments: &mut BTreeSet<FlatSegment>,
    low: Id,
    high: Id,
    parent_ids: &[Id],
) -> bool {
    if let [parent_id] = parent_ids {
        if *parent_id + 1 != low {
            return false;
        }
    } else {
        return false;
    }
    let upper_bound = FlatSegment {
        low,
        high: low,
        parents: Vec::new(),
    };
    if let Some(candidate) = segments.range(..=upper_bound).rev().next() {
        if candidate.high + 1 == low {
            // Merge
            let candidate = candidate.clone();
            let new_seg = FlatSegment {
                low: candidate.low,
                high,
                parents: candidate.parents.clone(),
            };
            segments.remove(&candidate);
            segments.insert(new_seg);
            return true;
        }
    }
    false
}

/// Check that segments are sorted and merged.
fn ensure_sorted_and_merged(segments: &BTreeSet<FlatSegment>) {
    let mut last_high = None;
    for seg in segments {
        // Sorted? No overlap?
        assert!(Some(seg.low) > last_high);
        // Merged?
        if let Some(last_high) = last_high {
            if seg.parents.len() == 1 && seg.parents[0] + 1 == seg.low {
                assert_ne!(last_high + 1, seg.low);
            }
        }
        last_high = Some(seg.high);
    }
}

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

    #[test]
    fn test_push_edge_out_of_order() {
        let mut segs = PreparedFlatSegments::default();
        segs.push_edge(Id(0), &[]);
        segs.push_edge(Id(50), &[]);
        segs.push_edge(Id(100), &[]);
        segs.push_edge(Id(1), &[Id(0)]);
        segs.push_edge(Id(51), &[Id(50)]);
        segs.push_edge(Id(101), &[Id(100)]);
        segs.push_edge(Id(2), &[]);
        segs.push_edge(Id(52), &[Id(51), Id(50)]);
        segs.push_edge(Id(102), &[Id(100)]);
        segs.push_edge(Id(103), &[Id(102)]);
        segs.push_edge(Id(53), &[Id(52)]);
        segs.push_edge(Id(105), &[Id(103)]);
        segs.push_edge(Id(106), &[Id(105)]);
        segs.push_edge(Id(104), &[Id(103)]);
        segs.push_edge(Id(3), &[Id(2)]);
        segs.push_edge(Id(4), &[Id(3)]);
        segs.push_edge(Id(54), &[Id(53)]);
        segs.push_edge(Id(49), &[Id(3)]);
        segs.push_edge(Id(107), &[Id(106)]);

        // Check that adjacent segments are merged.
        assert_eq!(
            dbg(&segs),
            [
                "FlatSegment { low: 0, high: 1, parents: [] }",
                "FlatSegment { low: 2, high: 4, parents: [] }",
                "FlatSegment { low: 49, high: 49, parents: [3] }",
                "FlatSegment { low: 50, high: 51, parents: [] }",
                "FlatSegment { low: 52, high: 54, parents: [51, 50] }",
                "FlatSegment { low: 100, high: 101, parents: [] }",
                "FlatSegment { low: 102, high: 104, parents: [100] }",
                "FlatSegment { low: 105, high: 107, parents: [103] }"
            ]
        );
    }

    #[test]
    fn test_push_segment() {
        let mut segs = PreparedFlatSegments::default();
        segs.push_segment(Id(10), Id(20), &[]);
        segs.push_segment(Id(40), Id(50), &[Id(15), Id(10), Id(12)]);
        segs.push_segment(Id(21), Id(30), &[Id(20)]);
        segs.push_segment(Id(31), Id(35), &[Id(20)]);
        assert_eq!(
            dbg(&segs),
            [
                "FlatSegment { low: 10, high: 30, parents: [] }",
                "FlatSegment { low: 31, high: 35, parents: [20] }",
                "FlatSegment { low: 40, high: 50, parents: [15, 10, 12] }",
            ]
        );
    }

    #[test]
    fn test_merge() {
        let mut segs1 = PreparedFlatSegments::default();
        segs1.push_edge(Id(10), &[]);
        segs1.push_edge(Id(11), &[Id(10)]);
        let mut segs2 = PreparedFlatSegments::default();
        segs2.push_edge(Id(12), &[Id(11)]);
        segs2.push_edge(Id(13), &[Id(12)]);
        segs2.push_edge(Id(14), &[Id(11)]);
        segs1.merge(segs2);
        assert_eq!(
            dbg(&segs1),
            [
                "FlatSegment { low: 10, high: 13, parents: [] }",
                "FlatSegment { low: 14, high: 14, parents: [11] }"
            ]
        );
    }

    fn dbg(segs: &PreparedFlatSegments) -> Vec<String> {
        segs.segments.iter().map(|s| format!("{:?}", s)).collect()
    }
}