pipecircuit 0.0.10

use alloc::vec::Vec;
use core::cmp::Ordering;
use core::marker::PhantomData;
use core::num::{NonZeroU32, NonZeroUsize};

use super::selection::SelectionMembership;
use super::{RangeIndexLookupOutcome, SecondaryIndexLookupOutcome, Selection};
use crate::storage::warp_table::secondary_index::RangeIndexOrderingLevel;
use crate::storage::warp_table::{WarpTable, WarpTableRow};

mod first_rows;

crate::value_wrapper! {
    #[derive(Clone, Copy)]
    pub struct OrderedRowLimit(NonZeroU32);
}

impl OrderedRowLimit {
    fn row_count(self) -> NonZeroUsize {
        NonZeroUsize::try_from(self.value()).unwrap_or(NonZeroUsize::MAX)
    }
}

#[must_use]
pub struct RangeIndexOrdering<'table, Row: WarpTableRow, Brand = Row> {
    first_level: RangeIndexOrderingLevel<'table>,
    later_levels: Vec<RangeIndexOrderingLevel<'table>>,
    first_sortable_word: fn(&Row::Ref<'table>) -> u64,
    later_sortable_words: Vec<fn(&Row::Ref<'table>) -> u64>,
    table: PhantomData<fn() -> &'table WarpTable<Row, Brand>>,
}

impl<'table, Row: WarpTableRow, Brand> RangeIndexOrdering<'table, Row, Brand> {
    pub fn then(mut self, next: Self) -> Self {
        self.later_levels.push(next.first_level);
        self.later_levels.extend(next.later_levels);
        self.later_sortable_words.push(next.first_sortable_word);
        self.later_sortable_words.extend(next.later_sortable_words);
        self
    }

    fn levels(
        &self,
    ) -> impl Iterator<
        Item = (
            &RangeIndexOrderingLevel<'table>,
            fn(&Row::Ref<'table>) -> u64,
        ),
    > {
        core::iter::once((&self.first_level, self.first_sortable_word)).chain(
            self.later_levels
                .iter()
                .zip(self.later_sortable_words.iter().copied()),
        )
    }

    fn level_count(&self) -> NonZeroUsize {
        NonZeroUsize::MIN.saturating_add(self.later_levels.len())
    }
}

enum OrderedCandidates<'table> {
    EveryRow { row_count: usize },
    Membership(SelectionMembership<'table>),
}

#[must_use]
pub struct OrderedSelection<'table, Row: WarpTableRow, Brand = Row> {
    table: &'table WarpTable<Row, Brand>,
    candidates: OrderedCandidates<'table>,
    ordering: RangeIndexOrdering<'table, Row, Brand>,
}

#[must_use]
pub struct OrderedSelectionComparingRows<'table, Row: WarpTableRow, Brand, Compare> {
    ordered: OrderedSelection<'table, Row, Brand>,
    compare: Compare,
}

impl<'table, Row: WarpTableRow, Brand> OrderedSelection<'table, Row, Brand> {
    pub fn then_compare_rows<Compare>(
        self,
        compare: Compare,
    ) -> OrderedSelectionComparingRows<'table, Row, Brand, Compare>
    where
        Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
    {
        OrderedSelectionComparingRows {
            ordered: self,
            compare,
        }
    }
}

#[must_use]
pub enum OrderedSelectionFirstCuttingTiesOutcome<RowRef> {
    Exact {
        rows: Vec<RowRef>,
    },
    TiesCut {
        rows: Vec<RowRef>,
        tied_rows_left_out: usize,
    },
}

#[must_use]
pub enum OrderedSelectionFirstKeepingTiesOutcome<RowRef> {
    Exact {
        rows: Vec<RowRef>,
    },
    TiesKept {
        leaders: Vec<RowRef>,
        tied: Vec<RowRef>,
    },
}

impl<'table, Row: WarpTableRow, Brand> Selection<'table, Row, Brand> {
    pub fn ordered_by(
        self,
        order: impl FnOnce(Row::Query<'table, Brand>) -> RangeIndexOrdering<'table, Row, Brand>,
    ) -> OrderedSelection<'table, Row, Brand>
    where
        Brand: 'table,
    {
        let table = self.table;
        OrderedSelection {
            table,
            candidates: OrderedCandidates::Membership(self.membership),
            ordering: order(Row::query(table)),
        }
    }
}

macro_rules! impl_lookup_ordering {
    ($outcome:ident) => {
        impl<'table, Row: WarpTableRow, Brand, Key> $outcome<'table, Row, Brand, Key> {
            pub fn ordered_by(
                self,
                order: impl FnOnce(
                    Row::Query<'table, Brand>,
                ) -> RangeIndexOrdering<'table, Row, Brand>,
            ) -> OrderedSelection<'table, Row, Brand>
            where
                Brand: 'table,
            {
                self.into_selection().ordered_by(order)
            }
        }
    };
}

impl_lookup_ordering!(SecondaryIndexLookupOutcome);
impl_lookup_ordering!(RangeIndexLookupOutcome);

#[doc(hidden)]
pub fn bind_range_ordering<'table, Row: WarpTableRow, Brand>(
    level: RangeIndexOrderingLevel<'table>,
    sortable_word: fn(&Row::Ref<'table>) -> u64,
) -> RangeIndexOrdering<'table, Row, Brand> {
    RangeIndexOrdering {
        first_level: level,
        later_levels: Vec::new(),
        first_sortable_word: sortable_word,
        later_sortable_words: Vec::new(),
        table: PhantomData,
    }
}

#[doc(hidden)]
pub fn bind_every_row_ordering<'table, Row: WarpTableRow, Brand>(
    table: &'table WarpTable<Row, Brand>,
    ordering: RangeIndexOrdering<'table, Row, Brand>,
) -> OrderedSelection<'table, Row, Brand> {
    OrderedSelection {
        table,
        candidates: OrderedCandidates::EveryRow {
            row_count: table.len(),
        },
        ordering,
    }
}

#[cfg(test)]
mod tests;