BinaryControlFlow.h

#ifndef ROSE_BinaryAnalysis_ControlFlow_H
#define ROSE_BinaryAnalysis_ControlFlow_H

#include <featureTests.h>
#ifdef ROSE_ENABLE_BINARY_ANALYSIS

#include "Map.h"
#include "WorkLists.h"
#include "SageBuilderAsm.h"

#include <boost/foreach.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/reverse_graph.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <Sawyer/GraphBoost.h>

class SgNode;
class SgAsmBlock;

namespace Rose {

namespace BinaryAnalysis {

class ControlFlow {
public:
    ControlFlow()
        : vertex_filter(NULL), edge_filter(NULL)
        {}


    typedef boost::adjacency_list<boost::setS,                                  /* edges of each vertex in std::list */
                                  boost::vecS,                                  /* vertices in std::vector */
                                  boost::bidirectionalS,
                                  boost::property<boost::vertex_name_t, SgAsmBlock*> > BlockGraph;

    typedef boost::adjacency_list<boost::setS,
                                  boost::vecS,
                                  boost::bidirectionalS,
                                  boost::property<boost::vertex_name_t, SgAsmInstruction*> > InsnGraph;

    typedef BlockGraph Graph;


    /**********************************************************************************************************************
     *                                      Filters
     **********************************************************************************************************************/
public:

    class VertexFilter {
    public:
        virtual ~VertexFilter() {}
        virtual bool operator()(ControlFlow*, SgAsmNode*) = 0;
    };

    class EdgeFilter {
    public:
        virtual ~EdgeFilter() {}
        virtual bool operator()(ControlFlow*, SgAsmNode *source, SgAsmNode *target) = 0;
    };

    void set_vertex_filter(VertexFilter *filter) { vertex_filter = filter; }
    VertexFilter *get_vertex_filter() const { return vertex_filter; }
    void set_edge_filter(EdgeFilter *filter) { edge_filter = filter; }
    EdgeFilter *get_edge_filter() const { return edge_filter; }
    bool is_vertex_filtered(SgAsmNode *bb_or_insn, VertexFilter *filter) { return filter && !(*filter)(this, bb_or_insn); }
    bool is_vertex_filtered(SgAsmNode *bb_or_insn) { return is_vertex_filtered(bb_or_insn, vertex_filter); }
    bool is_edge_filtered(SgAsmNode *src, SgAsmNode *dst, EdgeFilter *filter) {
        return filter && !(*filter)(this, src, dst);
    }
    bool is_edge_filtered(SgAsmNode *src, SgAsmNode *dst) {
        return is_edge_filtered(src, dst, edge_filter);
    }
protected:
    VertexFilter *vertex_filter;
    EdgeFilter *edge_filter;

    /**********************************************************************************************************************
     *                                      Methods that modify the AST
     **********************************************************************************************************************/
public:

    void clear_ast(SgNode *ast);

    template<class ControlFlowGraph>
    void apply_to_ast(const ControlFlowGraph&);

    template<class ControlFlowGraph>
    void cache_vertex_descriptors(const ControlFlowGraph&);

    /**********************************************************************************************************************
     *                                      Graph construction methods
     **********************************************************************************************************************/
public:

    template<class ControlFlowGraph>
    ControlFlowGraph build_block_cfg_from_ast(SgNode *root);

    template<class ControlFlowGraph>
    void build_block_cfg_from_ast(SgNode *root, ControlFlowGraph &cfg/*out*/);

    template<class ControlFlowGraph>
    ControlFlowGraph build_insn_cfg_from_ast(SgNode *root);

    template<class ControlFlowGraph>
    void build_insn_cfg_from_ast(SgNode *root, ControlFlowGraph &cfg/*out*/);
    template<class BlockCFG, class InsnCFG>
    void explode_blocks(const BlockCFG &cfgb, InsnCFG &cfgi/*out*/);
    template<class InsnCFG>
    void fixup_fcall_fret(InsnCFG &cfg/*in,out*/, bool preserve_call_fallthrough_edges);

    template<class ControlFlowGraph>
    ControlFlowGraph build_cg_from_ast(SgNode *root);

    template<class ControlFlowGraph>
    void build_cg_from_ast(SgNode *root, ControlFlowGraph &cfg/*out*/);
    template<class ControlFlowGraph>
    ControlFlowGraph copy(const ControlFlowGraph &src);

