armdb 0.11.2

use std::cmp::Ordering;
use std::collections::BinaryHeap;

use crate::Dir;

/// Lazily merges homogeneous sources that are already ordered for `dir`.
///
/// Heads that compare equal are emitted in ascending source index order.
pub struct MergeSorted<I, T> {
    dir: Dir,
    sources: Vec<I>,
    heads: BinaryHeap<Head<T>>,
    initialized: bool,
}

struct Head<T> {
    item: T,
    source: usize,
    dir: Dir,
}

impl<T: Ord> Ord for Head<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        let order = match self.dir {
            Dir::Asc => other.item.cmp(&self.item),
            Dir::Desc => self.item.cmp(&other.item),
        };
        order.then_with(|| other.source.cmp(&self.source))
    }
}

impl<T: Ord> PartialOrd for Head<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T: Ord> PartialEq for Head<T> {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

impl<T: Ord> Eq for Head<T> {}

/// Lazily merge sources that are already ordered for `dir`.
///
/// Heads that compare equal are emitted in ascending source index order.
pub fn merge_sorted<I, T>(dir: Dir, sources: Vec<I>) -> MergeSorted<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    // Initialization retains at most one head per source.
    let head_capacity = sources.len();
    MergeSorted {
        dir,
        sources,
        heads: BinaryHeap::with_capacity(head_capacity),
        initialized: false,
    }
}

impl<I, T> MergeSorted<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    /// Fold each adjacent run of equal merged values into its first value.
    ///
    /// For an equal run, the row from the lowest source index is the first
    /// output row and becomes the accumulator passed to `f`.
    ///
    /// The reducer must preserve every accumulator field that participates in
    /// [`Ord`]; it may merge only non-ordering payload. Each successor is
    /// compared with the current accumulator after every reduction.
    pub fn coalesce_equal<F>(self, f: F) -> CoalescedMerge<I, T, F>
    where
        F: FnMut(&mut T, T),
    {
        CoalescedMerge {
            rows: self,
            f,
            pending: None,
        }
    }

    fn initialize(&mut self) {
        if self.initialized {
            return;
        }
        self.initialized = true;
        for (source, rows) in self.sources.iter_mut().enumerate() {
            if let Some(item) = rows.next() {
                self.heads.push(Head {
                    item,
                    source,
                    dir: self.dir,
                });
            }
        }
    }
}

impl<I, T> Iterator for MergeSorted<I, T>
where
    I: Iterator<Item = T>,
    T: Ord,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        self.initialize();
        let head = self.heads.pop()?;
        if let Some(item) = self.sources[head.source].next() {
            debug_assert!(
                match self.dir {
                    Dir::Asc => head.item <= item,
                    Dir::Desc => head.item >= item,
                },
                "merge source must remain ordered",
            );
            self.heads.push(Head {
                item,
                source: head.source,
                dir: self.dir,
            });
        }
        Some(head.item)
    }
}

/// Iterator returned by [`MergeSorted::coalesce_equal`].
pub struct CoalescedMerge<I, T, F> {
    rows: MergeSorted<I, T>,
    f: F,
    pending: Option<T>,
}

impl<I, T, F> Iterator for CoalescedMerge<I, T, F>
where
    I: Iterator<Item = T>,
    T: Ord,
    F: FnMut(&mut T, T),
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        let mut acc = self.pending.take().or_else(|| self.rows.next())?;
        let dir = self.rows.dir;
        for next in self.rows.by_ref() {
            let order = acc.cmp(&next);
            // The merge emits rows in `dir` order, so a successor that sorts
            // behind the accumulator can only mean the reducer moved it.
            debug_assert!(
                match dir {
                    Dir::Asc => order != Ordering::Greater,
                    Dir::Desc => order != Ordering::Less,
                },
                "coalesce_equal reducer must preserve the accumulator's ordering fields",
            );
            if order == Ordering::Equal {
                (self.f)(&mut acc, next);
            } else {
                self.pending = Some(next);
                break;
            }
        }
        Some(acc)
    }
}

/// Iterator-erased rows for callers with heterogeneous source types.
pub type BoxedRows<'a, T> = Box<dyn Iterator<Item = T> + 'a>;

