pochoir-parser 0.12.2

use super::TreeRef;

pub struct TraverseBreadthIter<'a, 'b> {
    pub(super) children: Vec<TreeRef<'a, 'b>>,
    pub(super) index: usize,
}

impl<'a, 'b> Iterator for TraverseBreadthIter<'a, 'b> {
    type Item = TreeRef<'a, 'b>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.children.is_empty() {
            return None;
        }

        if let Some(tree_ref) = self.children.get(self.index) {
            self.index += 1;
            Some(*tree_ref)
        } else {
            // Merge all the children of all children
            let children = self.children.iter().flat_map(TreeRef::children).collect();

            self.children = children;
            self.index = 0;
            self.next()
        }
    }
}

pub struct TraverseDepthIter<'a, 'b> {
    pub(super) queue: Vec<TreeRef<'a, 'b>>,
}

impl<'a, 'b> Iterator for TraverseDepthIter<'a, 'b> {
    type Item = TreeRef<'a, 'b>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(tree_ref) = self.queue.pop() {
            let mut children: Vec<TreeRef> = tree_ref.children().collect();
            children.reverse();

            self.queue.extend(children);
            Some(tree_ref)
        } else {
            None
        }
    }
}