From ee5d708e555b56a760f3b18e43c1ae4661e170cb Mon Sep 17 00:00:00 2001
From: Zhongxing Xu <xuzhongxing@gmail.com>
Date: Thu, 23 Jul 2009 13:39:38 +0000
Subject: [PATCH] Add two nodes to the call graph:  - Root is the main function
 or 0.  - ExternalCallingNode has edges to all external functions.

llvm-svn: 76876
---
 clang/include/clang/Analysis/CallGraph.h | 15 ++++++++++-----
 clang/lib/Analysis/CallGraph.cpp         | 12 ++++++++++++
 clang/tools/wpa/clang-wpa.cpp            |  2 +-
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/clang/include/clang/Analysis/CallGraph.h b/clang/include/clang/Analysis/CallGraph.h
index 5725816f2441..8c437c5261b6 100644
--- a/clang/include/clang/Analysis/CallGraph.h
+++ b/clang/include/clang/Analysis/CallGraph.h
@@ -63,11 +63,14 @@ class CallGraph {
   /// CallerCtx maps a caller to its ASTContext.
   llvm::DenseMap<CallGraphNode *, ASTContext *> CallerCtx;
 
-  /// Entry node of the call graph.
-  // FIXME: find the entry of the graph.
-  CallGraphNode *Entry;
+  /// Root node is the 'main' function or 0.
+  CallGraphNode *Root;
+
+  /// ExternalCallingNode has edges to all external functions.
+  CallGraphNode *ExternalCallingNode;
 
 public:
+  CallGraph();
   ~CallGraph();
 
   typedef FunctionMapTy::iterator iterator;
@@ -78,7 +81,9 @@ public:
   const_iterator begin() const { return FunctionMap.begin(); }
   const_iterator end()   const { return FunctionMap.end();   }
 
-  CallGraphNode *getEntry() { return Entry; }
+  CallGraphNode *getRoot() { return Root; }
+
+  CallGraphNode *getExternalCallingNode() { return ExternalCallingNode; }
 
   void addTU(ASTUnit &AST);
 
@@ -106,7 +111,7 @@ template <> struct GraphTraits<clang::CallGraph> {
   typedef mapped_iterator<NodeType::iterator, CGNDerefFun> ChildIteratorType;
 
   static NodeType *getEntryNode(GraphType *CG) {
-    return CG->getEntry();
+    return CG->getExternalCallingNode();
   }
 
   static ChildIteratorType child_begin(NodeType *N) {
diff --git a/clang/lib/Analysis/CallGraph.cpp b/clang/lib/Analysis/CallGraph.cpp
index bfc2c0d2194f..07c2b35349f9 100644
--- a/clang/lib/Analysis/CallGraph.cpp
+++ b/clang/lib/Analysis/CallGraph.cpp
@@ -58,6 +58,10 @@ void CGBuilder::VisitCallExpr(CallExpr *CE) {
   }
 }
 
+CallGraph::CallGraph() : Root(0) {
+  ExternalCallingNode = getOrInsertFunction(Entity());
+}
+
 CallGraph::~CallGraph() {
   if (!FunctionMap.empty()) {
     for (FunctionMapTy::iterator I = FunctionMap.begin(), E = FunctionMap.end();
@@ -80,6 +84,14 @@ void CallGraph::addTU(ASTUnit &AST) {
         Entity Ent = Entity::get(FD, Prog);
         CallGraphNode *Node = getOrInsertFunction(Ent);
         CallerCtx[Node] = &Ctx;
+
+        // If this function has external linkage, anything could call it.
+        if (FD->isGlobal())
+          ExternalCallingNode->addCallee(idx::ASTLocation(), Node);
+
+        // Set root node to 'main' function.
+        if (FD->getNameAsString() == "main")
+          Root = Node;
         
         CGBuilder builder(*this, FD, Ent, Node);
         builder.Visit(FD->getBody());
diff --git a/clang/tools/wpa/clang-wpa.cpp b/clang/tools/wpa/clang-wpa.cpp
index dff8cb0b5249..e7515ecd2427 100644
--- a/clang/tools/wpa/clang-wpa.cpp
+++ b/clang/tools/wpa/clang-wpa.cpp
@@ -65,5 +65,5 @@ int main(int argc, char **argv) {
   for (unsigned i = 0, e = TUnits.size(); i != e; ++i)
     CG->addTU(*(TUnits[i]->AST));
 
-  CG->dump();
+  CG->ViewCallGraph();
 }