versatile-data 0.117.0

use std::{cmp::Ordering, fmt::Debug, num::NonZeroU32};

use idx_binary::{AvltrieeSearch, IdxBinary, IdxFileAvlTriee};

use crate::{Data, FieldName, RowSet};

pub trait CustomSort {
    fn compare(&self, a: NonZeroU32, b: NonZeroU32) -> Ordering;
    fn asc(&self) -> Vec<NonZeroU32>;
    fn desc(&self) -> Vec<NonZeroU32>;
}
impl Debug for dyn CustomSort {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "CustomSort")?;
        Ok(())
    }
}

pub struct NoCustomSort {}

impl CustomSort for NoCustomSort {
    fn compare(&self, _: NonZeroU32, _: NonZeroU32) -> Ordering {
        unreachable!()
    }

    fn asc(&self) -> Vec<NonZeroU32> {
        unreachable!()
    }

    fn desc(&self) -> Vec<NonZeroU32> {
        unreachable!()
    }
}

pub enum CustomOrderKey<C: CustomSort> {
    Serial,
    Row,
    TermBegin,
    TermEnd,
    LastUpdated,
    Field(FieldName),
    Custom(C),
}

pub type OrderKey = CustomOrderKey<NoCustomSort>;

pub enum Order<C: CustomSort> {
    Asc(CustomOrderKey<C>),
    Desc(CustomOrderKey<C>),
}

impl Data {
    /// Sort search results.
    pub fn sort<C: CustomSort>(&self, rows: &RowSet, orders: &[Order<C>]) -> Vec<NonZeroU32> {
        let sub_orders = &orders[1..];
        match &orders[0] {
            Order::Asc(key) => self.sort_with_key(rows, key, sub_orders),
            Order::Desc(key) => self.sort_with_key_desc(rows, key, sub_orders),
        }
    }

    /// Compares two rows by one key. Rows without a value come last
    /// in both ascending and descending order.
    fn compare_by_key<C: CustomSort>(
        &self,
        key: &CustomOrderKey<C>,
        a: NonZeroU32,
        b: NonZeroU32,
        desc: bool,
    ) -> Ordering {
        fn missing_last<T>(
            a: Option<T>,
            b: Option<T>,
            desc: bool,
            cmp: impl Fn(T, T) -> Ordering,
        ) -> Ordering {
            match (a, b) {
                (Some(a), Some(b)) => {
                    if desc {
                        cmp(b, a)
                    } else {
                        cmp(a, b)
                    }
                }
                (Some(_), None) => Ordering::Less,
                (None, Some(_)) => Ordering::Greater,
                (None, None) => Ordering::Equal,
            }
        }
        let index = |f: &Option<idx_binary::IdxFile<u64>>| {
            f.as_ref().map_or(Ordering::Equal, |f| {
                missing_last(f.value(a), f.value(b), desc, Ord::cmp)
            })
        };
        match key {
            CustomOrderKey::Serial => {
                missing_last(self.serial.value(a), self.serial.value(b), desc, Ord::cmp)
            }
            CustomOrderKey::Row => missing_last(Some(a), Some(b), desc, |a, b| a.cmp(&b)),
            CustomOrderKey::TermBegin => index(&self.term_begin),
            CustomOrderKey::TermEnd => index(&self.term_end),
            CustomOrderKey::LastUpdated => index(&self.last_updated),
            CustomOrderKey::Field(name) => self.fields.get(name).map_or(Ordering::Equal, |f| {
                missing_last(f.value(a), f.value(b), desc, IdxBinary::cmp)
            }),
            CustomOrderKey::Custom(custom_order) => {
                if desc {
                    custom_order.compare(b, a)
                } else {
                    custom_order.compare(a, b)
                }
            }
        }
    }

    fn subsort<C: CustomSort>(
        &self,
        mut tmp: Vec<NonZeroU32>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        tmp.sort_by(|a, b| {
            for order in sub_orders {
                let ord = match order {
                    Order::Asc(key) => self.compare_by_key(key, *a, *b, false),
                    Order::Desc(key) => self.compare_by_key(key, *a, *b, true),
                };
                if ord != Ordering::Equal {
                    return ord;
                }
            }
            Ordering::Equal
        });
        tmp
    }

