From 8c4f0868a42c9e28d6dcddd87690440d77454b61 Mon Sep 17 00:00:00 2001 From: Tom Alexander Date: Sun, 28 Jun 2026 23:06:38 -0400 Subject: [PATCH] Port reverse tree iterator to from python to rust. --- src/nix_util/tree_iter/reverse_tree_iter.rs | 49 ---- .../tree_iter/reverse_tree_iter/benchmarks.rs | 118 ++++++++ .../tree_iter/reverse_tree_iter/mod.rs | 220 +++++++++++++++ .../tree_iter/reverse_tree_iter/tests.rs | 266 ++++++++++++++++++ 4 files changed, 604 insertions(+), 49 deletions(-) delete mode 100644 src/nix_util/tree_iter/reverse_tree_iter.rs create mode 100644 src/nix_util/tree_iter/reverse_tree_iter/benchmarks.rs create mode 100644 src/nix_util/tree_iter/reverse_tree_iter/mod.rs create mode 100644 src/nix_util/tree_iter/reverse_tree_iter/tests.rs diff --git a/src/nix_util/tree_iter/reverse_tree_iter.rs b/src/nix_util/tree_iter/reverse_tree_iter.rs deleted file mode 100644 index 1f7b87a..0000000 --- a/src/nix_util/tree_iter/reverse_tree_iter.rs +++ /dev/null @@ -1,49 +0,0 @@ -use crate::nix_util::activity_tree::ActivityTree; -use crate::nix_util::activity_tree::ActivityTreeEntry; - -/// An iterator over the activity tree that will return nodes in the tree matching the supplied criteria. -/// -/// is_match: Return only nodes that pass this predicate and the parents of that node up to the root. -/// is_alive: If a node fails this predicate, do not return it and do not check any of its children. -/// is_transparent: Nodes passing this predicate will be returned but they will not be counted in the depth calculation. -/// -/// The elements returned from this iterator are a tuple containing: -/// - The depth -/// - Whether the node passed is_match -/// - Whether the node passed is_transparent -/// - The activity itself -/// -/// Using this information, you can draw a tree, using the depth information to glean which nodes are children of other nodes. -/// -/// This iterator returns the deepest/latest nodes first in a depth-first-search pattern. -/// -/// With the example tree: -/// -/// ```text -/// A - B - C -/// |\ \ -/// H G D - E -/// | \ -/// I F -/// | -/// J -/// ``` -/// -/// With B, E, and J matching -/// and D as transparent -/// you would get the following output: -/// -/// ```text -/// 4 True False J -/// 3 False False I -/// 2 False False H -/// 3 True False E -/// 3 False True D -/// 2 True False B -/// 1 False False A -/// ``` -pub(crate) struct ReverseTreeIter<'tree> { - activity_tree: &'tree ActivityTree, - origin: Option<&'tree ActivityTreeEntry>, - child_index: Vec, -} diff --git a/src/nix_util/tree_iter/reverse_tree_iter/benchmarks.rs b/src/nix_util/tree_iter/reverse_tree_iter/benchmarks.rs new file mode 100644 index 0000000..b757e73 --- /dev/null +++ b/src/nix_util/tree_iter/reverse_tree_iter/benchmarks.rs @@ -0,0 +1,118 @@ +use std::collections::BTreeSet; + +use crate::nix_util::activity::Activity; +use crate::nix_util::activity::ActivityState; +use crate::nix_util::activity::ActivityUnknown; + +use test::Bencher; + +use super::*; + +/// Build an ActivityTree matching the description at the top of the file. +fn build_example_tree( + num_levels: u64, + num_children_per_node: u64, +) -> Result> { + let mut activity_tree = ActivityTree::new(); + let mut next_index = 1; + let mut previous_level: Vec = Vec::new(); + + for _ in 0..num_levels { + if previous_level.is_empty() { + // first level, add num_children + for _ in 0..num_children_per_node { + activity_tree.add_activity( + next_index, + 0, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: format!("N{next_index}"), + }), + )?; + previous_level.push(next_index); + next_index += 1; + } + } else { + // Subsequent levels, add num_children to each node of the previous level. + let mut current_parents = Vec::new(); + std::mem::swap(&mut previous_level, &mut current_parents); + for parent in current_parents { + for _ in 0..num_children_per_node { + activity_tree.add_activity( + next_index, + parent, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: format!("N{next_index}"), + }), + )?; + previous_level.push(next_index); + next_index += 1; + } + } + } + } + + Ok(activity_tree) +} + +fn get_primes(limit: u64) -> BTreeSet { + primal::Primes::all() + .take_while(|p| *p <= (limit as usize)) + .map(|p| p as u64) + .collect() +} + +/// Test a tree with few levels but many nodes per level. +#[bench] +fn bench_shallow_wide_tree(b: &mut Bencher) -> Result<(), Box> { + let num_levels: u64 = 3; + let num_children_per_node: u64 = 27; + let total_nodes = num_children_per_node.pow(num_levels as u32); + + let primes = get_primes(total_nodes); + // match = power of 2 + // transparent = prime + let activity_tree = build_example_tree(num_levels, num_children_per_node)?; + + b.iter(|| { + let nodes = ReverseTreeIter::new( + &activity_tree, + None, + |entry| (*entry.get_activity_id().deref()).count_ones() == 1, + |entry| primes.contains(entry.get_activity_id().deref()), + |_entry| true, + ); + + for _ in nodes {} + }); + + Ok(()) +} + +/// Test a tree with many levels but few nodes per level. +#[bench] +fn bench_deep_narrow_tree(b: &mut Bencher) -> Result<(), Box> { + let num_levels: u64 = 9; + let num_children_per_node: u64 = 3; + let total_nodes = num_children_per_node.pow(num_levels as u32); + + let primes = get_primes(total_nodes); + // match = power of 2 + // transparent = prime + let activity_tree = build_example_tree(num_levels, num_children_per_node)?; + + b.iter(|| { + let nodes = ReverseTreeIter::new( + &activity_tree, + None, + |entry| (*entry.get_activity_id().deref()).count_ones() == 1, + |entry| primes.contains(entry.get_activity_id().deref()), + |_entry| true, + ); + + for _ in nodes {} + }); + + Ok(()) +} diff --git a/src/nix_util/tree_iter/reverse_tree_iter/mod.rs b/src/nix_util/tree_iter/reverse_tree_iter/mod.rs new file mode 100644 index 0000000..d482c37 --- /dev/null +++ b/src/nix_util/tree_iter/reverse_tree_iter/mod.rs @@ -0,0 +1,220 @@ +#[cfg(test)] +mod benchmarks; +#[cfg(test)] +mod tests; + +use std::ops::Deref; + +use crate::nix_util::activity_tree::ActivityTree; +use crate::nix_util::activity_tree::ActivityTreeEntry; + +use super::stack_entry::StackEntry; + +/// An iterator over the activity tree that will return nodes in the tree matching the supplied criteria. +/// +/// is_match: Return only nodes that pass this predicate and the parents of that node up to the root. +/// is_alive: If a node fails this predicate, do not return it and do not check any of its children. +/// is_transparent: Nodes passing this predicate will be returned but they will not be counted in the depth calculation. +/// +/// The elements returned from this iterator are a tuple containing: +/// - The depth +/// - Whether the node passed is_match +/// - Whether the node passed is_transparent +/// - The activity itself +/// +/// Using this information, you can draw a tree, using the depth information to glean which nodes are children of other nodes. +/// +/// This iterator returns the deepest/latest nodes first in a depth-first-search pattern. +/// +/// With the example tree: +/// +/// ```text +/// A - B - C +/// |\ \ +/// H G D - E +/// | \ +/// I F +/// | +/// J +/// ``` +/// +/// With B, E, and J matching +/// and D as transparent +/// you would get the following output: +/// +/// ```text +/// 4 True False J +/// 3 False False I +/// 2 False False H +/// 3 True False E +/// 3 False True D +/// 2 True False B +/// 1 False False A +/// ``` +pub(crate) struct ReverseTreeIter<'tree, MP, TP, AP> { + activity_tree: &'tree ActivityTree, + origin: &'tree ActivityTreeEntry, + is_match_predicate: MP, + is_transparent_predicate: TP, + is_alive_predicate: AP, + stack: Vec>, +} + +impl<'tree, MP, TP, AP> ReverseTreeIter<'tree, MP, TP, AP> +where + MP: FnMut(&'tree ActivityTreeEntry) -> bool, + TP: FnMut(&'tree ActivityTreeEntry) -> bool, + AP: FnMut(&'tree ActivityTreeEntry) -> bool, +{ + pub(crate) fn new( + activity_tree: &'tree ActivityTree, + origin: Option<&'tree ActivityTreeEntry>, + is_match_predicate: MP, + is_transparent_predicate: TP, + is_alive_predicate: AP, + ) -> ReverseTreeIter<'tree, MP, TP, AP> { + let origin_node = origin.unwrap_or_else(|| { + let root_id = activity_tree.get_root_id(); + activity_tree.get(&root_id) + }); + ReverseTreeIter { + activity_tree, + origin: origin_node, + is_match_predicate, + is_transparent_predicate, + is_alive_predicate, + stack: vec![StackEntry::new(origin_node)], + } + } + + /// Find the next StackEntry that should be returned from the iterator. + fn announce<'iter>( + &'iter mut self, + node: &'tree ActivityTreeEntry, + is_match: bool, + is_transparent: bool, + ) -> (usize, bool, bool, &'tree ActivityTreeEntry) { + let depth = self + .walk_up_stack(node) + .filter(|node| !node.is_transparent()) + .count(); + (depth, is_match, is_transparent, node) + } + + /// Walk up the stack, returning all visited nodes above the base. + /// + /// This should give you the path to base (starting from the end). + fn walk_up_stack<'iter>( + &'iter mut self, + base: &'tree ActivityTreeEntry, + ) -> impl Iterator> { + self.stack + .iter_mut() + .rev() + .skip_while(|node| !std::ptr::eq(node.get_node(), base)) + .skip(1) + .filter(|node| node.is_visited()) + } +} + +impl<'tree, MP, TP, AP> Iterator for ReverseTreeIter<'tree, MP, TP, AP> +where + MP: FnMut(&'tree ActivityTreeEntry) -> bool, + TP: FnMut(&'tree ActivityTreeEntry) -> bool, + AP: FnMut(&'tree ActivityTreeEntry) -> bool, +{ + type Item = (usize, bool, bool, &'tree ActivityTreeEntry); + + fn next(&mut self) -> Option { + // Walk through stack + loop { + let mut should_check_for_output = false; + + let (node, is_visited, should_announce, is_match, is_transparent) = { + let stack_entry = match self.stack.last_mut() { + Some(e) => e, + None => { + break; + } + }; + let node = stack_entry.get_node(); + let is_visited = stack_entry.is_visited(); + let should_announce = stack_entry.should_announce(); + let is_match = stack_entry.matches(); + let is_transparent = stack_entry.is_transparent(); + + // If we've already processed the children, remove it from the stack. + if stack_entry.is_visited() { + if stack_entry.should_announce() { + let output = Some(self.announce(node, is_match, is_transparent)); + self.stack.pop(); + return output; + } + self.stack.pop(); + continue; + } + + (node, is_visited, should_announce, is_match, is_transparent) + }; + + if is_visited { + if should_announce { + let output = Some(self.announce(node, is_match, is_transparent)); + self.stack.pop(); + return output; + } + self.stack.pop(); + continue; + } + + // Otherwise, check if it matches the predicate and mark all parents if it does. + if (self.is_match_predicate)(node) { + { + let stack_entry = self + .stack + .last_mut() + .expect("Loop condition guarantees this is Some."); + // TODO: Combine into single op. + stack_entry.set_should_announce(); + stack_entry.set_matches(); + should_check_for_output = true; + } + let origin_activity_id = *self.origin.get_activity_id().deref(); + for parent in self.walk_up_stack(node).filter(|entry| { + !entry.should_announce() + && *entry.get_node().get_activity_id().deref() != origin_activity_id + }) { + parent.set_should_announce(); + } + } + + let children = { + let stack_entry = self + .stack + .last_mut() + .expect("Loop condition guarantees this is Some."); + + // Mark the node as visited + stack_entry.set_visited(); + if (self.is_transparent_predicate)(node) { + stack_entry.set_transparent(); + } + + // Add all children to the stack if they pass the is_alive_predicate. + let children = node + .get_child_ids() + .iter() + .map(|id| self.activity_tree.get(id)) + .filter(|node| { + *node.get_activity_id().deref() != 0 && (self.is_alive_predicate)(node) + }) + .map(|node| StackEntry::new(node)); + children + }; + + self.stack.extend(children); + } + + None + } +} diff --git a/src/nix_util/tree_iter/reverse_tree_iter/tests.rs b/src/nix_util/tree_iter/reverse_tree_iter/tests.rs new file mode 100644 index 0000000..23d4d02 --- /dev/null +++ b/src/nix_util/tree_iter/reverse_tree_iter/tests.rs @@ -0,0 +1,266 @@ +use crate::nix_util::activity::Activity; +use crate::nix_util::activity::ActivityState; +use crate::nix_util::activity::ActivityUnknown; + +use super::*; + +/// Build an ActivityTree matching the description at the top of the file. +fn build_example_tree() -> Result> { + let mut activity_tree = ActivityTree::new(); + activity_tree.add_activity( + 1, + 0, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "A".to_owned(), + }), + )?; + activity_tree.add_activity( + 2, + 1, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "B".to_owned(), + }), + )?; + activity_tree.add_activity( + 3, + 2, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "C".to_owned(), + }), + )?; + activity_tree.add_activity( + 4, + 2, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "D".to_owned(), + }), + )?; + activity_tree.add_activity( + 5, + 4, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "E".to_owned(), + }), + )?; + activity_tree.add_activity( + 6, + 4, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "F".to_owned(), + }), + )?; + activity_tree.add_activity( + 7, + 1, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "G".to_owned(), + }), + )?; + activity_tree.add_activity( + 8, + 1, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "H".to_owned(), + }), + )?; + activity_tree.add_activity( + 9, + 8, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "I".to_owned(), + }), + )?; + activity_tree.add_activity( + 10, + 9, + Activity::Unknown(ActivityUnknown { + state: ActivityState::Stopped, + text: "J".to_owned(), + }), + )?; + + Ok(activity_tree) +} + +#[test] +fn reverse_tree_order() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ReverseTreeIter::new( + &activity_tree, + None, + |entry| match *entry.get_activity_id().deref() { + 2 | 5 | 10 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 4 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 => true, + _ => false, + }, + ); + + let output: Vec<_> = nodes + .map(|(depth, is_match, is_transparent, node)| { + ( + depth, + is_match, + is_transparent, + node.get_activity() + .display_name() + .expect("Should have a display_name"), + ) + }) + .collect(); + let expected = vec![ + (4, true, false, "Unknown(J)".into()), + (3, false, false, "Unknown(I)".into()), + (2, false, false, "Unknown(H)".into()), + (3, true, false, "Unknown(E)".into()), + (3, false, true, "Unknown(D)".into()), + (2, true, false, "Unknown(B)".into()), + (1, false, false, "Unknown(A)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_a() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ReverseTreeIter::new( + &activity_tree, + Some(activity_tree.get(&activity_tree.get_activity_id(1)?)), + |entry| match *entry.get_activity_id().deref() { + 2 | 5 | 10 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 4 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 => true, + _ => false, + }, + ); + + let output: Vec<_> = nodes + .map(|(depth, is_match, is_transparent, node)| { + ( + depth, + is_match, + is_transparent, + node.get_activity() + .display_name() + .expect("Should have a display_name"), + ) + }) + .collect(); + let expected = vec![ + (3, true, false, "Unknown(J)".into()), + (2, false, false, "Unknown(I)".into()), + (1, false, false, "Unknown(H)".into()), + (2, true, false, "Unknown(E)".into()), + (2, false, true, "Unknown(D)".into()), + (1, true, false, "Unknown(B)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_h() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ReverseTreeIter::new( + &activity_tree, + Some(activity_tree.get(&activity_tree.get_activity_id(8)?)), + |entry| match *entry.get_activity_id().deref() { + 2 | 5 | 10 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 4 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 => true, + _ => false, + }, + ); + + let output: Vec<_> = nodes + .map(|(depth, is_match, is_transparent, node)| { + ( + depth, + is_match, + is_transparent, + node.get_activity() + .display_name() + .expect("Should have a display_name"), + ) + }) + .collect(); + let expected = vec![ + (2, true, false, "Unknown(J)".into()), + (1, false, false, "Unknown(I)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_c() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ReverseTreeIter::new( + &activity_tree, + Some(activity_tree.get(&activity_tree.get_activity_id(3)?)), + |entry| match *entry.get_activity_id().deref() { + 2 | 5 | 10 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 4 => true, + _ => false, + }, + |entry| match *entry.get_activity_id().deref() { + 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 => true, + _ => false, + }, + ); + + let output: Vec<_> = nodes + .map(|(depth, is_match, is_transparent, node)| { + ( + depth, + is_match, + is_transparent, + node.get_activity() + .display_name() + .expect("Should have a display_name"), + ) + }) + .collect(); + let expected = vec![]; + + assert_eq!(output, expected); + + Ok(()) +}