llvm_ir_analysis/
control_flow_graph.rs

1use llvm_ir::{Function, Name, Terminator};
2use petgraph::prelude::{DiGraphMap, Direction};
3use std::fmt;
4
5/// The control flow graph for a particular function.
6///
7/// To construct a `ControlFlowGraph`, use
8/// [`FunctionAnalysis`](struct.FunctionAnalysis.html), which you can get
9/// from [`ModuleAnalysis`](struct.ModuleAnalysis.html).
10pub struct ControlFlowGraph<'m> {
11    /// The graph itself. Nodes are basic block names, and an edge from bbX to
12    /// bbY indicates that control may (immediately) flow from bbX to bbY
13    ///
14    /// Or, an edge from bbX to `Return` indicates that the function may return
15    /// from bbX
16    pub(crate) graph: DiGraphMap<CFGNode<'m>, ()>,
17
18    /// Entry node for the function
19    pub(crate) entry_node: CFGNode<'m>,
20}
21
22/// A CFGNode represents a basic block, or the special node `Return`
23#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
24pub enum CFGNode<'m> {
25    /// The block with the given `Name`
26    Block(&'m Name),
27    /// The special `Return` node indicating function return
28    Return,
29}
30
31impl<'m> fmt::Display for CFGNode<'m> {
32    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
33        match self {
34            CFGNode::Block(block) => write!(f, "{}", block),
35            CFGNode::Return => write!(f, "Return"),
36        }
37    }
38}
39
40impl<'m> ControlFlowGraph<'m> {
41    pub(crate) fn new(function: &'m Function) -> Self {
42        let mut graph: DiGraphMap<CFGNode<'m>, ()> = DiGraphMap::with_capacity(
43            function.basic_blocks.len() + 1,
44            2 * function.basic_blocks.len(), // arbitrary guess
45        );
46
47        for bb in &function.basic_blocks {
48            match &bb.term {
49                Terminator::Br(br) => {
50                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&br.dest), ());
51                }
52                Terminator::CondBr(condbr) => {
53                    graph.add_edge(
54                        CFGNode::Block(&bb.name),
55                        CFGNode::Block(&condbr.true_dest),
56                        (),
57                    );
58                    graph.add_edge(
59                        CFGNode::Block(&bb.name),
60                        CFGNode::Block(&condbr.false_dest),
61                        (),
62                    );
63                }
64                Terminator::IndirectBr(ibr) => {
65                    for dest in &ibr.possible_dests {
66                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
67                    }
68                }
69                Terminator::Switch(switch) => {
70                    graph.add_edge(
71                        CFGNode::Block(&bb.name),
72                        CFGNode::Block(&switch.default_dest),
73                        (),
74                    );
75                    for (_, dest) in &switch.dests {
76                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
77                    }
78                }
79                Terminator::Ret(_) | Terminator::Resume(_) => {
80                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
81                }
82                Terminator::Invoke(invoke) => {
83                    graph.add_edge(
84                        CFGNode::Block(&bb.name),
85                        CFGNode::Block(&invoke.return_label),
86                        (),
87                    );
88                    graph.add_edge(
89                        CFGNode::Block(&bb.name),
90                        CFGNode::Block(&invoke.exception_label),
91                        (),
92                    );
93                }
94                Terminator::CleanupRet(cleanupret) => {
95                    if let Some(dest) = &cleanupret.unwind_dest {
96                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
97                    } else {
98                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
99                    }
100                }
101                Terminator::CatchRet(catchret) => {
102                    // Despite its name, my reading of the LLVM 10 LangRef indicates that CatchRet cannot directly return from the function
103                    graph.add_edge(
104                        CFGNode::Block(&bb.name),
105                        CFGNode::Block(&catchret.successor),
106                        (),
107                    );
108                }
109                Terminator::CatchSwitch(catchswitch) => {
110                    if let Some(dest) = &catchswitch.default_unwind_dest {
111                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
112                    } else {
113                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
114                    }
115                    for handler in &catchswitch.catch_handlers {
116                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(handler), ());
117                    }
118                }
119                Terminator::CallBr(_) => unimplemented!("CallBr instruction"),
120                Terminator::Unreachable(_) => {
121                    // no successors
122                }
123            }
124        }
125
126        Self {
127            graph,
128            entry_node: CFGNode::Block(&function.basic_blocks[0].name),
129        }
130    }
131
132    /// Get the predecessors of the basic block with the given `Name`
133    pub fn preds<'s>(&'s self, block: &'m Name) -> impl Iterator<Item = &'m Name> + 's {
134        self.preds_of_cfgnode(CFGNode::Block(block))
135    }
136
137    /// Get the predecessors of the special `Return` node, i.e., get all blocks
138    /// which may directly return
139    pub fn preds_of_return<'s>(&'s self) -> impl Iterator<Item = &'m Name> + 's {
140        self.preds_of_cfgnode(CFGNode::Return)
141    }
142
143    pub(crate) fn preds_of_cfgnode<'s>(
144        &'s self,
145        node: CFGNode<'m>,
146    ) -> impl Iterator<Item = &'m Name> + 's {
147        self.preds_as_nodes(node).map(|cfgnode| match cfgnode {
148            CFGNode::Block(block) => block,
149            CFGNode::Return => panic!("Shouldn't have CFGNode::Return as a predecessor"), // perhaps you tried to call this on a reversed CFG? In-crate users can use `preds_as_nodes()` if they need to account for the possibility of a reversed CFG
150        })
151    }
152
153    pub(crate) fn preds_as_nodes<'s>(
154        &'s self,
155        node: CFGNode<'m>,
156    ) -> impl Iterator<Item = CFGNode<'m>> + 's {
157        self.graph.neighbors_directed(node, Direction::Incoming)
158    }
159
160    /// Get the successors of the basic block with the given `Name`.
161    /// Here, `CFGNode::Return` indicates that the function may directly return
162    /// from this basic block.
163    pub fn succs<'s>(&'s self, block: &'m Name) -> impl Iterator<Item = CFGNode<'m>> + 's {
164        self.graph
165            .neighbors_directed(CFGNode::Block(block), Direction::Outgoing)
166    }
167
168    /// Get the `Name` of the entry block for the function
169    pub fn entry(&self) -> &'m Name {
170        match self.entry_node {
171            CFGNode::Block(block) => block,
172            CFGNode::Return => panic!("Return node should not be entry"), // perhaps you tried to call this on a reversed CFG? In-crate users can use the `entry_node` field directly if they need to account for the possibility of a reversed CFG
173        }
174    }
175
176    /// Get the reversed CFG; i.e., the CFG where all edges have been reversed
177    pub(crate) fn reversed(&self) -> Self {
178        Self {
179            graph: DiGraphMap::from_edges(self.graph.all_edges().map(|(a, b, _)| (b, a, ()))),
180            entry_node: CFGNode::Return,
181        }
182    }
183}
llvm_ir_analysis/control_flow_graph.rs

llvm_ir_analysis/
control_flow_graph.rs
