[PATCH] IR: Refactor main detection

[David Majnemer]

Hi chandlerc, rnk,

Add Function::isMain and Module::getMainFunction, they do nothing but
standardize existing practices inside of LLVM.

http://llvm-reviews.chandlerc.com/D1668

Files:
  include/llvm/IR/Function.h
  include/llvm/IR/Module.h
  lib/IR/Function.cpp
  lib/IR/Module.cpp
  lib/Target/X86/X86ISelDAGToDAG.cpp
  lib/Target/X86/X86ISelLowering.cpp
  lib/Transforms/IPO/GlobalOpt.cpp
  lib/Transforms/Instrumentation/EdgeProfiling.cpp
  lib/Transforms/Instrumentation/OptimalEdgeProfiling.cpp
  lib/Transforms/Instrumentation/PathProfiling.cpp

Index: include/llvm/IR/Function.h
===================================================================
--- include/llvm/IR/Function.h
+++ include/llvm/IR/Function.h
@@ -155,6 +155,9 @@
   unsigned getIntrinsicID() const LLVM_READONLY;
   bool isIntrinsic() const { return getName().startswith("llvm."); }
 
+  // isMain - Returns true if this is main, false otherwise.
+  bool isMain() const;
+
   /// getCallingConv()/setCallingConv(CC) - These method get and set the
   /// calling convention of this function.  The enum values for the known
   /// calling conventions are defined in CallingConv.h.
Index: include/llvm/IR/Module.h
===================================================================
--- include/llvm/IR/Module.h
+++ include/llvm/IR/Module.h
@@ -344,6 +344,10 @@
   /// If it does not exist, return null.
   Function *getFunction(StringRef Name) const;
 
+  /// getMainFunction - Look up the main function in the module symbol table.
+  /// If it does not exist, return null.
+  Function *getMainFunction() const;
+
 /// @}
 /// @name Global Variable Accessors
 /// @{
Index: lib/IR/Function.cpp
===================================================================
--- lib/IR/Function.cpp
+++ lib/IR/Function.cpp
@@ -17,6 +17,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/Triple.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -720,3 +721,23 @@
 
   return false;
 }
+
+bool Function::isMain() const {
+  if (!hasExternalLinkage())
+    return false;
+
+  std::string FnName = getName();
+  if (FnName == "main")
+    return true;
+
+  // Using fortran? ... this kind of works
+  if (FnName == "MAIN__")
+    return true;
+
+  Triple TargetTriple(getParent()->getTargetTriple());
+  if (TargetTriple.getOS() == Triple::Win32)
+    if (FnName == "wmain")
+      return true;
+
+  return false;
+}
Index: lib/IR/Module.cpp
===================================================================
--- lib/IR/Module.cpp
+++ lib/IR/Module.cpp
@@ -17,6 +17,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/Triple.h"
 #include "llvm/GVMaterializer.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -222,6 +223,26 @@
   return dyn_cast_or_null<Function>(getNamedValue(Name));
 }
 
+Function *Module::getMainFunction() const {
+  if (Function *Main = dyn_cast_or_null<Function>(getNamedValue("main")))
+    if (Main->hasExternalLinkage())
+      return Main;
+
+  // Using fortran? ... this kind of works
+  if (Function *MainUnderUnder =
+          dyn_cast_or_null<Function>(getNamedValue("MAIN__")))
+    if (MainUnderUnder->hasExternalLinkage())
+      return MainUnderUnder;
+
+  Triple TargetTriple(getTargetTriple());
+  if (TargetTriple.getOS() == Triple::Win32)
+    if (Function *WMain = dyn_cast_or_null<Function>(getNamedValue("wmain")))
+      if (WMain->hasExternalLinkage())
+        return WMain;
+
+  return 0;
+}
+
 //===----------------------------------------------------------------------===//
 // Methods for easy access to the global variables in the module.
 //
Index: lib/Target/X86/X86ISelDAGToDAG.cpp
===================================================================
--- lib/Target/X86/X86ISelDAGToDAG.cpp
+++ lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -568,7 +568,7 @@
 void X86DAGToDAGISel::EmitFunctionEntryCode() {
   // If this is main, emit special code for main.
   if (const Function *Fn = MF->getFunction())
-    if (Fn->hasExternalLinkage() && Fn->getName() == "main")
+    if (Fn->isMain())
       EmitSpecialCodeForMain(MF->begin(), MF->getFrameInfo());
 }
 
