[LowerConstantIntrinsics] Fix heap-use-after-free bug in worklist

[dstenb]

Summary

This fixes PR51730, a heap-use-after-free bug in
replaceConditionalBranchesOnConstant().

With the attached reproducer we were left with a function looking
something like this after replaceAndRecursivelySimplify():

[...]

cont2.i:
  br i1 %.not1.i, label %handler.type_mismatch3.i, label %cont4.i

handler.type_mismatch3.i:
  %3 = phi i1 [ %2, %cont2.thread.i ], [ false, %cont2.i ]
  unreachable

cont4.i:
  unreachable

[...]

with both the branch instruction and PHI node being in the worklist. As
a result of replacing the branch instruction with an unconditional
branch, the PHI node in %handler.type_mismatch3.i would be removed. This
then resulted in a heap-use-after-free bug due to accessing that removed
PHI node in the next worklist iteration.

This is solved by using a value handle worklist. I am a unsure if this
is the most idiomatic solution. Another solution could have been to
produce a worklist just containing the interesting branch instructions,
but I thought that it perhaps was a bit cleaner to keep all worklist
filtering in the loop that does the rewrites.

[nickdesaulniers]

llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp

61–68

Why make a copy of the worklist? Should we just be using dyn_cast_or_null rather than dyn_cast (ie. without all of the other changes)?

[dstenb]

llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp

61–68

The extra code is needed because the worklist would not contain null pointers but rather stale pointers to deleted instructions.

With the attached reproducer we will end up with the following two instructions in the worklist:

* br i1 false, label %handler.type_mismatch3.i, label %cont4.i
  ptr=0x564e36fac680, parent=cont2.i

* %1 = phi i1 [ true, %cont2.thread.i ], [ false, %cont2.i ]
  ptr=0x564e36fac6d8, parent=handler.type_mismatch3.i

In the first iteration, the branch instruction is replaced with an
unconditional branch to %cont4.i. As a part of that, %cont2.i is removed as a
predecessor to %handler.type_mismatch3.i via removePredecessor(), which results
in the PHI node, which is in the worklist, being deleted. Before this patch, we then got a
heap-use-after-free bug in the next iteration when trying to use dyn_cast on the
deleted PHI.

[lebedev.ri]

LGTM
While you could use CallbackVH with a lambda that would erase from the set, the set contains instructions, not VH's, so this seems like the smallest fix.

[dstenb]

Thanks for the reviews! I'll land this shortly.

llvm/lib/Transforms/Scalar/LowerConstantIntrinsics.cpp

if (C->isManifestConstant())
       return ConstantInt::getTrue(II->getType());
   return ConstantInt::getFalse(II->getType());
 }
 
 static bool replaceConditionalBranchesOnConstant(Instruction *II,
                                                  Value *NewValue,
                                                  DomTreeUpdater *DTU) {
   bool HasDeadBlocks = false;
-  SmallSetVector<Instruction *, 8> Worklist;
+  SmallSetVector<Instruction *, 8> UnsimplifiedUsers;
   replaceAndRecursivelySimplify(II, NewValue, nullptr, nullptr, nullptr,
-                                &Worklist);
-  for (auto I : Worklist) {
-    BranchInst *BI = dyn_cast<BranchInst>(I);
+                                &UnsimplifiedUsers);
+  // UnsimplifiedUsers can contain PHI nodes that may be removed when
+  // replacing the branch instructions, so use a value handle worklist
+  // to handle those possibly removed instructions.
+  SmallVector<WeakVH, 8> Worklist(UnsimplifiedUsers.begin(),
+                                  UnsimplifiedUsers.end());
+
+  for (auto &VH : Worklist) {
+    BranchInst *BI = dyn_cast_or_null<BranchInst>(VH);
     if (!BI)
       continue;
     if (BI->isUnconditional())
       continue;
 
     BasicBlock *Target, *Other;
     if (match(BI->getOperand(0), m_Zero())) {
       Target = BI->getSuccessor(1);

llvm/test/Transforms/LowerConstantIntrinsics/stale-worklist-phi.ll

+; RUN: opt -lower-constant-intrinsics -S < %s | FileCheck %s
+
+; This is a reproducer for a heap-use-after-free bug that occured due to trying
+; to process a PHI node that was removed in a preceding worklist iteration. The
+; conditional branch in %cont2.i will be replaced with an unconditional branch
+; to %cont4.i. As a result of that, the PHI node in %handler.type_mismatch3.i
+; will be left with one predecessor and will thus be removed in that iteration.
+
+; CHECK-NOT: phi
+; CHECK: cont2.i:
+; CHECK-NEXT: br label %cont4.i
+; CHECK-NOT: phi
+
+%s = type { [2 x i16] }
+
+define fastcc void @foo(%s* %p) unnamed_addr {
+entry:
+  %0 = bitcast %s* %p to i8*
+  %1 = tail call i32 @llvm.objectsize.i32.p0i8(i8* %0, i1 false, i1 false, i1 false) #2
+  %2 = icmp ne i32 %1, 0
+  %.not1.i = icmp eq i32 %1, 0
+  br label %for.cond
+
+for.cond:                                         ; preds = %entry
+  br label %cont.i
+
+cont.i:                                           ; preds = %for.cond
+  br i1 undef, label %cont2.i, label %cont2.thread.i
+
+cont2.thread.i:                                   ; preds = %cont.i
+  br label %handler.type_mismatch3.i
+
+cont2.i:                                          ; preds = %cont.i
+  br i1 %.not1.i, label %handler.type_mismatch3.i, label %cont4.i
+
+handler.type_mismatch3.i:                         ; preds = %cont2.i, %cont2.thread.i
+  %3 = phi i1 [ %2, %cont2.thread.i ], [ false, %cont2.i ]
+  unreachable
+
+cont4.i:                                          ; preds = %cont2.i
+  unreachable
+}
+
+; Function Attrs: nofree nosync nounwind readnone speculatable willreturn
+declare i32 @llvm.objectsize.i32.p0i8(i8*, i1 immarg, i1 immarg, i1 immarg) #1
+
+attributes #1 = { nofree nosync nounwind readnone speculatable willreturn }
+attributes #2 = { nounwind }