intrex 0.2.0

use core::ops;

use crate::{
    bintree::{NodesCmpKey, NodesLink},
    node_data::NodesData,
};

use super::*;

/// # Iteration
impl<'tree, Nodes, Index> TreeAccessor<'tree, Nodes, Index>
where
    Nodes: NodesRb<Index>,
    Index: PartialEq + Clone,
{
    /// Iterate over `(Index, Nodes::Data)`.
    #[inline]
    pub fn iter(&self) -> Iter<'tree, &Nodes, Index> {
        Iter(self.raw.iter())
    }

    /// Iterate over `(Index, Nodes::Data<'_>)` in a given range.
    #[inline]
    pub fn iter_index_range(
        &self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Iter<'tree, &Nodes, Index> {
        Iter(self.raw.iter_index_range(range))
    }

    /// Iterate over `(Index, Nodes::Data)` in a given range, consuming `self`.
    #[inline]
    pub fn into_iter_index_range(
        self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Iter<'tree, Nodes, Index> {
        Iter(self.raw.into_iter_index_range(range))
    }

    /// Iterate over `Index`.
    #[inline]
    pub fn indices(&self) -> Indices<'tree, &Nodes, Index> {
        Indices(self.raw.indices())
    }

    /// Iterate over `Index`, consuming `self`.
    #[inline]
    pub fn into_indices(self) -> Indices<'tree, Nodes, Index> {
        Indices(self.raw.into_indices())
    }

    /// Iterate over `Index` in a given range.
    #[inline]
    pub fn indices_index_range(
        &self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Indices<'tree, &Nodes, Index> {
        Indices(self.raw.indices_index_range(range))
    }

    /// Iterate over `Index` in a given range, consuming `self`.
    #[inline]
    pub fn into_indices_index_range(
        self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Indices<'tree, Nodes, Index> {
        Indices(self.raw.into_indices_index_range(range))
    }

    /// Iterate over `Nodes::Data<'_>`.
    #[inline]
    pub fn values(&self) -> Values<'tree, &Nodes, Index> {
        Values(self.raw.values())
    }

    /// Iterate over `Nodes::Data`, consuming `self`.
    #[inline]
    pub fn into_values(self) -> Values<'tree, Nodes, Index> {
        Values(self.raw.into_values())
    }

    /// Iterate over `Nodes::Data<'_>` in a given range.
    #[inline]
    pub fn values_index_range(
        &self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Values<'tree, &Nodes, Index> {
        Values(self.raw.values_index_range(range))
    }

    /// Iterate over `Nodes::Data` in a given range, consuming `self`.
    #[inline]
    pub fn into_values_index_range(
        self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Values<'tree, Nodes, Index> {
        Values(self.raw.into_values_index_range(range))
    }
}

/// # BST Ranged Iteration
///
/// The following methods assume `self` points to a valid binary search tree.
impl<'tree, Nodes, Index> TreeAccessor<'tree, Nodes, Index>
where
    Nodes: NodesRb<Index>,
    Index: PartialEq + Clone,
{
    /// Iterate over `(Index, Nodes::Data<'_>)` in a given key range.
    #[inline]
    pub fn bst_iter_key_range<Key>(
        &self,
        range: impl ops::RangeBounds<Key>,
    ) -> Iter<'tree, &Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Iter(self.raw.bst_iter_key_range(range))
    }

    /// Iterate over `(Index, Nodes::Data)` in a given key range, consuming
    /// `self`.
    #[inline]
    pub fn bst_into_iter_key_range<Key>(
        self,
        range: impl ops::RangeBounds<Key>,
    ) -> Iter<'tree, Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Iter(self.raw.bst_into_iter_key_range(range))
    }

    /// Iterate over `Index` in a given key range.
    #[inline]
    pub fn bst_indices_key_range<Key>(
        &self,
        range: impl ops::RangeBounds<Key>,
    ) -> Indices<'tree, &Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Indices(self.raw.bst_indices_key_range(range))
    }

    /// Iterate over `Index` in a given key range, consuming `self`.
    #[inline]
    pub fn bst_into_indices_key_range<Key>(
        self,
        range: impl ops::RangeBounds<Key>,
    ) -> Indices<'tree, Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Indices(self.raw.bst_into_indices_key_range(range))
    }

    /// Iterate over `Nodes::Data<'_>` in a given key range.
    #[inline]
    pub fn bst_values_key_range<Key>(
        &self,
        range: impl ops::RangeBounds<Key>,
    ) -> Values<'tree, &Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Values(self.raw.bst_values_key_range(range))
    }

    /// Iterate over `Nodes::Data` in a given key range, consuming `self`.
    #[inline]
    pub fn bst_into_values_key_range<Key>(
        self,
        range: impl ops::RangeBounds<Key>,
    ) -> Values<'tree, Nodes, Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        Values(self.raw.bst_into_values_key_range(range))
    }
}

