nix_builder/src/nix_util/transparent_iter.rs

use super::activity_tree::ActivityTree;
use super::activity_tree::ActivityTreeEntry;

/// An iterator that returns either the original activity if it is not transparent, or if it is transparent, the earliest non-transparent children.
pub(crate) struct TransparentIter<'tree> {
    activity_tree: &'tree ActivityTree,
    origin: Option<&'tree ActivityTreeEntry>,
    child_index: Vec<usize>,
}

impl<'tree> Iterator for TransparentIter<'tree> {
    type Item = &'tree ActivityTreeEntry;

    fn next(&mut self) -> Option<Self::Item> {
        let origin = match self.origin {
            Some(o) => o,
            None => {
                return None;
            }
        };

        if !origin.get_activity().is_transparent() {
            let origin = self
                .origin
                .take()
                .expect("We would have returned early if origin was None.");
            return Some(origin);
        }

        if self.child_index.is_empty() {
            self.child_index.push(0);
        }

        loop {
            let current_entry = self.get_current_entry();
            if let Some(child) = current_entry {
                if child.get_activity().is_transparent() {
                    self.child_index.push(0);
                } else {
                    let last_index = self.child_index.last_mut().expect("Stack cannot be empty.");
                    *last_index += 1;
                    return Some(child);
                }
            } else {
                self.child_index.pop();
                if self.child_index.is_empty() {
                    self.origin.take();
                    return None;
                }
                let last_index = self.child_index.last_mut().expect("Stack cannot be empty.");
                *last_index += 1;
            }
        }
    }
}

impl<'tree> TransparentIter<'tree> {
    pub(crate) fn new(
        activity_tree: &'tree ActivityTree,
        origin: &'tree ActivityTreeEntry,
    ) -> TransparentIter<'tree> {
        TransparentIter {
            activity_tree,
            origin: Some(origin),
            child_index: Vec::new(),
        }
    }

    fn get_current_entry(&mut self) -> Option<&'tree ActivityTreeEntry> {
        let mut current_entry = match self.origin {
            Some(o) => o,
            None => {
                return None;
            }
        };
        for ind in self.child_index.iter() {
            if let Some(child) = current_entry
                .get_child_ids()
                .get(*ind)
                .map(|child_index| self.activity_tree.get(child_index))
            {
                current_entry = child;
                continue;
            } else {
                return None;
            }
        }
        Some(current_entry)
    }
}