/// Erase an iterator type for use as a [`BoxedRows`].
pub fn boxed<'a, T>(it: impl Iterator<Item = T> + 'a) -> BoxedRows<'a, T> {
    Box::new(it)
}

#[cfg(test)]
mod tests {
    use std::cell::{Cell, RefCell};
    use std::rc::Rc;

    use super::*;
    use crate::Dir;

    struct CountingIter<T> {
        rows: std::vec::IntoIter<T>,
        calls: Rc<Cell<usize>>,
    }

    impl<T> CountingIter<T> {
        fn new(rows: Vec<T>, calls: Rc<Cell<usize>>) -> Self {
            Self {
                rows: rows.into_iter(),
                calls,
            }
        }
    }

    impl<T> Iterator for CountingIter<T> {
        type Item = T;

        fn next(&mut self) -> Option<Self::Item> {
            self.calls.set(self.calls.get() + 1);
            self.rows.next()
        }
    }

    #[test]
    fn merge_is_lazy_and_stable_for_equal_heads() {
        let calls = Rc::new(Cell::new(0));
        let sources = vec![
            CountingIter::new(vec![(1, 0), (3, 0)], Rc::clone(&calls)),
            CountingIter::new(vec![(1, 1), (2, 1)], Rc::clone(&calls)),
        ];
        let mut rows = merge_sorted(Dir::Asc, sources);
        assert_eq!(calls.get(), 0, "construction must not pull sources");
        assert_eq!(rows.next(), Some((1, 0)), "lower source index wins");
        assert_eq!(rows.collect::<Vec<_>>(), vec![(1, 1), (2, 1), (3, 0)]);
    }

    #[test]
    fn merge_pulls_only_the_rows_a_short_take_needs() {
        const SOURCES: usize = 4;
        const PER_SOURCE: u64 = 500;
        const TAKE: usize = 20;

        let calls = Rc::new(Cell::new(0));
        let sources: Vec<_> = (0..SOURCES)
            .map(|source| {
                let rows: Vec<u64> = (0..PER_SOURCE)
                    .map(|i| i * SOURCES as u64 + source as u64)
                    .collect();
                CountingIter::new(rows, Rc::clone(&calls))
            })
            .collect();

        let rows = merge_sorted(Dir::Asc, sources)
            .take(TAKE)
            .collect::<Vec<_>>();

        assert_eq!(rows, (0..TAKE as u64).collect::<Vec<_>>());
        // One head per source, then one replacement per emitted row — not the
        // 2000 rows an eager collect-and-sort would have to read first.
        assert_eq!(calls.get(), SOURCES + TAKE);
    }

    #[derive(Debug)]
    struct Row {
        key: i32,
        payload: i32,
    }

    impl Row {
        fn new(key: i32, payload: i32) -> Self {
            Self { key, payload }
        }
    }

