1 files changed, 24 insertions, 17 deletions

diff --git a/lib/Transforms/Instrumentation/BlockProfiling.cpp b/lib/Transforms/Instrumentation/BlockProfiling.cpp
index 7b38d70d9e..8dacc92bc2 100644
--- a/lib/Transforms/Instrumentation/BlockProfiling.cpp
+++ b/lib/Transforms/Instrumentation/BlockProfiling.cpp
@@ -31,7 +31,7 @@ static void insertInitializationCall(Function *MainFn, const char *FnName,
   const Type *ArgVTy = PointerType::get(PointerType::get(Type::SByteTy));
   const Type *UIntPtr = PointerType::get(Type::UIntTy);
   Module &M = *MainFn->getParent();
-  Function *InitFn = M.getOrInsertFunction(FnName, Type::VoidTy, Type::IntTy,
+  Function *InitFn = M.getOrInsertFunction(FnName, Type::IntTy, Type::IntTy,
                                            ArgVTy, UIntPtr, Type::UIntTy, 0);
 
   // This could force argc and argv into programs that wouldn't otherwise have
@@ -45,38 +45,45 @@ static void insertInitializationCall(Function *MainFn, const char *FnName,
   BasicBlock::iterator InsertPos = Entry->begin();
   while (isa<AllocaInst>(InsertPos)) ++InsertPos;
 
+  ConstantPointerRef *ArrayCPR = ConstantPointerRef::get(Array);
+  std::vector<Constant*> GEPIndices(2, Constant::getNullValue(Type::LongTy));
+  Args[2] = ConstantExpr::getGetElementPtr(ArrayCPR, GEPIndices);
+  
+  unsigned NumElements =
+    cast<ArrayType>(Array->getType()->getElementType())->getNumElements();
+  Args[3] = ConstantUInt::get(Type::UIntTy, NumElements);
+  
+  Instruction *InitCall = new CallInst(InitFn, Args, "newargc", InsertPos);
+
+  // If argc or argv are not available in main, just pass null values in.
   Function::aiterator AI;
   switch (MainFn->asize()) {
   default:
   case 2:
     AI = MainFn->abegin(); ++AI;
     if (AI->getType() != ArgVTy) {
-      Args[1] = new CastInst(AI, ArgVTy, "argv.cast", InsertPos);
+      InitCall->setOperand(2, new CastInst(AI, ArgVTy, "argv.cast", InitCall));
     } else {
-      Args[1] = AI;
+      InitCall->setOperand(2, AI);
     }
 
   case 1:
     AI = MainFn->abegin();
+    // If the program looked at argc, have it look at the return value of the
+    // init call instead.
     if (AI->getType() != Type::IntTy) {
-      Args[0] = new CastInst(AI, Type::IntTy, "argc.cast", InsertPos);
+      if (!AI->use_empty())
+        AI->replaceAllUsesWith(new CastInst(InitCall, AI->getType(), "",
+                                            InsertPos));
+      InitCall->setOperand(1, new CastInst(AI, Type::IntTy, "argc.cast",
+                                           InitCall));
     } else {
-      Args[0] = AI;
+      AI->replaceAllUsesWith(InitCall);
+      InitCall->setOperand(1, AI);
     }
     
-  case 0:
-    break;
+  case 0: break;
   }
-
-  ConstantPointerRef *ArrayCPR = ConstantPointerRef::get(Array);
-  std::vector<Constant*> GEPIndices(2, Constant::getNullValue(Type::LongTy));
-  Args[2] = ConstantExpr::getGetElementPtr(ArrayCPR, GEPIndices);
-  
-  unsigned NumElements =
-    cast<ArrayType>(Array->getType()->getElementType())->getNumElements();
-  Args[3] = ConstantUInt::get(Type::UIntTy, NumElements);
-  
-  new CallInst(InitFn, Args, "", InsertPos);
 }
 
 static void IncrementCounterInBlock(BasicBlock *BB, unsigned CounterNum,