From 1b0c4ab66ce165da09fd1d060de6a48e9a1a228d Mon Sep 17 00:00:00 2001
From: Gor Nishanov <GorNishanov@gmail.com>
Date: Fri, 2 Jun 2017 02:18:36 +0000
Subject: [PATCH] [coroutines] PR33271: Remove stray coro.save intrinsics
 during CoroSplit

Summary:
Optimization passes may remove llvm.coro.suspend intrinsic while leaving matching llvm.coro.save intrinsic orphaned.
Make sure we clean up orphaned coro.saves.  The bug manifested with a crash similar to this:

```
    llvm_unreachable("Unknown type!");
    llvm::MVT::getVT (Ty=0x489518, HandleUnknown=false)
    llvm::EVT::getEVT
    llvm::TargetLoweringBase::getValueType
    llvm::ComputeValueVTs
    llvm::SelectionDAGBuilder::visitTargetIntrinsic
```

Reviewers: GorNishanov

Subscribers: EricWF, llvm-commits

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@304518 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Coroutines/Coroutines.cpp    | 12 ++++++++++++
 test/Transforms/Coroutines/coro-split-02.ll |  5 ++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/Transforms/Coroutines/Coroutines.cpp b/lib/Transforms/Coroutines/Coroutines.cpp
index ea48043f938..44e1f9b404e 100644
--- a/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/lib/Transforms/Coroutines/Coroutines.cpp
@@ -218,6 +218,8 @@ void coro::Shape::buildFrom(Function &F) {
   size_t FinalSuspendIndex = 0;
   clear(*this);
   SmallVector<CoroFrameInst *, 8> CoroFrames;
+  SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+
   for (Instruction &I : instructions(F)) {
     if (auto II = dyn_cast<IntrinsicInst>(&I)) {
       switch (II->getIntrinsicID()) {
@@ -229,6 +231,12 @@ void coro::Shape::buildFrom(Function &F) {
       case Intrinsic::coro_frame:
         CoroFrames.push_back(cast<CoroFrameInst>(II));
         break;
+      case Intrinsic::coro_save:
+        // After optimizations, coro_suspends using this coro_save might have
+        // been removed, remember orphaned coro_saves to remove them later.
+        if (II->use_empty())
+          UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));
+        break;
       case Intrinsic::coro_suspend:
         CoroSuspends.push_back(cast<CoroSuspendInst>(II));
         if (CoroSuspends.back()->isFinal()) {
@@ -311,4 +319,8 @@ void coro::Shape::buildFrom(Function &F) {
   if (HasFinalSuspend &&
       FinalSuspendIndex != CoroSuspends.size() - 1)
     std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
+
+  // Remove orphaned coro.saves.
+  for (CoroSaveInst *CoroSave : UnusedCoroSaves)
+    CoroSave->eraseFromParent();
 }
diff --git a/test/Transforms/Coroutines/coro-split-02.ll b/test/Transforms/Coroutines/coro-split-02.ll
index 953c2508865..4dc8921cd69 100644
--- a/test/Transforms/Coroutines/coro-split-02.ll
+++ b/test/Transforms/Coroutines/coro-split-02.ll
@@ -1,5 +1,6 @@
 ; Tests that coro-split can handle the case when a code after coro.suspend uses
 ; a value produces between coro.save and coro.suspend (%Result.i19)
+; and checks whether stray coro.saves are properly removed
 ; RUN: opt < %s -coro-split -S | FileCheck %s
 
 %"struct.std::coroutine_handle" = type { i8* }
@@ -24,9 +25,10 @@ entry:
     i8 1, label %exit
   ]
 await.ready:
+  %StrayCoroSave = call token @llvm.coro.save(i8* null)
   %val = load i32, i32* %Result.i19
   call void @print(i32 %val)
-  br label %exit  
+  br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
@@ -35,6 +37,7 @@ exit:
 ; CHECK-LABEL: @a.resume(
 ; CHECK:         getelementptr inbounds %a.Frame
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
+; CHECK-NOT:     call token @llvm.coro.save(i8* null)
 ; CHECK-NEXT:    %val = load i32, i32* %Result
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void
-- 
2.40.0