charon_lib/transform/add_missing_info/
sccs.rs

1use std::fmt::Debug;
2use std::mem;
3use std::vec::Vec;
4
5use itertools::Itertools;
6use petgraph::algo::scc::tarjan_scc;
7use petgraph::graphmap::NodeTrait;
8use petgraph::prelude::DiGraphMap;
9use petgraph::visit::{DfsPostOrder, Walker};
10
11use crate::ast::*;
12
13/// Compute the SCCs (Strongly Connected Components) of a set of identifiers, where the order of
14/// the SCCs and the order of the identifiers inside the SCCs attempt to respect as much as
15/// possible the sort order given. The provided `sort_by` should be cheap.
16///
17/// This is used to generate the translated definitions in a consistent and stable manner. For
18/// instance, let's say we extract 4 definitions  `f`, `g1`, `g2` and `h`, where:
19/// - `g1` and `g2` are mutually recursive
20/// - `h` calls `g1`
21///
22/// When translating those functions, we group together the mutually recursive
23/// functions, because they have to be extracted in one single group, and thus
24/// apply Tarjan's algorithm on the call graph to find out which functions are
25/// mutually recursive.
26/// The implementation of Tarjan's algorithm we use gives us the Strongly Connected
27/// SCCs of the call graph in an arbitrary order, so we can have:
28/// `[[f], [g1, g2], [h]]`, but also `[[h], [f], [g2, g1]]`, etc.
29///
30/// If the user defined those functions in the order: `f, g1, h, g2`, we want
31/// to reorder them into: `f, g1, g2, h`, so that we can translate the mutually
32/// recursive functions together, while performing a minimal amount of reordering.
33/// And if reordering is not needed, because the user defined those functions
34/// in the order `f, g1, g2, h`, or `g1, g2, f, h` or ... then we want to translate
35/// them in this exact order.
36///
37/// The order in which Tarjan's algorithm generates the SCCs is arbitrary, while we want to keep as
38/// much as possible the original order (the order in which the user generated the ids). For this,
39/// we iterate through the ordered ids and try to add the SCC containing the id to a new list of
40/// SCCs Of course, this SCC might have dependencies, so we need to add the dependencies first (in
41/// which case we have reordering to do). This is what this function does: we add an SCC and its
42/// dependencies to the list of reordered SCCs by doing a depth-first search.
43pub fn ordered_scc<Id: NodeTrait + Debug, O: Ord>(
44    graph: &DiGraphMap<Id, ()>,
45    sort_by: impl Fn(&Id) -> O,
46) -> Vec<Vec<Id>> {
47    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
48    struct SccId(usize);
49
50    let mut sccs = tarjan_scc(graph);
51
52    // Reorder the identifiers inside the SCCs.
53    for scc in sccs.iter_mut() {
54        scc.sort_by_key(&sort_by);
55    }
56
57    // Map the nodes to their SCC index. The map is sorted.
58    let id_to_scc: SeqHashMap<Id, SccId> = sccs
59        .iter()
60        .enumerate()
61        .map(|(i, nodes)| (SccId(i), nodes))
62        .flat_map(|(scc_id, scc_nodes)| scc_nodes.iter().map(move |node| (*node, scc_id)))
63        .sorted_by_key(|(node, _)| sort_by(node))
64        .collect();
65
66    // Compute the graph of the sccs, where there is an edge between sccs if there's an edge
67    // between some nodes of each scc.
68    let mut scc_graph: DiGraphMap<SccId, ()> = DiGraphMap::new();
69    // Add sccs in the order of their earliest element.
70    for (_, &scc_id) in &id_to_scc {
71        scc_graph.add_node(scc_id);
72    }
73    for (scc_id, scc_nodes) in sccs.iter().enumerate().map(|(i, nodes)| (SccId(i), nodes)) {
74        // Add the scc neighbors in the desired node order.
75        for neighbor_scc_id in scc_nodes
76            .iter()
77            .flat_map(|node| graph.neighbors(*node))
78            .sorted_by_key(&sort_by)
79            .map(|neighbor| *id_to_scc.get(&neighbor).unwrap())
80        {
81            scc_graph.add_edge(scc_id, neighbor_scc_id, ());
82        }
83    }
84
85    // Reorder the SCCs among themselves by a post-order visit of the graph: for each scc (in the
86    // order of their earliest node in the desired order), we list it and all the sccs it depends
87    // on we haven't yet included, in postorder. This guarantees that the final order is a
88    // topological sort for the scc DAG. Moreover this will explore other sccs in order too because
89    // we sorted the list of neighbors.
90    let mut reordered_sccs_ids: SeqHashSet<SccId> = SeqHashSet::new();
91    for scc_id in scc_graph.nodes() {
92        for scc_id in DfsPostOrder::new(&scc_graph, scc_id).iter(&scc_graph) {
93            reordered_sccs_ids.insert(scc_id);
94        }
95    }
96
97    // Output the fully reordered SCCs.
98    reordered_sccs_ids
99        .into_iter()
100        .map(|scc_id| mem::take(&mut sccs[scc_id.0]))
101        .collect()
102}
103
104#[cfg(test)]
105mod tests {
106    use petgraph::prelude::DiGraphMap;
107
108    #[test]
109    fn test_reorder_sccs1() {
110        let mut graph = DiGraphMap::new();
111
112        // First nontrivial cc
113        graph.add_edge(1, 2, ());
114        graph.add_edge(2, 1, ());
115        // Second nontrivial cc
116        graph.add_edge(3, 4, ());
117        graph.add_edge(4, 5, ());
118        graph.add_edge(5, 3, ());
119        // Other deps
120        graph.add_edge(1, 0, ());
121        graph.add_edge(0, 3, ());
122        graph.add_edge(1, 3, ());
123
124        let reordered = super::ordered_scc(&graph, |i| *i);
125
126        assert_eq!(reordered, vec![vec![3, 4, 5], vec![0], vec![1, 2],]);
127    }
128}
charon_lib/transform/add_missing_info/sccs.rs

charon_lib/transform/add_missing_info/
sccs.rs
