[llvm] [DropUnnecessaryAssumes] Make the ephemeral value check more precise (PR #160700)

[Nikita Popov via llvm-commits]

================
@@ -17,13 +18,44 @@ using namespace llvm;
 using namespace llvm::PatternMatch;
 
 static bool affectedValuesAreEphemeral(ArrayRef<Value *> Affected) {
-  // If all the affected uses have only one use (part of the assume), then
-  // the assume does not provide useful information. Note that additional
-  // users may appear as a result of inlining and CSE, so we should only
-  // make this assumption late in the optimization pipeline.
-  // TODO: Handle dead cyclic usages.
-  // TODO: Handle multiple dead assumes on the same value.
-  return all_of(Affected, match_fn(m_OneUse(m_Value())));
+  // Check whether all the uses are ephemeral, i.e. recursively only used
+  // by assumes. In that case, the assume does not provide useful information.
+  // Note that additional users may appear as a result of inlining and CSE,
+  // so we should only make this assumption late in the optimization pipeline.
+  SmallSetVector<User *, 16> Worklist;
+  auto AddUser = [&](User *U) {
+    // Bail out if we need to inspect too many users.
+    if (Worklist.size() >= 32)
+      return false;
+    Worklist.insert(U);
+    return true;
+  };
+
+  for (Value *V : Affected)
+    for (User *U : V->users())
+      if (!AddUser(U))
+        return false;
+
+  for (unsigned Idx = 0; Idx < Worklist.size(); ++Idx) {
+    auto *I = dyn_cast<Instruction>(Worklist[Idx]);
+    if (!I)
+      return false;
----------------
nikic wrote:

I'm generally not a fan of fully flattening users up-front, as it means we can't have early exits.

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