From 88690a965892e82cac05a162a9d10e2ce4e2355f Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban@azul.com>
Date: Mon, 14 Sep 2020 11:42:23 +0700
Subject: [PATCH] [CodeGenPrepare] Fix zapping dead operands of assume

This patch fixes a problem of the commit 52cc97a0.
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
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp                |  8 ++++---
 .../X86/recursively-delete-dead-instructions.ll    | 27 ++++++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/CodeGenPrepare/X86/recursively-delete-dead-instructions.ll

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 529975c..bb0bad7 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/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;
     }
 
diff --git a/llvm/test/Transforms/CodeGenPrepare/X86/recursively-delete-dead-instructions.ll b/llvm/test/Transforms/CodeGenPrepare/X86/recursively-delete-dead-instructions.ll
new file mode 100644
index 0000000..0366b7d
--- /dev/null
+++ b/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
+}
-- 
2.7.4