impl<'tree, Nodes, Index> IntoIterator for TreeAccessor<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index> + NodesData<Index>,
    Index: PartialEq + Clone,
{
    type Item = (Index, Nodes::Data);
    type IntoIter = Iter<'tree, Nodes, Index>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        Iter(self.raw.into_iter())
    }
}

impl<'a, 'tree, Nodes, Index> IntoIterator for &'a TreeAccessor<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index> + NodesDataLend<Index>,
    Index: PartialEq + Clone,
{
    type Item = (Index, <Nodes as NodesDataLendGat<'a, Index>>::Data);
    type IntoIter = Iter<'tree, &'a Nodes, Index>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        Iter((&self.raw).into_iter())
    }
}

/// An iterator over the indices and elements of [`TreeAccessor`].
#[derive(Debug, Clone)]
pub struct Iter<'tree, Nodes, Index>(crate::bintree::accessor::Iter<'tree, Nodes, Index>);

/// An iterator over the elements of [`TreeAccessor`].
#[derive(Debug, Clone)]
pub struct Values<'tree, Nodes, Index>(crate::bintree::accessor::Values<'tree, Nodes, Index>);

/// An iterator over the element indices of [`TreeAccessor`].
#[derive(Debug, Clone)]
pub struct Indices<'tree, Nodes, Index>(crate::bintree::accessor::Indices<'tree, Nodes, Index>);

dry::macro_for!($Ty in [Iter, Values, Indices] {
    impl<'tree, Nodes, Index> Default for $Ty<'tree, Nodes, Index>
    where
        Nodes: Default,
    {
        #[inline]
        fn default() -> Self {
            Self(Default::default())
        }
    }

    impl<'tree, Nodes, Index> Iterator for $Ty<'tree, Nodes, Index>
    where
        Nodes: NodesLink<Index> + NodesData<Index>,
        Index: PartialEq + Clone,
    {
        type Item = <crate::bintree::accessor::$Ty<'tree, Nodes, Index> as Iterator>::Item;

        #[inline]
        fn next(&mut self) -> Option<Self::Item> {
            self.0.next()
        }
    }

    impl<'tree, Nodes, Index> DoubleEndedIterator for $Ty<'tree, Nodes, Index>
    where
        Nodes: NodesLink<Index> + NodesData<Index>,
        Index: PartialEq + Clone,
    {
        #[inline]
        fn next_back(&mut self) -> Option<Self::Item> {
            self.0.next_back()
        }
    }

    impl<'tree, Nodes, Index> core::iter::FusedIterator for $Ty<'tree, Nodes, Index>
    where
        Nodes: NodesLink<Index> + NodesData<Index>,
        Index: PartialEq + Clone,
    {
    }
});

#[cfg(test)]
mod tests {
    use core::ops::RangeBounds as _;
    use std::vec::Vec;

    use crate::rbtree::tests::{Node, any_rbbst};

    use super::*;

    #[proptest::property_test]
    fn pt_bst_values_range(
        #[strategy = any_rbbst(true)] (nodes, root): (Vec<Node<u8>>, Tree),
        start: Bound<u8>,
        end: Bound<u8>,
    ) {
        let got_values: Vec<u8> = root
            .read(&nodes[..])
            .bst_values_key_range((start, end))
            .copied()
            .collect();

        let expected_values: Vec<u8> = root
            .read(&nodes[..])
            .values()
            .copied()
            .filter(|i| (start, end).contains(i))
            .collect();

        assert_eq!(got_values, expected_values);
    }
}