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

use std::cmp::Ordering;

use crate::errors::bug;
use crate::errors::programming;
use crate::id::Group;
use crate::id::Id;
use crate::segment::Segment;
use crate::segment::SegmentFlags;
use crate::spanset::Span;
use crate::IdSet;
use crate::Level;
use crate::Result;

mod in_process_store;

#[cfg(any(test, feature = "indexedlog-backend"))]
pub(crate) mod indexedlog_store;

pub(crate) use in_process_store::InProcessStore;
#[cfg(any(test, feature = "indexedlog-backend"))]
pub(crate) use indexedlog_store::IndexedLogStore;

pub trait IdDagStore: Send + Sync + 'static {
    /// Maximum level segment in the store
    fn max_level(&self) -> Result<Level>;

    /// Find segment by level and head.
    fn find_segment_by_head_and_level(&self, head: Id, level: u8) -> Result<Option<Segment>>;

    /// Find flat segment containing the given id.
    fn find_flat_segment_including_id(&self, id: Id) -> Result<Option<Segment>>;

    /// Add a new segment.
    ///
    /// For simplicity, it does not check if the new segment overlaps with
    /// an existing segment (which is a logic error). Those checks can be
    /// offline.
    fn insert(
        &mut self,
        flags: SegmentFlags,
        level: Level,
        low: Id,
        high: Id,
        parents: &[Id],
    ) -> Result<()> {
        let segment = Segment::new(flags, level, low, high, parents);
        self.insert_segment(segment)
    }

    fn insert_segment(&mut self, segment: Segment) -> Result<()>;

    /// Remove a flat segment. The segment cannot have descendants.
    fn remove_flat_segment(&mut self, segment: &Segment) -> Result<()> {
        self.resize_flat_segment(segment, None)
    }

    /// Remove or truncate a flat segment.
    /// - If `new_high` is None, remove the flat segment. The segment cannot
    ///   have descendants.
    /// - If `new_high` is set, trucnate the flat segment by resetting `high` to
    ///   the given value. If `new_high` is smaller than `high`, then
    ///   `descendants(new_high+1:high)` must be empty.
    fn resize_flat_segment(&mut self, segment: &Segment, new_high: Option<Id>) -> Result<()> {
        // `segment` must be a flat segment.
        if segment.level()? != 0 {
            return programming(format!(
                "resize_flat_segment requires a flat segment, got {:?}",
                segment
            ));
        }

        // Figure out deleted and inserted spans for checks.
        let span = segment.span()?;
        let (deleted_span, inserted_span) = get_deleted_inserted_spans(span, new_high);
        if deleted_span.is_none() && inserted_span.is_none() {
            // No need to do anything.
            return Ok(());
        }

        // `segment` needs to match an existing flat segment?
        if let Some(existing_segment) = self.find_flat_segment_including_id(span.high)? {
            if &existing_segment != segment {
                return programming(format!(
                    "resize_flat_segment got {:?} which does not match the existing segment {:?}",
                    segment, existing_segment
                ));
            }
        } else {
            return programming(format!(
                "resize_flat_segment got a non-existed segment {:?}",
                segment
            ));
        }

        // When inserting a new span, the space cannot be taken.
        if let Some(span) = inserted_span {
            let taken = self.all_ids_in_groups(&[span.high.group()])?;
            let overlap = taken.intersection(&span.into());
            if !overlap.is_empty() {
                return programming(format!(
                    "resize_flat_segment cannot overlap with existing segments (segment: {:?} new_high: {:?}, overlap: {:?})",
                    segment, new_high, overlap,
                ));
            }
        }

        // When deleting a span, it should have no descendants.
        if let Some(span) = deleted_span {
            if let Some(item) = self.iter_flat_segments_with_parent_span(span)?.next() {
                let (_, child) = item?;
                return programming(format!(
                    "resize_flat_segment requires a segment without descendants, got {:?} with child segment {:?}",
                    segment, child
                ));
            }
        }

        // Prepare the resized segment and check `new_high`.
        let new_segment = match new_high {
            Some(high) => Some(segment.with_high(high)?),
            None => None,
        };

        // Check passed, call the low-level function.
        self.remove_flat_segment_unchecked(segment)?;

        // Need to insert a new segment.
        if let Some(seg) = new_segment {
            self.insert_segment(seg)?;
        }

        Ok(())
    }

    /// Actual implementation of the segment removal part of
    /// `resize_flat_segment` without checks.
    /// This should only be called via `resize_flat_segment` for integrity.
    fn remove_flat_segment_unchecked(&mut self, segment: &Segment) -> Result<()>;

    /// Return all ids from given groups. This is useful to implement the
    /// `all()` operation.
    ///
    /// With discontinuous segments, this might return multiple spans for
    /// a single group.
    fn all_ids_in_groups(&self, groups: &[Group]) -> Result<IdSet>;

    /// Find all ids covered by a specific level of segments.
    ///
    /// This function assumes that segments are built in order,
    /// and higher level segments do not cover more than lower
    /// levels.
    ///
    /// That is, if range `Id(x)..Id(y)` is covered by segment
    /// level `n`. Then segment level `n+1` would cover `Id(x)..Id(p)`
    /// and not cover `Id(p)..Id(y)` (x <= p <= y). In other words,
    /// the following cases are forbidden:
    ///
    /// ```plain,ignore
    ///     level n     [-------covered-------]
    ///     level n+1   [covered] gap [covered]
    ///
    ///     level n     [---covered---]
    ///     level n+1   gap [covered]
    ///
    ///     level n     [covered] gap
    ///     level n+1   [---covered---]
    /// ```
    ///
    /// The following cases are okay:
    ///
    /// ```plain,ignore
    ///     level n     [---covered---]
    ///     level n+1   [covered] gap
    ///
    ///     level n     [---covered---]
    ///     level n+1   [---covered---]
    /// ```
    fn all_ids_in_segment_level(&self, level: Level) -> Result<IdSet> {
        let all_ids = self.all_ids_in_groups(&Group::ALL)?;
        if level == 0 {
            return Ok(all_ids);
        }

        let mut result = IdSet::empty();
        for span in all_ids.as_spans() {
            // In this span:
            //
            //      [---------span--------]
            //                 seg-]
            //
            // If we found the right side of a segment, then we can
            // assume the segments cover till the left side without
            // checking the actual segments:
            //
            //      [---------span--------]
            //      [seg][...][seg-]
            let seg = self.iter_segments_descending(span.high, level)?.next();
            if let Some(seg) = seg {
                let seg = seg?;
                let seg_span = seg.span()?;
                if span.contains(seg_span.high) {
                    // sanity check
                    if !span.contains(seg_span.low) {
                        return programming(format!(
                            "span {:?} from all_ids_in_groups should cover all segment {:?}",
                            span, seg
                        ));
                    }
                    result.push(span.low..=seg_span.high);
                }
            }
        }
        Ok(result)
    }

    /// Find segments that covers `id..` range at the given level, within a same group.
    fn next_segments(&self, id: Id, level: Level) -> Result<Vec<Segment>>;

    /// Find segments that fully covers the given range. Return segments in ascending order.
    fn segments_in_span_ascending(&self, span: Span, level: Level) -> Result<Vec<Segment>> {
        let mut iter = self.iter_segments_ascending(span.low, level)?;
        let mut result = Vec::new();
        while let Some(item) = iter.next() {
            let seg = item?;
            let seg_span = seg.span()?;
            if seg_span.low >= span.low && seg_span.high <= span.high {
                result.push(seg);
            }
            if seg_span.low > span.high {
                break;
            }
        }
        Ok(result)
    }

    /// Iterate through segments at the given level in descending order.
    fn iter_segments_descending<'a>(
        &'a self,
        max_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>>;

    /// Iterate through segments at the given level in ascending order.
    fn iter_segments_ascending<'a>(
        &'a self,
        min_high_id: Id,
        level: Level,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a + Send + Sync>>;

    /// Iterate through `(parent_id, segment)` for master flat segments
    /// that have a parent in the given span.
    ///
    /// The order of returned segments is implementation-specific.
    /// Different stores might return different order.
    fn iter_flat_segments_with_parent_span<'a>(
        &'a self,
        parent_span: Span,
    ) -> Result<Box<dyn Iterator<Item = Result<(Id, Segment)>> + 'a>>;

    /// Iterate through flat segments that have the given parent.
    ///
    /// The order of returned segments is implementation-specific.
    /// Different stores might return different order.
    fn iter_flat_segments_with_parent<'a>(
        &'a self,
        parent: Id,
    ) -> Result<Box<dyn Iterator<Item = Result<Segment>> + 'a>>;

    /// Remove all non master Group identifiers from the DAG.
    fn remove_non_master(&mut self) -> Result<()>;

    /// Attempt to merge the flat `segment` with the last flat segment to reduce
    /// fragmentation.
    ///
    /// ```plain,ignore
    /// [---last segment---] [---segment---]
    ///                    ^---- the only parent of segment
    /// [---merged segment-----------------]
    /// ```
    ///
    /// Return the merged segment if it's mergeable.
    fn maybe_merged_flat_segment(&self, segment: &Segment) -> Result<Option<Segment>> {
        let level = segment.level()?;
        if level != 0 {
            // Only applies to flat segments.
            return Ok(None);
        }
        if segment.has_root()? {
            // Cannot merge if segment has roots (implies no parent for a flat segment).
            return Ok(None);
        }
        let span = segment.span()?;
        let group = span.low.group();
        if group != Group::MASTER {
            // Do not merge non-master groups for simplicity.
            return Ok(None);
        }
        let parents = segment.parents()?;
        if parents.len() != 1 || parents[0] + 1 != span.low {
            // Cannot merge - span.low dos not have parent [low-1] (non linear).
            return Ok(None);
        }
        let last_segment = match self.iter_segments_descending(span.low, 0)?.next() {
            Some(Ok(s)) => s,
            _ => return Ok(None), // Cannot merge - No last flat segment.
        };
        let last_span = last_segment.span()?;
        if last_span.high + 1 != span.low {
            // Cannot merge - Two spans are not connected.
            return Ok(None);
        }

        // Can merge!

        // Sanity check: No high-level segments should cover "last_span".
        // This is because we intentionally dropped the last (incomplete)
        // high-level segment when building.
        for lv in 1..=self.max_level()? {
            if self
                .find_segment_by_head_and_level(last_span.high, lv)?
                .is_some()
            {
                return bug(format!(
                    "lv{} segment should not cover last flat segment {:?}! ({})",
                    lv, last_span, "check build_high_level_segments"
                ));
            }
        }

        // Calculate the merged segment.
        let merged = {
            let last_parents = last_segment.parents()?;
            let flags = {
                let last_flags = last_segment.flags()?;
                let flags = segment.flags()?;
                (flags & SegmentFlags::ONLY_HEAD) | (last_flags & SegmentFlags::HAS_ROOT)
            };
            Segment::new(flags, level, last_span.low, span.high, &last_parents)
        };

        tracing::debug!(
            "merge flat segments {:?} + {:?} => {:?}",
            &last_segment,
            &segment,
            &merged
        );

        Ok(Some(merged))
    }
}

