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

Port forward tree iterator from python to rust.

Browse Source
This commit is contained in:
Tom Alexander
2026-06-27 19:31:59 -04:00
parent fba8203cda
commit 1b09c77492
15 changed files with 862 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/error/error.rs
View File

@@ -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)

4
src/main.rs
View File

@@ -1,3 +1,7 @@
#![feature(test)]
extern crate test;
use std::process::ExitCode;
use clap::Parser;

1
src/nix_util/mod.rs
View File

@@ -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::*;

34
src/nix_util/running_update.rs
View File

@@ -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<Item = &ActivityTreeEntry> {
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))
}
}

118
src/nix_util/tree_iter/forward_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 = 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<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 = 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(())
}

235
src/nix_util/tree_iter/forward_tree_iter/mod.rs Normal file
View File

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

266
src/nix_util/tree_iter/forward_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 forward_tree_order() -> Result<(), Box<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
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(())
}

6
src/nix_util/tree_iter/mod.rs Normal file
View File

@@ -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;

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

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

63
src/nix_util/tree_iter/stack_entry.rs Normal file
View File

@@ -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;
}
}