    template<class ControlFlowGraph>
    void copy(const ControlFlowGraph &src, ControlFlowGraph &dst/*out*/);
    /***********************************************************************************************************************
     *                                      Graph output
     ***********************************************************************************************************************/

    template<class CFG>
    struct FunctionSubgraphInfo {
        std::vector<typename boost::graph_traits<CFG>::vertex_descriptor> vertices;
        std::vector<typename boost::graph_traits<CFG>::edge_descriptor> edges;
    };

    template<class CFG>
    struct DefaultVertexPropertyWriter {
        void operator()(std::ostream &o, typename boost::graph_traits<CFG>::vertex_descriptor vertex) const {}
    };

    template<class CFG>
    struct DefaultEdgePropertyWriter {
        void operator()(std::ostream &o, typename boost::graph_traits<CFG>::edge_descriptor vertex) const {}
    };

    template<typename CFG, class VertexPropertyWriter, class EdgePropertyWriter>
    void write_graphviz(std::ostream&, const CFG&, const VertexPropertyWriter&, const EdgePropertyWriter&);

    template<typename CFG>
    void write_graphviz(std::ostream &out, const CFG &cfg) {
        write_graphviz(out, cfg, DefaultVertexPropertyWriter<CFG>(), DefaultEdgePropertyWriter<CFG>());
    }

    template<typename CFG, class VertexPropertyWriter>
    void write_graphviz(std::ostream &out, const CFG &cfg, const VertexPropertyWriter &vpw) {
        write_graphviz(out, cfg, vpw, DefaultEdgePropertyWriter<CFG>());
    }
    /**********************************************************************************************************************
     *                                      Miscellaneous members
     **********************************************************************************************************************/

private:
    /* Visitor used by flow_order().  Declaring this in function scope results in boost errors (boost-1.42, 2011-05). */
    template<class ControlFlowGraph>
    struct FlowOrder: public boost::default_dfs_visitor {
        typedef typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor Vertex;
        typedef std::vector<Vertex> VertexList;
        typedef std::vector<size_t> ReverseVertexList;
        VertexList *forward_order;
        FlowOrder(VertexList *forward_order): forward_order(forward_order) {}
        void compute(const ControlFlowGraph &g, Vertex v0, ReverseVertexList *reverse_order);
        void finish_vertex(Vertex v, ControlFlowGraph g);
    };

    /* Helper class for build_block_cfg_from_ast().  Adds vertices to its 'cfg' member. Vertices are any SgAsmBlock that
     * contains at least one SgAsmInstruction. */
    template<class ControlFlowGraph>
    class VertexInserter: public AstSimpleProcessing {
    public:
        ControlFlow *analyzer;
        ControlFlowGraph &cfg;
        typedef typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor Vertex;
        typedef Map<SgAsmBlock*, Vertex> BlockVertexMap;
        BlockVertexMap &bv_map;
        VertexInserter(ControlFlow *analyzer, ControlFlowGraph &cfg, BlockVertexMap &bv_map)
            : analyzer(analyzer), cfg(cfg), bv_map(bv_map)
            {}
        // Add basic block to graph if it hasn't been added already.
        void conditionally_add_vertex(SgAsmBlock *block);