/// Used by `resize_flat_segment` functions.
pub(crate) fn get_deleted_inserted_spans(
    span: Span,
    new_high: Option<Id>,
) -> (Option<Span>, Option<Span>) {
    match new_high {
        Some(new_high) => match new_high.cmp(&span.high) {
            Ordering::Less => (Some(Span::from(new_high + 1..=span.high)), None),
            Ordering::Equal => (None, None),
            Ordering::Greater => (None, Some(Span::from(span.high + 1..=new_high))),
        },
        None => (Some(span), None),
    }
}

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

    use once_cell::sync::Lazy;

    use super::*;

    fn nid(id: u64) -> Id {
        Group::NON_MASTER.min_id() + id
    }
    //  0--1--2--3--4--5--10--11--12--13--N0--N1--N2--N5--N6
    //         \-6-7-8--9-/-----------------\-N3--N4--/
    static LEVEL0_HEAD2: Lazy<Segment> =
        Lazy::new(|| Segment::new(SegmentFlags::HAS_ROOT, 0 as Level, Id(0), Id(2), &[]));
    static LEVEL0_HEAD5: Lazy<Segment> =
        Lazy::new(|| Segment::new(SegmentFlags::ONLY_HEAD, 0 as Level, Id(3), Id(5), &[Id(2)]));
    static LEVEL0_HEAD9: Lazy<Segment> =
        Lazy::new(|| Segment::new(SegmentFlags::empty(), 0 as Level, Id(6), Id(9), &[Id(2)]));
    static LEVEL0_HEAD13: Lazy<Segment> = Lazy::new(|| {
        Segment::new(
            SegmentFlags::empty(),
            0 as Level,
            Id(10),
            Id(13),
            &[Id(5), Id(9)],
        )
    });

    static MERGED_LEVEL0_HEAD5: Lazy<Segment> = Lazy::new(|| {
        Segment::new(
            SegmentFlags::HAS_ROOT | SegmentFlags::ONLY_HEAD,
            0 as Level,
            Id(0),
            Id(5),
            &[],
        )
    });

    static LEVEL0_HEADN2: Lazy<Segment> =
        Lazy::new(|| Segment::new(SegmentFlags::empty(), 0 as Level, nid(0), nid(2), &[Id(13)]));
    static LEVEL0_HEADN4: Lazy<Segment> = Lazy::new(|| {
        Segment::new(
            SegmentFlags::empty(),
            0 as Level,
            nid(3),
            nid(4),
            &[nid(0), Id(9)],
        )
    });
    static LEVEL0_HEADN6: Lazy<Segment> = Lazy::new(|| {
        Segment::new(
            SegmentFlags::empty(),
            0 as Level,
            nid(5),
            nid(6),
            &[nid(2), nid(4)],
        )
    });

    static LEVEL1_HEAD13: Lazy<Segment> =
        Lazy::new(|| Segment::new(SegmentFlags::HAS_ROOT, 1 as Level, Id(0), Id(13), &[]));
    static LEVEL1_HEADN6: Lazy<Segment> = Lazy::new(|| {
        Segment::new(
            SegmentFlags::HAS_ROOT,
            1 as Level,
            nid(0),
            nid(6),
            &[Id(13)],
        )
    });

    // Helpers
    const ROOT: SegmentFlags = SegmentFlags::HAS_ROOT;
    const EMPTY: SegmentFlags = SegmentFlags::empty();

    const M: Group = Group::MASTER;
    const N: Group = Group::NON_MASTER;

    fn seg(flags: SegmentFlags, group: Group, low: u64, high: u64, parents: &[u64]) -> Segment {
        Segment::new(
            flags,
            0,
            group.min_id() + low,
            group.min_id() + high,
            &parents.iter().copied().map(Id).collect::<Vec<_>>(),
        )
    }

    /// High-level segment.
    fn hseg(
        level: Level,
        flags: SegmentFlags,
        group: Group,
        low: u64,
        high: u64,
        parents: &[u64],
    ) -> Segment {
        Segment::new(
            flags,
            level,
            group.min_id() + low,
            group.min_id() + high,
            &parents.iter().copied().map(Id).collect::<Vec<_>>(),
        )
    }

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

    fn fmt_iter<T: fmt::Debug>(iter: impl Iterator<Item = Result<T>>) -> Vec<String> {
        iter.map(|i| fmt(i.unwrap())).collect()
    }

    fn insert_segments(store: &mut dyn IdDagStore, segments: Vec<&Segment>) {
        for segment in segments {
            store.insert_segment(segment.clone()).unwrap();
        }
    }

    fn get_segments() -> Vec<&'static Segment> {
        vec![
            &LEVEL0_HEAD2,
            &LEVEL0_HEAD5,
            &LEVEL0_HEAD9,
            &LEVEL0_HEAD13,
            &LEVEL1_HEAD13,
            &LEVEL0_HEADN2,
            &LEVEL0_HEADN4,
            &LEVEL0_HEADN6,
            &LEVEL1_HEADN6,
        ]
    }

    fn segments_to_owned(segments: &[&Segment]) -> Vec<Segment> {
        segments.into_iter().cloned().cloned().collect()
    }

    fn test_find_segment_by_head_and_level(store: &dyn IdDagStore) {
        let segment = store
            .find_segment_by_head_and_level(Id(13), 1 as Level)
            .unwrap()
            .unwrap();
        assert_eq!(&segment, LEVEL1_HEAD13.deref());

        let opt_segment = store
            .find_segment_by_head_and_level(Id(2), 0 as Level)
            .unwrap();
        assert!(opt_segment.is_none());

        let segment = store
            .find_segment_by_head_and_level(Id(5), 0 as Level)
            .unwrap()
            .unwrap();
        assert_eq!(&segment, MERGED_LEVEL0_HEAD5.deref());

        let segment = store
            .find_segment_by_head_and_level(nid(2), 0 as Level)
            .unwrap()
            .unwrap();
        assert_eq!(&segment, LEVEL0_HEADN2.deref());
    }

    fn test_find_flat_segment_including_id(store: &dyn IdDagStore) {
        let segment = store
            .find_flat_segment_including_id(Id(10))
            .unwrap()
            .unwrap();
        assert_eq!(&segment, LEVEL0_HEAD13.deref());

        let segment = store
            .find_flat_segment_including_id(Id(0))
            .unwrap()
            .unwrap();
        assert_eq!(&segment, MERGED_LEVEL0_HEAD5.deref());

        let segment = store
            .find_flat_segment_including_id(Id(2))
            .unwrap()
            .unwrap();
        assert_eq!(&segment, MERGED_LEVEL0_HEAD5.deref());

        let segment = store
            .find_flat_segment_including_id(Id(5))
            .unwrap()
            .unwrap();
        assert_eq!(&segment, MERGED_LEVEL0_HEAD5.deref());

        let segment = store
            .find_flat_segment_including_id(nid(1))
            .unwrap()
            .unwrap();
        assert_eq!(&segment, LEVEL0_HEADN2.deref());
    }

    fn test_all_ids_in_groups(store: &mut dyn IdDagStore) {
        let all_id_str = |store: &dyn IdDagStore, groups| {
            format!("{:?}", store.all_ids_in_groups(groups).unwrap())
        };

        // Insert some discontinuous segments. Then query all_ids_in_groups.
        store.insert_segment(seg(ROOT, M, 10, 20, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=20");

        store.insert_segment(seg(ROOT, M, 30, 40, &[])).unwrap();
        store.insert_segment(seg(ROOT, M, 50, 60, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=20 30..=40 50..=60");

        // Insert adjacent segments and check that spans are merged.
        store.insert_segment(seg(EMPTY, M, 41, 45, &[40])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=20 30..=45 50..=60");

        store.insert_segment(seg(EMPTY, M, 46, 49, &[45])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=20 30..=60");

        store.insert_segment(seg(EMPTY, M, 61, 70, &[60])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=20 30..=70");

        store.insert_segment(seg(ROOT, M, 21, 29, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "10..=70");

        store.insert_segment(seg(ROOT, M, 0, 5, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "0..=5 10..=70");

        store.insert_segment(seg(ROOT, M, 6, 9, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "0..=70");

        // Spans in the non-master group.
        store.insert_segment(seg(EMPTY, N, 0, 10, &[])).unwrap();
        store.insert_segment(seg(EMPTY, N, 20, 30, &[])).unwrap();
        assert_eq!(all_id_str(store, &[N]), "N0..=N10 N20..=N30");
        store.insert_segment(seg(EMPTY, N, 11, 15, &[])).unwrap();
        assert_eq!(all_id_str(store, &[N]), "N0..=N15 N20..=N30");
        store.insert_segment(seg(EMPTY, N, 17, 19, &[])).unwrap();
        assert_eq!(all_id_str(store, &[N]), "N0..=N15 N17..=N30");
        store.insert_segment(seg(EMPTY, N, 16, 16, &[])).unwrap();
        assert_eq!(all_id_str(store, &[M]), "0..=70");
        assert_eq!(all_id_str(store, &[N]), "N0..=N30");
        assert_eq!(all_id_str(store, &[M, N]), "0..=70 N0..=N30");

        store.remove_non_master().unwrap();
        assert_eq!(all_id_str(store, &[N]), "");
        assert_eq!(all_id_str(store, &[M, N]), "0..=70");
    }

    fn test_all_ids_in_segment_level(store: &mut dyn IdDagStore) {
        let level_id_str = |store: &dyn IdDagStore, level| {
            format!("{:?}", store.all_ids_in_segment_level(level).unwrap())
        };

        // Insert some discontinuous segments. Then query all_ids_in_groups.
        insert_segments(
            store,
            vec![
                &seg(ROOT, M, 0, 10, &[]),
                &seg(EMPTY, M, 11, 20, &[9]),
                &seg(EMPTY, M, 21, 30, &[15]),
                &seg(ROOT, M, 50, 60, &[]),
                &seg(EMPTY, M, 61, 70, &[51]),
                &seg(EMPTY, M, 71, 75, &[51]),
                &seg(EMPTY, M, 76, 80, &[51]),
                &seg(EMPTY, M, 81, 85, &[51]),
                &seg(ROOT, M, 100, 110, &[]),
                &seg(EMPTY, M, 111, 120, &[105]),
                &seg(EMPTY, M, 121, 130, &[115]),
                &seg(ROOT, N, 0, 10, &[]),
                &seg(EMPTY, N, 11, 20, &[9]),
                &seg(EMPTY, N, 21, 30, &[15]),
                &hseg(1, ROOT, M, 0, 10, &[]),
                &hseg(1, EMPTY, M, 11, 20, &[9]),
                &hseg(1, ROOT, M, 50, 70, &[]),
                &hseg(1, EMPTY, M, 71, 80, &[51]),
                &hseg(1, ROOT, M, 100, 120, &[]),
                &hseg(1, ROOT, N, 0, 30, &[]),
                &hseg(2, ROOT, M, 50, 80, &[]),
                &hseg(2, ROOT, M, 100, 120, &[]),
            ],
        );

        assert_eq!(level_id_str(store, 0), "0..=30 50..=85 100..=130 N0..=N30");
        assert_eq!(level_id_str(store, 1), "0..=20 50..=80 100..=120 N0..=N30");
        assert_eq!(level_id_str(store, 2), "50..=80 100..=120");
        assert_eq!(level_id_str(store, 3), "");
    }

    fn test_discontinuous_merges(store: &mut dyn IdDagStore) {
        insert_segments(
            store,
            vec![
                &seg(ROOT, M, 0, 10, &[]),
                &seg(EMPTY, M, 20, 30, &[5]),
                &seg(EMPTY, M, 11, 15, &[10]),
                &seg(EMPTY, M, 31, 35, &[30]),
            ],
        );

        let iter = store.iter_segments_descending(Id(25), 0).unwrap();
        assert_eq!(fmt_iter(iter), ["R0-15[]"]);

        // 0-10 and 11-15 are merged.
        let seg = store.find_segment_by_head_and_level(Id(10), 0).unwrap();
        assert_eq!(fmt(seg), "None");
        let seg = store.find_segment_by_head_and_level(Id(15), 0).unwrap();
        assert_eq!(fmt(seg), "Some(R0-15[])");

        // 20-30 and 31-35 are merged.
        let seg = store.find_segment_by_head_and_level(Id(30), 0).unwrap();
        assert_eq!(fmt(seg), "None");
        let seg = store.find_segment_by_head_and_level(Id(35), 0).unwrap();
        assert_eq!(fmt(seg), "Some(20-35[5])");

        // 0-10 and 11-15 are merged.
        let seg = store.find_flat_segment_including_id(Id(9)).unwrap();
        assert_eq!(fmt(seg), "Some(R0-15[])");
        let seg = store.find_flat_segment_including_id(Id(14)).unwrap();
        assert_eq!(fmt(seg), "Some(R0-15[])");
        let seg = store.find_flat_segment_including_id(Id(16)).unwrap();
        assert_eq!(fmt(seg), "None");

        // 20-30 and 31-35 are merged.
        let seg = store.find_flat_segment_including_id(Id(35)).unwrap();
        assert_eq!(fmt(seg), "Some(20-35[5])");
        let seg = store.find_flat_segment_including_id(Id(36)).unwrap();
        assert_eq!(fmt(seg), "None");

        // Parent lookup.
        let iter = store.iter_flat_segments_with_parent(Id(5)).unwrap();
        assert_eq!(fmt_iter(iter), ["20-35[5]"]);
        let iter = store.iter_flat_segments_with_parent(Id(10)).unwrap();
        assert_eq!(fmt_iter(iter), [] as [String; 0]);
        let iter = store.iter_flat_segments_with_parent(Id(30)).unwrap();
        assert_eq!(fmt_iter(iter), [] as [String; 0]);
    }

    fn test_next_segments(store: &dyn IdDagStore) {
        let segments = store.next_segments(Id(4), 0 as Level).unwrap();
        let expected = segments_to_owned(&[&MERGED_LEVEL0_HEAD5, &LEVEL0_HEAD9, &LEVEL0_HEAD13]);
        assert_eq!(segments, expected);

        let segments = store.next_segments(Id(14), 0 as Level).unwrap();
        assert!(segments.is_empty());

        let segments = store.next_segments(Id(0), 1 as Level).unwrap();
        let expected = segments_to_owned(&[&LEVEL1_HEAD13]);
        assert_eq!(segments, expected);

        let segments = store.next_segments(Id(0), 2 as Level).unwrap();
        assert!(segments.is_empty());
    }

    fn test_max_level(store: &dyn IdDagStore) {
        assert_eq!(store.max_level().unwrap(), 1);
    }

    fn test_empty_store_max_level(store: &dyn IdDagStore) {
        assert_eq!(store.max_level().unwrap(), 0);
    }

    fn test_iter_segments_descending(store: &dyn IdDagStore) {
        let answer = store
            .iter_segments_descending(Id(12), 0)
            .unwrap()
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let expected = segments_to_owned(&[&LEVEL0_HEAD9, &MERGED_LEVEL0_HEAD5]);
        assert_eq!(answer, expected);

        let mut answer = store.iter_segments_descending(Id(1), 0).unwrap();
        assert!(answer.next().is_none());

        let answer = store
            .iter_segments_descending(Id(13), 1)
            .unwrap()
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let expected = segments_to_owned(&[&LEVEL1_HEAD13]);
        assert_eq!(answer, expected);

        let mut answer = store.iter_segments_descending(Id(5), 2).unwrap();
        assert!(answer.next().is_none());
    }

    fn test_iter_segments_ascending(store: &dyn IdDagStore) {
        let answer = store
            .iter_segments_ascending(Id(12), 0)
            .unwrap()
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let expected = segments_to_owned(&[
            &LEVEL0_HEAD13,
            &LEVEL0_HEADN2,
            &LEVEL0_HEADN4,
            &LEVEL0_HEADN6,
        ]);
        assert_eq!(answer, expected);

        let answer = store
            .iter_segments_ascending(Id(14), 0)
            .unwrap()
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let expected = segments_to_owned(&[&LEVEL0_HEADN2, &LEVEL0_HEADN4, &LEVEL0_HEADN6]);
        assert_eq!(answer, expected);

        let mut answer = store.iter_segments_ascending(nid(7), 0).unwrap();
        assert!(answer.next().is_none());

        let answer = store
            .iter_segments_ascending(nid(3), 1)
            .unwrap()
            .collect::<Result<Vec<_>>>()
            .unwrap();
        let expected = segments_to_owned(&[&LEVEL1_HEADN6]);
        assert_eq!(answer, expected);

        let mut answer = store.iter_segments_ascending(Id(5), 2).unwrap();
        assert!(answer.next().is_none());
    }

    fn test_store_iter_flat_segments_with_parent_span(store: &dyn IdDagStore) {
        let query = |span: Span| -> String {
            let iter = store.iter_flat_segments_with_parent_span(span).unwrap();
            let mut answer_str_list: Vec<String> =
                iter.map(|s| format!("{:?}", s.unwrap())).collect();
            answer_str_list.sort_unstable();
            answer_str_list.join(" ")
        };

        assert_eq!(query(Id(2).into()), "(2, 6-9[2])");
        assert_eq!(query((Id(0)..=Id(3)).into()), "(2, 6-9[2])");
        assert_eq!(query(Id(13).into()), "(13, N0-N2[13])");
        assert_eq!(query(Id(4).into()), "");
        assert_eq!(query(nid(2).into()), "(N2, N5-N6[N2, N4])");
    }

    fn test_store_iter_flat_segments_with_parent(store: &dyn IdDagStore) {
        let lookup = |id: Id| -> Vec<Segment> {
            let mut list = store
                .iter_flat_segments_with_parent(id)
                .unwrap()
                .collect::<Result<Vec<_>>>()
                .unwrap();
            list.sort_unstable_by_key(|seg| seg.low().unwrap());
            list
        };

        let answer = lookup(Id(2));
        // LEVEL0_HEAD5 is not in answer because it was merged into MERGED_LEVEL0_HEAD5
        // and MERGED_LEVEL0_HEAD5 no longer has parent 2.
        let expected = segments_to_owned(&[&LEVEL0_HEAD9]);
        assert_eq!(answer, expected);

        let answer = lookup(Id(13));
        let expected = segments_to_owned(&[&LEVEL0_HEADN2]);
        assert_eq!(answer, expected);

        let answer = lookup(Id(4));
        assert!(answer.is_empty());

        let answer = lookup(nid(2));
        let expected = segments_to_owned(&[&LEVEL0_HEADN6]);
        assert_eq!(answer, expected);

        let answer = lookup(Id(9));
        let expected = segments_to_owned(&[&LEVEL0_HEAD13, &LEVEL0_HEADN4]);
        assert_eq!(answer, expected);
    }

    fn test_remove_non_master(store: &mut dyn IdDagStore) {
        store.remove_non_master().unwrap();

        assert!(
            store
                .find_segment_by_head_and_level(nid(2), 0 as Level)
                .unwrap()
                .is_none()
        );
        assert!(
            store
                .find_flat_segment_including_id(nid(1))
                .unwrap()
                .is_none()
        );
        assert!(
            store
                .iter_flat_segments_with_parent_span(nid(2).into())
                .unwrap()
                .next()
                .is_none()
        );
    }

    pub(crate) fn test_remove_segment(store: &mut dyn IdDagStore) {
        // Prepare segments, 3 segments per group.
        let parents_nid_3_4 = [nid(4), nid(3)];
        let parents_id_12_nid_4 = [Id(12), nid(4)];
        let segs: Vec<(Id, Id, &[Id])> = vec![
            (Id(0), Id(5), &[]),
            (Id(6), Id(10), &[Id(4), Id(3)]),
            (Id(11), Id(15), &[Id(4)]),
            (nid(0), nid(5), &[]),
            (nid(6), nid(10), &parents_nid_3_4),
            (nid(11), nid(15), &parents_id_12_nid_4),
        ];
        let mut segs: Vec<Segment> = segs
            .into_iter()
            .map(|(low, high, parents)| {
                let flags = if parents.is_empty() { ROOT } else { EMPTY };
                Segment::new(flags, 0, low, high, parents)
            })
            .collect();
        // Also insert high-level segments.
        segs.push(Segment::new(ROOT, 1, Id(0), Id(10), &[]));
        segs.push(Segment::new(ROOT, 1, nid(0), nid(10), &[]));
        for seg in segs.clone() {
            store.insert_segment(seg).unwrap();
        }

        // Cannot remove segment with descendants, or high-level segments.
        for i in [0, 2, 3, 6, 7] {
            assert!(store.remove_flat_segment(&segs[i]).is_err());
        }

        let all_before_delete = store.all_ids_in_groups(&Group::ALL).unwrap();
        assert_eq!(format!("{:?}", &all_before_delete), "0..=15 N0..=N15");
        assert_eq!(
            dump_store_state(store, &all_before_delete),
            r#"
Lv0: R0-5[], 6-10[4, 3], 11-15[4], RN0-N5[], N6-N10[N4, N3], N11-N15[12, N4]
Lv1: R0-10[], RN0-N10[]
P->C: 3->6, 4->6, 4->11, 12->N11, N3->N6, N4->N6, N4->N11"#
        );

        // Remove the "middle" segment per group.
        for i in [1, 4] {
            store.remove_flat_segment(&segs[i]).unwrap();
        }

        // Check that the removed segments are actually removed.
        let all_after_delete = store.all_ids_in_groups(&Group::ALL).unwrap();
        assert_eq!(
            format!("{:?}", &all_after_delete),
            "0..=5 11..=15 N0..=N5 N11..=N15"
        );
        assert_eq!(
            dump_store_state(store, &all_before_delete),
            r#"
Lv0: R0-5[], 11-15[4], RN0-N5[], N11-N15[12, N4]
P->C: 4->11, 12->N11, N4->N11"#
        );
        let deleted_ids = IdSet::from_spans(vec![Id(6)..=Id(10), nid(6)..=nid(10)]);
        assert_eq!(dump_store_state(store, &deleted_ids), "");
    }

    pub(crate) fn test_resize_segment(store: &mut dyn IdDagStore) {
        // Prepare segments, 3 segments per group.
        let segs: Vec<(Id, Id, &[Id])> = vec![
            (Id(0), Id(100), &[]),
            (Id(200), Id(200), &[]),
            (nid(100), nid(200), &[Id(50)]),
            (nid(300), nid(400), &[Id(50)]),
        ];
        let mut segs: Vec<Segment> = segs
            .into_iter()
            .map(|(low, high, parents)| {
                let flags = if parents.is_empty() { ROOT } else { EMPTY };
                Segment::new(flags, 0, low, high, parents)
            })
            .collect();
        // Also insert high-level segments.
        segs.push(Segment::new(ROOT, 1, Id(0), Id(100), &[]));
        segs.push(Segment::new(ROOT, 1, nid(100), nid(200), &[Id(50)]));
        for seg in segs.clone() {
            store.insert_segment(seg).unwrap();
        }

        let all_before_resize = store.all_ids_in_groups(&Group::ALL).unwrap();
        assert_eq!(
            format!("{:?}", &all_before_resize),
            "0..=100 200 N100..=N200 N300..=N400"
        );
        assert_eq!(
            dump_store_state(store, &all_before_resize),
            r#"
Lv0: R0-100[], R200-200[], N100-N200[50], N300-N400[50]
Lv1: R0-100[], RN100-N200[50]
P->C: 50->N100, 50->N300"#
        );

        // Check error cases.
        let mut e = |i, id| -> String {
            store
                .resize_flat_segment(&segs[i], Some(id))
                .unwrap_err()
                .to_string()
        };

        // Cannot resize because of descendants.
        assert_eq!(
            e(0, Id(49)),
            "ProgrammingError: resize_flat_segment requires a segment without descendants, got R0-100[] with child segment N100-N200[50]"
        );

        // Cannot resize because of overlap.
        assert_eq!(
            e(0, Id(200)),
            "ProgrammingError: resize_flat_segment cannot overlap with existing segments (segment: R0-100[] new_high: Some(200), overlap: 200)"
        );

        // Cannot resize because of high < low.
        assert_eq!(
            e(1, Id(199)),
            "ProgrammingError: with_high got invalid input (segment: R200-200[] new_high: 199)"
        );

        // Cannot resize because of high and new_high are in different groups.
        assert_eq!(
            e(1, nid(0)),
            "ProgrammingError: with_high got invalid input (segment: R200-200[] new_high: N0)"
        );

        // Do resize.
        let mut resize = |i, id| store.resize_flat_segment(&segs[i], Some(id)).unwrap();

        // Shrink.
        resize(0, Id(50));
        resize(3, nid(350));

        // Grow.
        resize(1, Id(250));
        resize(2, nid(250));

        // Check state after resize.
        let all_after_resize = store.all_ids_in_groups(&Group::ALL).unwrap();
        assert_eq!(
            format!("{:?}", &all_after_resize),
            "0..=50 200..=250 N100..=N250 N300..=N350"
        );
        let all = all_after_resize.union(&all_before_resize);
        assert_eq!(
            dump_store_state(store, &all),
            r#"
Lv0: R0-50[], R200-250[], N100-N250[50], N300-N350[50]
P->C: 50->N100, 50->N300"#
        );
    }

    /// Dump the store state in the given `id_set` as a string for testing.
    pub(crate) fn dump_store_state(store: &dyn IdDagStore, id_set: &IdSet) -> String {
        let mut output = Vec::new();
        let max_level = store.max_level().unwrap();
        // Segments per level. Exercises the "head" index.
        for level in 0..=max_level {
            let mut level_segments = Vec::new();
            for &span in id_set.iter_span_asc() {
                let segs = store.segments_in_span_ascending(span, level).unwrap();
                for seg in segs {
                    if seg.level().unwrap() == level {
                        level_segments.push(format!("{:?}", seg));
                    }
                }
            }
            if !level_segments.is_empty() {
                output.push(format!("\nLv{}: {}", level, level_segments.join(", ")));
            }
        }
        // Parent indexes in the id_set. Exercises the "parent->child" index.
        let mut parent_child_relations = Vec::new();
        for &span in id_set.iter_span_asc() {
            let parent_child_segs = store
                .iter_flat_segments_with_parent_span(span)
                .unwrap()
                .collect::<Result<Vec<_>>>()
                .unwrap();
            let mut relations = parent_child_segs
                .into_iter()
                .map(|(parent_id, child_seg)| (parent_id, child_seg.low().unwrap()))
                .collect::<Vec<_>>();
            relations.sort_unstable();
            let mut relations = relations
                .into_iter()
                .map(|(parent, child)| format!("{:?}->{:?}", parent, child))
                .collect::<Vec<_>>();
            parent_child_relations.append(&mut relations);
        }
        if !parent_child_relations.is_empty() {
            output.push(format!("\nP->C: {}", parent_child_relations.join(", ")));
        }
        output.concat()
    }

    fn for_each_empty_store(f: impl Fn(&mut dyn IdDagStore)) {
        let mut store = InProcessStore::new();
        tracing::debug!("testing InProcessStore");
        f(&mut store);

        #[cfg(feature = "indexedlog-backend")]
        {
            let dir = tempfile::tempdir().unwrap();
            let mut store = IndexedLogStore::open(&dir.path()).unwrap();
            tracing::debug!("testing IndexedLogStore");
            f(&mut store);
        }
    }

    fn for_each_store(f: impl Fn(&mut dyn IdDagStore)) {
        for_each_empty_store(|store| {
            insert_segments(store, get_segments());
            f(store);
        })
    }

    #[test]
    fn test_multi_stores_insert() {
        // `for_each_store` does inserts, we care that nothings panics.
        for_each_store(|_store| ())
    }

    #[test]
    fn test_multi_stores_find_segment_by_head_and_level() {
        for_each_store(|store| test_find_segment_by_head_and_level(store));
    }

    #[test]
    fn test_multi_stores_find_flat_segment_including_id() {
        for_each_store(|store| test_find_flat_segment_including_id(store));
    }

    #[test]
    fn test_multi_stores_all_ids_in_groups() {
        for_each_empty_store(|store| {
            test_all_ids_in_groups(store);
        })
    }

    #[test]
    fn test_multi_stores_all_ids_in_segment_level() {
        for_each_empty_store(|store| {
            test_all_ids_in_segment_level(store);
        })
    }

    #[test]
    fn test_multi_stores_next_segments() {
        for_each_store(|store| test_next_segments(store));
    }

    #[test]
    fn test_multi_stores_max_level() {
        for_each_empty_store(|store| test_empty_store_max_level(store));
    }

    #[test]
    fn test_multi_stores_iter_segments_descending() {
        for_each_store(|store| test_iter_segments_descending(store));
    }

    #[test]
    fn test_multi_stores_iter_segments_ascending() {
        for_each_store(|store| test_iter_segments_ascending(store));
    }

    #[test]
    fn test_multi_stores_iter_flat_segments_with_parent_span() {
        for_each_store(|store| test_store_iter_flat_segments_with_parent_span(store));
    }

    #[test]
    fn test_multi_stores_iter_flat_segments_with_parent() {
        for_each_store(|store| test_store_iter_flat_segments_with_parent(store));
    }

    #[test]
    fn test_multi_stores_remove_non_master() {
        for_each_store(|store| test_remove_non_master(store));
    }

    #[test]
    fn test_multi_stores_discontinuous_merges() {
        for_each_empty_store(|store| test_discontinuous_merges(store));
    }

    #[test]
    fn test_multi_stores_remove_segment() {
        for_each_empty_store(|store| test_remove_segment(store));
    }

    #[test]
    fn test_multi_stores_resize_segment() {
        for_each_empty_store(|store| test_resize_segment(store));
    }
}


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

//! # idmap
//!
//! See [`IdMap`] for the main structure.

use std::borrow::Cow;

use crate::errors::bug;
use crate::id::Group;
use crate::id::Id;
use crate::id::VertexName;
use crate::ops::IdConvert;
use crate::ops::Parents;
use crate::segment::PreparedFlatSegments;
use crate::types_ext::PreparedFlatSegmentsExt;
use crate::Error;
use crate::IdSet;
use crate::Result;

#[cfg(any(test, feature = "indexedlog-backend"))]
mod indexedlog_idmap;
mod mem_idmap;

#[cfg(any(test, feature = "indexedlog-backend"))]
pub use indexedlog_idmap::IdMap;
pub(crate) use mem_idmap::CoreMemIdMap;
pub use mem_idmap::MemIdMap;

/// DAG-aware write operations.
#[async_trait::async_trait]
pub trait IdMapAssignHead: IdConvert + IdMapWrite {
    /// Assign an id for a head in a DAG. This implies ancestors of the
    /// head will also have ids assigned.
    ///
    /// This function is incremental. If the head or any of its ancestors
    /// already have an id stored in this map, the existing ids will be
    /// reused.
    ///
    /// This function needs roughly `O(N)` heap memory. `N` is the number of
    /// ids to assign. When `N` is very large, try assigning ids to a known
    /// ancestor first.
    ///
    /// New `id`s inserted by this function will have the specified `group`.
    /// Existing `id`s that are ancestors of `head` will get re-assigned
    /// if they have a higher `group`.
    ///
    /// `covered_ids` specifies what ranges of `Id`s are already covered.
    /// This is usually obtained from `IdDag::all_ids_in_groups(&Group::ALL)`.
    /// `IdMap` itself might not be able to provide that information
    /// efficiently because it might be lazy. `covered_ids` will be updated
    /// to cover newly inserted `Id`s.
    ///
    /// `reserved_ids` specifies what ranges are reserved for future growth
    /// of other important heads (usually a couple of mainline branches that
    /// are long-lived, growing, and used by many people). This is useful
    /// to reduce fragmentation.
    async fn assign_head(
        &mut self,
        head: VertexName,
        parents_by_name: &dyn Parents,
        group: Group,
        covered_ids: &mut IdSet,
        reserved_ids: &IdSet,
    ) -> Result<PreparedFlatSegments> {
        // There are some interesting cases to optimize the numbers:
        //
        // C     For a merge C, it has choice to assign numbers to A or B
        // |\    first (A and B are abstract branches that have many nodes).
        // A B   Suppose branch A is linear and B have merges, and D is
        // |/    (::A & ::B). Then:
        // D
        //
        // - If `D` is empty or already assigned, it's better to assign A last.
        //   This is because (A+C) can then always form a segment regardless of
        //   the complexity of B:
        //
        //      B   A   C       vs.        A   B   C
        //     ~~~  ^^^^^                     ~~~
        //     xxxxxx                          *****
        //                                 xxxxx
        //
        //   [~]: Might be complex (ex. many segments)
        //   [^]: Can always form a segment. (better)
        //   [*]: Can only be a segment if segment size is large enough.
        //   [x]: Cannot form a segment.
        //
        // - If `D` is not empty (and not assigned), it _might_ be better to
        //   assign D and A first. This provides benefits for A and D to be
        //   continuous, with the downside that A and C are not continuous.
        //
        //   A typical pattern is one branch continuously merges into the other
        //   (see also segmented-changelog.pdf, page 19):
        //
        //        B---D---F
        //         \   \   \
        //      A---C---E---G
        //
        // The code below is optimized for cases where p1 branch is linear,
        // but p2 branch is not.
        //
        // However, the above visit order (first parent last) is not optimal
        // for incremental build case with pushrebase branches. Because
        // pushrebase uses the first parent as the mainline. For example:
        //
        //    A---B---C-...-D---M  (parents(M) = [D, F])
        //                     /
        //              E-...-F
        //
        // The A ... M branch is the mainline. The head of the mainline
        // was A ... D, then M. An incremental build job might have built up
        // A, B, ..., D before it sees M. In this case it's better to make
        // the incremental build finish the A ... D part before jumping to
        // E ... F.
        //
        // We choose first parent last order if `covered` is empty, or when
        // visiting ancestors of non-first parents.
        let mut outcome = PreparedFlatSegments::default();

        #[derive(Copy, Clone, Debug)]
        enum VisitOrder {
            /// Visit the first parent first.
            FirstFirst,
            /// Visit the first parent last.
            FirstLast,
        }

        // Emulate the stack in heap to avoid overflow.
        #[derive(Debug)]
        enum Todo {
            /// Visit parents. Finally assign self. This will eventually turn into AssignedId.
            Visit {
                head: VertexName,

                /// The `Id` in `IdMap` that is known assigned to the `head`.
                /// This can be non-`None` if `IdMap` has more entries than `IdDag`.
                known_id: Option<Id>,

                order: VisitOrder,
            },

            /// Assign an `Id` if not assigned. Their parents are prepared in the
            /// `parent_ids` stack. `Assign` `head` and `Visit` `head`'s parents
            /// are pushed together so the `Visit` entries can turn into `Id`s in
            /// the `parent_ids` stack.
            Assign {
                /// The vertex to assign. Its parents are already visited and assigned.
                head: VertexName,

                /// The `Id` in `IdMap` that is known assigned to the `head`.
                /// This can be non-`None` if `IdMap` has more entries than `IdDag`.
                known_id: Option<Id>,

                /// The number of parents, at the end of the `parent_ids`.
                parent_len: usize,

                /// The order of parents if extracted from `parent_ids`.
                order: VisitOrder,
            },

            /// Assigned Id. Will be picked by and pushed to the current `parent_ids` stack.
            AssignedId { id: Id },
        }
        use Todo::Assign;
        use Todo::AssignedId;
        use Todo::Visit;
        let mut parent_ids: Vec<Id> = Vec::new();

        let mut todo_stack: Vec<Todo> = {
            let order = if covered_ids.is_empty() {
                // Assume re-building from scratch.
                VisitOrder::FirstLast
            } else {
                // Assume incremental updates with pushrebase.
                VisitOrder::FirstFirst
            };
            vec![Visit {
                head: head.clone(),
                known_id: None,
                order,
            }]
        };
        while let Some(todo) = todo_stack.pop() {
            tracing::trace!(target: "dag::assign", "todo: {:?}", &todo);
            match todo {
                Visit {
                    head,
                    known_id,
                    order,
                } => {
                    // If the id was not assigned, or was assigned to a higher group,
                    // (re-)assign it to this group.
                    //
                    // PERF: This might trigger remote fetch too frequently.
                    let known_id = match known_id {
                        Some(id) => Some(id),
                        None => self.vertex_id_with_max_group(&head, group).await?,
                    };
                    match known_id {
                        Some(id) if covered_ids.contains(id) => todo_stack.push(AssignedId { id }),
                        _ => {
                            let parents = parents_by_name.parent_names(head.clone()).await?;
                            tracing::trace!(target: "dag::assign", "visit {:?} ({:?}) with parents {:?}", &head, known_id, &parents);
                            todo_stack.push(Assign {
                                head,
                                known_id,
                                parent_len: parents.len(),
                                order,
                            });
                            let mut visit = parents;
                            match order {
                                VisitOrder::FirstFirst => {}
                                VisitOrder::FirstLast => visit.reverse(),
                            }
                            for (i, p) in visit.into_iter().enumerate() {
                                // If the parent was not assigned, or was assigned to a higher group,
                                // (re-)assign the parent to this group.
                                match self.vertex_id_with_max_group(&p, group).await {
                                    Ok(Some(id)) if covered_ids.contains(id) => {
                                        todo_stack.push(AssignedId { id })
                                    }
                                    // Go deeper if IdMap has the entry but IdDag misses it.
                                    Ok(Some(id)) => todo_stack.push(Visit {
                                        head: p,
                                        known_id: Some(id),
                                        order,
                                    }),
                                    Ok(None) => {
                                        let parent_order = match (order, i) {
                                            (VisitOrder::FirstFirst, 0) => VisitOrder::FirstFirst,
                                            _ => VisitOrder::FirstLast,
                                        };
                                        todo_stack.push(Visit {
                                            head: p,
                                            known_id: None,
                                            order: parent_order,
                                        })
                                    }
                                    Err(e) => return Err(e),
                                }
                            }
                        }
                    }
                }
                Assign {
                    head,
                    known_id,
                    parent_len,
                    order,
                } => {
                    let parent_start = parent_ids.len() - parent_len;
                    let known_id = match known_id {
                        Some(id) => Some(id),
                        None => self.vertex_id_with_max_group(&head, group).await?,
                    };
                    let id = match known_id {
                        Some(id) if covered_ids.contains(id) => id,
                        _ => {
                            let parents = match order {
                                VisitOrder::FirstLast => Cow::Borrowed(&parent_ids[parent_start..]),
                                VisitOrder::FirstFirst => Cow::Owned(
                                    parent_ids[parent_start..]
                                        .iter()
                                        .cloned()
                                        .rev()
                                        .collect::<Vec<_>>(),
                                ),
                            };
                            let id = match known_id {
                                Some(id) => id,
                                None => {
                                    let candidate_id = match parents.iter().max() {
                                        Some(&max_parent_id) => {
                                            (max_parent_id + 1).max(group.min_id())
                                        }
                                        None => group.min_id(),
                                    };
                                    adjust_candidate_id(
                                        self,
                                        covered_ids,
                                        reserved_ids,
                                        candidate_id,
                                    )
                                    .await?
                                }
                            };
                            if id.group() != group {
                                return Err(Error::IdOverflow(group));
                            }
                            covered_ids.push(id);
                            if known_id.is_none() {
                                tracing::trace!(target: "dag::assign", "assign {:?} = {:?}", &head, id);
                                self.insert(id, head.as_ref()).await?;
                            } else {
                                tracing::trace!(target: "dag::assign", "assign {:?} = {:?} (known)", &head, id);
                            }
                            if parents.iter().any(|&p| p >= id) {
                                return bug(format!(
                                    "IdMap Ids are not topo-sorted: {:?} ({:?}) has parent ids {:?}",
                                    id, head, &parents,
                                ));
                            }
                            outcome.push_edge(id, &parents);
                            id
                        }
                    };
                    parent_ids.truncate(parent_start);
                    todo_stack.push(AssignedId { id });
                }
                AssignedId { id } => {
                    if !covered_ids.contains(id) {
                        return bug(format!(
                            concat!(
                                "attempted to assign id with wrong order: {:?} ",
                                "is being pushed as parent id but it cannot be used ",
                                "because it is not yet covered by IdDag",
                            ),
                            &id
                        ));
                    }
                    parent_ids.push(id);
                }
            }
        }

        Ok(outcome)
    }
}

/// Pick a minimal `n`, so `candidate_id + n` is an `Id` that is not "covered",
/// not "reserved", and not in the "map".  Return the picked `Id`.
async fn adjust_candidate_id(
    map: &(impl IdConvert + ?Sized),
    covered_ids: &IdSet,
    reserved_ids: &IdSet,
    mut candidate_id: Id,
) -> Result<Id> {
    loop {
        // (Fast) test using covered_ids + reserved_ids.
        loop {
            if let Some(span) = covered_ids.span_contains(candidate_id) {
                candidate_id = span.high + 1;
                continue;
            }
            if let Some(span) = reserved_ids.span_contains(candidate_id) {
                candidate_id = span.high + 1;
                continue;
            }
            break;
        }
        // (Slow) test using the IdMap.
        // PERF: This lacks of batching if it forms a loop. But it
        // is also expected to be rare - only when the server
        // tailer (assuming only one tailer is writing globally) is
        // killed abnormally, *and* the branch being assigned has
        // non-fast-forward move, this code path becomes useful.
        //
        // Technically, not using `locally` is more correct in a
        // lazy `IdMap`. However, lazy `IdMap` is only used by
        // client (local) dag, which ensures `IdMap` and `IdDag`
        // are in-sync, meaning that the above `covered_ids` check
        // is sufficient. So this is really only protecting the
        // server's out-of-sync `IdMap` use-case, where the
        // `locally` variant is the same as the non-`locally`,
        // since the server has a non-lazy `IdMap`.
        if let [true] = &map.contains_vertex_id_locally(&[candidate_id]).await?[..] {
            candidate_id = candidate_id + 1;
        } else {
            break;
        }
    }
    Ok(candidate_id)
}

impl<T> IdMapAssignHead for T where T: IdConvert + IdMapWrite {}

/// Write operations for IdMap.
#[async_trait::async_trait]
pub trait IdMapWrite {
    /// Insert a new `(id, name)` pair to the map.
    ///
    /// The `id` and `name` mapping should be unique, it's an error to map an id
    /// to multiple names, or map a name to multiple ids. Note: older versions
    /// of `IdMap` allowed mapping a name to a non-master Id, then a master Id
    /// (in this order), and the master Id is used for lookups. This is no
    /// longer permitted.
    async fn insert(&mut self, id: Id, name: &[u8]) -> Result<()>;
    /// Remove ids in the range `low..=high` and their associated names.
    /// Return removed names.
    async fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>>;
}

#[cfg(test)]
mod tests {
    use nonblocking::non_blocking_result as r;
    #[cfg(feature = "indexedlog-backend")]
    use tempfile::tempdir;

    use super::*;
    #[cfg(feature = "indexedlog-backend")]
    use crate::ops::Persist;
    #[cfg(feature = "indexedlog-backend")]
    use crate::ops::PrefixLookup;

    #[cfg(feature = "indexedlog-backend")]
    #[test]
    fn test_basic_operations() {
        let dir = tempdir().unwrap();
        let mut map = IdMap::open(dir.path()).unwrap();
        let lock = map.lock().unwrap();
        map.reload(&lock).unwrap();
        map.insert(Id(1), b"abc").unwrap();
        map.insert(Id(2), b"def").unwrap();
        map.insert(Id(10), b"ghi").unwrap();
        map.insert(Id(11), b"ghi").unwrap_err(); // ghi maps to 10
        map.insert(Id(10), b"ghi2").unwrap_err(); // 10 maps to ghi

        // Test another group.
        let id = Group::NON_MASTER.min_id();
        map.insert(id, b"jkl").unwrap();
        map.insert(id, b"jkl").unwrap();
        map.insert(id, b"jkl2").unwrap_err(); // id maps to jkl
        map.insert(id + 1, b"jkl2").unwrap();
        map.insert(id + 2, b"jkl2").unwrap_err(); // jkl2 maps to id + 1
        map.insert(Id(15), b"jkl2").unwrap_err(); // reassign jkl2 to master group - error.
        map.insert(id + 3, b"abc").unwrap_err(); // reassign abc to non-master group - error.

        // Test hex lookup.
        assert_eq!(0x6a, b'j');
        assert_eq!(
            r(map.vertexes_by_hex_prefix(b"6a", 3)).unwrap(),
            [
                VertexName::from(&b"jkl"[..]),
                VertexName::from(&b"jkl2"[..])
            ]
        );
        assert_eq!(
            r(map.vertexes_by_hex_prefix(b"6a", 1)).unwrap(),
            [VertexName::from(&b"jkl"[..])]
        );
        assert!(r(map.vertexes_by_hex_prefix(b"6b", 1)).unwrap().is_empty());

        for _ in 0..=1 {
            assert_eq!(map.find_name_by_id(Id(1)).unwrap().unwrap(), b"abc");
            assert_eq!(map.find_name_by_id(Id(2)).unwrap().unwrap(), b"def");
            assert!(map.find_name_by_id(Id(3)).unwrap().is_none());
            assert_eq!(map.find_name_by_id(Id(10)).unwrap().unwrap(), b"ghi");

            assert_eq!(map.find_id_by_name(b"abc").unwrap().unwrap().0, 1);
            assert_eq!(map.find_id_by_name(b"def").unwrap().unwrap().0, 2);
            assert_eq!(map.find_id_by_name(b"ghi").unwrap().unwrap().0, 10);
            assert_eq!(map.find_id_by_name(b"jkl").unwrap().unwrap(), id);
            assert_eq!(
                format!("{:?}", map.find_id_by_name(b"jkl2").unwrap().unwrap()),
                "N1"
            );
            assert!(map.find_id_by_name(b"jkl3").unwrap().is_none());
        }

        // Test Debug
        assert_eq!(
            format!("{:?}", &map),
            r#"IdMap {
  abc: 1,
  def: 2,
  ghi: 10,
  jkl: N0,
  jkl2: N1,
}
"#
        );
    }

    #[test]
    fn test_remove_range() {
        let map = MemIdMap::new();
        check_remove_range(map);

        #[cfg(feature = "indexedlog-backend")]
        {
            let dir = tempdir().unwrap();
            let path = dir.path();
            let map = IdMap::open(path).unwrap();
            check_remove_range(map);
        }
    }

    fn check_remove_range(mut map: impl IdConvert + IdMapWrite) {
        let items: &[(Id, &[u8])] = &[
            (Id(0), b"z"),
            (Id(1), b"a"),
            (Id(2), b"bbb"),
            (Id(3), b"bb"),
            (Id(4), b"cc"),
            (Id(5), b"ccc"),
            (Id(9), b"ddd"),
            (Id(11), b"e"),
            (Id(13), b"ff"),
            (nid(0), b"n"),
            (nid(1), b"n1"),
            (nid(2), b"n2"),
            (nid(3), b"n3"),
            (nid(4), b"n4"),
            (nid(5), b"n5"),
            (nid(12), b"n12"),
            (nid(20), b"n20"),
        ];
        for (id, name) in items {
            r(map.insert(*id, name)).unwrap();
        }

        // deleted ids in a string, with extra consistency checks.
        let deleted = |map: &dyn IdConvert| -> String {
            let mut deleted_ids = Vec::new();
            for (id, name) in items {
                let name = VertexName::copy_from(name);
                let id = *id;
                let has_id = r(map.contains_vertex_id_locally(&[id])).unwrap()[0];
                let lookup_id = r(map.vertex_id_optional(&name)).unwrap();
                let lookup_name = if has_id {
                    Some(r(map.vertex_name(id)).unwrap())
                } else {
                    None
                };

                match (lookup_id, lookup_name) {
                    (None, None) => deleted_ids.push(id),
                    (None, Some(_)) => {
                        panic!("name->id deleted but not id->name: ({:?} {:?})", id, name)
                    }
                    (Some(_), None) => {
                        panic!("id->name deleted but not name->id: ({:?} {:?})", id, name)
                    }
                    (Some(lid), Some(lname)) => {
                        assert_eq!(lid, id);
                        assert_eq!(lname, name);
                    }
                }
            }
            format!("{:?}", deleted_ids)
        };

        let f = |vs: Vec<VertexName>| -> String {
            let mut vs = vs;
            vs.sort_unstable();
            format!("{:?}", vs)
        };

        let removed = r(map.remove_range(Id(1), Id(3))).unwrap();
        assert_eq!(f(removed), "[a, bb, bbb]");
        assert_eq!(deleted(&map), "[1, 2, 3]");

        let removed = r(map.remove_range(Id(8), Id(12))).unwrap();
        assert_eq!(f(removed), "[ddd, e]");
        assert_eq!(deleted(&map), "[1, 2, 3, 9, 11]");

        let removed = r(map.remove_range(nid(2), nid(4))).unwrap();
        assert_eq!(f(removed), "[n2, n3, n4]");
        assert_eq!(deleted(&map), "[1, 2, 3, 9, 11, N2, N3, N4]");

        let removed = r(map.remove_range(nid(20), nid(10000))).unwrap();
        assert_eq!(f(removed), "[n20]");
        assert_eq!(deleted(&map), "[1, 2, 3, 9, 11, N2, N3, N4, N20]");
    }

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


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

//! Integrity checks.

use std::collections::BTreeSet;

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

use crate::iddag::IdDag;
use crate::iddagstore::IdDagStore;
use crate::idmap::IdMapAssignHead;
use crate::namedag::AbstractNameDag;
use crate::nameset::NameSet;
use crate::ops::CheckIntegrity;
use crate::ops::DagAlgorithm;
use crate::ops::IdConvert;
use crate::ops::Persist;
use crate::ops::TryClone;
use crate::segment::SegmentFlags;
use crate::Group;
use crate::Id;
use crate::Result;
use crate::VertexName;

#[async_trait::async_trait]
impl<IS, M, P, S> CheckIntegrity for AbstractNameDag<IdDag<IS>, M, P, S>
where
    IS: IdDagStore + Persist + 'static,
    IdDag<IS>: TryClone,
    M: TryClone + IdMapAssignHead + Persist + Send + Sync + 'static,
    P: TryClone + Send + Sync + 'static,
    S: TryClone + Persist + Send + Sync + 'static,
{
    async fn check_universal_ids(&self) -> Result<Vec<Id>> {
        let universal_ids: Vec<Id> = self.dag.universal_ids()?.into_iter().collect();
        tracing::debug!("{} universally known vertexes", universal_ids.len());
        let exists = self.map.contains_vertex_id_locally(&universal_ids).await?;
        let missing_ids = universal_ids
            .into_iter()
            .zip(exists)
            .filter_map(|(id, b)| if b { None } else { Some(id) })
            .collect();
        Ok(missing_ids)
    }

    async fn check_segments(&self) -> Result<Vec<String>> {
        // Detected problems.
        let mut problems = Vec::new();

        // Track heads and roots in the graph.
        let mut heads: BTreeSet<Id> = Default::default();
        let mut roots: BTreeSet<Id> = Default::default();

        for level in 0..=self.dag.max_level()? {
            let mut expected_low = Id::MIN;

            // Check all levels.
            for seg in self.dag.iter_segments_ascending(Id::MIN, level)? {
                let seg = seg?;
                let span = seg.span()?;
                let mut add_problem =
                    |msg| problems.push(format!("Level {} segment {:?} {}", level, &seg, msg));

                // Spans need to be sorted and non-overlapping within a group.
                if span.low.group() > expected_low.group() {
                    expected_low = span.low.group().min_id();
                }
                if span.low > span.high || span.low.group() != span.high.group() {
                    add_problem(format!("has invalid span {:?}", span));
                }
                if span.low < expected_low {
                    add_problem(format!(
                        "has unexpected span ({:?}), expected low ({:?})",
                        span, expected_low
                    ));
                }
                expected_low = span.high + 1;

                // Parents should be < low to avoid cycles, and should not have duplicates.
                let mut parents = seg.parents()?;
                let orig_parents_len = parents.len();
                parents.sort_unstable();
                parents.dedup();
                if parents.len() < orig_parents_len {
                    add_problem("has duplicated parents".to_string());
                }
                if parents.iter().any(|&p| p >= span.low) {
                    add_problem("has parents that might cause cycles".to_string());
                }

                // Maintain heads and roots. This can only be calculated from flat segments.
                if level == 0 {
                    let previous_head = heads.iter().rev().next().cloned();
                    if let Some(head) = previous_head {
                        if span.low <= head {
                            add_problem(format!(
                                "overlapped segments: {:?} with previous head {:?}",
                                span, head
                            ));
                        }
                    }
                    for p in &parents {
                        heads.remove(p);
                    }
                    heads.insert(span.high);
                    if parents.is_empty() {
                        roots.insert(span.low);
                    }
                }

                // Check flags. min: must have flags, max: all possible flags.
                let mut expected_flags_max = SegmentFlags::empty();
                let mut expected_flags_min = SegmentFlags::empty();
                if roots.range(span.low..=span.high).take(1).count() > 0 {
                    expected_flags_min |= SegmentFlags::HAS_ROOT;
                    expected_flags_max |= SegmentFlags::HAS_ROOT;
                }
                if level == 0 && heads.len() == 1 && span.high.group() == Group::MASTER {
                    // ONLY_HEAD is optional.
                    expected_flags_max |= SegmentFlags::ONLY_HEAD;
                }
                let flags = seg.flags()?;
                if !flags.contains(expected_flags_min) || !expected_flags_max.contains(flags) {
                    add_problem(format!(
                        "has unexpected flags: {:?} (expected: min: {:?}, max: {:?})",
                        flags, expected_flags_min, expected_flags_max
                    ));
                }

                // Check parents for high-level segments.
                if level > 0 {
                    let mut expected_parents = Vec::with_capacity(parents.len());
                    for seg in self.dag.iter_segments_ascending(span.low, level - 1)? {
                        let seg = seg?;
                        let subspan = seg.span()?;
                        if subspan.high > span.high && subspan.low <= span.high {
                            add_problem(format!(
                                "does not align with low-level segment {:?}",
                                &seg
                            ));
                        }
                        if subspan.low > span.high {
                            break;
                        }
                        for p in seg.parents()? {
                            if p < span.low {
                                expected_parents.push(p);
                            }
                        }
                    }
                    expected_parents.sort_unstable();
                    expected_parents.dedup();
                    if parents != expected_parents {
                        add_problem(format!(
                            "has unexpected parents (expected: {:?})",
                            expected_parents
                        ));
                    }
                }
            }
        }
        Ok(problems)
    }

    async fn check_isomorphic_graph(
        &self,
        other: &dyn DagAlgorithm,
        heads: NameSet,
    ) -> Result<Vec<String>> {
        let mut problems = Vec::new();

        // Prefetch merges and their parents in both graphs' master group.
        // This reduces round-trips.
        tracing::debug!("prefetching merges and parents");
        for graph in [self, other] as [&dyn DagAlgorithm; 2] {
            if !graph.is_vertex_lazy() {
                continue;
            }
            let related = graph.master_group().await? & graph.ancestors(heads.clone()).await?;
            let merges = graph.merges(related).await?;
            let parents = graph.parents(merges.clone()).await?;
            let prefetch = merges | parents;
            let mut iter = prefetch.iter().await?;
            while let Some(_) = iter.next().await {}
        }
        tracing::trace!("prefetched merges and parents");

        // To verify two graphs are isomorphic. Start from some heads,
        // check and compare their linear ancestors recursively.
        //
        // For example, starting from "to_check" having just "E":
        //
        //      A--C--D--E      B--C--D--E
        //        /               /
        //       B               A
        //
        // 1. Figure out the linear portion (C--D--E).
        // 2. Check the linear portion is the same in both graphs.
        //    (only its root C can be a merge and C has the same parents)
        // 3. Remove the head (E) from "to_check", insert root (C)'s parents
        //    to "to_check".
        // 4. Repeat from 1 until "to_check" is empty.
        let mut to_check: Vec<VertexName> = heads.iter().await?.try_collect().await?;
        let mut visited: BTreeSet<VertexName> = Default::default();
        while let Some(head) = to_check.pop() {
            if !visited.insert(head.clone()) {
                continue;
            }

            // Use flat segment to figure out the linear portion.
            let head_id = self.vertex_id(head.clone()).await?;
            let seg = match self.dag.find_flat_segment_including_id(head_id)? {
                Some(seg) => seg,
                None => {
                    problems.push(format!(
                        "head_id {:?} should be covered by a flat segment",
                        head_id
                    ));
                    continue;
                }
            };
            let span = seg.span()?;
            let root_id = span.low;
            let root = self.vertex_name(root_id).await?;
            let parents = self.parent_names(root.clone()).await?;
            let mut add_problem = |msg| {
                problems.push(format!(
                    "range {:?}::{:?} with parents {:?}: {}",
                    &root, &head, &parents, msg
                ));
            };
            tracing::trace!("checking range {:?}::{:?}", &root, &head);

            // Check against the other graph for various properties.
            // Check vertex count in the range.
            let this_count = head_id.0 - root_id.0 + 1;
            let set = match other.range(root.clone().into(), head.clone().into()).await {
                Ok(set) => set,
                Err(e) => {
                    add_problem(format!("cannot resolve range on the other graph: {:?}", e));
                    continue;
                }
            };
            let other_count = set.count().await? as u64;
            if other_count != this_count {
                add_problem(format!(
                    "length mismatch: {} != {}",
                    this_count, other_count
                ));
            }

            // Check that merge can only be at most 1 (`root`).
            let other_merges = other.merges(set).await?.count().await?;
            let this_merges = if parents.len() > 1 { 1 } else { 0 };
            if other_merges != this_merges {
                add_problem(format!(
                    "merge mismatch: {} != {}",
                    this_merges, other_merges
                ));
            }

            // Check parents of root.
            let other_parents = match other.parent_names(root.clone()).await {
                Ok(ps) => ps,
                Err(e) => {
                    add_problem(format!(
                        "cannot get parents of {:?} on the other graph: {:?}",
                        &root, e
                    ));
                    continue;
                }
            };
            if other_parents != parents {
                add_problem(format!(
                    "parents mismatch: {:?} != {:?}",
                    &parents, other_parents
                ));
            }

            // Check parents recursively.
            to_check.extend(parents);
        }

        Ok(problems)
    }
}


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

#![allow(dead_code)]
#![allow(clippy::iter_nth_zero, clippy::for_loops_over_fallibles)]

//! # dag
//!
//! Building blocks for the commit graph used by source control.

mod bsearch;
mod default_impl;
mod delegate;
pub mod errors;
mod fmt;
pub mod iddag;
pub mod iddagstore;
pub mod idmap;
mod integrity;
pub mod namedag;
pub mod nameset;
pub mod ops;
pub mod protocol;
#[cfg(any(test, feature = "render"))]
pub mod render;
pub mod segment;
mod spanset;
pub(crate) mod types_ext;
pub mod utils;
mod verlink;
mod vertex_options;

pub use dag_types::clone;
pub use dag_types::id;
pub use dag_types::CloneData;
pub use dag_types::Group;
pub use dag_types::Id;
pub use dag_types::Location;
pub use dag_types::VertexName;
pub use iddag::FirstAncestorConstraint;
pub use iddag::IdDag;
pub use iddag::IdDagAlgorithm;
pub use iddagstore::IdDagStore;
#[cfg(any(test, feature = "indexedlog-backend"))]
pub use idmap::IdMap;
#[cfg(any(test, feature = "indexedlog-backend"))]
pub use namedag::NameDag;
pub use namedag::NameDagBuilder;
pub use nameset::NameSet;
pub use ops::DagAlgorithm;
pub use segment::FlatSegment;
pub use segment::IdSegment;
pub use segment::PreparedFlatSegments;
pub use verlink::VerLink;
pub use vertex_options::VertexListWithOptions;
pub use vertex_options::VertexOptions;

pub type Level = u8;
pub type InProcessIdDag = IdDag<iddagstore::InProcessStore>;
#[cfg(any(test, feature = "indexedlog-backend"))]
pub type OnDiskIdDag = IdDag<iddagstore::IndexedLogStore>;

// Short aliases for main public types.
#[cfg(any(test, feature = "indexedlog-backend"))]
pub type Dag = NameDag;
pub type Set = NameSet;
pub type IdSet = spanset::SpanSet;
pub type IdSetIter<T> = spanset::SpanSetIter<T>;
pub type IdSpan = spanset::Span;
pub use namedag::MemNameDag as MemDag;
pub use nameset::NameIter as SetIter;
pub type Vertex = VertexName;

#[cfg(feature = "indexedlog-backend")]
pub use iddagstore::indexedlog_store::describe_indexedlog_entry;

#[cfg(any(test, feature = "indexedlog-backend"))]
pub mod tests;

pub use errors::DagError as Error;
pub type Result<T> = std::result::Result<T, Error>;

// Re-export
#[cfg(feature = "indexedlog-backend")]
pub use indexedlog::Repair;
pub use nonblocking;

#[macro_export]
macro_rules! failpoint {
    ($name:literal) => {
        ::fail::fail_point!($name, |_| {
            let msg = format!("failpoint injected by FAILPOINTS: {}", $name);
            Err($crate::errors::DagError::from(
                $crate::errors::BackendError::Generic(msg),
            ))
        })
    };
}

/// Whether running inside a test.
pub(crate) fn is_testing() -> bool {
    std::env::var("TESTTMP").is_ok()
}

#[cfg(test)]
dev_logger::init!();