From 1b09c7749283438fe0b0e8a317dcbfcbcc0fe3d5 Mon Sep 17 00:00:00 2001 From: Tom Alexander Date: Sat, 27 Jun 2026 19:31:59 -0400 Subject: [PATCH] Port forward tree iterator from python to rust. --- Cargo.lock | 53 ++++ Cargo.toml | 3 + flake.lock | 12 +- .../flip_exponent_for_tests/flip.py | 39 +++ rust-toolchain.toml | 2 +- src/error/error.rs | 2 + src/main.rs | 4 + src/nix_util/mod.rs | 1 + src/nix_util/running_update.rs | 34 ++- .../tree_iter/forward_tree_iter/benchmarks.rs | 118 ++++++++ .../tree_iter/forward_tree_iter/mod.rs | 235 ++++++++++++++++ .../tree_iter/forward_tree_iter/tests.rs | 266 ++++++++++++++++++ src/nix_util/tree_iter/mod.rs | 6 + src/nix_util/tree_iter/reverse_tree_iter.rs | 49 ++++ src/nix_util/tree_iter/stack_entry.rs | 63 +++++ 15 files changed, 862 insertions(+), 25 deletions(-) create mode 100755 python_experiments/flip_exponent_for_tests/flip.py create mode 100644 src/nix_util/tree_iter/forward_tree_iter/benchmarks.rs create mode 100644 src/nix_util/tree_iter/forward_tree_iter/mod.rs create mode 100644 src/nix_util/tree_iter/forward_tree_iter/tests.rs create mode 100644 src/nix_util/tree_iter/mod.rs create mode 100644 src/nix_util/tree_iter/reverse_tree_iter.rs create mode 100644 src/nix_util/tree_iter/stack_entry.rs diff --git a/Cargo.lock b/Cargo.lock index 80d9985..5392a52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -609,6 +609,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "hamming" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65043da274378d68241eb9a8f8f8aa54e349136f7b8e12f63e3ef44043cc30e1" + [[package]] name = "hashbrown" version = "0.12.3" @@ -1047,6 +1053,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions", + "primal", "serde", "serde_json", "sha2", @@ -1339,6 +1346,52 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "primal" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1e5f354948532e6017fc91f9a5ff5ba1be0dabd3a0c9e9c417969cd4c1ad6e8" +dependencies = [ + "primal-check", + "primal-estimate", + "primal-sieve", +] + +[[package]] +name = "primal-bit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252429dbb8aeacc3233df500dc3a6a367bf28eb3a711272884d7540a7b636055" +dependencies = [ + "hamming", +] + +[[package]] +name = "primal-check" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08" +dependencies = [ + "num-integer", +] + +[[package]] +name = "primal-estimate" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a432100a0b3a61085e75b5f89e9f42de73c0acb7dea5038b893697918105d822" + +[[package]] +name = "primal-sieve" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e982796d82203351983d3602a8d6372d1d7894e86960047ba0d4b7426a5edd3" +dependencies = [ + "primal-bit", + "primal-estimate", + "smallvec", +] + [[package]] name = "proc-macro2" version = "1.0.105" diff --git a/Cargo.toml b/Cargo.toml index 552eed7..9bca627 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,9 @@ tracing-opentelemetry = { version = "0.20.0", optional = true } tracing-subscriber = { version = "0.3.17", optional = true, features = ["env-filter"] } url = { version = "2.5.8", default-features = false, features = ["std"] } +[dev-dependencies] +primal = { version = "0.3.3", default-features = false, features = [] } + # Optimized build for any sort of release. [profile.release-lto] inherits = "release" diff --git a/flake.lock b/flake.lock index edbd142..1e0861f 100644 --- a/flake.lock +++ b/flake.lock @@ -2,11 +2,11 @@ "nodes": { "nixpkgs": { "locked": { - "lastModified": 1771008912, - "narHash": "sha256-gf2AmWVTs8lEq7z/3ZAsgnZDhWIckkb+ZnAo5RzSxJg=", + "lastModified": 1778869304, + "narHash": "sha256-30sZNZoA1cqF5JNO9fVX+wgiQYjB7HJqqJ4ztCDeBZE=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "a82ccc39b39b621151d6732718e3e250109076fa", + "rev": "d233902339c02a9c334e7e593de68855ad26c4cb", "type": "github" }, "original": { @@ -29,11 +29,11 @@ ] }, "locked": { - "lastModified": 1771211437, - "narHash": "sha256-lcNK438i4DGtyA+bPXXyVLHVmJjYpVKmpux9WASa3ro=", + "lastModified": 1779506148, + "narHash": "sha256-OffE5yuMrSW71zIJNLvtl9MCO6WTAHEjlFPUCvkd/QM=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "c62195b3d6e1bb11e0c2fb2a494117d3b55d410f", + "rev": "d9973e2ab49747fada06ebbe26cda27eb0220cf1", "type": "github" }, "original": { diff --git a/python_experiments/flip_exponent_for_tests/flip.py b/python_experiments/flip_exponent_for_tests/flip.py new file mode 100755 index 0000000..1297f41 --- /dev/null +++ b/python_experiments/flip_exponent_for_tests/flip.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python +# +from math import log10 + + +def main(): + for x, y, z, num_nodes in sorted(list(get_matches(90, 90)), key=lambda val: val[3]): + if num_nodes < 1000000: + print(f"{x}**{y} == {y}**{z} == {num_nodes}") + + +def get_matches(max_x: int, max_y: int): + for x in range(2, max_x): + for y in range(2, max_y): + diff, y_exp = check(x, y) + if diff == 0: + yield (x, y, y_exp, x**y) + + +def check(x: int, y: int): + num_nodes = x**y + # y**z = x**y + # zlogy = ylogx + # z = ylogx/logy + + # flipped: + exp_for_y = y * log10(x) / log10(y) + # Round the exp for y to be a plain integer + exp_for_y = round(exp_for_y) + flipped_num_nodes = y**exp_for_y + + # compare the num nodes after rounding to find a good fit with integers. + diff = abs(num_nodes - flipped_num_nodes) + + return diff, exp_for_y + + +if __name__ == "__main__": + main() diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 8427192..251338f 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,4 +1,4 @@ [toolchain] -channel = "stable" +channel = "nightly-2026-05-23" profile = "default" components = ["clippy", "rustfmt"] diff --git a/src/error/error.rs b/src/error/error.rs index bea2db9..2c1e344 100644 --- a/src/error/error.rs +++ b/src/error/error.rs @@ -28,6 +28,8 @@ pub(crate) enum CustomError { SystemTime(#[allow(dead_code)] SystemTimeError), } +impl std::error::Error for CustomError {} + impl Display for CustomError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{:?}", self) diff --git a/src/main.rs b/src/main.rs index cd44eab..b184aee 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,7 @@ +#![feature(test)] + +extern crate test; + use std::process::ExitCode; use clap::Parser; diff --git a/src/nix_util/mod.rs b/src/nix_util/mod.rs index f5b6bf2..e249b12 100644 --- a/src/nix_util/mod.rs +++ b/src/nix_util/mod.rs @@ -7,6 +7,7 @@ mod output_stream; mod running_build; mod running_update; mod transparent_iter; +mod tree_iter; pub(crate) use activity_tree::ActivityIdAlreadyInTreeError; pub(crate) use activity_tree::ActivityIdNotInTreeError; pub(crate) use high_level::*; diff --git a/src/nix_util/running_update.rs b/src/nix_util/running_update.rs index f48f300..03d0b78 100644 --- a/src/nix_util/running_update.rs +++ b/src/nix_util/running_update.rs @@ -1,3 +1,5 @@ +use std::borrow::Cow; + use tokio::process::Child; use tracing::error; @@ -5,6 +7,7 @@ use crate::Result; use crate::nix_util::NixOutputStream; use crate::nix_util::nix_output_stream::NixAction; use crate::nix_util::output_stream::OutputStream; +use crate::nix_util::tree_iter::ForwardTreeIter; use super::activity_tree::ActivityId; use super::activity_tree::ActivityTreeEntry; @@ -113,26 +116,21 @@ impl RunningUpdate { fn maybe_print_current_status(&mut self) -> () {} fn print_current_status(&mut self) -> () { - let nodes = self.get_children_in_order(self.activity_tree.get_tree().get_root_id()); + let nodes = ForwardTreeIter::new( + self.activity_tree.get_tree(), + None, + |entry| true, + |entry| false, + |entry| true, + ); println!("\n\n\nvvvvvv\n\n"); - for n in nodes { - let activity = n.get_activity(); - // if activity.is_active() { - println!("{:?}", activity.display_name()); - // } + for (depth, is_match, is_transparent, node) in nodes { + let name = node + .get_activity() + .display_name() + .unwrap_or(Cow::Borrowed("null")); + println!("{depth}\t{is_match}\t{is_transparent}\t{name}"); } println!("\n\n\n^^^^^^\n\n"); } - - fn get_children_in_order( - &self, - parent_id: ActivityId, - ) -> impl Iterator { - let parent = self.activity_tree.get_tree().get(&parent_id); - parent - .get_child_ids() - .iter() - .map(|child_id| self.activity_tree.get_tree().get(child_id)) - .flat_map(|child| TransparentIter::new(self.activity_tree.get_tree(), child)) - } } diff --git a/src/nix_util/tree_iter/forward_tree_iter/benchmarks.rs b/src/nix_util/tree_iter/forward_tree_iter/benchmarks.rs new file mode 100644 index 0000000..970a92f --- /dev/null +++ b/src/nix_util/tree_iter/forward_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 = ForwardTreeIter::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 = ForwardTreeIter::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/forward_tree_iter/mod.rs b/src/nix_util/tree_iter/forward_tree_iter/mod.rs new file mode 100644 index 0000000..14abf22 --- /dev/null +++ b/src/nix_util/tree_iter/forward_tree_iter/mod.rs @@ -0,0 +1,235 @@ +#[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 earliest/highest 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 +/// 1 False False A +/// 2 True False B +/// 3 False True D +/// 3 True False E +/// 2 False False H +/// 3 False False I +/// 4 True False J +/// ``` +pub(crate) struct ForwardTreeIter<'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> ForwardTreeIter<'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, + ) -> ForwardTreeIter<'tree, MP, TP, AP> { + let origin_node = origin.unwrap_or_else(|| { + let root_id = activity_tree.get_root_id(); + activity_tree.get(&root_id) + }); + ForwardTreeIter { + 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 get_next_announce<'iter>( + &'iter mut self, + ) -> Option<(usize, bool, bool, &'tree ActivityTreeEntry)> { + let (activity, is_match, is_transparent) = match self + .stack + .iter_mut() + .filter(|entry| entry.should_announce() && !entry.is_announced()) + .next() + { + Some(entry) => { + entry.set_announced(); + let is_match = entry.matches(); + let is_transparent = entry.is_transparent(); + (entry.get_node(), is_match, is_transparent) + } + None => { + return None; + } + }; + + let depth = self + .walk_up_stack(activity) + .filter(|node| !node.is_transparent()) + .count(); + + Some((depth, is_match, is_transparent, activity)) + } + + /// 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 ForwardTreeIter<'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 { + // If any entries in the stack are waiting to be returned, output them. + match self.get_next_announce() { + out @ Some(_) => { + return out; + } + None => {} + }; + + // Walk through stack + loop { + let mut should_check_for_output = false; + + let node = { + let stack_entry = match self.stack.last_mut() { + Some(e) => e, + None => { + break; + } + }; + let node = stack_entry.get_node(); + + // If we've already processed the children, remove it from the stack. + if stack_entry.is_visited() { + self.stack.pop(); + continue; + } + + node + }; + + // 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() + .rev() + .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); + + // If any entries in the stack are waiting to be returned, output them. + if should_check_for_output { + match self.get_next_announce() { + out @ Some(_) => { + return out; + } + None => {} + }; + } + } + + None + } +} diff --git a/src/nix_util/tree_iter/forward_tree_iter/tests.rs b/src/nix_util/tree_iter/forward_tree_iter/tests.rs new file mode 100644 index 0000000..9cda1fa --- /dev/null +++ b/src/nix_util/tree_iter/forward_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 forward_tree_order() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ForwardTreeIter::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![ + (1, false, false, "Unknown(A)".into()), + (2, true, false, "Unknown(B)".into()), + (3, false, true, "Unknown(D)".into()), + (3, true, false, "Unknown(E)".into()), + (2, false, false, "Unknown(H)".into()), + (3, false, false, "Unknown(I)".into()), + (4, true, false, "Unknown(J)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_a() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ForwardTreeIter::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![ + (1, true, false, "Unknown(B)".into()), + (2, false, true, "Unknown(D)".into()), + (2, true, false, "Unknown(E)".into()), + (1, false, false, "Unknown(H)".into()), + (2, false, false, "Unknown(I)".into()), + (3, true, false, "Unknown(J)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_h() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ForwardTreeIter::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![ + (1, false, false, "Unknown(I)".into()), + (2, true, false, "Unknown(J)".into()), + ]; + + assert_eq!(output, expected); + + Ok(()) +} + +#[test] +fn start_from_c() -> Result<(), Box> { + let activity_tree = build_example_tree()?; + let nodes = ForwardTreeIter::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(()) +} diff --git a/src/nix_util/tree_iter/mod.rs b/src/nix_util/tree_iter/mod.rs new file mode 100644 index 0000000..b1aa94f --- /dev/null +++ b/src/nix_util/tree_iter/mod.rs @@ -0,0 +1,6 @@ +mod forward_tree_iter; +mod reverse_tree_iter; +mod stack_entry; + +pub(crate) use forward_tree_iter::ForwardTreeIter; +pub(crate) use reverse_tree_iter::ReverseTreeIter; diff --git a/src/nix_util/tree_iter/reverse_tree_iter.rs b/src/nix_util/tree_iter/reverse_tree_iter.rs new file mode 100644 index 0000000..1f7b87a --- /dev/null +++ b/src/nix_util/tree_iter/reverse_tree_iter.rs @@ -0,0 +1,49 @@ +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/stack_entry.rs b/src/nix_util/tree_iter/stack_entry.rs new file mode 100644 index 0000000..7b13584 --- /dev/null +++ b/src/nix_util/tree_iter/stack_entry.rs @@ -0,0 +1,63 @@ +use crate::nix_util::activity_tree::ActivityTreeEntry; + +const FLAG_ANNOUNCED: u8 = 1 << 0; +const FLAG_SHOULD_ANNOUNCE: u8 = 1 << 1; +const FLAG_VISITED: u8 = 1 << 2; +const FLAG_MATCHED: u8 = 1 << 3; +const FLAG_TRANSPARENT: u8 = 1 << 4; + +pub(crate) struct StackEntry<'tree> { + // TODO: node_id? node itself? + node: &'tree ActivityTreeEntry, + flags: u8, +} + +impl<'tree> StackEntry<'tree> { + pub(crate) fn new(node: &'tree ActivityTreeEntry) -> StackEntry<'tree> { + StackEntry { node, flags: 0 } + } + + pub(crate) fn get_node(&self) -> &'tree ActivityTreeEntry { + self.node + } + + pub(crate) fn is_announced(&self) -> bool { + self.flags & FLAG_ANNOUNCED != 0 + } + + pub(crate) fn set_announced(&mut self) -> () { + self.flags = self.flags | FLAG_ANNOUNCED; + } + + pub(crate) fn should_announce(&self) -> bool { + self.flags & FLAG_SHOULD_ANNOUNCE != 0 + } + + pub(crate) fn set_should_announce(&mut self) -> () { + self.flags = self.flags | FLAG_SHOULD_ANNOUNCE; + } + + pub(crate) fn is_visited(&self) -> bool { + self.flags & FLAG_VISITED != 0 + } + + pub(crate) fn set_visited(&mut self) -> () { + self.flags = self.flags | FLAG_VISITED; + } + + pub(crate) fn matches(&self) -> bool { + self.flags & FLAG_MATCHED != 0 + } + + pub(crate) fn set_matches(&mut self) -> () { + self.flags = self.flags | FLAG_MATCHED; + } + + pub(crate) fn is_transparent(&self) -> bool { + self.flags & FLAG_TRANSPARENT != 0 + } + + pub(crate) fn set_transparent(&mut self) -> () { + self.flags = self.flags | FLAG_TRANSPARENT; + } +}