cairo_lang_sierra/algorithm/topological_order.rs
1use crate::program::StatementIdx;
2
3/// The status of a node during the topological ordering finding algorithm.
4#[derive(Clone, Debug)]
5enum TopologicalOrderStatus {
6 /// The computation for that statement did not start.
7 NotStarted,
8 /// The computation is in progress.
9 InProgress,
10 /// The computation was completed, and all the children were visited.
11 Done,
12}
13
14/// Returns the reverse topological ordering.
15///
16/// `detect_cycles` - if true, the function will return an error if a cycle is detected, else will
17/// not detect the cycle, and ordering within cycles won't be topological.
18/// `roots` - the roots of the graph.
19/// `node_count` - the number of nodes in the graph.
20/// `get_children` - a function that returns the children of a node.
21/// `out_of_bounds_err` - a function that returns an error for a node is out of bounds.
22/// `cycle_err` - a function that returns an error for a node that is part of a cycle.
23/// Note: Will only work properly if the nodes are in the range [0, node_count).
24pub fn reverse_topological_ordering<E>(
25 detect_cycles: bool,
26 roots: impl Iterator<Item = StatementIdx>,
27 node_count: usize,
28 get_children: impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>,
29 cycle_err: impl Fn(StatementIdx) -> E,
30) -> Result<Vec<StatementIdx>, E> {
31 let mut ordering = vec![];
32 let mut status = vec![TopologicalOrderStatus::NotStarted; node_count];
33 for root in roots {
34 calculate_reverse_topological_ordering(
35 detect_cycles,
36 &mut ordering,
37 &mut status,
38 root,
39 &get_children,
40 &cycle_err,
41 )?;
42 }
43 Ok(ordering)
44}
45
46/// Calculates the reverse topological ordering starting from `root`. For more info see
47/// `reverse_topological_ordering`.
48fn calculate_reverse_topological_ordering<E>(
49 detect_cycles: bool,
50 ordering: &mut Vec<StatementIdx>,
51 status: &mut [TopologicalOrderStatus],
52 root: StatementIdx,
53 get_children: &impl Fn(StatementIdx) -> Result<Vec<StatementIdx>, E>,
54 cycle_err: &impl Fn(StatementIdx) -> E,
55) -> Result<(), E> {
56 // A stack of statements to visit.
57 let mut stack = vec![root];
58
59 while let Some(idx) = stack.pop() {
60 match status[idx.0] {
61 TopologicalOrderStatus::NotStarted => {
62 // Mark the statement as `InProgress`.
63 status[idx.0] = TopologicalOrderStatus::InProgress;
64
65 // Push the statement back to the stack, so that after visiting all
66 // of its children, we would add it to the ordering.
67 // Add the missing children on top of it.
68 stack.push(idx);
69 for child in get_children(idx)? {
70 match status[child.0] {
71 TopologicalOrderStatus::InProgress if detect_cycles => {
72 return Err(cycle_err(child));
73 }
74 TopologicalOrderStatus::NotStarted => stack.push(child),
75 TopologicalOrderStatus::Done | TopologicalOrderStatus::InProgress => {}
76 }
77 }
78 }
79 TopologicalOrderStatus::InProgress => {
80 // Mark the statement as `Done`.
81 status[idx.0] = TopologicalOrderStatus::Done;
82
83 // Add the element to the ordering after visiting all its children.
84 // This gives us reverse topological ordering.
85 ordering.push(idx);
86 }
87 TopologicalOrderStatus::Done => {}
88 }
89 }
90
91 Ok(())
92}