[coroutines] Self-review: two combined fixes

[GorNishanov]

Summary

TODO: add unit tests

34267

https://bugs.llvm.org/show_bug.cgi?id=34267

[GorNishanov]

This got fixed with:

https://reviews.llvm.org/D39768

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@317981 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Coroutines/CoroFrame.cpp

 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/circular_raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/PromoteMemToReg.h"
 
 using namespace llvm;
 
 // The "coro-suspend-crossing" flag is very noisy. There is another debug type,
 // "coro-frame", which results in leaner debug spew.
 #define DEBUG_TYPE "coro-suspend-crossing"
 
 enum { SmallVectorThreshold = 32 };

 static bool materializable(Instruction &V) {
   return isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
          isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V);
 }
 
 // Check for structural coroutine intrinsics that should not be spilled into
 // the coroutine frame.
 static bool isCoroutineStructureIntrinsic(Instruction &I) {
-  return isa<CoroIdInst>(&I) || isa<CoroBeginInst>(&I) ||
+  return isa<CoroIdInst>(&I) ||
          isa<CoroSaveInst>(&I) || isa<CoroSuspendInst>(&I);
 }
 
 // For every use of the value that is across suspend point, recreate that value
 // after a suspend point.
 static void rewriteMaterializableInstructions(IRBuilder<> &IRB,
                                               SpillInfo const &Spills) {
   BasicBlock *CurrentBlock = nullptr;
if (&BB->front() == I) {
     if (BB->getSinglePredecessor()) {
       BB->setName(Name);
       return BB;
     }
   }
   return BB->splitBasicBlock(I, Name);
 }
 
+static bool onlyStoresConstants(AllocaInst *AI) {
+  if (!AI->getAllocatedType()->isIntegerTy())
+    return false;
+
+  for (const User *U : AI->users())
+    if (const StoreInst *SI = dyn_cast<StoreInst>(U))
+      if (!isa<ConstantInt>(SI->getOperand(0)))
+        return false;
+
+  return true;
+}
+
+// Allocas that follow the pattern of cleanup.dest.slot, namely, they are
+// integer allocas that don't escape and only store constants. We don't want
+// them to be saved into the coroutine frame (as some of the cleanup code will
+// access them after coroutine frame destroyed).
+static void promoteCleanupSlotLikeAllocasToRegisters(Function &F, AllocaInst* AI) {
+  SmallVector<AllocaInst *, 2> Allocas;
+  BasicBlock &BB = F.getEntryBlock(); // Get the entry node for the function
+
+  // Find allocas that are safe to promote, by looking at all instructions in
+  // the entry node
+  for (BasicBlock::iterator I = BB.begin(), E = --BB.end(); I != E; ++I)
+    if (AllocaInst *AI = dyn_cast<AllocaInst>(I)) // Is it an alloca?
+      if (isAllocaPromotable(AI) && onlyStoresConstants(AI))
+        Allocas.push_back(AI);
+
+  if (Allocas.empty())
+      return;
+
+  DominatorTree DT(F);
+  PromoteMemToReg(Allocas, DT);
+}
+
 // Split above and below a particular instruction so that it
 // will be all alone by itself in a block.
 static void splitAround(Instruction *I, const Twine &Name) {
   splitBlockIfNotFirst(I, Name);
   splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
 }
 
 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   // Make sure that all coro.save, coro.suspend and the fallthrough coro.end
   // intrinsics are in their own blocks to simplify the logic of building up
   // SuspendCrossing data.
   for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
     splitAround(CSI->getCoroSave(), "CoroSave");
     splitAround(CSI, "CoroSuspend");
   }
 
+  // Promote cleanup.slot like allocas to registers
+  promoteCleanupSlotLikeAllocasToRegisters(F, Shape.PromiseAlloca);
+
   // Put CoroEnds into their own blocks.
   for (CoroEndInst *CE : Shape.CoroEnds)
     splitAround(CE, "CoroEnd");
 
   // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
   // never has its definition separated from the PHI by the suspend point.
   rewritePHIs(F);
 
void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   for (Argument &A : F.args())
     for (User *U : A.users())
       if (Checker.isDefinitionAcrossSuspend(A, U))
         Spills.emplace_back(&A, U);
 
   for (Instruction &I : instructions(F)) {
     // Values returned from coroutine structure intrinsics should not be part
     // of the Coroutine Frame.
-    if (isCoroutineStructureIntrinsic(I))
+    if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
       continue;
     // The Coroutine Promise always included into coroutine frame, no need to
     // check for suspend crossing.
     if (Shape.PromiseAlloca == &I)
       continue;
 
     for (User *U : I.users())
       if (Checker.isDefinitionAcrossSuspend(I, U)) {