Index: lib/Target/X86/X86ISelLowering.cpp
===================================================================
--- lib/Target/X86/X86ISelLowering.cpp
+++ lib/Target/X86/X86ISelLowering.cpp
@@ -2134,10 +2134,8 @@
   MachineFunction &MF = DAG.getMachineFunction();
   X86MachineFunctionInfo *FuncInfo = MF.getInfo<X86MachineFunctionInfo>();
 
-  const Function* Fn = MF.getFunction();
-  if (Fn->hasExternalLinkage() &&
-      Subtarget->isTargetCygMing() &&
-      Fn->getName() == "main")
+  const Function *Fn = MF.getFunction();
+  if (Subtarget->isTargetCygMing() && Fn->isMain())
     FuncInfo->setForceFramePointer(true);
 
   MachineFrameInfo *MFI = MF.getFrameInfo();
Index: lib/Transforms/IPO/GlobalOpt.cpp
===================================================================
--- lib/Transforms/IPO/GlobalOpt.cpp
+++ lib/Transforms/IPO/GlobalOpt.cpp
@@ -1946,12 +1946,10 @@
   // are just replacing static memory to stack memory.
   //
   // If the global is in different address space, don't bring it to stack.
-  if (!GS.HasMultipleAccessingFunctions &&
-      GS.AccessingFunction && !GS.HasNonInstructionUser &&
+  if (!GS.HasMultipleAccessingFunctions && GS.AccessingFunction &&
+      !GS.HasNonInstructionUser &&
       GV->getType()->getElementType()->isSingleValueType() &&
-      GS.AccessingFunction->getName() == "main" &&
-      GS.AccessingFunction->hasExternalLinkage() &&
-      GV->getType()->getAddressSpace() == 0) {
+      GS.AccessingFunction->isMain() && GV->getType()->getAddressSpace() == 0) {
     DEBUG(dbgs() << "LOCALIZING GLOBAL: " << *GV);
     Instruction &FirstI = const_cast<Instruction&>(*GS.AccessingFunction
                                                    ->getEntryBlock().begin());
Index: lib/Transforms/Instrumentation/EdgeProfiling.cpp
===================================================================
--- lib/Transforms/Instrumentation/EdgeProfiling.cpp
+++ lib/Transforms/Instrumentation/EdgeProfiling.cpp
@@ -52,7 +52,7 @@
 ModulePass *llvm::createEdgeProfilerPass() { return new EdgeProfiler(); }
 
 bool EdgeProfiler::runOnModule(Module &M) {
-  Function *Main = M.getFunction("main");
+  Function *Main = M.getMainFunction();
   if (Main == 0) {
     errs() << "WARNING: cannot insert edge profiling into a module"
            << " with no main function!\n";
Index: lib/Transforms/Instrumentation/OptimalEdgeProfiling.cpp
===================================================================
--- lib/Transforms/Instrumentation/OptimalEdgeProfiling.cpp
+++ lib/Transforms/Instrumentation/OptimalEdgeProfiling.cpp
@@ -73,7 +73,7 @@
 }
 
 bool OptimalEdgeProfiler::runOnModule(Module &M) {
-  Function *Main = M.getFunction("main");
+  Function *Main = M.getMainFunction();
   if (Main == 0) {
     errs() << "WARNING: cannot insert edge profiling into a module"
            << " with no main function!\n";
Index: lib/Transforms/Instrumentation/PathProfiling.cpp
===================================================================
--- lib/Transforms/Instrumentation/PathProfiling.cpp
+++ lib/Transforms/Instrumentation/PathProfiling.cpp
@@ -1338,11 +1338,7 @@
         << "****************************************\n");
 
   // No main, no instrumentation!
-  Function *Main = M.getFunction("main");
-
-  // Using fortran? ... this kind of works
-  if (!Main)
-    Main = M.getFunction("MAIN__");
+  Function *Main = M.getMainFunction();
 
   if (!Main) {
     errs() << "WARNING: cannot insert path profiling into a module"
-------------- next part --------------
A non-text attachment was scrubbed...
Name: D1668.1.patch
Type: text/x-patch
Size: 7223 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20130912/84b31310/attachment.bin>