[Coroutines] Don't transform cmpinst prematurely in simplifyTerminatorLeadingToRet

Previously, we would try to transform cmpinst in
simplifyTerminatorLeadingToRet if we found it was a constant. However,
this is incorrect.

Since the resolved constants in simplifyTerminatorLeadingToRet are not
truely constants. They are basically constants along cerntain code
paths.

In this way, it is clearly incorrect to transform the compare
instruction to a constant.

It will cause confusing miscompilations. This patch tries to fix this.

Author: ChuanqiXu9

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -1249,8 +1249,11 @@ scanPHIsAndUpdateValueMap(Instruction *Prev, BasicBlock *NewBlock,
 // instruction. Suspend instruction represented by a switch, track the PHI
 // values and select the correct case successor when possible.
 static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
+  // There is nothing to simplify.
+  if (isa<ReturnInst>(InitialInst))
+    return false;
+
   DenseMap<Value *, Value *> ResolvedValues;
-  BasicBlock *UnconditionalSucc = nullptr;
   assert(InitialInst->getModule());
   const DataLayout &DL = InitialInst->getModule()->getDataLayout();
 
@@ -1280,39 +1283,35 @@ static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
   Instruction *I = InitialInst;
   while (I->isTerminator() || isa<CmpInst>(I)) {
     if (isa<ReturnInst>(I)) {
-      if (I != InitialInst) {
-        // If InitialInst is an unconditional branch,
-        // remove PHI values that come from basic block of InitialInst
-        if (UnconditionalSucc)
-          UnconditionalSucc->removePredecessor(InitialInst->getParent(), true);
-        ReplaceInstWithInst(InitialInst, I->clone());
-      }
+      ReplaceInstWithInst(InitialInst, I->clone());
       return true;
     }
+
     if (auto *BR = dyn_cast<BranchInst>(I)) {
-      if (BR->isUnconditional()) {
-        BasicBlock *Succ = BR->getSuccessor(0);
-        if (I == InitialInst)
-          UnconditionalSucc = Succ;
-        scanPHIsAndUpdateValueMap(I, Succ, ResolvedValues);
-        I = GetFirstValidInstruction(Succ->getFirstNonPHIOrDbgOrLifetime());
-        continue;
+      unsigned SuccIndex = 0;
+      if (BR->isConditional()) {
+        // Handle the case the condition of the conditional branch is constant.
+        // e.g.,
+        //
+        //     br i1 false, label %cleanup, label %CoroEnd
+        //
+        // It is possible during the transformation. We could continue the
+        // simplifying in this case.
+        ConstantInt *Cond = TryResolveConstant(BR->getCondition());
+        if (!Cond)
+          return false;
+
+        SuccIndex = Cond->isOne() ? 0 : 1;
       }
 
-      BasicBlock *BB = BR->getParent();
-      // Handle the case the condition of the conditional branch is constant.
-      // e.g.,
-      //
-      //     br i1 false, label %cleanup, label %CoroEnd
-      //
-      // It is possible during the transformation. We could continue the
-      // simplifying in this case.
-      if (ConstantFoldTerminator(BB, /*DeleteDeadConditions=*/true)) {
-        // Handle this branch in next iteration.
-        I = BB->getTerminator();
-        continue;
-      }
-    } else if (auto *CondCmp = dyn_cast<CmpInst>(I)) {
+      BasicBlock *Succ = BR->getSuccessor(SuccIndex);
+      scanPHIsAndUpdateValueMap(I, Succ, ResolvedValues);
+      I = GetFirstValidInstruction(Succ->getFirstNonPHIOrDbgOrLifetime());
+
+      continue;
+    }
+
+    if (auto *CondCmp = dyn_cast<CmpInst>(I)) {
       // If the case number of suspended switch instruction is reduced to
       // 1, then it is simplified to CmpInst in llvm::ConstantFoldTerminator.
       auto *BR = dyn_cast<BranchInst>(
@@ -1336,13 +1335,14 @@ static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
       if (!ConstResult)
         return false;
 
-      CondCmp->replaceAllUsesWith(ConstResult);
-      CondCmp->eraseFromParent();
+      ResolvedValues[BR->getCondition()] = ConstResult;
 
       // Handle this branch in next iteration.
       I = BR;
       continue;
-    } else if (auto *SI = dyn_cast<SwitchInst>(I)) {
+    }
+
+    if (auto *SI = dyn_cast<SwitchInst>(I)) {
       ConstantInt *Cond = TryResolveConstant(SI->getCondition());
       if (!Cond)
         return false;
@@ -1352,9 +1352,8 @@ static bool simplifyTerminatorLeadingToRet(Instruction *InitialInst) {
       I = GetFirstValidInstruction(BB->getFirstNonPHIOrDbgOrLifetime());
       continue;
     }
-
-    return false;
   }
+
   return false;
 }
 

llvm/test/Transforms/Coroutines/coro-split-musttail12.ll

@@ -0,0 +1,83 @@
+; Tests that coro-split won't convert the cmp instruction prematurely.
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+declare void @fakeresume1(ptr)
+declare void @print()
+
+define void @f(i1 %cond) #0 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %alloc = call ptr @malloc(i64 16) #3
+  %vFrame = call noalias nonnull ptr @llvm.coro.begin(token %id, ptr %alloc)
+
+  %save = call token @llvm.coro.save(ptr null)
+
+  %init_suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %init_suspend, label %coro.end [
+    i8 0, label %await.ready
+    i8 1, label %coro.end
+  ]
+await.ready:
+  %save2 = call token @llvm.coro.save(ptr null)
+  br i1 %cond, label %then, label %else
+
+then:
+  call fastcc void @fakeresume1(ptr align 8 null)
+  br label %merge
+
+else:
+  br label %merge
+
+merge:
+  %v0 = phi i1 [0, %then], [1, %else]
+  br label %compare
+
+compare:
+  %cond.cmp = icmp eq i1 %v0, 0
+  br i1 %cond.cmp, label %ready, label %prepare
+
+prepare:
+  call void @print()
+  br label %ready
+
+ready:
+  %suspend = call i8 @llvm.coro.suspend(token %save2, i1 true)
+  %switch = icmp ult i8 %suspend, 2
+  br i1 %switch, label %cleanup, label %coro.end
+
+cleanup:
+  %free.handle = call ptr @llvm.coro.free(token %id, ptr %vFrame)
+  %.not = icmp eq ptr %free.handle, null
+  br i1 %.not, label %coro.end, label %coro.free
+
+coro.free:
+  call void @delete(ptr nonnull %free.handle) #2
+  br label %coro.end
+
+coro.end:
+  call i1 @llvm.coro.end(ptr null, i1 false)
+  ret void
+}
+
+; CHECK-LABEL: @f.resume(
+; CHECK-NOT:      }
+; CHECK:          call void @print()
+
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #1
+declare i1 @llvm.coro.alloc(token) #2
+declare i64 @llvm.coro.size.i64() #3
+declare ptr @llvm.coro.begin(token, ptr writeonly) #2
+declare token @llvm.coro.save(ptr) #2
+declare ptr @llvm.coro.frame() #3
+declare i8 @llvm.coro.suspend(token, i1) #2
+declare ptr @llvm.coro.free(token, ptr nocapture readonly) #1
+declare i1 @llvm.coro.end(ptr, i1) #2
+declare ptr @llvm.coro.subfn.addr(ptr nocapture readonly, i8) #1
+declare ptr @malloc(i64)
+declare void @delete(ptr nonnull) #2
+
+attributes #0 = { presplitcoroutine }
+attributes #1 = { argmemonly nounwind readonly }
+attributes #2 = { nounwind }
+attributes #3 = { nounwind readnone }