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}