[VPlan] Account for removed users in replaceAllUsesWith.

Make sure we do not iterate using an invalid iterator.

Another small fix/step towards traversing the def-use chains in VPlan.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -884,10 +884,18 @@ void VPWidenCanonicalIVRecipe::print(raw_ostream &O, const Twine &Indent,
 template void DomTreeBuilder::Calculate<VPDominatorTree>(VPDominatorTree &DT);
 
 void VPValue::replaceAllUsesWith(VPValue *New) {
-  for (VPUser *User : users())
+  for (unsigned J = 0; J < getNumUsers();) {
+    VPUser *User = Users[J];
+    unsigned NumUsers = getNumUsers();
     for (unsigned I = 0, E = User->getNumOperands(); I < E; ++I)
       if (User->getOperand(I) == this)
         User->setOperand(I, New);
+    // If a user got removed after updating the current user, the next user to
+    // update will be moved to the current position, so we only need to
+    // increment the index if the number of users did not change.
+    if (NumUsers == getNumUsers())
+      J++;
+  }
 }
 
 void VPValue::printAsOperand(raw_ostream &OS, VPSlotTracker &Tracker) const {

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

@@ -167,7 +167,13 @@ TEST(VPInstructionTest, replaceAllUsesWith) {
   EXPECT_EQ(0u, VPV2->getNumUsers());
   EXPECT_EQ(0u, VPV3->getNumUsers());
 
+  VPInstruction *I2 = new VPInstruction(0, {VPV1, VPV2});
+  EXPECT_EQ(3u, VPV1->getNumUsers());
+  VPV1->replaceAllUsesWith(VPV3);
+  EXPECT_EQ(3u, VPV3->getNumUsers());
+
   delete I1;
+  delete I2;
   delete VPV1;
   delete VPV2;
   delete VPV3;