[llvm] [ValueTracking] Support multiple predecessors in willNotFreeBetween (PR #223580)

[Antonio Frighetto via llvm-commits]

================
@@ -733,50 +733,85 @@ bool llvm::isValidAssumeForContext(const Instruction *Inv,
   return false;
 }
 
-bool llvm::willNotFreeBetween(const Instruction *Assume,
-                              const Instruction *CtxI) {
-  // Helper to check if there are any calls in the range that may free memory.
-  unsigned NumChecked = 0;
-  auto hasNoFreeInRange = [&NumChecked](auto Range) {
-    for (const Instruction &I : Range) {
-      if (NumChecked++ > MaxInstrsToCheckForFree)
-        return false;
-
-      if (auto *CB = dyn_cast<CallBase>(&I)) {
-        if (!CB->hasFnAttr(Attribute::NoFree))
-          return false;
-      } else if (I.maySynchronize())
+static bool hasNoFreeInRange(BasicBlock::const_iterator Begin,
+                             BasicBlock::const_iterator End,
+                             unsigned &NumChecked) {
+  for (const Instruction &I : make_range(Begin, End)) {
+    if (NumChecked++ > MaxInstrsToCheckForFree)
+      return false;
+    if (auto *CB = dyn_cast<CallBase>(&I)) {
+      if (!CB->hasFnAttr(Attribute::NoFree))
         return false;
+    } else if (I.maySynchronize()) {
+      return false;
     }
-    return true;
-  };
+  }
+  return true;
+}
+
+bool llvm::willNotFreeBetween(const Instruction *Assume,
+                              const Instruction *CtxI,
+                              const DominatorTree *DT) {
 
   const BasicBlock *CtxBB = CtxI->getParent();
   const BasicBlock *AssumeBB = Assume->getParent();
+  unsigned NumChecked = 0;
   BasicBlock::const_iterator CtxIter = CtxI->getIterator();
+
   if (CtxBB == AssumeBB) {
-    // Same block case: check that Assume comes before CtxI.
     if (Assume != CtxI && !Assume->comesBefore(CtxI))
       return false;
-    return hasNoFreeInRange(make_range(Assume->getIterator(), CtxIter));
+    return hasNoFreeInRange(Assume->getIterator(), CtxIter, NumChecked);
   }
 
-  // Handle chain of single-predecessor blocks.
-  const BasicBlock *CurBB = CtxBB;
-  while (true) {
-    if (CurBB == AssumeBB)
-      return hasNoFreeInRange(
-          make_range(Assume->getIterator(), AssumeBB->end()));
+  if (DT && !DT->dominates(Assume, CtxI))
+    return false;
+
+  if (!hasNoFreeInRange(CtxBB->begin(), CtxIter, NumChecked))
+    return false;
+  if (pred_empty(CtxBB))
+    return false;
+
+  SmallVector<const BasicBlock *, 16> Worklist;
+  SmallPtrSet<const BasicBlock *, 16> Visited;
+  // Enqueue all predecessors of CtxBB.
+  // Note: CtxBB is NOT pre-inserted to ensure that loop
+  // backedges returning to CtxBB are enqueued and checked correclty.
+  for (const BasicBlock *Pred : predecessors(CtxBB)) {
+    if (Visited.insert(Pred).second)
+      Worklist.push_back(Pred);
+  }
+  while (!Worklist.empty()) {
+    const BasicBlock *CurBB = Worklist.pop_back_val();
 
-    const BasicBlock *PredBB = CurBB->getSinglePredecessor();
-    if (!PredBB)
+    if (CurBB == AssumeBB) {
+      if (!hasNoFreeInRange(Assume->getIterator(), AssumeBB->end(), NumChecked))
+        return false;
+      continue;
+    }
+
+    if (DT && !DT->dominates(AssumeBB, CurBB))
       return false;
 
-    if (!hasNoFreeInRange(make_range(CurBB->begin(),
-                                     CurBB == CtxBB ? CtxIter : CurBB->end())))
+    if (pred_empty(CurBB))
       return false;
-    CurBB = PredBB;
+    // If CurBB == CtxBB (due to a loop backedge targeting CtxBB), check
+    // instructions from CtxIter to the end of CtxBB (instructions before
+    // CtxIter were checked above). Otherwise, check the entire block.
+    BasicBlock::const_iterator Start =
+        (CurBB == CtxBB) ? CtxIter : CurBB->begin();
----------------
antoniofrighetto wrote:

```suggestion
    auto StartIt = (CurBB == CtxBB) ? CtxIter : CurBB->begin();
```

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