cleanup as per Duncan's review

[oota-llvm.git] / lib / Transforms / IPO / IPConstantPropagation.cpp

diff --git a/lib/Transforms/IPO/IPConstantPropagation.cpp b/lib/Transforms/IPO/IPConstantPropagation.cpp
index 0e654a50eefbf34c7e33656752c89818671eb4a7..42c02e6a45c45907960061c5f72c7ec8fc896832 100644 (file)
--- a/lib/Transforms/IPO/IPConstantPropagation.cpp
+++ b/lib/Transforms/IPO/IPConstantPropagation.cpp
@@ -145,6 +145,10 @@ bool IPCP::PropagateConstantsIntoArguments(Function &F) {
 // all callers that use those return values with the constant value. This will
 // leave in the actual return values and instructions, but deadargelim will
 // clean that up.
+//
+// Additionally if a function always returns one of its arguments directly,
+// callers will be updated to use the value they pass in directly instead of
+// using the return value.
 bool IPCP::PropagateConstantReturn(Function &F) {
   if (F.getReturnType() == Type::VoidTy)
     return false; // No return value.
@@ -188,8 +192,8 @@ bool IPCP::PropagateConstantReturn(Function &F) {
           if (isa<UndefValue>(V))
             continue;
           
-          // Try to see if all the rets return the same constant.
-          if (isa<Constant>(V)) {
+          // Try to see if all the rets return the same constant or argument.
+          if (isa<Constant>(V) || isa<Argument>(V)) {
             if (isa<UndefValue>(RV)) {
               // No value found yet? Try the current one.
               RetVals[i] = V;
@@ -214,20 +218,27 @@ bool IPCP::PropagateConstantReturn(Function &F) {
   // constant.
   bool MadeChange = false;
   for (Value::use_iterator UI = F.use_begin(), E = F.use_end(); UI != E; ++UI) {
-    // Make sure this is an invoke or call and that the use is for the callee.
-    if (!(isa<InvokeInst>(*UI) || isa<CallInst>(*UI)) ||
-        UI.getOperandNo() != 0) {
+    CallSite CS = CallSite::get(*UI);
+    Instruction* Call = CS.getInstruction();
+
+    // Not a call instruction or a call instruction that's not calling F
+    // directly?
+    if (!Call || UI.getOperandNo() != 0)
       continue;
-    }
     
-    Instruction *Call = cast<Instruction>(*UI);
+    // Call result not used?
     if (Call->use_empty())
       continue;
 
     MadeChange = true;
 
     if (STy == 0) {
-      Call->replaceAllUsesWith(RetVals[0]);
+      Value* New = RetVals[0];
+      if (Argument *A = dyn_cast<Argument>(New))
+        // Was an argument returned? Then find the corresponding argument in
+        // the call instruction and use that.
+        New = CS.getArgument(A->getArgNo());
+      Call->replaceAllUsesWith(New);
       continue;
     }
    
@@ -244,9 +255,7 @@ bool IPCP::PropagateConstantReturn(Function &F) {
 
       // Find the index of the retval to replace with
       int index = -1;
-      if (GetResultInst *GR = dyn_cast<GetResultInst>(Ins))
-        index = GR->getIndex();
-      else if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(Ins))
+      if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(Ins))
         if (EV->hasIndices())
           index = *EV->idx_begin();
 
@@ -255,6 +264,10 @@ bool IPCP::PropagateConstantReturn(Function &F) {
       if (index != -1) {
         Value *New = RetVals[index];
         if (New) {
+          if (Argument *A = dyn_cast<Argument>(New))
+            // Was an argument returned? Then find the corresponding argument in
+            // the call instruction and use that.
+            New = CS.getArgument(A->getArgNo());
           Ins->replaceAllUsesWith(New);
           Ins->eraseFromParent();
         }