smt_scope/analysis/graph/
subgraph.rs

1use std::ops::Deref;
2
3use fxhash::FxHashSet;
4#[cfg(feature = "mem_dbg")]
5use mem_dbg::{MemDbg, MemSize};
6use petgraph::{
7    graph::{DiGraph, EdgeIndex, IndexType, Neighbors, NodeIndex},
8    Directed,
9    Direction::{self, Incoming, Outgoing},
10    Undirected,
11};
12
13use super::{
14    raw::{RawInstGraph, RawIx},
15    RawNodeIndex,
16};
17
18#[cfg_attr(feature = "mem_dbg", derive(MemSize, MemDbg))]
19#[derive(Debug)]
20pub struct Subgraphs {
21    pub subgraphs: TiVec<GraphIdx, Subgraph>,
22    pub singletons: Vec<RawNodeIndex>,
23}
24
25impl Subgraphs {
26    pub fn topo_node_indices(&self) -> impl Iterator<Item = RawNodeIndex> + '_ {
27        self.in_subgraphs().chain(self.singletons())
28    }
29    pub fn in_subgraphs(&self) -> impl Iterator<Item = RawNodeIndex> + '_ {
30        self.subgraphs.iter().flat_map(|s| s.nodes.iter().copied())
31    }
32    pub fn singletons(&self) -> impl Iterator<Item = RawNodeIndex> + '_ {
33        self.singletons.iter().copied()
34    }
35}
36
37impl Deref for Subgraphs {
38    type Target = TiVec<GraphIdx, Subgraph>;
39    fn deref(&self) -> &Self::Target {
40        &self.subgraphs
41    }
42}
43
44impl RawInstGraph {
45    pub fn partition(&mut self) -> Result<Subgraphs> {
46        let mut subgraphs = TiVec::default();
47        let mut discovered = VisitBox {
48            dfs: self.graph.visit_map(),
49        };
50        for node in self.node_indices() {
51            let has_parents = self
52                .graph
53                .neighbors_directed(node.0, Incoming)
54                .next()
55                .is_some();
56            if has_parents {
57                continue;
58            }
59            let has_children = self
60                .graph
61                .neighbors_directed(node.0, Outgoing)
62                .next()
63                .is_some();
64            if !has_children {
65                continue;
66            }
67            // We've skipped all singleton nodes and all non-root nodes.
68
69            if self.graph[node.0].subgraph.is_some() {
70                continue;
71            }
72
73            // Construct subgraph
74            let idx = subgraphs.next_key();
75            let (subgraph, discovered_) =
76                Subgraph::new(node, &mut self.graph, discovered, |node, i| {
77                    node.subgraph = Some((idx, i))
78                })?;
79            discovered = discovered_;
80            subgraphs.raw.try_reserve(1)?;
81            subgraphs.push(subgraph);
82        }
83
84        // Collect all missing singleton nodes
85        discovered.dfs.toggle_range(..);
86        debug_assert!(discovered.visited().all(|a| {
87            let no_parents = self.graph.edges_directed(a.0, Incoming).next().is_none();
88            let no_children = self.graph.edges_directed(a.0, Outgoing).next().is_none();
89            no_parents && no_children
90        }));
91        let singletons = discovered.visited().collect();
92        Ok(Subgraphs {
93            singletons,
94            subgraphs,
95        })
96    }
97}
98
99#[cfg_attr(feature = "mem_dbg", derive(MemSize, MemDbg))]
100#[derive(Debug)]
101pub struct Subgraph {
102    pub(super) nodes: Vec<RawNodeIndex>,
103}
104
105pub struct VisitBox<D: VisitMap<NodeIndex<RawIx>>> {
106    pub dfs: D,
107}
108
109impl VisitBox<<petgraph::graph::DiGraph<(), ()> as petgraph::visit::Visitable>::Map> {
110    pub fn visited(&self) -> impl Iterator<Item = RawNodeIndex> + '_ {
111        self.dfs.ones().map(NodeIndex::new).map(RawNodeIndex)
112    }
113}
114
115impl Subgraph {
116    pub fn new<N, E, D: VisitMap<NodeIndex<RawIx>>>(
117        node: RawNodeIndex,
118        graph: &mut DiGraph<N, E, RawIx>,
119        mut visit: VisitBox<D>,
120        mut f: impl FnMut(&mut N, u32),
121    ) -> Result<(Self, VisitBox<D>)> {
122        let mut start_nodes = Vec::new();
123
124        let mut un_graph = std::mem::replace(graph, DiGraph::<N, E, RawIx>::with_capacity(0, 0))
125            .into_edge_type::<Undirected>();
126        let mut dfs: Dfs<NodeIndex<RawIx>, _> =
127            petgraph::visit::Dfs::from_parts(Vec::new(), visit.dfs);
128        dfs.move_to(node.0);
129        while let Some(node) = dfs.next(&un_graph) {
130            let di_graph = un_graph.into_edge_type::<Directed>();
131            let has_parents = di_graph.neighbors_directed(node, Incoming).next().is_some();
132            un_graph = di_graph.into_edge_type();
133            if !has_parents {
134                start_nodes.push(node);
135            }
136        }
137        visit.dfs = dfs.discovered;
138        *graph = un_graph.into_edge_type();
139
140        // OPTIMISATION: use a `VisitMap` from `VisitBox` to avoid allocating a
141        // `FxHashSet` here (as well as the need for `SubgraphStartNodes`).
142        let mut topo = petgraph::visit::Topo::new(&SubgraphStartNodes {
143            start_nodes: &start_nodes,
144            graph,
145        });
146        let mut nodes = Vec::new();
147        let mut count = 0_u32;
148        while let Some(node) = topo.next(&SubgraphStartNodes {
149            start_nodes: &start_nodes,
150            graph,
151        }) {
152            let n = &mut graph[node];
153            f(n, count);
154            count += 1;
155            nodes.try_reserve(1)?;
156            nodes.push(RawNodeIndex(node));
157        }
158
159        Ok((Self { nodes }, visit))
160    }
161}
162
163// Graph wrapper for Topo walk
164
165pub(super) struct SubgraphStartNodes<'g, N, E, Ix: IndexType> {
166    pub(super) start_nodes: &'g Vec<NodeIndex<Ix>>,
167    pub(super) graph: &'g DiGraph<N, E, Ix>,
168}
169use petgraph::visit::*;
170
171use crate::{items::GraphIdx, Result, TiVec};
172
173impl<N, E, Ix: IndexType> GraphBase for SubgraphStartNodes<'_, N, E, Ix> {
174    type NodeId = NodeIndex<Ix>;
175    type EdgeId = EdgeIndex;
176}
177impl<'a, N, E, Ix: IndexType> IntoNodeIdentifiers for &'a SubgraphStartNodes<'_, N, E, Ix> {
178    type NodeIdentifiers = std::iter::Copied<std::slice::Iter<'a, NodeIndex<Ix>>>;
179    fn node_identifiers(self) -> Self::NodeIdentifiers {
180        self.start_nodes.iter().copied()
181    }
182}
183impl<'a, N, E, Ix: IndexType> IntoNeighbors for &'a SubgraphStartNodes<'_, N, E, Ix> {
184    type Neighbors = Neighbors<'a, E, Ix>;
185    fn neighbors(self, n: Self::NodeId) -> Self::Neighbors {
186        self.graph.neighbors(n)
187    }
188}
189impl<'a, N, E, Ix: IndexType> IntoNeighborsDirected for &'a SubgraphStartNodes<'_, N, E, Ix> {
190    type NeighborsDirected = Neighbors<'a, E, Ix>;
191    fn neighbors_directed(self, n: Self::NodeId, d: Direction) -> Self::NeighborsDirected {
192        self.graph.neighbors_directed(n, d)
193    }
194}
195impl<'a, N, E, Ix: IndexType> Visitable for &'a SubgraphStartNodes<'_, N, E, Ix> {
196    type Map = FxHashSet<NodeIndex<Ix>>;
197    fn visit_map(&self) -> Self::Map {
198        FxHashSet::default()
199    }
200    fn reset_map(&self, map: &mut Self::Map) {
201        map.clear()
202    }
203}