[CodeGenPrepare] Fix zapping dead operands of assume

This patch fixes a problem of the commit 52cc97a.
A test case is created to demonstrate the crash caused by
the instruction iterator invalidated by the recursive
removal of dead operands of assume. The solution restarts
from the blocks's first instruction in case CurInstIterator
is invalidated by RecursivelyDeleteTriviallyDeadInstructions().

Reviewed By: bkramer

Differential Revision: https://reviews.llvm.org/D87434

Author: Yevgeny Rouban

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -2047,9 +2047,11 @@ bool CodeGenPrepare::optimizeCallInst(CallInst *CI, bool &ModifiedDT) {
       Value *Operand = II->getOperand(0);
       II->eraseFromParent();
       // Prune the operand, it's most likely dead.
-      RecursivelyDeleteTriviallyDeadInstructions(
-          Operand, TLInfo, nullptr,
-          [&](Value *V) { removeAllAssertingVHReferences(V); });
+      resetIteratorIfInvalidatedWhileCalling(BB, [&]() {
+        RecursivelyDeleteTriviallyDeadInstructions(
+            Operand, TLInfo, nullptr,
+            [&](Value *V) { removeAllAssertingVHReferences(V); });
+      });
       return true;
     }
 

llvm/test/Transforms/CodeGenPrepare/X86/recursively-delete-dead-instructions.ll

@@ -0,0 +1,27 @@
+; RUN: opt -codegenprepare -S -mtriple=x86_64-linux < %s | FileCheck %s
+
+declare void @llvm.assume(i1 noundef) nounwind willreturn
+
+; Recursively deleting dead operands of assume() may result in its next
+; instruction deleted and the iterator pointing to the next instruction
+; invalidated. This prevents the following simple loop in
+; CodeGenPrepare::optimizeBlock() unless CurInstIterator is fixed:
+;
+;   CurInstIterator = BB.begin();
+;   while (CurInstIterator != BB.end())
+;     optimizeInst(&*CurInstIterator++, ModifiedDT);
+;
+define i32 @test_assume_in_loop(i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_assume_in_loop(
+; CHECK-NEXT:  entry:
+entry:
+  br label %loop
+
+; CHECK: loop:
+; CHECK-NEXT:  br label %loop
+loop:
+  %cond3 = phi i1 [%cond1, %entry], [%cond4, %loop]
+  call void @llvm.assume(i1 %cond3)
+  %cond4 = icmp ult i1 %cond1, %cond2
+  br label %loop
+}