[llvm] [InstCombine] Collect users iteratively (PR #131956)

[Matt Arsenault via llvm-commits]

================
@@ -270,77 +270,87 @@ class PointerReplacer {
 } // end anonymous namespace
 
 bool PointerReplacer::collectUsers() {
-  if (!collectUsersRecursive(Root))
-    return false;
+  std::stack<User *> Worklist;
+  SmallSetVector<Instruction *, 4> ValuesToRevisit;
 
-  // Ensure that all outstanding (indirect) users of I
-  // are inserted into the Worklist. Return false
-  // otherwise.
-  return llvm::set_is_subset(ValuesToRevisit, Worklist);
-}
+  auto PushUsersToWorklist = [&](Instruction *Inst) {
+    for (auto *U : Inst->users())
+      if (isa<Instruction>(U) && !isAvailable(cast<Instruction>(U)))
+        Worklist.push(U);
+  };
 
-bool PointerReplacer::collectUsersRecursive(Instruction &I) {
-  for (auto *U : I.users()) {
-    auto *Inst = cast<Instruction>(&*U);
+  PushUsersToWorklist(&Root);
+  while (!Worklist.empty()) {
+    auto *Inst = dyn_cast<Instruction>(Worklist.top());
+    Worklist.pop();
+    if (!Inst)
+      return false;
+    if (isAvailable(Inst))
+      continue;
     if (auto *Load = dyn_cast<LoadInst>(Inst)) {
       if (Load->isVolatile())
         return false;
-      Worklist.insert(Load);
+      UsersToReplace.insert(Load);
     } else if (auto *PHI = dyn_cast<PHINode>(Inst)) {
       // All incoming values must be instructions for replacability
       if (any_of(PHI->incoming_values(),
                  [](Value *V) { return !isa<Instruction>(V); }))
         return false;
 
-      // If at least one incoming value of the PHI is not in Worklist,
-      // store the PHI for revisiting and skip this iteration of the
-      // loop.
-      if (any_of(PHI->incoming_values(), [this](Value *V) {
-            return !isAvailable(cast<Instruction>(V));
-          })) {
-        ValuesToRevisit.insert(Inst);
+      // If all incoming values are available, mark this PHI as
+      // replacable and push it's users into the worklist.
+      if (all_of(PHI->incoming_values(),
+                 [&](Value *V) { return isAvailable(cast<Instruction>(V)); })) {
+        UsersToReplace.insert(PHI);
+        PushUsersToWorklist(PHI);
         continue;
       }
 
-      Worklist.insert(PHI);
-      if (!collectUsersRecursive(*PHI))
-        return false;
-    } else if (auto *SI = dyn_cast<SelectInst>(Inst)) {
-      if (!isa<Instruction>(SI->getTrueValue()) ||
-          !isa<Instruction>(SI->getFalseValue()))
+      // Not all incoming values are available. If this PHI was already
+      // visited prior to this iteration, return false.
+      if (ValuesToRevisit.contains(PHI))
         return false;
 
-      if (!isAvailable(cast<Instruction>(SI->getTrueValue())) ||
-          !isAvailable(cast<Instruction>(SI->getFalseValue()))) {
-        ValuesToRevisit.insert(Inst);
-        continue;
+      // Revisit PHI, after all it's incoming values have been visited.
+      Worklist.push(PHI);
+      ValuesToRevisit.insert(PHI);
+      for (unsigned Idx = 0; Idx < PHI->getNumIncomingValues(); ++Idx) {
+        auto *IncomingValue = cast<Instruction>(PHI->getIncomingValue(Idx));
+        if (UsersToReplace.contains(IncomingValue))
+          continue;
+        if (ValuesToRevisit.contains(IncomingValue))
+          return false;
+        ValuesToRevisit.insert(IncomingValue);
----------------
arsenm wrote:

Check if the insert succeeded instead of double map operation 

https://github.com/llvm/llvm-project/pull/131956