    /// Sorts rows in the order of the index. Rows that are not in the index
    /// (rows without a value for the key) come last.
    fn sort_with_triee_inner<T: PartialEq, I: ?Sized, C: CustomSort>(
        &self,
        rows: &RowSet,
        triee: &IdxFileAvlTriee<T, I>,
        iter: impl Iterator<Item = NonZeroU32>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        let sorted_group = |tmp: Vec<NonZeroU32>| {
            if tmp.len() <= 1 || sub_orders.is_empty() {
                tmp
            } else {
                self.subsort(tmp, sub_orders)
            }
        };

        let mut ret = Vec::with_capacity(rows.len());
        let mut before: Option<&T> = None;
        let mut tmp: Vec<NonZeroU32> = Vec::new();
        for r in iter {
            if rows.contains(&r) {
                let value: &T = triee.node(r).unwrap();
                if before.is_some_and(|before| before.ne(value)) {
                    ret.extend(sorted_group(std::mem::take(&mut tmp)));
                }
                tmp.push(r);
                before = Some(value);
            }
        }
        ret.extend(sorted_group(tmp));

        if ret.len() < rows.len() {
            let in_index: RowSet = ret.iter().copied().collect();
            let missing = rows
                .iter()
                .copied()
                .filter(|r| !in_index.contains(r))
                .collect();
            ret.extend(sorted_group(missing));
        }
        ret
    }

    fn sort_with_triee<T: PartialEq, I: ?Sized, C: CustomSort>(
        &self,
        rows: &RowSet,
        triee: &IdxFileAvlTriee<T, I>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        self.sort_with_triee_inner(rows, triee, triee.iter(), sub_orders)
    }

    fn sort_with_triee_desc<T: PartialEq, I: ?Sized, C: CustomSort>(
        &self,
        rows: &RowSet,
        index: &IdxFileAvlTriee<T, I>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        self.sort_with_triee_inner(rows, index, index.desc_iter(), sub_orders)
    }

    fn sort_with_key<C: CustomSort>(
        &self,
        rows: &RowSet,
        key: &CustomOrderKey<C>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        match key {
            CustomOrderKey::Serial => self.sort_with_triee::<u32, u32, C>(rows, &self.serial, &[]),
            CustomOrderKey::Row => rows.iter().cloned().collect(),
            CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
                || rows.iter().cloned().collect(),
                |f| self.sort_with_triee(rows, f, sub_orders),
            ),
            CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
                || rows.iter().cloned().collect(),
                |f| self.sort_with_triee(rows, f, sub_orders),
            ),
            CustomOrderKey::LastUpdated => self.last_updated.as_ref().map_or_else(
                || rows.iter().cloned().collect(),
                |f| self.sort_with_triee(rows, f, sub_orders),
            ),
            CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
                || rows.iter().cloned().collect(),
                |f| self.sort_with_triee(rows, f.as_ref(), sub_orders),
            ),
            CustomOrderKey::Custom(custom_order) => custom_order.asc(),
        }
    }

    fn sort_with_key_desc<C: CustomSort>(
        &self,
        rows: &RowSet,
        key: &CustomOrderKey<C>,
        sub_orders: &[Order<C>],
    ) -> Vec<NonZeroU32> {
        match key {
            CustomOrderKey::Serial => {
                self.sort_with_triee_desc::<u32, u32, C>(rows, &self.serial, &[])
            }
            CustomOrderKey::Row => rows.iter().rev().cloned().collect(),
            CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
                || rows.iter().rev().cloned().collect(),
                |f| self.sort_with_triee_desc(rows, f, sub_orders),
            ),
            CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
                || rows.iter().rev().cloned().collect(),
                |f| self.sort_with_triee_desc(rows, f, sub_orders),
            ),
            CustomOrderKey::LastUpdated => self.last_updated.as_ref().map_or_else(
                || rows.iter().rev().cloned().collect(),
                |f| self.sort_with_triee_desc(rows, f, sub_orders),
            ),
            CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
                || rows.iter().rev().cloned().collect(),
                |f| self.sort_with_triee_desc(rows, f.as_ref(), sub_orders),
            ),
            CustomOrderKey::Custom(custom_order) => custom_order.desc(),
        }
    }
}