1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use llvm_ir::{Function, Name, Terminator};
use petgraph::prelude::{DiGraphMap, Direction};
use std::fmt;

/// The control flow graph for a particular function.
///
/// To construct a `ControlFlowGraph`, use
/// [`FunctionAnalysis`](struct.FunctionAnalysis.html), which you can get
/// from [`ModuleAnalysis`](struct.ModuleAnalysis.html).
pub struct ControlFlowGraph<'m> {
    /// The graph itself. Nodes are basic block names, and an edge from bbX to
    /// bbY indicates that control may (immediately) flow from bbX to bbY
    ///
    /// Or, an edge from bbX to `Return` indicates that the function may return
    /// from bbX
    pub(crate) graph: DiGraphMap<CFGNode<'m>, ()>,

    /// Entry node for the function
    pub(crate) entry_node: CFGNode<'m>,
}

/// A CFGNode represents a basic block, or the special node `Return`
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
pub enum CFGNode<'m> {
    /// The block with the given `Name`
    Block(&'m Name),
    /// The special `Return` node indicating function return
    Return,
}

impl<'m> fmt::Display for CFGNode<'m> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            CFGNode::Block(block) => write!(f, "{}", block),
            CFGNode::Return => write!(f, "Return"),
        }
    }
}

impl<'m> ControlFlowGraph<'m> {
    pub(crate) fn new(function: &'m Function) -> Self {
        let mut graph: DiGraphMap<CFGNode<'m>, ()> = DiGraphMap::with_capacity(
            function.basic_blocks.len() + 1,
            2 * function.basic_blocks.len(), // arbitrary guess
        );

        for bb in &function.basic_blocks {
            match &bb.term {
                Terminator::Br(br) => {
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&br.dest), ());
                },
                Terminator::CondBr(condbr) => {
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&condbr.true_dest), ());
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&condbr.false_dest), ());
                },
                Terminator::IndirectBr(ibr) => {
                    for dest in &ibr.possible_dests {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
                    }
                },
                Terminator::Switch(switch) => {
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&switch.default_dest), ());
                    for (_, dest) in &switch.dests {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
                    }
                },
                Terminator::Ret(_) | Terminator::Resume(_) => {
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
                }
                Terminator::Invoke(invoke) => {
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&invoke.return_label), ());
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&invoke.exception_label), ());
                },
                Terminator::CleanupRet(cleanupret) => {
                    if let Some(dest) = &cleanupret.unwind_dest {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
                    } else {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
                    }
                },
                Terminator::CatchRet(catchret) => {
                    // Despite its name, my reading of the LLVM 10 LangRef indicates that CatchRet cannot directly return from the function
                    graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(&catchret.successor), ());
                },
                Terminator::CatchSwitch(catchswitch) => {
                    if let Some(dest) = &catchswitch.default_unwind_dest {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(dest), ());
                    } else {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Return, ());
                    }
                    for handler in &catchswitch.catch_handlers {
                        graph.add_edge(CFGNode::Block(&bb.name), CFGNode::Block(handler), ());
                    }
                },
                #[cfg(not(feature = "llvm-8"))]
                Terminator::CallBr(_) => unimplemented!("CallBr instruction"),
                Terminator::Unreachable(_) => {
                    // no successors
                },
            }
        }

        Self {
            graph,
            entry_node: CFGNode::Block(&function.basic_blocks[0].name),
        }
    }

    /// Get the predecessors of the basic block with the given `Name`
    pub fn preds<'s>(&'s self, block: &'m Name) -> impl Iterator<Item = &'m Name> + 's {
        self.preds_of_cfgnode(CFGNode::Block(block))
    }

    /// Get the predecessors of the special `Return` node, i.e., get all blocks
    /// which may directly return
    pub fn preds_of_return<'s>(&'s self) -> impl Iterator<Item = &'m Name> + 's {
        self.preds_of_cfgnode(CFGNode::Return)
    }

    pub(crate) fn preds_of_cfgnode<'s>(&'s self, node: CFGNode<'m>) -> impl Iterator<Item = &'m Name> + 's {
        self.preds_as_nodes(node).map(|cfgnode| match cfgnode {
            CFGNode::Block(block) => block,
            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
        })
    }

    pub(crate) fn preds_as_nodes<'s>(&'s self, node: CFGNode<'m>) -> impl Iterator<Item = CFGNode<'m>> + 's {
        self.graph.neighbors_directed(node, Direction::Incoming)
    }

    /// Get the successors of the basic block with the given `Name`.
    /// Here, `CFGNode::Return` indicates that the function may directly return
    /// from this basic block.
    pub fn succs<'s>(&'s self, block: &'m Name) -> impl Iterator<Item = CFGNode<'m>> + 's {
        self.graph.neighbors_directed(CFGNode::Block(block), Direction::Outgoing)
    }

    /// Get the `Name` of the entry block for the function
    pub fn entry(&self) -> &'m Name {
        match self.entry_node {
            CFGNode::Block(block) => block,
            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
        }
    }

    /// Get the reversed CFG; i.e., the CFG where all edges have been reversed
    pub(crate) fn reversed(&self) -> Self {
        Self {
            graph: DiGraphMap::from_edges(
                self.graph.all_edges().map(|(a, b, _)| (b, a, ()))
            ),
            entry_node: CFGNode::Return,
        }
    }
}