talexander/nix_builder
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Port reverse tree iterator to from python to rust.

Browse Source
This commit is contained in:
Tom Alexander
2026-06-28 23:06:38 -04:00
parent 1b09c77492
commit 8c4f0868a4
4 changed files with 604 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
src/nix_util/tree_iter/reverse_tree_iter.rs
View File

@@ -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<usize>,
}

118
src/nix_util/tree_iter/reverse_tree_iter/benchmarks.rs Normal file
View File

@@ -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<ActivityTree, Box<dyn std::error::Error>> {
let mut activity_tree = ActivityTree::new();
let mut next_index = 1;
let mut previous_level: Vec<u64> = 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<u64> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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(())
}

220
src/nix_util/tree_iter/reverse_tree_iter/mod.rs Normal file
View File

@@ -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<StackEntry<'tree>>,
}
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<Item = &'iter mut StackEntry<'tree>> {
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<Self::Item> {
// 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
}
}

266
src/nix_util/tree_iter/reverse_tree_iter/tests.rs Normal file
View File

@@ -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<ActivityTree, Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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(())
}