        void visit(SgNode *node) {
            if (isSgAsmFunction(node)) {
                // Add the function entry block before the other blocks of the function.  This ensures that the entry block
                // of a function has a lower vertex number than the other blocks of the function (the traversal is not
                // guaranteed to visit the function basic blocks in that order).
                conditionally_add_vertex(isSgAsmFunction(node)->get_entry_block());
            } else {
                conditionally_add_vertex(isSgAsmBlock(node));
            }
        }
    };

public:
    template<class ControlFlowGraph>
    std::vector<typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor>
    flow_order(const ControlFlowGraph&,
               typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor start,
               std::vector<size_t> *reverse_order=NULL);

private:
    /* Visitor used by return_blocks(). Declaring this in function scope results in boost errors (boost-1.42, 2011-05). */
    template<class ControlFlowGraph>
    struct ReturnBlocks: public boost::default_dfs_visitor {
        typedef typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor Vertex;
        typedef std::vector<Vertex> Vector;
        Vector &blocks;
        ReturnBlocks(Vector &blocks): blocks(blocks) {}
        void finish_vertex(Vertex v, ControlFlowGraph g);
    };

public:
    template<class ControlFlowGraph>
    std::vector<typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor>
    return_blocks(const ControlFlowGraph &cfg,
                  typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor start);
};


/*******************************************************************************************************************************
 *                                      Functions
 *******************************************************************************************************************************/

template<class V, class E>
typename Sawyer::Container::Graph<V, E>::VertexValue
get_ast_node(const Sawyer::Container::Graph<V, E> &cfg, size_t vertexId) {
    typedef typename Sawyer::Container::Graph<V, E> CFG;
    typename CFG::ConstVertexValueIterator iter = cfg.findVertex(vertexId);
    ASSERT_forbid2(iter==cfg.vertices().end(), "invalid vertex ID " + StringUtility::numberToString(vertexId));
    return *iter;
}

template<class V, class E, class AstNode>
void
put_ast_node(Sawyer::Container::Graph<V, E> &cfg, size_t vertexId, AstNode *astNode) {
    typedef typename Sawyer::Container::Graph<V, E> CFG;
    typename CFG::VertexValueIterator iter = cfg.findVertex(vertexId);
    ASSERT_forbid2(iter==cfg.vertices().end(), "invalid vertex ID " + StringUtility::numberToString(vertexId));
    *iter = astNode;
}

// Sorry about this mess!  The goal is to match only boost::adjacency_list graphs.
template<class A, class B, class C, class D, class E, class F, class G>
typename boost::property_traits<typename boost::property_map<boost::adjacency_list<A, B, C, D, E, F, G>,
                                                             boost::vertex_name_t>::type>::value_type
get_ast_node(const boost::adjacency_list<A, B, C, D, E, F, G> &cfg,
             typename boost::graph_traits<boost::adjacency_list<A, B, C, D, E, F, G> >::vertex_descriptor vertex) {
    return boost::get(boost::vertex_name, cfg, vertex);
}

// Sorry about this mess!  The goal is to match only boost::adjacency_list graphs.
template<class A, class B, class C, class D, class E, class F, class G>
void
put_ast_node(boost::adjacency_list<A, B, C, D, E, F, G> &cfg,
             typename boost::graph_traits<boost::adjacency_list<A, B, C, D, E, F, G> >::vertex_descriptor vertex,
             typename boost::property_traits<
                 typename boost::property_map<boost::adjacency_list<A, B, C, D, E, F, G>, boost::vertex_name_t>::type
                 >::value_type ast_node) {
    boost::put(boost::vertex_name, cfg, vertex, ast_node);
}

/******************************************************************************************************************************
 *                                      Function template definitions
 ******************************************************************************************************************************/

template<class ControlFlowGraph>
void
ControlFlow::apply_to_ast(const ControlFlowGraph &cfg)
{
    typename boost::graph_traits<ControlFlowGraph>::vertex_iterator vi, vi_end;
    for (boost::tie(vi, vi_end)=boost::vertices(cfg); vi!=vi_end; ++vi) {
        SgAsmBlock *block = get_ast_node(cfg, *vi); // FIXME: Instruction CFGs not supported yet
        if (!block || is_vertex_filtered(block))
            continue;

        /* Delete old targets */
        const SgAsmIntegerValuePtrList &targets = block->get_successors();
        for (SgAsmIntegerValuePtrList::const_iterator ti=targets.begin(); ti!=targets.end(); ++ti)
            delete *ti;

        /* Add new targets */
        block->set_successors_complete(true);
        block->get_successors().clear();
        typename boost::graph_traits<ControlFlowGraph>::out_edge_iterator ei, ei_end;
        for (boost::tie(ei, ei_end)=boost::out_edges(*vi, cfg); ei!=ei_end; ++ei) {
            SgAsmBlock *target_block = get_ast_node(cfg, boost::target(*ei, cfg));
            if (target_block && !is_edge_filtered(block, target_block)) {
                SgAsmIntegerValueExpression *target = SageBuilderAsm::buildValueU64(target_block->get_address());
                target->makeRelativeTo(target_block);
                target->set_parent(block);
                block->get_successors().push_back(target);
            }
        }
    }
}

template<class ControlFlowGraph>
void
ControlFlow::cache_vertex_descriptors(const ControlFlowGraph &cfg)
{
    typename boost::graph_traits<ControlFlowGraph>::vertex_iterator vi, vi_end;
    for (boost::tie(vi, vi_end)=boost::vertices(cfg); vi!=vi_end; ++vi) {
        SgAsmBlock *block = get_ast_node(cfg, *vi); // FIXME: Instruction CFGs not supported yet
        if (block && !is_vertex_filtered(block))
            block->set_cached_vertex(*vi);
    }
}

template<class ControlFlowGraph>
void
ControlFlow::VertexInserter<ControlFlowGraph>::conditionally_add_vertex(SgAsmBlock *block)
{
    if (block && block->has_instructions() && !analyzer->is_vertex_filtered(block) && !bv_map.exists(block)) {
        Vertex vertex = boost::add_vertex(cfg);
        bv_map[block] = vertex;
        put_ast_node(cfg, vertex, block);
    }
}

template<class ControlFlowGraph>
void
ControlFlow::build_block_cfg_from_ast(SgNode *root, ControlFlowGraph &cfg)
{
    typedef typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor Vertex;
    Vertex NO_VERTEX = boost::graph_traits<ControlFlowGraph>::null_vertex();
    typedef Map<SgAsmBlock*, Vertex> BlockVertexMap;
    BlockVertexMap bv_map;

    // Add the vertices
    cfg.clear();
    VertexInserter<ControlFlowGraph>(this, cfg, bv_map).traverse(root, preorder);

    // Mapping from block entry address to CFG vertex
    Map<rose_addr_t, Vertex> addrToVertex;
    for (typename BlockVertexMap::iterator bvi=bv_map.begin(); bvi!=bv_map.end(); ++bvi)
        addrToVertex[bvi->first->get_address()] = bvi->second;

    // Add the edges
    BOOST_FOREACH (Vertex sourceVertex, boost::vertices(cfg)) {
        SgAsmBlock *sourceBlock = get_ast_node(cfg, sourceVertex);
        BOOST_FOREACH (SgAsmIntegerValueExpression *integerValue, sourceBlock->get_successors()) {
            Vertex targetVertex = addrToVertex.get_value_or(integerValue->get_absoluteValue(), NO_VERTEX);
            if (targetVertex!=NO_VERTEX) {
                SgAsmBlock *targetBlock = get_ast_node(cfg, targetVertex);
                assert(targetBlock!=NULL); // since we have a vertex, there must be an SgAsmBlock!
                if (!is_edge_filtered(sourceBlock, targetBlock))
                    boost::add_edge(sourceVertex, targetVertex, cfg);
            }
        }
    }
}

template<class ControlFlowGraph>
void
ControlFlow::build_insn_cfg_from_ast(SgNode *root, ControlFlowGraph &cfg)
{
    BlockGraph cfgb;
    build_block_cfg_from_ast(root, cfgb);
    explode_blocks(cfgb, cfg);
    bool preserve_call_fallthrough_edges = false;
    fixup_fcall_fret(cfg, preserve_call_fallthrough_edges);
}

template<class ControlFlowGraph>
void
ControlFlow::build_cg_from_ast(SgNode *root, ControlFlowGraph &cfg/*out*/)
{
    struct T1: public EdgeFilter {
        EdgeFilter *parent;
        T1(EdgeFilter *parent): parent(parent) {}
        bool operator()(ControlFlow *analyzer, SgAsmNode *src, SgAsmNode *dst) {
            SgAsmFunction *src_func = SageInterface::getEnclosingNode<SgAsmFunction>(src, true);
            SgAsmBlock *dst_block = SageInterface::getEnclosingNode<SgAsmBlock>(dst, true);
            SgAsmFunction *dst_func = SageInterface::getEnclosingNode<SgAsmFunction>(dst_block);
            if (!src_func || !dst_func || dst_block!=dst_func->get_entry_block()) {
                return false;
            } else if (src_func!=dst_func) {
                // inter-function call, not a return edge
            } else {
                // FIXME: this might not actually be a recursive call [Robb P. Matzke 2013-09-05]
            }
            return parent ? (*parent)(analyzer, src, dst) : true;
        }
    };

    EdgeFilter *parent = get_edge_filter();
    T1 edge_filter(parent);
    try {
        set_edge_filter(&edge_filter);
        build_block_cfg_from_ast(root, cfg);
    } catch (...) {
        set_edge_filter(parent);
        throw;
    }
    set_edge_filter(parent);
}

template<class ControlFlowGraph>
void
ControlFlow::copy(const ControlFlowGraph &src, ControlFlowGraph &dst/*out*/)
{
    typedef typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor Vertex;
    Vertex NO_VERTEX = boost::graph_traits<ControlFlowGraph>::null_vertex();

    dst.clear();
    std::vector<Vertex> src_to_dst(boost::num_vertices(src), NO_VERTEX);

    typename boost::graph_traits<const ControlFlowGraph>::vertex_iterator vi, vi_end;
    for (boost::tie(vi, vi_end)=boost::vertices(src); vi!=vi_end; ++vi) {
        SgAsmNode *node = get_ast_node(src, *vi);
        if (!is_vertex_filtered(node)) {
            src_to_dst[*vi] = boost::add_vertex(dst);
            put_ast_node(dst, src_to_dst[*vi], get_ast_node(src, *vi));
        }
    }

    typename boost::graph_traits<const ControlFlowGraph>::edge_iterator ei, ei_end;
    for (boost::tie(ei, ei_end)=boost::edges(src); ei!=ei_end; ++ei) {
        if (NO_VERTEX!=src_to_dst[boost::source(*ei, src)] && NO_VERTEX!=src_to_dst[boost::target(*ei, src)]) {
            SgAsmNode *node1 = get_ast_node(src, boost::source(*ei, src));
            SgAsmNode *node2 = get_ast_node(src, boost::target(*ei, src));
            if (!is_edge_filtered(node1, node2))
                boost::add_edge(src_to_dst[boost::source(*ei, src)], src_to_dst[boost::target(*ei, src)], dst);
        }
    }
}

template<class ControlFlowGraph>
ControlFlowGraph
ControlFlow::copy(const ControlFlowGraph &src)
{
    ControlFlowGraph dst;
    copy(src, dst);
    return dst;
}

template<class BlockCFG, class InsnCFG>
void
ControlFlow::explode_blocks(const BlockCFG &cfgb, InsnCFG &cfgi/*out*/)
{
    // BlockCFG is the basic-block binary control flow graph
    typedef typename boost::graph_traits<const BlockCFG>::vertex_descriptor BlockCFG_Vertex;
    typedef typename boost::graph_traits<const BlockCFG>::vertex_iterator BlockCFG_VertexIterator;
    typedef typename boost::graph_traits<const BlockCFG>::edge_iterator BlockCFG_EdgeIterator;

    // InsnCFG is the instruction binary control flow graph--it points to instructions rather than basic blocks, and changes
    // some edges regarding function calls.
    typedef typename boost::graph_traits<InsnCFG>::vertex_descriptor InsnCFG_Vertex;
    typedef std::pair<InsnCFG_Vertex, InsnCFG_Vertex> InsnCFG_VertexPair;

    // Expand the cfgb basic blocks to create a cfgi that has instructions instead of blocks, and add the intra-block edges
    cfgi.clear();
    Map<BlockCFG_Vertex, InsnCFG_VertexPair> vertex_translation; // enter and leave instructions for each of the blocks in cfgb
    {
        BlockCFG_VertexIterator vi, vi_end;
        for (boost::tie(vi, vi_end)=boost::vertices(cfgb); vi!=vi_end; ++vi) {
            SgAsmBlock *blk = get_ast_node(cfgb, *vi);
            const SgAsmStatementPtrList &insns = blk->get_statementList();
            assert(!insns.empty());
            InsnCFG_Vertex enter_vertex = boost::graph_traits<InsnCFG>::null_vertex();
            InsnCFG_Vertex prev_vertex = boost::graph_traits<InsnCFG>::null_vertex();
            for (SgAsmStatementPtrList::const_iterator ii=insns.begin(); ii!=insns.end(); ++ii) {
                SgAsmInstruction *insn = isSgAsmInstruction(*ii);
                assert(insn!=NULL); // basic blocks contain only instructions, no other type of asm statement
                InsnCFG_Vertex vertex = boost::add_vertex(cfgi);
                put_ast_node(cfgi, vertex, insn);
                if (ii==insns.begin()) {
                    enter_vertex = vertex;
                } else {
                    boost::add_edge(prev_vertex, vertex, cfgi);
                }
                prev_vertex = vertex;
            }
            assert(prev_vertex!=boost::graph_traits<InsnCFG>::null_vertex()); // basic block had no instructions but was in CFG!
            vertex_translation[*vi] = InsnCFG_VertexPair(enter_vertex, prev_vertex);
        }
    }
    
    // Insert the edges from cfgb.  The corresponding edge in cfgi must emanate from the final instruction of the source basic
    // block and enter at the first instruction of the target basic block.
    {
        BlockCFG_EdgeIterator ei, ei_end;
        for (boost::tie(ei, ei_end)=boost::edges(cfgb); ei!=ei_end; ++ei) {
            InsnCFG_Vertex src_leave_vertex = vertex_translation.get_one(boost::source(*ei, cfgb)).second;
            InsnCFG_Vertex dst_enter_vertex = vertex_translation.get_one(boost::target(*ei, cfgb)).first;
            assert(src_leave_vertex!=boost::graph_traits<InsnCFG>::null_vertex());
            assert(dst_enter_vertex!=boost::graph_traits<InsnCFG>::null_vertex());
            boost::add_edge(src_leave_vertex, dst_enter_vertex, cfgi);
        }
    }
}

template<class InsnCFG>
void
ControlFlow::fixup_fcall_fret(InsnCFG &cfg, bool preserve_call_fallthrough_edges)
{
    typedef typename boost::graph_traits<InsnCFG>::vertex_descriptor CFG_Vertex;
    typedef typename boost::graph_traits<InsnCFG>::vertex_iterator CFG_VertexIterator;
    typedef typename boost::graph_traits<InsnCFG>::in_edge_iterator CFG_InEdgeIterator;
    typedef std::pair<CFG_Vertex, CFG_Vertex> CFG_VertexPair;
    typedef Map<SgAsmInstruction*, CFG_Vertex> InsnToVertex;
    CFG_Vertex NO_VERTEX = boost::graph_traits<InsnCFG>::null_vertex();

    // Build mappings needed later and find the function return points.  We just look for the x86
    // RET instruction for now and assume that each one we find is a return if it has no control flow successors.  They have no
    // successors at this point because CFG1 didn't have any.
    InstructionMap insns;
    InsnToVertex insn_to_vertex;
    std::vector<bool> isret(boost::num_vertices(cfg), false);
    {
        CFG_VertexIterator vi, vi_end;
        for (boost::tie(vi, vi_end)=boost::vertices(cfg); vi!=vi_end; ++vi) {
            SgAsmInstruction *insn = get_ast_node(cfg, *vi);
            insns[insn->get_address()] = insn;
            insn_to_vertex[insn] = *vi;

            if (0==boost::out_degree(*vi, cfg)) {
                // FIXME: Architecture-specific code here
                if (SgAsmX86Instruction *insn_x86 = isSgAsmX86Instruction(insn)) {
                    isret[*vi] = x86_ret==insn_x86->get_kind();
                }
            }
        }
    }

    // Return the entry vertex for a function that owns the indicated instruction
    struct FunctionEntryVertex {
        const InsnToVertex &insn_to_vertex;
        const InstructionMap &imap;
        FunctionEntryVertex(const InsnToVertex &insn_to_vertex, const InstructionMap &imap)
            : insn_to_vertex(insn_to_vertex), imap(imap) {}
        CFG_Vertex operator()(SgAsmInstruction *insn) {
            SgAsmFunction *func = SageInterface::getEnclosingNode<SgAsmFunction>(insn, true);
            SgAsmInstruction *entry_insn = imap.get_one(func->get_entry_va());
            CFG_Vertex entry_vertex = insn_to_vertex.get_one(entry_insn);
            return entry_vertex;
        }
    } function_entry_vertex(insn_to_vertex, insns);
    
    // Process each return site in order to add edges from the return site to the vertex representing the return address
    std::vector<CFG_VertexPair> edges_to_insert, edges_to_erase;
    {
        CFG_VertexIterator vi, vi_end;
        for (boost::tie(vi, vi_end)=boost::vertices(cfg); vi!=vi_end; ++vi) {
            CFG_Vertex returner_vertex = *vi;
            if (!isret[returner_vertex])
                continue;
            SgAsmInstruction *returner_insn = get_ast_node(cfg, returner_vertex);

            // Find all of the true call sites for the function that owns the returner instruction (e.g., RET) by recursively
            // following inter-function CFG edges until we find the true calls (those edges that follow CALL semantics).
            // Inter-function CFG edges can represent true calls or simply inter-function branches such as thunks.  We have to
            // gather up the information without adding it to the CFG yet (can't add while we're iterating)
            std::vector<bool> seen(boost::num_vertices(cfg), false);
            WorkList<CFG_Vertex> worklist; // targets of inter-function CFG edges; function callees
            worklist.push(function_entry_vertex(returner_insn));
            while (!worklist.empty()) {
                CFG_Vertex callee_vertex = worklist.shift();
                CFG_InEdgeIterator ei, ei_end;
                for (boost::tie(ei, ei_end)=boost::in_edges(callee_vertex, cfg); ei!=ei_end; ++ei) {
                    CFG_Vertex caller_vertex = boost::source(*ei, cfg); // caller is a inter-function call or branch site
                    if (!seen[caller_vertex]) {
                        seen[caller_vertex] = true;
                        SgAsmInstruction *caller_insn = get_ast_node(cfg, caller_vertex);
                        SgAsmBlock *caller_block = SageInterface::getEnclosingNode<SgAsmBlock>(caller_insn);
                        assert(caller_block!=NULL);
                        rose_addr_t target_va, returnee_va; // returnee_va is usually the call's fall-through address
                        if (caller_block->is_function_call(target_va/*out*/, returnee_va/*out*/)) {
                            // This is a true call, so we need to add a return edge from the return instruction (the
                            // "returner") to what is probably the fall-through address of the call site (the returnee).
                            SgAsmInstruction *returnee_insn = insns.get_value_or(returnee_va, NULL);
                            CFG_Vertex returnee_vertex = insn_to_vertex.get_value_or(returnee_insn, NO_VERTEX);
                            if (returnee_vertex!=NO_VERTEX) {
                                edges_to_insert.push_back(CFG_VertexPair(returner_vertex, returnee_vertex));
                                edges_to_erase.push_back(CFG_VertexPair(caller_vertex, returnee_vertex));
                            }
                        } else {
                            // This is a non-call inter-function edge; probably a thunk. We need to find its call sites and add
                            // the returnee addresses (call fall throughs) to the returnee addresses of the RET we're
                            // processing.
                            worklist.push(function_entry_vertex(caller_insn));
                        }
                    }
                }
            }
        }
    }
    
    // Erase and insert edges now that we're done iterating.
    if (!preserve_call_fallthrough_edges) {
        for (size_t i=0; i<edges_to_erase.size(); ++i)
            boost::remove_edge(edges_to_erase[i].first, edges_to_erase[i].second, cfg);
    }
    for (size_t i=0; i<edges_to_insert.size(); ++i)
        boost::add_edge(edges_to_insert[i].first, edges_to_insert[i].second, cfg);
}

template<class ControlFlowGraph>
void
ControlFlow::FlowOrder<ControlFlowGraph>::compute(const ControlFlowGraph &g, Vertex v0,
                                                  ReverseVertexList *reverse_order) {
    forward_order->clear();
    std::vector<boost::default_color_type> colors(boost::num_vertices(g), boost::white_color);
    boost::depth_first_visit(g, v0, *this, &(colors[0]));
    assert(!forward_order->empty()); /* it should at least contain v0 */
    std::reverse(forward_order->begin(), forward_order->end());
    if (reverse_order) {
        reverse_order->clear();
        reverse_order->resize(boost::num_vertices(g), INVALID_INDEX);
        for (size_t i=0; i<forward_order->size(); i++)
            (*reverse_order)[(*forward_order)[i]] = i;
    }
}

template<class ControlFlowGraph>
void
ControlFlow::FlowOrder<ControlFlowGraph>::finish_vertex(Vertex v, ControlFlowGraph g) {
    forward_order->push_back(v);
}

template<class ControlFlowGraph>
std::vector<typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor>
ControlFlow::flow_order(const ControlFlowGraph &cfg,
                        typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor start,
                        std::vector<size_t> *reverse_order/*=NULL*/)
{
    std::vector<typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor> forward_order;
    FlowOrder<ControlFlowGraph>(&forward_order).compute(cfg, start, reverse_order);
    return forward_order;
}

template<class ControlFlowGraph>
void
ControlFlow::ReturnBlocks<ControlFlowGraph>::finish_vertex(Vertex v, ControlFlowGraph g)
{
    typename boost::graph_traits<ControlFlowGraph>::out_edge_iterator ei, ei_end;
    boost::tie(ei, ei_end) = boost::out_edges(v, g);
    if (ei==ei_end)
        blocks.push_back(v);
}

template<class ControlFlowGraph>
std::vector<typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor>
ControlFlow::return_blocks(const ControlFlowGraph &cfg,
                           typename boost::graph_traits<ControlFlowGraph>::vertex_descriptor start)
{
    typename ReturnBlocks<ControlFlowGraph>::Vector result;
    ReturnBlocks<ControlFlowGraph> visitor(result);
    std::vector<boost::default_color_type> colors(boost::num_vertices(cfg), boost::white_color);
    boost::depth_first_visit(cfg, start, visitor, &(colors[0]));
    return result;
}

template<class ControlFlowGraph>
ControlFlowGraph
ControlFlow::build_block_cfg_from_ast(SgNode *root)
{
    ControlFlowGraph cfg;
    build_block_cfg_from_ast(root, cfg);
    return cfg;
}

template<class ControlFlowGraph>
ControlFlowGraph
ControlFlow::build_insn_cfg_from_ast(SgNode *root)
{
    ControlFlowGraph cfg;
    build_insn_cfg_from_ast(root, cfg);
    return cfg;
}

template<class ControlFlowGraph>
ControlFlowGraph
ControlFlow::build_cg_from_ast(SgNode *root)
{
    ControlFlowGraph cfg;
    build_cg_from_ast(root, cfg);
    return cfg;
}

template<typename CFG, class VertexPropertyWriter, class EdgePropertyWriter>
void
ControlFlow::write_graphviz(std::ostream &out, const CFG &cfg,
                            const VertexPropertyWriter &vpw, const EdgePropertyWriter &epw)
{
 // typedef typename boost::graph_traits<CFG>::vertex_descriptor CFG_Vertex;
    typedef typename boost::graph_traits<CFG>::edge_descriptor CFG_Edge;
    typedef typename boost::graph_traits<CFG>::vertex_iterator CFG_VertexIterator;
    typedef typename boost::graph_traits<CFG>::out_edge_iterator CFG_OutEdgeIterator;

    // Partition the graph into functions and inter-function edges
    typedef Map<SgAsmFunction*, FunctionSubgraphInfo<CFG> > Functions;
    Functions funcs;
    std::vector<CFG_Edge> interfunc_edges;
    CFG_VertexIterator vi, vi_end;
    for (boost::tie(vi, vi_end)=boost::vertices(cfg); vi!=vi_end; ++vi) {
        SgAsmFunction *func = SageInterface::getEnclosingNode<SgAsmFunction>(get_ast_node(cfg, *vi), true);
        FunctionSubgraphInfo<CFG> &f = funcs[func];
        f.vertices.push_back(*vi);
        CFG_OutEdgeIterator ei, ei_end;
        for (boost::tie(ei, ei_end)=boost::out_edges(*vi, cfg); ei!=ei_end; ++ei) {
            SgNode *tgt_node = get_ast_node(cfg, boost::target(*ei, cfg));
            SgAsmFunction *tgt_func = SageInterface::getEnclosingNode<SgAsmFunction>(tgt_node, true);
            if (tgt_func==func) {
                f.edges.push_back(*ei);
            } else {
                interfunc_edges.push_back(*ei);
            }
        }
    }

    // Output subgraph info, each function in its own cluster
    out <<"digraph G {\n";
    for (typename Functions::iterator fi=funcs.begin(); fi!=funcs.end(); ++fi) {
        FunctionSubgraphInfo<CFG> &f = fi->second;
        if (!f.vertices.empty() || !f.edges.empty()) {
            SgNode *node = get_ast_node(cfg, f.vertices.front());
            SgAsmFunction *func = SageInterface::getEnclosingNode<SgAsmFunction>(node, true);
            char cluster_name[64];
            sprintf(cluster_name, "cluster_F%" PRIx64, func->get_entry_va());
            out <<"  subgraph " <<cluster_name <<" {\n"
                <<"    style=filled;\n"
                <<"    color=lightgrey;\n"
                <<"    label=\"Function " <<StringUtility::addrToString(func->get_entry_va())
                <<(func->get_name().empty()?std::string(""):(" <"+func->get_name()+">")) <<"\";\n";
            for (size_t i=0; i<f.vertices.size(); ++i) {
                out <<"    " <<f.vertices[i];
                vpw(out, f.vertices[i]);
                out <<";\n";
            }
            for (size_t i=0; i<f.edges.size(); ++i) {
                out <<"    " <<boost::source(f.edges[i], cfg) <<"->" <<boost::target(f.edges[i], cfg);
                epw(out, f.edges[i]);
                out <<";\n";
            }
            out <<"  }\n"; // subgraph
        }
    }

    // Inter-function edges
    for (size_t i=0; i<interfunc_edges.size(); ++i) {
        out <<"  " <<boost::source(interfunc_edges[i], cfg) <<"->" <<boost::target(interfunc_edges[i], cfg);
        epw(out, interfunc_edges[i]);
        out <<";\n";
    }
    out <<"}\n"; // digraph
}

} // namespace
} // namespace

#endif
#endif