Use the new getMainFunction() method to efficiently locate main

llvm-svn: 4630

1 files changed, 5 insertions, 10 deletions

diff --git a/llvm/lib/Transforms/IPO/Internalize.cpp b/llvm/lib/Transforms/IPO/Internalize.cpp
index dfd4d218641..910e1d4322d 100644
--- a/llvm/lib/Transforms/IPO/Internalize.cpp
+++ b/llvm/lib/Transforms/IPO/Internalize.cpp
@@ -17,21 +17,16 @@ namespace {
 
   class InternalizePass : public Pass {
     virtual bool run(Module &M) {
-      bool FoundMain = false;   // Look for a function named main...
-      for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I)
-        if (I->getName() == "main" && !I->isExternal() &&
-            I->hasExternalLinkage()) {
-          FoundMain = true;
-          break;
-        }
-      
-      if (!FoundMain) return false;  // No main found, must be a library...
+      Function *MainFunc = M.getMainFunction();
+
+      if (MainFunc == 0 || MainFunc->isExternal())
+        return false;  // No main found, must be a library...
       
       bool Changed = false;
       
       // Found a main function, mark all functions not named main as internal.
       for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I)
-        if (I->getName() != "main" &&   // Leave the main function external
+        if (&*I != MainFunc &&          // Leave the main function external
             !I->isExternal() &&         // Function must be defined here
             !I->hasInternalLinkage()) { // Can't already have internal linkage
           I->setInternalLinkage(true);