Fix a use-after-free in the coro.alloca treatment.

llvm-svn: 368792

1 files changed, 10 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
index 0d0ca511ca2..52c167c9f71 100644
--- a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1042,7 +1042,8 @@ static bool localAllocaNeedsStackSave(CoroAllocaAllocInst *AI) {
 
 /// Turn each of the given local allocas into a normal (dynamic) alloca
 /// instruction.
-static void lowerLocalAllocas(ArrayRef<CoroAllocaAllocInst*> LocalAllocas) {
+static void lowerLocalAllocas(ArrayRef<CoroAllocaAllocInst*> LocalAllocas,
+                              SmallVectorImpl<Instruction*> &DeadInsts) {
   for (auto AI : LocalAllocas) {
     auto M = AI->getModule();
     IRBuilder<> Builder(AI);
@@ -1075,10 +1076,10 @@ static void lowerLocalAllocas(ArrayRef<CoroAllocaAllocInst*> LocalAllocas) {
                              StackSave);
         }
       }
-      cast<Instruction>(U)->eraseFromParent();
+      DeadInsts.push_back(cast<Instruction>(U));
     }
 
-    AI->eraseFromParent();
+    DeadInsts.push_back(AI);
   }
 }
 
@@ -1201,6 +1202,11 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
       continue;
     }
 
+    // Ignore alloca.get; we process this as part of coro.alloca.alloc.
+    if (isa<CoroAllocaGetInst>(I)) {
+      continue;
+    }
+
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {
         // We cannot spill a token.
@@ -1214,7 +1220,7 @@ void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   moveSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, Spills);
   Shape.FramePtr = insertSpills(Spills, Shape);
-  lowerLocalAllocas(LocalAllocas);
+  lowerLocalAllocas(LocalAllocas, DeadInstructions);
 
   for (auto I : DeadInstructions)
     I->eraseFromParent();