cursive-tree 0.0.9

use super::{super::backend::*, list::*, node::*};

//
// NodeIterator
//

/// Tree node iterator.
///
/// Iterates the nodes in visual order from top to bottom.
pub struct NodeIterator<'model, BackendT>
where
    BackendT: TreeBackend,
{
    current_node_list: &'model NodeList<BackendT>,
    next_index: usize,
    only_expanded: bool,
    stack: Vec<(&'model NodeList<BackendT>, usize)>,
}

impl<'model, BackendT> NodeIterator<'model, BackendT>
where
    BackendT: TreeBackend,
{
    /// Constructor.
    ///
    /// When only_expanded is true will skip the children of collapsed branches.
    pub fn new(roots: &'model NodeList<BackendT>, only_expanded: bool) -> Self {
        Self { current_node_list: roots, next_index: 0, only_expanded, stack: Default::default() }
    }
}

impl<'model, BackendT> Iterator for NodeIterator<'model, BackendT>
where
    BackendT: TreeBackend,
{
    type Item = &'model Node<BackendT>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.current_node_list.0.get(self.next_index) {
            Some(node) => {
                self.next_index += 1;

                if node.kind.is_branch()
                    && (!self.only_expanded || node.branch_state.is_expanded())
                    && let Some(children) = &node.children
                {
                    // Remember our place for later
                    self.stack.push((self.current_node_list, self.next_index));

                    // Continue into children
                    self.current_node_list = children;
                    self.next_index = 0;
                }

                Some(node)
            }

            None => match self.stack.pop() {
                Some((current_nodes, next_index)) => {
                    // Go back to where we were before
                    self.current_node_list = current_nodes;
                    self.next_index = next_index;
                    self.next()
                }

                None => None,
            },
        }
    }
}