intrex 0.2.0

//! Iterators for [`TreeAccessor`]
use core::ops;

use crate::node_data::NodesData;

use super::*;

/// # In-Order Iteration
impl<'tree, Nodes, Index> TreeAccessor<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index>,
    Index: PartialEq + Clone,
{
    /// Perform in-order iteration over `(Index, Nodes::Data)`.
    #[inline]
    pub fn iter(&self) -> Iter<'tree, &Nodes, Index> {
        Iter {
            indices: self.indices(),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.indices_index_range(range),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.into_indices_index_range(range),
        }
    }

    /// Perform in-order iteration over `Index`.
    #[inline]
    pub fn indices(&self) -> Indices<'tree, &Nodes, Index> {
        Indices {
            range: Option::zip(self.front_index(), self.back_index())
                .map(|(start, end)| [start, end]),
            tree: self.tree,
            nodes: &self.nodes,
        }
    }

    /// Perform in-order iteration over `Index`, consuming `self`.
    #[inline]
    pub fn into_indices(self) -> Indices<'tree, Nodes, Index> {
        Indices {
            range: Option::zip(self.front_index(), self.back_index())
                .map(|(start, end)| [start, end]),
            tree: self.tree,
            nodes: self.nodes,
        }
    }

    /// Perform in-order iteration over `Index` in a given range.
    #[inline]
    pub fn indices_index_range(
        &self,
        range: impl ops::RangeBounds<Option<Index>>,
    ) -> Indices<'tree, &Nodes, Index> {
        let (range, _) = self.resolve_index_range(range);
        Indices {
            tree: self.tree,
            range,
            nodes: &self.nodes,
        }
    }

    /// Perform in-order iteration 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> {
        let (range, _) = self.resolve_index_range(range);
        Indices {
            tree: self.tree,
            range,
            nodes: self.nodes,
        }
    }

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

    /// Perform in-order iteration over `Nodes::Data`, consuming `self`.
    #[inline]
    pub fn into_values(self) -> Values<'tree, Nodes, Index> {
        Values {
            indices: self.into_indices(),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.indices_index_range(range),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.into_indices_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: NodesLink<Index>,
    Index: PartialEq + Clone,
{
    /// Perform in-order iteration 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 {
            indices: self.bst_indices_key_range(range),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.bst_into_indices_key_range(range),
        }
    }

    /// Perform in-order iteration 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>,
    {
        let range = self.bst_resolve_key_range(range);
        Indices {
            tree: self.tree,
            range,
            nodes: &self.nodes,
        }
    }

    /// Perform in-order iteration 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>,
    {
        let range = self.bst_resolve_key_range(range);
        Indices {
            tree: self.tree,
            range,
            nodes: self.nodes,
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.bst_indices_key_range(range),
        }
    }

    /// Perform in-order iteration 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 {
            indices: self.bst_into_indices_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 {
            indices: self.into_indices(),
        }
    }
}

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 {
        self.iter()
    }
}

/// An iterator over the indices and elements of [`TreeAccessor`].
#[derive(Debug)]
pub struct Iter<'tree, Nodes, Index> {
    indices: Indices<'tree, Nodes, Index>,
}

impl<'tree, Nodes, Index> Clone for Iter<'tree, Nodes, Index>
where
    Nodes: Clone,
    Index: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            indices: self.indices.clone(),
        }
    }
}

impl<'tree, Nodes, Index> Default for Iter<'tree, Nodes, Index>
where
    Nodes: Default,
{
    #[inline]
    fn default() -> Self {
        Self {
            indices: <_>::default(),
        }
    }
}

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

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.indices
            .next()
            .map(|i| (i.clone(), self.indices.nodes.node_data(i)))
    }
}

impl<'tree, Nodes, Index> DoubleEndedIterator for Iter<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index> + NodesData<Index>,
    Index: PartialEq + Clone,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        self.indices
            .next_back()
            .map(|i| (i.clone(), self.indices.nodes.node_data(i)))
    }
}

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

/// An iterator over the elements of [`TreeAccessor`].
#[derive(Debug)]
pub struct Values<'tree, Nodes, Index> {
    indices: Indices<'tree, Nodes, Index>,
}

