rbt-rs 0.1.1

use std::{iter::FusedIterator, mem::ManuallyDrop};

use crate::{
    Node,
    NodeKind::{self},
    NodePtr, RedBlackTree,
};

/// An iterator over shared references of the nodes of a [RedBlackTree].
///
/// This `struct` is created by the [`iter_node`] method on [RedBlackTree] or [Node].
///
/// [`iter_node`]: RedBlackTree::iter_node
#[derive(Debug)]
pub struct IterNode<'a, T> {
    pub(crate) stack: Vec<(&'a Node<T>, bool)>,
}

impl<'a, T> Iterator for IterNode<'a, T> {
    type Item = &'a Node<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((top, recursed)) = self.stack.pop() {
            let top = if !recursed {
                self.stack.push((top, true));

                let mut x = top;
                while let NodeKind::Node(ref left) = x.left {
                    x = left;
                    self.stack.push((left, true));
                }

                self.stack.pop().unwrap().0
            } else {
                top
            };

            if let NodeKind::Node(ref right) = top.right {
                self.stack.push((right, false));
            }
            return Some(top);
        }
        None
    }
}

impl<'a, T> FusedIterator for IterNode<'a, T> {}

/// An iterator over shared references of the values of a [RedBlackTree].
///
/// This `struct` is created by the [`iter`] method on [RedBlackTree] or [Node].
///
/// [`iter`]: RedBlackTree::iter
#[derive(Debug)]
pub struct Iter<'a, T> {
    pub(crate) iter_node: IterNode<'a, T>,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter_node.next().map(|n| &n.t)
    }
}

impl<T> FusedIterator for Iter<'_, T> {}

impl<'a, T> IntoIterator for &'a Node<T> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a, T> IntoIterator for &'a RedBlackTree<T> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

/// An iterator over owned values of the values of a [RedBlackTree].
///
/// This `struct` is created by the [`into_iter`] method on [RedBlackTree] (provided by the [`IntoIterator`] trait).
///
/// [`into_iter`]: IntoIterator::into_iter
#[derive(Debug)]
pub struct IntoIter<T> {
    stack: Vec<(NodePtr<T>, bool)>,
    // ManuallyDrop stops RedBlackTree::drop as well as the drop-glue for RedBlackTree from being called.
    // Since RedBlackTree contains only stack allocations apart from the nodes, this is fine.
    _tree: ManuallyDrop<RedBlackTree<T>>,
}

// SAFETY: We don't access the T itself.
unsafe impl<#[may_dangle] T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        // Simply exhaust the iterator, which drops the unvisted allocated nodes.
        while self.next().is_some() {}
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((top, recursed)) = self.stack.pop() {
            let top = if !recursed {
                self.stack.push((top.clone(), true));

                let mut x = top;
                while let NodeKind::Node(left) = x.left.clone() {
                    x = left.clone();
                    self.stack.push((left, true));
                }

                self.stack.pop().unwrap().0
            } else {
                top
            };

            let Node { right, t, .. } = Box::into_inner(unsafe { Box::from_raw(top.0.as_ptr()) });
            if let Some(right) = right.node() {
                self.stack.push((right, false));
            }
            return Some(t);
        }
        None
    }
}

impl<T> FusedIterator for IntoIter<T> {}

impl<T> IntoIterator for RedBlackTree<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        if let Some(root) = self.root.clone().node() {
            IntoIter {
                stack: vec![(root, false)],
                _tree: ManuallyDrop::new(self),
            }
        } else {
            IntoIter {
                stack: vec![],
                _tree: ManuallyDrop::new(self),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn works_iter() {
        let tree = RedBlackTree::from_iter([243, 116, 212, 255, 177]);
        assert_eq!(
            tree.iter().map(|n| n).cloned().collect::<Vec<i32>>(),
            &[116, 177, 212, 243, 255]
        );
    }

    #[test]
    fn works_into_iter() {
        let tree = RedBlackTree::from_iter([15, 6, 3, 2, 4, 18]);
        assert_eq!(
            tree.into_iter().map(|n| n).collect::<Vec<i32>>(),
            &[2, 3, 4, 6, 15, 18]
        );
    }

    #[test]
    fn no_double_drop_into_iter() {
        let mut iter = RedBlackTree::from_iter([243, 116, 212, 255, 177]).into_iter();
        iter.next();
        iter.next();
    }
}