call_graph.h

// This file if part of the llir-opt project.
// Licensing information can be found in the LICENSE file.
// (C) 2018 Nandor Licker. All rights reserved.
 
#pragma once
 
#include <llvm/ADT/GraphTraits.h>
#include <llvm/ADT/PointerUnion.h>
#include <llvm/Support/DOTGraphTraits.h>
 
#include "core/func.h"
#include "core/inst.h"
#include "core/prog.h"
 
 
 
class CallGraph final {
public:
  class Node {
  public:
    class iterator : public std::iterator<std::forward_iterator_tag, const Node *> {
    public:
      iterator(const Node *node, Inst *start);
      iterator(const Node *node, Func *func);
      iterator() : it_(static_cast<Inst *>(nullptr)) {}
 
      bool operator!=(const iterator &that) const { return !(*this == that); }
      bool operator==(const iterator &that) const;
 
      iterator &operator++();
 
      iterator operator++(int)
      {
        iterator it(*this);
        ++*this;
        return it;
      }
 
      const Node *operator*() const;
 
    private:
      const Node *node_;
      llvm::PointerUnion<Inst *, Func *> it_;
    };
 
  public:
 
    Node(const CallGraph *graph, Prog *prog);
    Node(const CallGraph *graph, Func *caller);
 
    iterator begin() const;
    iterator end() const { return iterator(); }
 
    Func *GetCaller() const;
 
    bool IsRecursive() const;
 
  private:
    friend class iterator;
    const CallGraph *graph_;
    llvm::PointerUnion<Func *, Prog *> node_;
  };
 
private:
  using NodeMap = std::unordered_map<const Func *, std::unique_ptr<Node>>;
 
public:
  struct const_node_iterator : llvm::iterator_adaptor_base
      < const_node_iterator
      , NodeMap::const_iterator
      , std::random_access_iterator_tag
      , const Node *
      >
  {
    explicit const_node_iterator(NodeMap::const_iterator it)
      : iterator_adaptor_base(it)
    {
    }
 
    const Node *operator*() const { return I->second.get(); }
    const Node *operator->() const { return I->second.get(); }
  };
 
public:
  CallGraph(Prog &p);
 
  ~CallGraph();
 
  const Node *Entry() const { return &entry_; }
  const Node *operator[](Func *f) const;
 
  const_node_iterator begin() const { return const_node_iterator(nodes_.begin()); }
  const_node_iterator end() const { return const_node_iterator(nodes_.end()); }
 
private:
  friend class Node::iterator;
  Node entry_;
  mutable NodeMap nodes_;
};
 
namespace llvm {
 
template <>
struct GraphTraits<CallGraph::Node *> {
  using NodeRef = const CallGraph::Node *;
 
  using ChildIteratorType = CallGraph::Node::iterator;
  static ChildIteratorType child_begin(NodeRef n) { return n->begin(); }
  static ChildIteratorType child_end(NodeRef n) { return n->end(); }
};
 
template <>
struct GraphTraits<CallGraph *> : public GraphTraits<CallGraph::Node *> {
  static NodeRef getEntryNode(const CallGraph *g) { return g->Entry(); }
 
  using nodes_iterator = CallGraph::const_node_iterator;
  static nodes_iterator nodes_begin(CallGraph *g) { return g->begin(); }
  static nodes_iterator nodes_end(CallGraph *g) { return g->end(); }
};
 
template<>
struct DOTGraphTraits<CallGraph*> : public DefaultDOTGraphTraits {
  DOTGraphTraits(bool isSimple = false) : DefaultDOTGraphTraits(isSimple) {}
 
  static std::string getNodeLabel(
      const CallGraph::Node *n,
      const CallGraph *g)
  {
    if (auto *f = n->GetCaller()) {
      return std::string(f->GetName());
    } else {
      return "root";
    }
  }
};
 
}