    impl Ord for Row {
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
            self.key.cmp(&other.key)
        }
    }

    impl PartialOrd for Row {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
            Some(self.cmp(other))
        }
    }

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

    impl Eq for Row {}

    #[test]
    fn coalesce_equal_never_receives_unequal_neighbours() {
        let rows = merge_sorted(
            Dir::Asc,
            vec![
                vec![Row::new(1, 1), Row::new(3, 3), Row::new(5, 5)].into_iter(),
                vec![Row::new(1, 1), Row::new(2, 2), Row::new(5, 5)].into_iter(),
            ],
        )
        .coalesce_equal(|acc, next| acc.payload += next.payload)
        .map(|row| (row.key, row.payload))
        .collect::<Vec<_>>();
        assert_eq!(rows, vec![(1, 2), (2, 2), (3, 3), (5, 10)]);
    }

    #[test]
    fn merge_orders_the_same_logical_data_in_both_directions() {
        let asc_sources = vec![
            vec![1, 4, 7].into_iter(),
            vec![2, 5, 8].into_iter(),
            vec![3, 6, 9].into_iter(),
        ];
        let desc_sources = vec![
            vec![7, 4, 1].into_iter(),
            vec![8, 5, 2].into_iter(),
            vec![9, 6, 3].into_iter(),
        ];

        assert_eq!(
            merge_sorted(Dir::Asc, asc_sources).collect::<Vec<_>>(),
            vec![1, 2, 3, 4, 5, 6, 7, 8, 9],
        );
        assert_eq!(
            merge_sorted(Dir::Desc, desc_sources).collect::<Vec<_>>(),
            vec![9, 8, 7, 6, 5, 4, 3, 2, 1],
        );
    }

    #[test]
    fn merge_handles_zero_and_one_source() {
        assert_eq!(
            merge_sorted::<std::vec::IntoIter<i32>, _>(Dir::Asc, vec![]).collect::<Vec<_>>(),
            Vec::<i32>::new(),
        );
        assert_eq!(
            merge_sorted(Dir::Asc, vec![vec![1, 2, 3].into_iter()]).collect::<Vec<_>>(),
            vec![1, 2, 3],
        );
    }

    #[test]
    fn merge_skips_initially_empty_sources_with_unequal_lengths() {
        assert_eq!(
            merge_sorted(
                Dir::Asc,
                vec![
                    vec![].into_iter(),
                    vec![1, 4, 6].into_iter(),
                    vec![2].into_iter(),
                    vec![].into_iter(),
                ],
            )
            .collect::<Vec<_>>(),
            vec![1, 2, 4, 6],
        );
    }

    #[test]
    fn boxed_preserves_the_input_iterator() {
        assert_eq!(
            boxed(vec![1, 2, 3].into_iter()).collect::<Vec<_>>(),
            vec![1, 2, 3]
        );
    }

    #[test]
    fn coalesce_equal_keeps_separate_equal_runs() {
        let calls = Rc::new(RefCell::new(Vec::new()));
        let seen = Rc::clone(&calls);
        let rows = merge_sorted(
            Dir::Asc,
            vec![
                vec![Row::new(1, 1), Row::new(3, 3), Row::new(5, 5)].into_iter(),
                vec![Row::new(1, 1), Row::new(2, 2), Row::new(5, 5)].into_iter(),
            ],
        )
        .coalesce_equal(move |acc, next| {
            seen.borrow_mut().push((acc.key, next.key));
            acc.payload += next.payload;
        })
        .map(|row| (row.key, row.payload))
        .collect::<Vec<_>>();
        assert_eq!(rows, vec![(1, 2), (2, 2), (3, 3), (5, 10)]);
        assert_eq!(&*calls.borrow(), &[(1, 1), (5, 5)]);
    }

    #[test]
    fn coalesce_equal_uses_lower_source_row_as_accumulator() {
        let mut calls = Vec::new();
        let rows = merge_sorted(
            Dir::Asc,
            vec![
                vec![Row::new(1, 10)].into_iter(),
                vec![Row::new(1, 20)].into_iter(),
            ],
        )
        .coalesce_equal(|acc, next| {
            calls.push((acc.payload, next.payload));
            acc.payload += next.payload;
        })
        .map(|row| (row.key, row.payload))
        .collect::<Vec<_>>();

        assert_eq!(calls, vec![(10, 20)]);
        assert_eq!(rows, vec![(1, 30)]);
    }

    #[derive(Eq, PartialEq)]
    struct NotClone(i32);

    impl Ord for NotClone {
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
            self.0.cmp(&other.0)
        }
    }

    impl PartialOrd for NotClone {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
            Some(self.cmp(other))
        }
    }

    #[test]
    fn merge_does_not_require_clone() {
        let rows = merge_sorted(
            Dir::Asc,
            vec![
                vec![NotClone(1), NotClone(3)].into_iter(),
                vec![NotClone(2)].into_iter(),
            ],
        )
        .map(|item| item.0)
        .collect::<Vec<_>>();
        assert_eq!(rows, vec![1, 2, 3]);
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "merge source must remain ordered")]
    fn merge_rejects_an_unordered_replacement_in_debug_builds() {
        let _ = merge_sorted(Dir::Asc, vec![vec![2, 1].into_iter()]).collect::<Vec<_>>();
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "coalesce_equal reducer must preserve")]
    fn coalesce_equal_rejects_a_key_moving_reducer_in_debug_builds() {
        // The source stays ordered, so only the reducer can move the accumulator
        // past its successor.
        let _ = merge_sorted(Dir::Asc, vec![vec![1u64, 1, 2].into_iter()])
            .coalesce_equal(|acc, _| *acc = 9)
            .collect::<Vec<_>>();
    }
}