[llvm-commits] [llvm] r58822 - /llvm/trunk/lib/CodeGen/StackProtector.cpp

[Bill Wendling]

Author: void
Date: Thu Nov  6 17:55:49 2008
New Revision: 58822

URL: http://llvm.org/viewvc/llvm-project?rev=58822&view=rev
Log:
Don't build a vector of returns. Just modify the Function in the loop.

Modified:
    llvm/trunk/lib/CodeGen/StackProtector.cpp

Modified: llvm/trunk/lib/CodeGen/StackProtector.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/CodeGen/StackProtector.cpp?rev=58822&r1=58821&r2=58822&view=diff

==============================================================================
--- llvm/trunk/lib/CodeGen/StackProtector.cpp (original)
+++ llvm/trunk/lib/CodeGen/StackProtector.cpp Thu Nov  6 17:55:49 2008
@@ -98,31 +98,8 @@
 ///  - The epilogue checks the value stored in the prologue against the original
 ///    value. It calls __stack_chk_fail if they differ.
 bool StackProtector::InsertStackProtectors() {
-  std::vector<BasicBlock*> ReturnBBs;
-
-  for (Function::iterator I = F->begin(); I != F->end(); ++I)
-    if (isa<ReturnInst>(I->getTerminator()))
-      ReturnBBs.push_back(I);
-
-  // If this function doesn't return, don't bother with stack protectors.
-  if (ReturnBBs.empty()) return false;
-
-  // Insert code into the entry block that stores the __stack_chk_guard variable
-  // onto the stack.
-  BasicBlock &Entry = F->getEntryBlock();
-  Instruction *InsertPt = &Entry.front();
-
-  const PointerType *GuardTy = PointerType::getUnqual(Type::Int8Ty);
-
-  // The global variable for the stack guard.
-  Constant *StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", GuardTy);
-  LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
-  CallInst::
-    Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_create),
-           LI, "", InsertPt);
-
-  // Create the basic block to jump to when the guard check fails.
-  BasicBlock *FailBB = CreateFailBB();
+  Constant *StackGuardVar = 0;  // The global variable for the stack guard.
+  BasicBlock *FailBB = 0;       // The basic block to jump to if check fails.
 
   // Loop through the basic blocks that have return instructions. Convert this:
   //
@@ -146,29 +123,55 @@
   //     call void @__stack_chk_fail()
   //     unreachable
   //
-  for (std::vector<BasicBlock*>::iterator
-         I = ReturnBBs.begin(), E = ReturnBBs.end(); I != E; ++I) {
-    BasicBlock *BB = *I;
-    ReturnInst *RI = cast<ReturnInst>(BB->getTerminator());
-    Function::iterator InsPt = BB; ++InsPt; // Insertion point for new BB.
-
-    // Split the basic block before the return instruction.
-    BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
-
-    // Move the newly created basic block to the point right after the old basic
-    // block so that it's in the "fall through" position.
-    NewBB->removeFromParent();
-    F->getBasicBlockList().insert(InsPt, NewBB);
-
-    // Generate the stack protector instructions in the old basic block.
-    LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
-    CallInst *CI = CallInst::
-      Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_check),
-             "", BB);
-    ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
-    BranchInst::Create(NewBB, FailBB, Cmp, BB);
+  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I) {
+    BasicBlock *BB = I;
+
+    if (isa<ReturnInst>(BB->getTerminator())) {
+      // Create the basic block to jump to when the guard check fails.
+      if (!FailBB)
+        FailBB = CreateFailBB();
+
+      if (!StackGuardVar)
+        StackGuardVar =
+          M->getOrInsertGlobal("__stack_chk_guard",
+                               PointerType::getUnqual(Type::Int8Ty));
+
+      ReturnInst *RI = cast<ReturnInst>(BB->getTerminator());
+      Function::iterator InsPt = BB; ++InsPt; // Insertion point for new BB.
+      ++I;
+
+      // Split the basic block before the return instruction.
+      BasicBlock *NewBB = BB->splitBasicBlock(RI, "SP_return");
+
+      // Move the newly created basic block to the point right after the old basic
+      // block so that it's in the "fall through" position.
+      NewBB->removeFromParent();
+      F->getBasicBlockList().insert(InsPt, NewBB);
+
+      // Generate the stack protector instructions in the old basic block.
+      LoadInst *LI1 = new LoadInst(StackGuardVar, "", false, BB);
+      CallInst *CI = CallInst::
+        Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_check),
+               "", BB);
+      ICmpInst *Cmp = new ICmpInst(CmpInst::ICMP_EQ, CI, LI1, "", BB);
+      BranchInst::Create(NewBB, FailBB, Cmp, BB);
+    }
   }
 
+  // Return if we didn't modify any basic blocks. I.e., there are no return
+  // statements in the function.
+  if (!FailBB) return false;
+
+  // Insert code into the entry block that stores the __stack_chk_guard variable
+  // onto the stack.
+  BasicBlock &Entry = F->getEntryBlock();
+  Instruction *InsertPt = &Entry.front();
+
+  LoadInst *LI = new LoadInst(StackGuardVar, "StackGuard", false, InsertPt);
+  CallInst::
+    Create(Intrinsic::getDeclaration(M, Intrinsic::stackprotector_create),
+           LI, "", InsertPt);
+
   return true;
 }