impl<'tree, Nodes, Index> Clone for Values<'tree, Nodes, Index>
where
    Nodes: Clone,
    Index: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            indices: self.indices.clone(),
        }
    }
}

impl<'tree, Nodes, Index> Default for Values<'tree, Nodes, Index>
where
    Nodes: Default,
{
    #[inline]
    fn default() -> Self {
        Self {
            indices: <_>::default(),
        }
    }
}

impl<'tree, Nodes, Index> Iterator for Values<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index> + NodesData<Index>,
    Index: PartialEq + Clone,
{
    type Item = Nodes::Data;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.indices.next().map(|i| self.indices.nodes.node_data(i))
    }
}

impl<'tree, Nodes, Index> DoubleEndedIterator for Values<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index> + NodesData<Index>,
    Index: PartialEq + Clone,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        self.indices
            .next_back()
            .map(|i| self.indices.nodes.node_data(i))
    }
}

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

/// An iterator over the element indices of [`TreeAccessor`].
#[derive(Debug)]
pub struct Indices<'tree, Nodes, Index> {
    tree: &'tree Tree<Index>,
    nodes: Nodes,
    range: Option<[Index; 2]>,
}

impl<'tree, Nodes, Index> Clone for Indices<'tree, Nodes, Index>
where
    Nodes: Clone,
    Index: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            tree: self.tree,
            nodes: self.nodes.clone(),
            range: self.range.clone(),
        }
    }
}

impl<'tree, Nodes, Index> Default for Indices<'tree, Nodes, Index>
where
    Nodes: Default,
{
    #[inline]
    fn default() -> Self {
        Self {
            tree: &const { Tree { root: None } },
            nodes: Nodes::default(),
            range: None,
        }
    }
}

impl<'tree, Nodes, Index> Iterator for Indices<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index>,
    Index: PartialEq + Clone,
{
    type Item = Index;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let [start, end] = self.range.take()?;
        if start != end
            && let Some(new_start) = next_index(&self.nodes, start.clone())
        {
            self.range = Some([new_start, end]);
        }

        Some(start)
    }
}

impl<'tree, Nodes, Index> DoubleEndedIterator for Indices<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index>,
    Index: PartialEq + Clone,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        let [start, end] = self.range.take()?;
        if start != end
            && let Some(new_end) = prev_index(&self.nodes, end.clone())
        {
            self.range = Some([start, new_end]);
        }

        Some(end)
    }
}

impl<'tree, Nodes, Index> core::iter::FusedIterator for Indices<'tree, Nodes, Index>
where
    Nodes: NodesLink<Index>,
    Index: PartialEq + Clone,
{
}

#[cfg(test)]
mod tests {
    use std::{prelude::rust_2024::*, vec};

    use crate::bintree::tests::tree_from_str;

    #[test]
    fn values() {
        let (nodes, tree) = tree_from_str("_");
        assert_eq!(tree.read(&nodes[..]).values().next(), None);
        assert_eq!(tree.read(&nodes[..]).values().next_back(), None);

        let (nodes, tree) = tree_from_str("(4 _ _)");
        assert_eq!(
            tree.read(&nodes[..]).values().copied().collect::<Vec<_>>(),
            vec![4]
        );
        assert_eq!(
            tree.read(&nodes[..])
                .values()
                .rev()
                .copied()
                .collect::<Vec<_>>(),
            vec![4]
        );

        let (nodes, tree) = tree_from_str("(2 (1 _ _) (3 _ _))");
        assert_eq!(
            tree.read(&nodes[..]).values().copied().collect::<Vec<_>>(),
            vec![1, 2, 3]
        );
        assert_eq!(
            tree.read(&nodes[..])
                .values()
                .rev()
                .copied()
                .collect::<Vec<_>>(),
            vec![3, 2, 1]
        );

        let (nodes, tree) = tree_from_str("(3 (1 _ (2 _ _)) (4 _ _))");
        assert_eq!(
            tree.read(&nodes[..]).values().copied().collect::<Vec<_>>(),
            vec![1, 2, 3, 4]
        );
        assert_eq!(
            tree.read(&nodes[..])
                .values()
                .rev()
                .copied()
                .collect::<Vec<_>>(),
            vec![4, 3, 2, 1]
        );
    }
}