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

llvm-svn: 368792

Author: rjmccall

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();

llvm/test/Transforms/Coroutines/coro-retcon-alloca.ll

@@ -202,6 +202,34 @@ loop2:
 ; CHECK-NEXT:    br i1 %cmp, label %loop2,
 ; CHECK-NEXT:  }
 
+declare {i8*, i32} @prototype_j(i8*)
+define {i8*, i32} @j(i8* %buffer, i32 %n) {
+entry:
+  %id = call token @llvm.coro.id.retcon(i32 1024, i32 8, i8* %buffer, i8* bitcast ({i8*, i32} (i8*)* @prototype_j to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
+  br label %forward
+
+back:
+  ; We should encounter this 'get' before we encounter the 'alloc'.
+  %ptr = call i8* @llvm.coro.alloca.get(token %alloca)
+  call void @use(i8* %ptr)
+  call void @llvm.coro.alloca.free(token %alloca)
+  %k = add i32 %n.val, 1
+  %cmp = icmp ugt i32 %k, 128
+  br i1 %cmp, label %forward, label %end
+
+forward:
+  %n.val = phi i32 [ %n, %entry ], [ %k, %back ]
+  call void (...) @llvm.coro.suspend.retcon.isVoid(i32 %n.val)
+  %alloca = call token @llvm.coro.alloca.alloc.i32(i32 %n.val, i32 8)
+  %inc = add i32 %n.val, 1
+  br label %back
+
+end:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  unreachable
+}
+
 declare token @llvm.coro.id.retcon(i32, i32, i8*, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.suspend.retcon.i1(...)