geometric-traits 0.2.3

//! Iterator of the sparse coordinates of the CSR2D matrix.
use core::cmp::Ordering;

use num_traits::{AsPrimitive, One, Zero};

use crate::prelude::*;

/// Iterator of the sparse coordinates of the CSR2D matrix.
pub struct CSR2DColumns<'a, CSR: SparseMatrix2D> {
    /// The CSR matrix.
    csr2d: &'a CSR,
    /// The row index.
    next_row: CSR::RowIndex,
    /// The end row index.
    back_row: CSR::RowIndex,
    /// The row associated with the index at the beginning of the iteration.
    next: Option<CSR::SparseRow<'a>>,
    /// The row associated with the index at the end of the iteration.
    back: Option<CSR::SparseRow<'a>>,
}

impl<CSR: SparseMatrix2D> Iterator for CSR2DColumns<'_, CSR> {
    type Item = CSR::ColumnIndex;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(column) = self.next.as_mut()?.next() {
                return Some(column);
            }

            match self.next_row.cmp(&self.back_row) {
                Ordering::Less => {
                    self.next_row += CSR::RowIndex::one();
                    if self.next_row == self.back_row {
                        self.next = self.back.take();
                    } else {
                        self.next = Some(self.csr2d.sparse_row(self.next_row));
                    }
                }
                Ordering::Equal | Ordering::Greater => {
                    return self.back.as_mut().and_then(Iterator::next);
                }
            }
        }
    }
}

impl<'matrix, CSR: SizedSparseMatrix2D> ExactSizeIterator for CSR2DColumns<'matrix, CSR>
where
    CSR::SparseRow<'matrix>: ExactSizeIterator,
{
    #[inline]
    fn len(&self) -> usize {
        if self.next.is_none() {
            return 0;
        }
        if self.back.is_none() {
            return self.next.as_ref().map_or(0, ExactSizeIterator::len);
        }

        let next_row_rank = self.csr2d.rank_row(self.next_row).as_();
        let already_observed_in_next_row =
            self.csr2d.number_of_defined_values_in_row(self.next_row).as_()
                - self.next.as_ref().map_or(0, ExactSizeIterator::len);
        let back_row_rank = self.csr2d.rank_row(self.back_row).as_();
        let already_observed_in_back_row = self.back.as_ref().map_or(0, |back| {
            self.csr2d.number_of_defined_values_in_row(self.back_row).as_() - back.len()
        });
        back_row_rank - next_row_rank - already_observed_in_next_row - already_observed_in_back_row
    }
}

impl<CSR: SparseMatrix2D> DoubleEndedIterator for CSR2DColumns<'_, CSR> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(back) = self.back.as_mut() {
                if let Some(column) = back.next_back() {
                    return Some(column);
                }

                if self.back_row > self.next_row {
                    self.back_row -= CSR::RowIndex::one();
                    if self.back_row == self.next_row {
                        self.back = None;
                    } else {
                        self.back = Some(self.csr2d.sparse_row(self.back_row));
                    }
                } else {
                    self.back = None;
                }
            } else {
                return self.next.as_mut().and_then(DoubleEndedIterator::next_back);
            }
        }
    }
}

impl<'a, CSR: SparseMatrix2D> From<&'a CSR> for CSR2DColumns<'a, CSR> {
    #[inline]
    fn from(csr2d: &'a CSR) -> Self {
        let next_row = CSR::RowIndex::zero();
        let mut back_row = CSR::RowIndex::zero();
        let has_rows = next_row < csr2d.number_of_rows();
        let next = has_rows.then(|| csr2d.sparse_row(next_row));
        let back = if has_rows {
            back_row = csr2d.number_of_rows() - CSR::RowIndex::one();
            (next_row < back_row).then(|| csr2d.sparse_row(back_row))
        } else {
            None
        };
        Self { csr2d, next_row, back_row, next, back }
    }
}

#[cfg(all(test, feature = "alloc"))]
mod tests {
    use super::*;
    use crate::traits::{MatrixMut, SparseMatrixMut};

    type TestCSR2D = CSR2D<usize, usize, usize>;

    fn sample_csr() -> TestCSR2D {
        let mut csr = TestCSR2D::with_sparse_shape((3, 4));
        MatrixMut::add(&mut csr, (0, 1)).expect("insert (0,1)");
        MatrixMut::add(&mut csr, (0, 3)).expect("insert (0,3)");
        MatrixMut::add(&mut csr, (2, 0)).expect("insert (2,0)");
        csr
    }

    #[test]
    fn test_columns_forward_backward_and_len() {
        let csr = sample_csr();
        let mut columns = CSR2DColumns::from(&csr);

        assert_eq!(columns.len(), 2);
        assert_eq!(columns.next(), Some(1));
        assert_eq!(columns.next_back(), Some(0));
        assert_eq!(columns.next(), Some(3));
        assert_eq!(columns.next(), None);
        assert_eq!(columns.next_back(), None);
    }

    #[test]
    fn test_columns_single_row_back_falls_through_to_front_iter() {
        let mut csr = TestCSR2D::with_sparse_shape((1, 5));
        MatrixMut::add(&mut csr, (0, 2)).expect("insert (0,2)");
        MatrixMut::add(&mut csr, (0, 4)).expect("insert (0,4)");

        let mut columns = CSR2DColumns::from(&csr);
        assert_eq!(columns.len(), 2);
        assert_eq!(columns.next_back(), Some(4));
        assert_eq!(columns.len(), 1);
        assert_eq!(columns.next(), Some(2));
        assert_eq!(columns.next_back(), None);
    }

    #[test]
    fn test_columns_empty_matrix() {
        let csr = TestCSR2D::with_sparse_shape((0, 0));
        let mut columns = CSR2DColumns::from(&csr);

        assert_eq!(columns.len(), 0);
        assert_eq!(columns.next(), None);
        assert_eq!(columns.next_back(), None);
    }

    #[test]
    fn test_columns_next_on_single_entry_exhausts_via_equal_branch() {
        let mut csr = TestCSR2D::with_sparse_shape((1, 3));
        MatrixMut::add(&mut csr, (0, 1)).expect("insert (0,1)");

        let mut columns = CSR2DColumns::from(&csr);
        assert_eq!(columns.next(), Some(1));
        assert_eq!(columns.next(), None);
    }

    #[test]
    fn test_columns_manual_state_hits_back_none_branch() {
        let mut csr = TestCSR2D::with_sparse_shape((1, 1));
        MatrixMut::add(&mut csr, (0, 0)).expect("insert (0,0)");

        let mut columns = CSR2DColumns {
            csr2d: &csr,
            next_row: 1,
            back_row: 0,
            next: Some(csr.sparse_row(0)),
            back: Some(csr.sparse_row(0)),
        };
        columns.back.as_mut().expect("back iterator exists").next();

        assert_eq!(columns.next_back(), Some(0));
        assert_eq!(columns.next_back(), None);
        assert!(columns.back.is_none());
    }
}