Make DeadArgElim keep liveness of the return value as a whole in addition to

only the liveness of partial return values (for functions returning a struct).
This is more explicit to prevent unwanted changes in the return value. 

In particular, deadargelim now canonicalizes a function returning {i32} to
returning i32 and {} to void, if the struct returned is not used in its
entirety, but only the single element is used.


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@53606 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 124 insertions, 126 deletions

diff --git a/lib/Transforms/IPO/DeadArgumentElimination.cpp b/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 19be0eeb3c..40012d1ad7 100644
--- a/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -47,12 +47,14 @@ namespace {
 
     /// Struct that represents (part of) either a return value or a function
     /// argument.  Used so that arguments and return values can be used
-    /// interchangably.
+    /// interchangably. Idx == -1 means the entire return value, while other
+    /// indices mean the corresponding element in the struct return type (if
+    /// any).
     struct RetOrArg {
-      RetOrArg(const Function* F, unsigned Idx, bool IsArg) : F(F), Idx(Idx),
+      RetOrArg(const Function* F, signed Idx, bool IsArg) : F(F), Idx(Idx),
                IsArg(IsArg) {}
       const Function *F;
-      unsigned Idx;
+      signed Idx;
       bool IsArg;
 
       /// Make RetOrArg comparable, so we can put it into a map.
@@ -71,8 +73,10 @@ namespace {
       }
 
       std::string getDescription() const {
-        return std::string((IsArg ? "Argument #" : "Return value #")) 
-               + utostr(Idx) + " of function " + F->getName();
+        return std::string((!IsArg && Idx != -1 ? "partial " : "")) 
+               + (IsArg ? "argument #" : "return value") 
+               + (!IsArg && Idx == -1 ? "" : " #" + utostr(Idx)) 
+               + " of function " + F->getName();
       }
     };
 
@@ -84,11 +88,11 @@ namespace {
     enum Liveness { Live, MaybeLive };
 
     /// Convenience wrapper
-    RetOrArg CreateRet(const Function *F, unsigned Idx) {
+    RetOrArg CreateRet(const Function *F, signed Idx) {
       return RetOrArg(F, Idx, false);
     }
     /// Convenience wrapper
-    RetOrArg CreateArg(const Function *F, unsigned Idx) {
+    RetOrArg CreateArg(const Function *F, signed Idx) {
       return RetOrArg(F, Idx, true);
     }
 
@@ -129,7 +133,7 @@ namespace {
   private:
     Liveness MarkIfNotLive(RetOrArg Use, UseVector &MaybeLiveUses);
     Liveness SurveyUse(Value::use_iterator U, UseVector &MaybeLiveUses,
-                       unsigned RetValNum = 0);
+                       signed RetValNum = -1);
     Liveness SurveyUses(Value *V, UseVector &MaybeLiveUses);
 
     void SurveyFunction(Function &F);
@@ -278,18 +282,6 @@ bool DAE::DeleteDeadVarargs(Function &Fn) {
   return true;
 }
 
-/// Convenience function that returns the number of return values. It returns 0
-/// for void functions and 1 for functions not returning a struct. It returns
-/// the number of struct elements for functions returning a struct.
-static unsigned NumRetVals(const Function *F) {
-  if (F->getReturnType() == Type::VoidTy)
-    return 0;
-  else if (const StructType *STy = dyn_cast<StructType>(F->getReturnType()))
-    return STy->getNumElements();
-  else
-    return 1;
-}
-
 /// MarkIfNotLive - This checks Use for liveness in LiveValues. If Use is not
 /// live, it adds Use to the MaybeLiveUses argument. Returns the determined
 /// liveness of Use.
@@ -304,16 +296,15 @@ DAE::Liveness DAE::MarkIfNotLive(RetOrArg Use, UseVector &MaybeLiveUses) {
   return MaybeLive;
 }
 
-
 /// SurveyUse - This looks at a single use of an argument or return value
 /// and determines if it should be alive or not. Adds this use to MaybeLiveUses
 /// if it causes the used value to become MaybeAlive.
 ///
 /// RetValNum is the return value number to use when this use is used in a
 /// return instruction. This is used in the recursion, you should always leave
-/// it at 0.
+/// it at -1.
 DAE::Liveness DAE::SurveyUse(Value::use_iterator U, UseVector &MaybeLiveUses,
-                             unsigned RetValNum) {
+                             signed RetValNum) {
     Value *V = *U;
     if (ReturnInst *RI = dyn_cast<ReturnInst>(V)) {
       // The value is returned from a function. It's only live when the
@@ -402,16 +393,23 @@ DAE::Liveness DAE::SurveyUses(Value *V, UseVector &MaybeLiveUses) {
 // well as arguments to functions which have their "address taken".
 //
 void DAE::SurveyFunction(Function &F) {
-  unsigned RetCount = NumRetVals(&F);
+  const StructType *STy = dyn_cast<StructType>(F.getFunctionType()->getReturnType());
+  // Store the number of partial return values, which we only have if the return
+  // type is a struct.
+  unsigned PartialRetVals = (STy ? STy->getNumElements() : 0);
   // Assume all return values are dead
   typedef SmallVector<Liveness, 5> RetVals;
-  RetVals RetValLiveness(RetCount, MaybeLive);
+  // Allocate one slot for the entire return value (index 0) and one for each
+  // partial return value.
+  RetVals RetValLiveness(PartialRetVals + 1, MaybeLive);
 
   typedef SmallVector<UseVector, 5> RetUses;
   // These vectors map each return value to the uses that make it MaybeLive, so
   // we can add those to the Uses map if the return value really turns out to be
-  // MaybeLive. Initialized to a list of RetCount empty lists.
-  RetUses MaybeLiveRetUses(RetCount);
+  // MaybeLive. 
+  // Allocate one slot for the entire return value (index 0) and one for each
+  // partial return value.
+  RetUses MaybeLiveRetUses(PartialRetVals + 1);
 
   for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB)
     if (ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator()))
@@ -428,10 +426,6 @@ void DAE::SurveyFunction(Function &F) {
   }
 
   DOUT << "DAE - Inspecting callers for fn: " << F.getName() << "\n";
-  // Keep track of the number of live retvals, so we can skip checks once all
-  // of them turn out to be live.
-  unsigned NumLiveRetVals = 0;
-  const Type *STy = dyn_cast<StructType>(F.getReturnType());
   // Loop all uses of the function.
   for (Value::use_iterator I = F.use_begin(), E = F.use_end(); I != E; ++I) {
     // If the function is PASSED IN as an argument, its address has been
@@ -452,43 +446,42 @@ void DAE::SurveyFunction(Function &F) {
     // If we end up here, we are looking at a direct call to our function.
 
     // Now, check how our return value(s) is/are used in this caller. Don't
-    // bother checking return values if all of them are live already.
-    if (NumLiveRetVals != RetCount) {
-      if (STy) {
-        // Check all uses of the return value.
-        for (Value::use_iterator I = TheCall->use_begin(),
-             E = TheCall->use_end(); I != E; ++I) {
-          ExtractValueInst *Ext = dyn_cast<ExtractValueInst>(*I);
-          if (Ext && Ext->hasIndices()) {
-            // This use uses a part of our return value, survey the uses of
-            // that part and store the results for this index only.
-            unsigned Idx = *Ext->idx_begin();
-            if (RetValLiveness[Idx] != Live) {
-              RetValLiveness[Idx] = SurveyUses(Ext, MaybeLiveRetUses[Idx]);
-              if (RetValLiveness[Idx] == Live)
-                NumLiveRetVals++;
-            }
-          } else {
-            // Used by something else than extractvalue. Mark all return
-            // values as live.
-            for (unsigned i = 0; i != RetCount; ++i )
-              RetValLiveness[i] = Live;
-            NumLiveRetVals = RetCount;
-            break;
+    // bother checking return values if the entire value is live already.
+    if (RetValLiveness[0] != Live) {
+      // Check all uses of the return value.
+      for (Value::use_iterator I = TheCall->use_begin(),
+           E = TheCall->use_end(); I != E; ++I) {
+        ExtractValueInst *Ext = dyn_cast<ExtractValueInst>(*I);
+        if (Ext && Ext->hasIndices()) {
+          // This use uses a part of our return value, survey the uses of
+          // that part and store the results for this index only.
+          unsigned Idx = *Ext->idx_begin();
+          // + 1 to skip the "entire retval" index (0)
+          Liveness &IsLive = RetValLiveness[Idx + 1];
+          if (IsLive != Live) {
+            // Don't bother checking if this retval was already live
+            // Survey all uses of the extractvalue
+            // + 1 to skip the "entire retval" index (0)
+            IsLive = SurveyUses(Ext, MaybeLiveRetUses[Idx + 1]);
           }
+        } else {
+          // Survey this use
+          RetValLiveness[0] = SurveyUse(I, MaybeLiveRetUses[0]);
+          if (RetValLiveness[0] == Live)
+            break;
         }
-      } else {
-        // Single return value
-        RetValLiveness[0] = SurveyUses(TheCall, MaybeLiveRetUses[0]);
-        if (RetValLiveness[0] == Live)
-          NumLiveRetVals = RetCount;
       }
     }
   }
 
   // Now we've inspected all callers, record the liveness of our return values.
-  for (unsigned i = 0; i != RetCount; ++i)
-    MarkValue(CreateRet(&F, i), RetValLiveness[i], MaybeLiveRetUses[i]);
+  MarkValue(CreateRet(&F, -1), RetValLiveness[0], MaybeLiveRetUses[0]);
+  if (RetValLiveness[0] != Live)
+    // If the entire retval (0) is Live, MarkValue will have marked all other retvals live
+    // as well, so we can skip this.
+    for (unsigned i = 0; i != PartialRetVals; ++i)
+      // + 1 to skip the "entire retval" index (0)
+      MarkValue(CreateRet(&F, i), RetValLiveness[i + 1], MaybeLiveRetUses[i + 1]);
 
   DOUT << "DAE - Inspecting args for fn: " << F.getName() << "\n";
 
@@ -539,8 +532,11 @@ void DAE::MarkLive(const Function &F) {
     for (unsigned i = 0, e = F.arg_size(); i != e; ++i)
       PropagateLiveness(CreateArg(&F, i));
     // Mark all return values as live.
-    for (unsigned i = 0, e = NumRetVals(&F); i != e; ++i)
-      PropagateLiveness(CreateRet(&F, i));
+    const Type *RTy = F.getFunctionType()->getReturnType();
+    if (const StructType *STy = dyn_cast<StructType>(RTy))
+      for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i)
+        PropagateLiveness(CreateRet(&F, i));
+    PropagateLiveness(CreateRet(&F, -1));
 }
 
 /// MarkLive - Mark the given return value or argument as live. Additionally,
@@ -555,6 +551,16 @@ void DAE::MarkLive(const RetOrArg &RA) {
 
   DOUT << "DAE - Marking " << RA.getDescription() << " live\n";
   PropagateLiveness(RA);
+
+  if (!RA.IsArg && RA.Idx == -1) {
+    // Entire return value live? 
+    const Type *RTy = RA.F->getFunctionType()->getReturnType();
+    if (const StructType *STy = dyn_cast<StructType>(RTy))
+      // And the function returns a struct?
+      for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i)
+        // Mark all partial return values live as well then
+        MarkLive(CreateRet(RA.F, i));
+  }
 }
 
 /// PropagateLiveness - Given that RA is a live value, propagate it's liveness
@@ -599,20 +605,20 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
   // Find out the new return value.
 
   const Type *RetTy = FTy->getReturnType();
+  const StructType *STy = dyn_cast<StructType>(RetTy);
+  const unsigned PartialRetVals = (STy ? STy->getNumElements() : 0);
   const Type *NRetTy = NULL;
-  unsigned RetCount = NumRetVals(F);
-  // Explicitly track if anything changed, for debugging.
-  bool Changed = false;
-  // -1 means unused, other numbers are the new index
-  SmallVector<int, 5> NewRetIdxs(RetCount, -1);
+  // -1 means unused, other numbers are the new index. Initialized to -1 for
+  // every partial return value we have.
+  SmallVector<int, 5> NewRetIdxs(PartialRetVals, -1);
   std::vector<const Type*> RetTypes;
-  if (RetTy == Type::VoidTy) {
-    NRetTy = Type::VoidTy;
+  if (LiveValues.count(CreateRet(F, -1))) {
+    // If the entire return value is live, leave it unchanged.
+    NRetTy = RetTy;
   } else {
-    const StructType *STy = dyn_cast<StructType>(RetTy);
-    if (STy)
-      // Look at each of the original return values individually.
-      for (unsigned i = 0; i != RetCount; ++i) {
+    if (STy) {
+      // Look at each of the partial return values individually.
+      for (unsigned i = 0; i != PartialRetVals; ++i) {
         RetOrArg Ret = CreateRet(F, i);
         if (LiveValues.erase(Ret)) {
           RetTypes.push_back(STy->getElementType(i));
@@ -621,25 +627,20 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
           ++NumRetValsEliminated;
           DOUT << "DAE - Removing return value " << i << " from "
                << F->getNameStart() << "\n";
-          Changed = true;
+          // We remove the value by not adding anything to RetTypes.
         }
       }
-    else
-      // We used to return a single value.
-      if (LiveValues.erase(CreateRet(F, 0))) {
-        RetTypes.push_back(RetTy);
-        NewRetIdxs[0] = 0;
-      } else {
-        DOUT << "DAE - Removing return value from " << F->getNameStart()
-             << "\n";
-        ++NumRetValsEliminated;
-        Changed = true;
-      }
-    if (RetTypes.size() > 1 || (STy && STy->getNumElements()==RetTypes.size()))
-      // More than one return type? Return a struct with them. Also, if we used
-      // to return a struct and didn't change the number of return values,
-      // return a struct again. This prevents changing {something} into
-      // something and {} into void.
+    } else if (RetTy != Type::VoidTy) {
+      // We used to return a single value, which is now dead (already checked in
+      // the if above)
+      DOUT << "DAE - Removing return value from " << F->getNameStart()
+           << "\n";
+      ++NumRetValsEliminated;
+      // We remove the value by not adding anything to RetTypes.
+    }
+
+    if (RetTypes.size() > 1)
+      // More than one return type? Return a struct with them. 
       // Make the new struct packed if we used to return a packed struct
       // already.
       NRetTy = StructType::get(RetTypes, STy->isPacked());
@@ -648,7 +649,7 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
       // return a struct with that simple value before.
       NRetTy = RetTypes.front();
     else if (RetTypes.size() == 0)
-      // No return types? Make it void, but only if we didn't use to return {}.
+      // No return types? Make it void.
       NRetTy = Type::VoidTy;
   }
 
@@ -688,7 +689,6 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
       ++NumArgumentsEliminated;
       DOUT << "DAE - Removing argument " << i << " (" << I->getNameStart()
            << ") from " << F->getNameStart() << "\n";
-      Changed = true;
     }
   }
 
@@ -714,11 +714,6 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
   if (NFTy == FTy)
     return false;
 
-  // The function type is only allowed to be different if we actually left out
-  // an argument or return value.
-  assert(Changed && "Function type changed while no arguments or return values"
-                    "were removed!");
-
   // Create the new function body and insert it into the module...
   Function *NF = Function::Create(NFTy, F->getLinkage());
   NF->copyAttributesFrom(F);
@@ -803,39 +798,42 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
                                          "void. The old return type must have"
                                          "been a struct!");
         // The original return value was a struct, update all uses (which are
-        // all extractvalue instructions).
+        // all extractvalue instructions, or uses that are unused themselves).
         for (Value::use_iterator I = Call->use_begin(), E = Call->use_end();
              I != E;) {
-          assert(isa<ExtractValueInst>(*I) && "Return value not only used by"
-                                              "extractvalue?");
-          ExtractValueInst *EV = cast<ExtractValueInst>(*I);
-          // Increment now, since we're about to throw away this use.
-          ++I;
-          assert(EV->hasIndices() && "Return value used by extractvalue without"
-                                     "indices?");
-          unsigned Idx = *EV->idx_begin();
-          if (NewRetIdxs[Idx] != -1) {
-            if (RetTypes.size() > 1) {
-              // We're still returning a struct, create a new extractvalue
-              // instruction with the first index updated
-              std::vector<unsigned> NewIdxs(EV->idx_begin(), EV->idx_end());
-              NewIdxs[0] = NewRetIdxs[Idx];
-              Value *NEV = ExtractValueInst::Create(New, NewIdxs.begin(),
-                                                    NewIdxs.end(), "retval",
-                                                    EV);
-              EV->replaceAllUsesWith(NEV);
-              EV->eraseFromParent();
+          if (ExtractValueInst *EV = dyn_cast<ExtractValueInst>(*I)) {
+            // Increment now, since we're about to throw away this use.
+            ++I;
+            assert(EV->hasIndices() && "Return value used by extractvalue without"
+                                       "indices?");
+            unsigned Idx = *EV->idx_begin();
+            if (NewRetIdxs[Idx] != -1) {
+              if (RetTypes.size() > 1) {
+                // We're still returning a struct, create a new extractvalue
+                // instruction with the first index updated
+                std::vector<unsigned> NewIdxs(EV->idx_begin(), EV->idx_end());
+                NewIdxs[0] = NewRetIdxs[Idx];
+                Value *NEV = ExtractValueInst::Create(New, NewIdxs.begin(),
+                                                      NewIdxs.end(), "retval",
+                                                      EV);
+                EV->replaceAllUsesWith(NEV);
+                EV->eraseFromParent();
+              } else {
+                // We are now only returning a simple value, remove the
+                // extractvalue.
+                EV->replaceAllUsesWith(New);
+                EV->eraseFromParent();
+              }
             } else {
-              // We are now only returning a simple value, remove the
-              // extractvalue.
-              EV->replaceAllUsesWith(New);
+              // Value unused, replace uses by null for now, they will get removed
+              // later on.
+              EV->replaceAllUsesWith(Constant::getNullValue(EV->getType()));
               EV->eraseFromParent();
             }
           } else {
-            // Value unused, replace uses by null for now, they will get removed
-            // later on.
-            EV->replaceAllUsesWith(Constant::getNullValue(EV->getType()));
-            EV->eraseFromParent();
+            // Not an extractvalue, so this use will become dead soon. Just
+            // replace it with null.
+            I.getUse().set(Constant::getNullValue(I.getUse().get()->getType()));
           }
         }
         New->takeName(Call);
@@ -888,7 +886,7 @@ bool DAE::RemoveDeadStuffFromFunction(Function *F) {
           Value *OldRet = RI->getOperand(0);
           // Start out building up our return value from undef
           RetVal = llvm::UndefValue::get(NRetTy);
-          for (unsigned i = 0; i != RetCount; ++i)
+          for (unsigned i = 0; i != PartialRetVals; ++i)
             if (NewRetIdxs[i] != -1) {
               ExtractValueInst *EV = ExtractValueInst::Create(OldRet, i,
                                                               "oldret", RI);