forest/
forest.rs

1use generic_cursors::with_data::{MoveDecision, MutRefStackWithData};
2
3pub struct Forest<T> {
4    roots: Vec<ForestNode<T>>,
5}
6
7struct ForestNode<T> {
8    data: T,
9    children: Vec<ForestNode<T>>,
10}
11
12pub fn preorder_traverse<T, F: FnMut(&mut T, usize)>(tree: &mut Forest<T>, mut callback: F) {
13    struct TraversalState {
14        next_index: usize,
15        depth: usize,
16    }
17    let mut cursor = MutRefStackWithData::new(
18        &mut *tree.roots,
19        TraversalState {
20            next_index: 0,
21            depth: 0,
22        },
23    );
24    loop {
25        let Ok(_) = cursor.move_with(|element, state| {
26            if state.next_index >= element.len() {
27                MoveDecision::Ascend
28            } else {
29                callback(&mut element[state.next_index].data, state.depth);
30                let decision = MoveDecision::Descend(
31                    &mut *element[state.next_index].children,
32                    TraversalState {
33                        next_index: 0,
34                        depth: state.depth + 1,
35                    },
36                );
37                state.next_index += 1;
38                decision
39            }
40        }) else {
41            break;
42        };
43    }
44}
45
46fn main() {
47    let mut forest = Forest {
48        roots: vec![
49            ForestNode {
50                data: 0u32,
51                children: vec![
52                    ForestNode {
53                        data: 1u32,
54                        children: vec![],
55                    },
56                    ForestNode {
57                        data: 2u32,
58                        children: vec![],
59                    },
60                    ForestNode {
61                        data: 3u32,
62                        children: vec![],
63                    },
64                ],
65            },
66            ForestNode {
67                data: 4u32,
68                children: vec![],
69            },
70            ForestNode {
71                data: 5u32,
72                children: vec![],
73            },
74            ForestNode {
75                data: 6u32,
76                children: vec![ForestNode {
77                    data: 7u32,
78                    children: vec![ForestNode {
79                        data: 8u32,
80                        children: vec![ForestNode {
81                            data: 9u32,
82                            children: vec![],
83                        }],
84                    }],
85                }],
86            },
87        ],
88    };
89    preorder_traverse(&mut forest, |t, depth| {
90        println!("{:depth$}{t}", "");
91        *t *= *t;
92    });
93    println!();
94    preorder_traverse(&mut forest, |t, depth| {
95        println!("{:depth$}{t}", "");
96    });
97}