[Coroutines] Refactor Spill and Alloca information

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

do {
     }
   } while (Changed);
   LLVM_DEBUG(dump());
 }
 
 #undef DEBUG_TYPE // "coro-suspend-crossing"
 #define DEBUG_TYPE "coro-frame"
 
-// We build up the list of spills for every case where a use is separated
-// from the definition by a suspend point.
-
-static const unsigned InvalidFieldIndex = ~0U;
-
 namespace {
-class Spill {
-  Value *Def = nullptr;
-  Instruction *User = nullptr;
-  unsigned FieldNo = InvalidFieldIndex;
-
-public:
-  Spill(Value *Def, llvm::User *U) : Def(Def), User(cast<Instruction>(U)) {}
+class FrameTypeBuilder;
+// Mapping from the to-be-spilled value to all the users that need reload.
+using SpillInfo = SmallMapVector<Value *, SmallVector<Instruction *, 2>, 8>;
+struct FrameDataInfo {
+  // All the values (that are not allocas) that needs to be spilled to the
+  // frame.
+  SpillInfo Spills;
+  // Allocas contains all values defined as allocas that need to live in the
+  // frame.
+  SmallVector<AllocaInst *, 8> Allocas;
 
-  Value *def() const { return Def; }
-  Instruction *user() const { return User; }
-  BasicBlock *userBlock() const { return User->getParent(); }
-
-  // Note that field index is stored in the first SpillEntry for a particular
-  // definition. Subsequent mentions of a defintion do not have fieldNo
-  // assigned. This works out fine as the users of Spills capture the info about
-  // the definition the first time they encounter it. Consider refactoring
-  // SpillInfo into two arrays to normalize the spill representation.
-  unsigned fieldIndex() const {
-    assert(FieldNo != InvalidFieldIndex && "Accessing unassigned field");
-    return FieldNo;
-  }
-  void setFieldIndex(unsigned FieldNumber) {
-    assert(FieldNo == InvalidFieldIndex && "Reassigning field number");
-    FieldNo = FieldNumber;
+  SmallVector<Value *, 8> getAllDefs() const {
+    SmallVector<Value *, 8> Defs;
+    for (auto &P : Spills)
+      Defs.push_back(P.first);
+    for (auto *A : Allocas)
+      Defs.push_back(A);
+    return Defs;
+  }
+
+  uint32_t getFieldIndex(Value *V) const {
+    auto Itr = FieldIndexMap.find(V);
+    assert(Itr != FieldIndexMap.end() &&
+           "Value does not have a frame field index");
+    return Itr->second;
+  }
+
+  void setFieldIndex(Value *V, uint32_t Index) {
+    assert((LayoutIndexUpdateStarted || FieldIndexMap.count(V) == 0) &&
+           "Cannot set the index for the same field twice.");
+    FieldIndexMap[V] = Index;
   }
+
+  // Remap the index of every field in the frame, using the final layout index.
+  void updateLayoutIndex(FrameTypeBuilder &B);
+
+private:
+  // LayoutIndexUpdateStarted is used to avoid updating the index of any field
+  // twice by mistake.
+  bool LayoutIndexUpdateStarted = false;
+  // Map from values to their slot indexes on the frame. They will be first set
+  // with their original insertion field index. After the frame is built, their
+  // indexes will be updated into the final layout index.
+  DenseMap<Value *, uint32_t> FieldIndexMap;
 };
 } // namespace
 
-// Note that there may be more than one record with the same value of Def in
-// the SpillInfo vector.
-using SpillInfo = SmallVector<Spill, 8>;
-
 #ifndef NDEBUG
-static void dump(StringRef Title, SpillInfo const &Spills) {
+static void dumpSpills(StringRef Title, const SpillInfo &Spills) {
   dbgs() << "------------- " << Title << "--------------\n";
-  Value *CurrentValue = nullptr;
-  for (auto const &E : Spills) {
-    if (CurrentValue != E.def()) {
-      CurrentValue = E.def();
-      CurrentValue->dump();
-    }
+  for (const auto &E : Spills) {
+    E.first->dump();
     dbgs() << "   user: ";
-    E.user()->dump();
+    for (auto *I : E.second)
+      I->dump();
   }
 }
+
+static void dumpAllocas(const SmallVectorImpl<AllocaInst *> &Allocas) {
+  dbgs() << "------------- Allocas --------------\n";
+  for (auto *A : Allocas)
+    A->dump();
+}
 #endif
 
 namespace {
+using FieldIDType = size_t;
 // We cannot rely solely on natural alignment of a type when building a
 // coroutine frame and if the alignment specified on the Alloca instruction
 // differs from the natural alignment of the alloca type we will need to insert
 // padding.
 class FrameTypeBuilder {
-public:
-  using ForSpillType = SmallVector<Spill *, 8>;
-
 private:
   struct Field {
     uint64_t Size;
     uint64_t Offset;
-    ForSpillType ForSpill;
     Type *Ty;
-    unsigned FieldIndex;
+    FieldIDType LayoutFieldIndex;
     Align Alignment;
     Align TyAlignment;
   };
 
   const DataLayout &DL;
   LLVMContext &Context;
   uint64_t StructSize = 0;
   Align StructAlign;
   bool IsFinished = false;
 
   SmallVector<Field, 8> Fields;
   DenseMap<Value*, unsigned> FieldIndexByKey;
 
 public:
   FrameTypeBuilder(LLVMContext &Context, DataLayout const &DL)
       : DL(DL), Context(Context) {}
 
-  class FieldId {
-    size_t Value;
-    explicit FieldId(size_t Value) : Value(Value) {}
-
-    friend class FrameTypeBuilder;
-  };
-
   /// Add a field to this structure for the storage of an `alloca`
   /// instruction.
-  FieldId addFieldForAlloca(AllocaInst *AI, ForSpillType ForSpill = {},
+  LLVM_NODISCARD FieldIDType addFieldForAlloca(AllocaInst *AI,
                                                bool IsHeader = false) {
     Type *Ty = AI->getAllocatedType();
 
     // Make an array type if this is a static array allocation.
     if (AI->isArrayAllocation()) {
       if (auto *CI = dyn_cast<ConstantInt>(AI->getArraySize()))
         Ty = ArrayType::get(Ty, CI->getValue().getZExtValue());
       else
         report_fatal_error("Coroutines cannot handle non static allocas yet");
     }
 
-    return addField(Ty, AI->getAlign(), ForSpill, IsHeader);
+    return addField(Ty, AI->getAlign(), IsHeader);
   }
 
   /// We want to put the allocas whose lifetime-ranges are not overlapped
   /// into one slot of coroutine frame.
   /// Consider the example at:https://bugs.llvm.org/show_bug.cgi?id=45566
   ///
   ///     cppcoro::task<void> alternative_paths(bool cond) {
   ///         if (cond) {
public:
   /// This function use StackLifetime algorithm to partition the AllocaInsts in
   /// Spills to non-overlapped sets in order to put Alloca in the same
   /// non-overlapped set into the same slot in the Coroutine Frame. Then add
   /// field for the allocas in the same non-overlapped set by using the largest
   /// type as the field type.
   ///
   /// Side Effects: Because We sort the allocas, the order of allocas in the
   /// frame may be different with the order in the source code.
-  void addFieldForAllocas(const Function &F, SpillInfo &Spills,
+  void addFieldForAllocas(const Function &F, FrameDataInfo &FrameData,
                           coro::Shape &Shape);
 
   /// Add a field to this structure.
-  FieldId addField(Type *Ty, MaybeAlign FieldAlignment,
-                   ForSpillType ForSpill = {}, bool IsHeader = false) {
+  LLVM_NODISCARD FieldIDType addField(Type *Ty, MaybeAlign FieldAlignment,
+                                      bool IsHeader = false) {
     assert(!IsFinished && "adding fields to a finished builder");
     assert(Ty && "must provide a type for a field");
 
     // The field size is always the alloc size of the type.
     uint64_t FieldSize = DL.getTypeAllocSize(Ty);
 
     // The field alignment might not be the type alignment, but we need
     // to remember the type alignment anyway to build the type.
     Align TyAlignment = DL.getABITypeAlign(Ty);
     if (!FieldAlignment) FieldAlignment = TyAlignment;
 
     // Lay out header fields immediately.
     uint64_t Offset;
     if (IsHeader) {
       Offset = alignTo(StructSize, FieldAlignment);
       StructSize = Offset + FieldSize;
 
     // Everything else has a flexible offset.
     } else {
       Offset = OptimizedStructLayoutField::FlexibleOffset;
     }
 
-    Fields.push_back({FieldSize, Offset, ForSpill, Ty, 0,
-                      *FieldAlignment, TyAlignment});
-    return FieldId(Fields.size() - 1);
+    Fields.push_back({FieldSize, Offset, Ty, 0, *FieldAlignment, TyAlignment});
+    return Fields.size() - 1;
   }
 
   /// Finish the layout and set the body on the given type.
   void finish(StructType *Ty);
 
   uint64_t getStructSize() const {
     assert(IsFinished && "not yet finished!");
     return StructSize;
   }
 
   Align getStructAlign() const {
     assert(IsFinished && "not yet finished!");
     return StructAlign;
   }
 
-  unsigned getFieldIndex(FieldId Id) const {
+  FieldIDType getLayoutFieldIndex(FieldIDType Id) const {
     assert(IsFinished && "not yet finished!");
-    return Fields[Id.Value].FieldIndex;
+    return Fields[Id].LayoutFieldIndex;
   }
 };
 } // namespace
 
-void FrameTypeBuilder::addFieldForAllocas(const Function &F, SpillInfo &Spills,
+void FrameDataInfo::updateLayoutIndex(FrameTypeBuilder &B) {
+  auto Updater = [&](Value *I) {
+    setFieldIndex(I, B.getLayoutFieldIndex(getFieldIndex(I)));
+  };
+  LayoutIndexUpdateStarted = true;
+  for (auto &S : Spills)
+    Updater(S.first);
+  for (auto *A : Allocas)
+    Updater(A);
+  LayoutIndexUpdateStarted = false;
+}
+
+void FrameTypeBuilder::addFieldForAllocas(const Function &F,
+                                          FrameDataInfo &FrameData,
                                           coro::Shape &Shape) {
   DenseMap<AllocaInst *, unsigned int> AllocaIndex;
-  SmallVector<AllocaInst *, 8> Allocas;
-  DenseMap<AllocaInst *, Spill *> SpillOfAllocas;
   using AllocaSetType = SmallVector<AllocaInst *, 4>;
   SmallVector<AllocaSetType, 4> NonOverlapedAllocas;
 
   // We need to add field for allocas at the end of this function. However, this
   // function has multiple exits, so we use this helper to avoid redundant code.
   struct RTTIHelper {
     std::function<void()> func;
     RTTIHelper(std::function<void()> &&func) : func(func) {}
     ~RTTIHelper() { func(); }
   } Helper([&]() {
-    for (auto AllocaSet : NonOverlapedAllocas) {
-      ForSpillType ForSpills;
-      for (auto Alloca : AllocaSet)
-        ForSpills.push_back(SpillOfAllocas[Alloca]);
-      auto *LargestAI = *AllocaSet.begin();
-      addFieldForAlloca(LargestAI, ForSpills);
+    for (auto AllocaList : NonOverlapedAllocas) {
+      auto *LargestAI = *AllocaList.begin();
+      FieldIDType Id = addFieldForAlloca(LargestAI);
+      for (auto *Alloca : AllocaList)
+        FrameData.setFieldIndex(Alloca, Id);
     }
   });
 
-  for (auto &Spill : Spills)
-    if (AllocaInst *AI = dyn_cast<AllocaInst>(Spill.def()))
-      if (find(Allocas, AI) == Allocas.end()) {
-        SpillOfAllocas[AI] = &Spill;
-        Allocas.emplace_back(AI);
-      }
-
   if (!Shape.ReuseFrameSlot && !EnableReuseStorageInFrame) {
-    for (auto Alloca : Allocas) {
+    for (auto Alloca : FrameData.Allocas) {
       AllocaIndex[Alloca] = NonOverlapedAllocas.size();
       NonOverlapedAllocas.emplace_back(AllocaSetType(1, Alloca));
     }
     return;
   }
 
   // Because there are pathes from the lifetime.start to coro.end
   // for each alloca, the liferanges for every alloca is overlaped
for (auto U : CoroSuspendInst->users()) {
       if (auto *ConstSWI = dyn_cast<SwitchInst>(U)) {
         auto *SWI = const_cast<SwitchInst *>(ConstSWI);
         DefaultSuspendDest[SWI] = SWI->getDefaultDest();
         SWI->setDefaultDest(SWI->getSuccessor(1));
       }
     }
   }
 
-  StackLifetime StackLifetimeAnalyzer(F, Allocas,
+  StackLifetime StackLifetimeAnalyzer(F, FrameData.Allocas,
                                       StackLifetime::LivenessType::May);
   StackLifetimeAnalyzer.run();
   auto IsAllocaInferenre = [&](const AllocaInst *AI1, const AllocaInst *AI2) {
     return StackLifetimeAnalyzer.getLiveRange(AI1).overlaps(
         StackLifetimeAnalyzer.getLiveRange(AI2));
   };
   auto GetAllocaSize = [&](const AllocaInst *AI) {
     Optional<uint64_t> RetSize = AI->getAllocationSizeInBits(DL);
     assert(RetSize && "We can't handle scalable type now.\n");
     return RetSize.getValue();
   };
   // Put larger allocas in the front. So the larger allocas have higher
   // priority to merge, which can save more space potentially. Also each
   // AllocaSet would be ordered. So we can get the largest Alloca in one
   // AllocaSet easily.
-  sort(Allocas, [&](auto Iter1, auto Iter2) {
+  sort(FrameData.Allocas, [&](auto Iter1, auto Iter2) {
     return GetAllocaSize(Iter1) > GetAllocaSize(Iter2);
   });
-  for (auto Alloca : Allocas) {
+  for (auto Alloca : FrameData.Allocas) {
     bool Merged = false;
     // Try to find if the Alloca is not inferenced with any existing
     // NonOverlappedAllocaSet. If it is true, insert the alloca to that
     // NonOverlappedAllocaSet.
     for (auto &AllocaSet : NonOverlapedAllocas) {
       assert(!AllocaSet.empty() && "Processing Alloca Set is not empty.\n");
       bool CouldMerge = none_of(AllocaSet, [&](auto Iter) {
         return IsAllocaInferenre(Alloca, Iter);
for (auto &LayoutField : LayoutFields) {
     // get from aligning to the field type's natural alignment.
     assert(Offset >= LastOffset);
     if (Offset != LastOffset) {
       if (Packed || alignTo(LastOffset, F.TyAlignment) != Offset)
         FieldTypes.push_back(ArrayType::get(Type::getInt8Ty(Context),
                                             Offset - LastOffset));
     }
 
-    // Record the layout information into both the Field and the
-    // original Spill, if there is one.
     F.Offset = Offset;
-    F.FieldIndex = FieldTypes.size();
-    for (auto Spill : F.ForSpill) {
-      Spill->setFieldIndex(F.FieldIndex);
-    }
+    F.LayoutFieldIndex = FieldTypes.size();
 
     FieldTypes.push_back(F.Ty);
     LastOffset = Offset + F.Size;
   }
 
   Ty->setBody(FieldTypes, Packed);
 
 #ifndef NDEBUG
   // Check that the IR layout matches the offsets we expect.
   auto Layout = DL.getStructLayout(Ty);
   for (auto &F : Fields) {
-    assert(Ty->getElementType(F.FieldIndex) == F.Ty);
-    assert(Layout->getElementOffset(F.FieldIndex) == F.Offset);
+    assert(Ty->getElementType(F.LayoutFieldIndex) == F.Ty);
+    assert(Layout->getElementOffset(F.LayoutFieldIndex) == F.Offset);
   }
 #endif
 
   IsFinished = true;
 }
 
 // Build a struct that will keep state for an active coroutine.
 //   struct f.frame {
 //     ResumeFnTy ResumeFnAddr;
 //     ResumeFnTy DestroyFnAddr;
 //     int ResumeIndex;
 //     ... promise (if present) ...
 //     ... spills ...
 //   };
 static StructType *buildFrameType(Function &F, coro::Shape &Shape,
-                                  SpillInfo &Spills) {
+                                  FrameDataInfo &FrameData) {
   LLVMContext &C = F.getContext();
   const DataLayout &DL = F.getParent()->getDataLayout();
   StructType *FrameTy = [&] {
     SmallString<32> Name(F.getName());
     Name.append(".Frame");
     return StructType::create(C, Name);
   }();
 
   FrameTypeBuilder B(C, DL);
 
   AllocaInst *PromiseAlloca = Shape.getPromiseAlloca();
-  Optional<FrameTypeBuilder::FieldId> PromiseFieldId;
-  Optional<FrameTypeBuilder::FieldId> SwitchIndexFieldId;
+  Optional<FieldIDType> SwitchIndexFieldId;
 
   if (Shape.ABI == coro::ABI::Switch) {
     auto *FramePtrTy = FrameTy->getPointerTo();
     auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
                                    /*IsVarArg=*/false);
     auto *FnPtrTy = FnTy->getPointerTo();
 
     // Add header fields for the resume and destroy functions.
     // We can rely on these being perfectly packed.
-    B.addField(FnPtrTy, None, {}, /*header*/ true);
-    B.addField(FnPtrTy, None, {}, /*header*/ true);
+    (void)B.addField(FnPtrTy, None, /*header*/ true);
+    (void)B.addField(FnPtrTy, None, /*header*/ true);
 
-    // Add a header field for the promise if there is one.
-    if (PromiseAlloca) {
-      PromiseFieldId = B.addFieldForAlloca(PromiseAlloca, {}, /*header*/ true);
-    }
+    // PromiseAlloca field needs to be explicitly added here because it's
+    // not in FrameData.Allocas.
+    if (PromiseAlloca)
+      FrameData.setFieldIndex(
+          PromiseAlloca, B.addFieldForAlloca(PromiseAlloca, /*header*/ true));
 
     // Add a field to store the suspend index.  This doesn't need to
     // be in the header.
     unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
     Type *IndexType = Type::getIntNTy(C, IndexBits);
 
     SwitchIndexFieldId = B.addField(IndexType, None);
   } else {
     assert(PromiseAlloca == nullptr && "lowering doesn't support promises");
   }
 
   // Because multiple allocas may own the same field slot,
   // we add allocas to field here.
-  B.addFieldForAllocas(F, Spills, Shape);
-  Value *CurrentDef = nullptr;
-  // Create an entry for every spilled value which is not an AllocaInst.
-  for (auto &S : Spills) {
-    // We can have multiple entries in Spills for a single value, but
-    // they should form a contiguous run.  Ignore all but the first.
-    if (CurrentDef == S.def())
-      continue;
-
-    CurrentDef = S.def();
-
-    assert(CurrentDef != PromiseAlloca &&
-           "recorded spill use of promise alloca?");
-
-    if (!isa<AllocaInst>(CurrentDef)) {
-      Type *Ty = CurrentDef->getType();
-      B.addField(Ty, None, {&S});
-    }
+  B.addFieldForAllocas(F, FrameData, Shape);
+  // Create an entry for every spilled value.
+  for (auto &S : FrameData.Spills) {
+    FieldIDType Id = B.addField(S.first->getType(), None);
+    FrameData.setFieldIndex(S.first, Id);
   }
 
   B.finish(FrameTy);
+  FrameData.updateLayoutIndex(B);
   Shape.FrameAlign = B.getStructAlign();
   Shape.FrameSize = B.getStructSize();
 
   switch (Shape.ABI) {
-  // In the switch ABI, remember the field indices for the promise and
-  // switch-index fields.
   case coro::ABI::Switch:
+    // In the switch ABI, remember the field indices for the switch-index field.
     Shape.SwitchLowering.IndexField =
-      B.getFieldIndex(*SwitchIndexFieldId);
-    Shape.SwitchLowering.PromiseField =
-      (PromiseAlloca ? B.getFieldIndex(*PromiseFieldId) : 0);
+        B.getLayoutFieldIndex(*SwitchIndexFieldId);
 
     // Also round the frame size up to a multiple of its alignment, as is
     // generally expected in C/C++.
     Shape.FrameSize = alignTo(Shape.FrameSize, Shape.FrameAlign);
     break;
 
   // In the retcon ABI, remember whether the frame is inline in the storage.
   case coro::ABI::Retcon:

 //  AllocaSpillBB:
 //    ; geps corresponding to allocas that were moved to coroutine frame
 //    br label PostSpill
 //
 //  PostSpill:
 //    whatever
 //
 //
-static Instruction *insertSpills(const SpillInfo &Spills, coro::Shape &Shape) {
+static Instruction *insertSpills(const FrameDataInfo &FrameData,
+                                 coro::Shape &Shape) {
   auto *CB = Shape.CoroBegin;
   LLVMContext &C = CB->getContext();
   IRBuilder<> Builder(CB->getNextNode());
   StructType *FrameTy = Shape.FrameTy;
   PointerType *FramePtrTy = FrameTy->getPointerTo();
   auto *FramePtr =
       cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
   DominatorTree DT(*CB->getFunction());
 
-  Value *CurrentValue = nullptr;
-  BasicBlock *CurrentBlock = nullptr;
-  Value *CurrentReload = nullptr;
-
-  // Proper field number will be read from field definition.
-  unsigned Index = InvalidFieldIndex;
-
-  // We need to keep track of any allocas that need "spilling"
-  // since they will live in the coroutine frame now, all access to them
-  // need to be changed, not just the access across suspend points
-  // we remember allocas and their indices to be handled once we processed
-  // all the spills.
-  SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
-
-  // Promise alloca (if present) doesn't show in the spills and has a
-  // special field number.
-  if (auto *PromiseAlloca = Shape.getPromiseAlloca()) {
-    assert(Shape.ABI == coro::ABI::Switch);
-    Allocas.emplace_back(PromiseAlloca, Shape.getPromiseField());
-  }
-
   // Create a GEP with the given index into the coroutine frame for the original
   // value Orig. Appends an extra 0 index for array-allocas, preserving the
   // original type.
-  auto GetFramePointer = [&](uint32_t Index, Value *Orig) -> Value * {
+  auto GetFramePointer = [&](Value *Orig) -> Value * {
+    FieldIDType Index = FrameData.getFieldIndex(Orig);
     SmallVector<Value *, 3> Indices = {
         ConstantInt::get(Type::getInt32Ty(C), 0),
         ConstantInt::get(Type::getInt32Ty(C), Index),
     };
 
     if (auto *AI = dyn_cast<AllocaInst>(Orig)) {
       if (auto *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
         auto Count = CI->getValue().getZExtValue();
if (isa<AllocaInst>(Orig)) {
       // AllocaInst. So we cast the GEP to the type of AllocaInst.
       if (GEP->getResultElementType() != Orig->getType())
         return Builder.CreateBitCast(GEP, Orig->getType(),
                                      Orig->getName() + Twine(".cast"));
     }
     return GEP;
   };
 
-  // Create a load instruction to reload the spilled value from the coroutine
-  // frame. Populates the Value pointer reference provided with the frame GEP.
-  auto CreateReload = [&](Instruction *InsertBefore, Value *&G) {
-    assert(Index != InvalidFieldIndex && "accessing unassigned field number");
-    Builder.SetInsertPoint(InsertBefore);
-
-    G = GetFramePointer(Index, CurrentValue);
-    G->setName(CurrentValue->getName() + Twine(".reload.addr"));
-
-    return isa<AllocaInst>(CurrentValue)
-               ? G
-               : Builder.CreateLoad(FrameTy->getElementType(Index), G,
-                                    CurrentValue->getName() + Twine(".reload"));
-  };
-
-  Value *GEP = nullptr, *CurrentGEP = nullptr;
-  for (auto const &E : Spills) {
-    // If we have not seen the value, generate a spill.
-    if (CurrentValue != E.def()) {
-      CurrentValue = E.def();
-      CurrentBlock = nullptr;
-      CurrentReload = nullptr;
-
-      Index = E.fieldIndex();
-
-      if (auto *AI = dyn_cast<AllocaInst>(CurrentValue)) {
-        // Spilled AllocaInst will be replaced with GEP from the coroutine frame
-        // there is no spill required.
-        Allocas.emplace_back(AI, Index);
-        if (!AI->isStaticAlloca())
-          report_fatal_error("Coroutines cannot handle non static allocas yet");
-      } else {
-        // Otherwise, create a store instruction storing the value into the
-        // coroutine frame.
-
-        Instruction *InsertPt = nullptr;
-        if (auto Arg = dyn_cast<Argument>(CurrentValue)) {
-          // For arguments, we will place the store instruction right after
-          // the coroutine frame pointer instruction, i.e. bitcast of
-          // coro.begin from i8* to %f.frame*.
-          InsertPt = FramePtr->getNextNode();
+  for (auto const &E : FrameData.Spills) {
+    Value *Def = E.first;
+    // Create a store instruction storing the value into the
+    if (auto Arg = dyn_cast<Argument>(Def)) {
 
-          // If we're spilling an Argument, make sure we clear 'nocapture'
-          // from the coroutine function.
-          Arg->getParent()->removeParamAttr(Arg->getArgNo(),
-                                            Attribute::NoCapture);
+      Arg->getParent()->removeParamAttr(Arg->getArgNo(), Attribute::NoCapture);
 
-        } else if (auto CSI = dyn_cast<AnyCoroSuspendInst>(CurrentValue)) {
-          // Don't spill immediately after a suspend; splitting assumes
-          // that the suspend will be followed by a branch.
-          InsertPt = CSI->getParent()->getSingleSuccessor()->getFirstNonPHI();
-        } else {
-          auto *I = cast<Instruction>(CurrentValue);
-          if (!DT.dominates(CB, I)) {
-            // If it is not dominated by CoroBegin, then spill should be
-            // inserted immediately after CoroFrame is computed.
-            InsertPt = FramePtr->getNextNode();
-          } else if (auto *II = dyn_cast<InvokeInst>(I)) {
-            // If we are spilling the result of the invoke instruction, split
-            // the normal edge and insert the spill in the new block.
-            auto *NewBB = SplitEdge(II->getParent(), II->getNormalDest());
-            InsertPt = NewBB->getTerminator();
-          } else if (isa<PHINode>(I)) {
-            // Skip the PHINodes and EH pads instructions.
-            BasicBlock *DefBlock = I->getParent();
-            if (auto *CSI =
-                    dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
-              InsertPt = splitBeforeCatchSwitch(CSI);
-            else
-              InsertPt = &*DefBlock->getFirstInsertionPt();
-          } else {
-            assert(!I->isTerminator() && "unexpected terminator");
-            // For all other values, the spill is placed immediately after
-            // the definition.
-            InsertPt = I->getNextNode();
-          }
-        }
+    } else if (auto CSI = dyn_cast<AnyCoroSuspendInst>(Def)) {
+      auto *I = cast<Instruction>(Def);
+        if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
 
-        Builder.SetInsertPoint(InsertPt);
-        auto *G = Builder.CreateConstInBoundsGEP2_32(
-            FrameTy, FramePtr, 0, Index,
-            CurrentValue->getName() + Twine(".spill.addr"));
-        Builder.CreateStore(CurrentValue, G);
-      }
-    }
+    auto Index = FrameData.getFieldIndex(Def);
+        FrameTy, FramePtr, 0, Index, Def->getName() + Twine(".spill.addr"));
+    Builder.CreateStore(Def, G);
 
-    // If we have not seen the use block, generate a reload in it.
-    if (CurrentBlock != E.userBlock()) {
-      CurrentBlock = E.userBlock();
-      CurrentReload = CreateReload(&*CurrentBlock->getFirstInsertionPt(), GEP);
-    }
-
-    // If we have a single edge PHINode, remove it and replace it with a reload
-    // from the coroutine frame. (We already took care of multi edge PHINodes
-    // by rewriting them in the rewritePHIs function).
-    if (auto *PN = dyn_cast<PHINode>(E.user())) {
-      assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
-                                                "values in the PHINode");
+    BasicBlock *CurrentBlock = nullptr;
+    Value *CurrentReload = nullptr;
+    for (auto *U : E.second) {
+      // If we have not seen the use block, create a load instruction to reload
+      // the spilled value from the coroutine frame. Populates the Value pointer
+      // reference provided with the frame GEP.
+      if (CurrentBlock != U->getParent()) {
+        CurrentBlock = U->getParent();
+        Builder.SetInsertPoint(&*CurrentBlock->getFirstInsertionPt());
+
+        auto GEP = GetFramePointer(E.first);
+        GEP->setName(E.first->getName() + Twine(".reload.addr"));
+        CurrentReload = Builder.CreateLoad(
+            FrameTy->getElementType(FrameData.getFieldIndex(E.first)), GEP,
+            E.first->getName() + Twine(".reload"));
+      }
+
+      // If we have a single edge PHINode, remove it and replace it with a
+      // reload from the coroutine frame. (We already took care of multi edge
+      // PHINodes by rewriting them in the rewritePHIs function).
+      if (auto *PN = dyn_cast<PHINode>(U)) {
+        assert(PN->getNumIncomingValues() == 1 &&
+               "unexpected number of incoming "
         PN->replaceAllUsesWith(CurrentReload);
         PN->eraseFromParent();
         continue;
       }
 
-    // If we have not seen this GEP instruction, migrate any dbg.declare from
-    // the alloca to it.
-    if (CurrentGEP != GEP) {
-      CurrentGEP = GEP;
-      TinyPtrVector<DbgDeclareInst *> DIs = FindDbgDeclareUses(CurrentValue);
-      if (!DIs.empty())
-        DIBuilder(*CurrentBlock->getParent()->getParent(),
-                  /*AllowUnresolved*/ false)
-            .insertDeclare(CurrentGEP, DIs.front()->getVariable(),
-                           DIs.front()->getExpression(),
-                           DIs.front()->getDebugLoc(), DIs.front());
+      // Replace all uses of CurrentValue in the current instruction with
+      // reload.
+      U->replaceUsesOfWith(Def, CurrentReload);
     }
-
-    // Replace all uses of CurrentValue in the current instruction with reload.
-    E.user()->replaceUsesOfWith(CurrentValue, CurrentReload);
   }
 
   BasicBlock *FramePtrBB = FramePtr->getParent();
 
   auto SpillBlock =
-    FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");      
+      FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
   SpillBlock->splitBasicBlock(&SpillBlock->front(), "PostSpill");
   Shape.AllocaSpillBlock = SpillBlock;
 
   // retcon and retcon.once lowering assumes all uses have been sunk.
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce) {
     // If we found any allocas, replace all of their remaining uses with Geps.
     Builder.SetInsertPoint(&SpillBlock->front());
-    for (auto &P : Allocas) {
-      auto *G = GetFramePointer(P.second, P.first);
+    for (auto &P : FrameData.Allocas) {
+      auto *G = GetFramePointer(P);
 
       // We are not using ReplaceInstWithInst(P.first, cast<Instruction>(G))
       // here, as we are changing location of the instruction.
-      G->takeName(P.first);
-      P.first->replaceAllUsesWith(G);
-      P.first->eraseFromParent();
+      G->takeName(P);
+      P->replaceAllUsesWith(G);
+      P->eraseFromParent();
     }
     return FramePtr;
   }
 
   // If we found any alloca, replace all of their remaining uses with GEP
   // instructions. Because new dbg.declare have been created for these alloca,
   // we also delete the original dbg.declare and replace other uses with undef.
   // Note: We cannot replace the alloca with GEP instructions indiscriminately,
   // as some of the uses may not be dominated by CoroBegin.
   bool MightNeedToCopy = false;
   Builder.SetInsertPoint(&Shape.AllocaSpillBlock->front());
   SmallVector<Instruction *, 4> UsersToUpdate;
-  for (auto &P : Allocas) {
-    AllocaInst *const A = P.first;
-
-    for (auto *DI : FindDbgDeclareUses(A))
-      DI->eraseFromParent();
-    replaceDbgUsesWithUndef(A);
-
+  for (AllocaInst *A : FrameData.Allocas) {
     UsersToUpdate.clear();
     for (User *U : A->users()) {
       auto *I = cast<Instruction>(U);
       if (DT.dominates(CB, I))
         UsersToUpdate.push_back(I);
       else
         MightNeedToCopy = true;
     }
     if (!UsersToUpdate.empty()) {
-      auto *G = GetFramePointer(P.second, A);
-      G->takeName(A);
+      auto *G = GetFramePointer(A);
+      G->setName(A->getName() + Twine(".reload.addr"));
+      TinyPtrVector<DbgDeclareInst *> DIs = FindDbgDeclareUses(A);
+      if (!DIs.empty())
+        DIBuilder(*A->getModule(),
+                  /*AllowUnresolved*/ false)
+            .insertDeclare(G, DIs.front()->getVariable(),
+                           DIs.front()->getExpression(),
+                           DIs.front()->getDebugLoc(), DIs.front());
+      for (auto *DI : FindDbgDeclareUses(A))
+        DI->eraseFromParent();
+      replaceDbgUsesWithUndef(A);
+
       for (Instruction *I : UsersToUpdate)
         I->replaceUsesOfWith(A, G);
     }
   }
   // If we discovered such uses not dominated by CoroBegin, see if any of them
   // preceed coro begin and have instructions that can modify the
   // value of the alloca and therefore would require a copying the value into
   // the spill slot in the coroutine frame.
   if (MightNeedToCopy) {
     Builder.SetInsertPoint(FramePtr->getNextNode());
 
-    for (auto &P : Allocas) {
-      AllocaInst *const A = P.first;
+    for (AllocaInst *A : FrameData.Allocas) {
       AllocaUseVisitor Visitor(A->getModule()->getDataLayout(), DT, *CB);
       auto PtrI = Visitor.visitPtr(*A);
       assert(!PtrI.isAborted());
       if (PtrI.isEscaped()) {
         // isEscaped really means potentially modified before CoroBegin.
         if (A->isArrayAllocation())
           report_fatal_error(
               "Coroutines cannot handle copying of array allocas yet");
 
-        auto *G = GetFramePointer(P.second, A);
+        auto *G = GetFramePointer(A);
         auto *Value = Builder.CreateLoad(A->getAllocatedType(), A);
         Builder.CreateStore(Value, G);
       }
       // For each alias to Alloca created before CoroBegin but used after
       // CoroBegin, we recreate them after CoroBegin by appplying the offset
       // to the pointer in the frame.
       for (const auto &Alias : Visitor.getAliases()) {
-        auto *FramePtr = GetFramePointer(P.second, A);
+        auto *FramePtr = GetFramePointer(A);
         auto *FramePtrRaw =
             Builder.CreateBitCast(FramePtr, Type::getInt8PtrTy(C));
         auto *AliasPtr = Builder.CreateGEP(
             FramePtrRaw, ConstantInt::get(Type::getInt64Ty(C), Alias.second));
         auto *AliasPtrTyped =
             Builder.CreateBitCast(AliasPtr, Alias.first->getType());
         Alias.first->replaceUsesWithIf(
             AliasPtrTyped, [&](Use &U) { return DT.dominates(CB, U); });

 static bool isCoroutineStructureIntrinsic(Instruction &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) {
+                                              const SpillInfo &Spills) {
+  for (const auto &E : Spills) {
+    Value *Def = E.first;
     BasicBlock *CurrentBlock = nullptr;
     Instruction *CurrentMaterialization = nullptr;
-  Instruction *CurrentDef = nullptr;
-
-  for (auto const &E : Spills) {
-    // If it is a new definition, update CurrentXXX variables.
-    if (CurrentDef != E.def()) {
-      CurrentDef = cast<Instruction>(E.def());
-      CurrentBlock = nullptr;
-      CurrentMaterialization = nullptr;
-    }
-
+    for (Instruction *U : E.second) {
       // If we have not seen this block, materialize the value.
-    if (CurrentBlock != E.userBlock()) {
-      CurrentBlock = E.userBlock();
-      CurrentMaterialization = cast<Instruction>(CurrentDef)->clone();
-      CurrentMaterialization->setName(CurrentDef->getName());
+      if (CurrentBlock != U->getParent()) {
+        CurrentBlock = U->getParent();
+        CurrentMaterialization = cast<Instruction>(Def)->clone();
+        CurrentMaterialization->setName(Def->getName());
         CurrentMaterialization->insertBefore(
             &*CurrentBlock->getFirstInsertionPt());
       }
-
-    if (auto *PN = dyn_cast<PHINode>(E.user())) {
-      assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
-                                                "values in the PHINode");
+      if (auto *PN = dyn_cast<PHINode>(U)) {
+        assert(PN->getNumIncomingValues() == 1 &&
+               "unexpected number of incoming "
         PN->replaceAllUsesWith(CurrentMaterialization);
         PN->eraseFromParent();
         continue;
       }
-
-    // Replace all uses of CurrentDef in the current instruction with the
+      // Replace all uses of Def in the current instruction with the
       // CurrentMaterialization for the block.
-    E.user()->replaceUsesOfWith(CurrentDef, CurrentMaterialization);
+      U->replaceUsesOfWith(Def, CurrentMaterialization);
+    }
   }
 }
 
 // Splits the block at a particular instruction unless it is the first
 // instruction in the block with a single predecessor.
 static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
   auto *BB = I->getParent();
   if (&BB->front() == I) {
static void eliminateSwiftError(Function &F, coro::Shape &Shape) {
   if (!AllocasToPromote.empty()) {
     DominatorTree DT(F);
     PromoteMemToReg(AllocasToPromote, DT);
   }
 }
 
 /// retcon and retcon.once conventions assume that all spill uses can be sunk
 /// after the coro.begin intrinsic.
-static void sinkSpillUsesAfterCoroBegin(Function &F, const SpillInfo &Spills,
+static void sinkSpillUsesAfterCoroBegin(Function &F,
+                                        const FrameDataInfo &FrameData,
                                         CoroBeginInst *CoroBegin) {
   DominatorTree Dom(F);
 
   SmallSetVector<Instruction *, 32> ToMove;
   SmallVector<Instruction *, 32> Worklist;
 
   // Collect all users that precede coro.begin.
-  for (auto const &Entry : Spills) {
-    auto *SpillDef = Entry.def();
-    for (User *U : SpillDef->users()) {
+  for (auto *Def : FrameData.getAllDefs()) {
+    for (User *U : Def->users()) {
       auto Inst = cast<Instruction>(U);
       if (Inst->getParent() != CoroBegin->getParent() ||
           Dom.dominates(CoroBegin, Inst))
         continue;
       if (ToMove.insert(Inst))
         Worklist.push_back(Inst);
     }
   }
for (BasicBlock *DomBB : DomSet) {
           S->eraseFromParent();
 
         break;
       }
     }
   }
 }
 
+static void collectFrameAllocas(Function &F, coro::Shape &Shape,
+                                SuspendCrossingInfo &Checker,
+                                SmallVectorImpl<AllocaInst *> &Allocas) {
+  // Collect lifetime.start info for each alloca.
+  using LifetimeStart = SmallPtrSet<Instruction *, 2>;
+  llvm::DenseMap<AllocaInst *, std::unique_ptr<LifetimeStart>> LifetimeMap;
+  for (Instruction &I : instructions(F)) {
+    auto *II = dyn_cast<IntrinsicInst>(&I);
+    if (!II || II->getIntrinsicID() != Intrinsic::lifetime_start)
+      continue;
+
+    if (auto *OpInst = dyn_cast<Instruction>(II->getOperand(1))) {
+      if (auto *AI = dyn_cast<AllocaInst>(OpInst->stripPointerCasts())) {
+
+        if (LifetimeMap.find(AI) == LifetimeMap.end())
+          LifetimeMap[AI] = std::make_unique<LifetimeStart>();
+        LifetimeMap[AI]->insert(isa<AllocaInst>(OpInst) ? II : OpInst);
+      }
+    }
+  }
+
+  for (Instruction &I : instructions(F)) {
+    auto *AI = dyn_cast<AllocaInst>(&I);
+    if (!AI)
+      continue;
+    // The PromiseAlloca will be specially handled since it needs to be in a
+    // fixed position in the frame.
+    if (AI == Shape.SwitchLowering.PromiseAlloca) {
+      continue;
+    }
+    auto Iter = LifetimeMap.find(AI);
+    for (User *U : I.users()) {
+      bool ShouldLiveOnFrame = false;
+
+      // Check against lifetime.start if the instruction has the info.
+      if (Iter != LifetimeMap.end())
+        for (auto *S : *Iter->second) {
+          if ((ShouldLiveOnFrame = Checker.isDefinitionAcrossSuspend(*S, U)))
+            break;
+        }
+      else
+        ShouldLiveOnFrame = Checker.isDefinitionAcrossSuspend(I, U);
+
+      if (ShouldLiveOnFrame) {
+        Allocas.push_back(AI);
+        break;
+      }
+    }
+  }
+}
+
 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   eliminateSwiftError(F, Shape);
 
   if (Shape.ABI == coro::ABI::Switch &&
       Shape.SwitchLowering.PromiseAlloca) {
     Shape.getSwitchCoroId()->clearPromise();
   }
 
void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   // 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);
 
   // Build suspend crossing info.
   SuspendCrossingInfo Checker(F, Shape);
 
   IRBuilder<> Builder(F.getContext());
-  SpillInfo Spills;
+  FrameDataInfo FrameData;
   SmallVector<CoroAllocaAllocInst*, 4> LocalAllocas;
   SmallVector<Instruction*, 4> DeadInstructions;
 
+  {
+    SpillInfo Spills;
     for (int Repeat = 0; Repeat < 4; ++Repeat) {
       // See if there are materializable instructions across suspend points.
       for (Instruction &I : instructions(F))
         if (materializable(I))
           for (User *U : I.users())
             if (Checker.isDefinitionAcrossSuspend(I, U))
-            Spills.emplace_back(&I, U);
+              Spills[&I].push_back(cast<Instruction>(U));
 
       if (Spills.empty())
         break;
 
-    // Rewrite materializable instructions to be materialized at the use point.
-    LLVM_DEBUG(dump("Materializations", Spills));
+      // Rewrite materializable instructions to be materialized at the use
+      // point.
+      LLVM_DEBUG(dumpSpills("Materializations", Spills));
       rewriteMaterializableInstructions(Builder, Spills);
       Spills.clear();
     }
+  }
 
   sinkLifetimeStartMarkers(F, Shape, Checker);
-  // Collect lifetime.start info for each alloca.
-  using LifetimeStart = SmallPtrSet<Instruction *, 2>;
-  llvm::DenseMap<Instruction *, std::unique_ptr<LifetimeStart>> LifetimeMap;
-  for (Instruction &I : instructions(F)) {
-    auto *II = dyn_cast<IntrinsicInst>(&I);
-    if (!II || II->getIntrinsicID() != Intrinsic::lifetime_start)
-      continue;
-
-    if (auto *OpInst = dyn_cast<Instruction>(II->getOperand(1))) {
-      if (auto *AI = dyn_cast<AllocaInst>(OpInst->stripPointerCasts())) {
-
-        if (LifetimeMap.find(AI) == LifetimeMap.end())
-          LifetimeMap[AI] = std::make_unique<LifetimeStart>();
-        LifetimeMap[AI]->insert(isa<AllocaInst>(OpInst) ? II : OpInst);
-      }
-    }
-  }
+  collectFrameAllocas(F, Shape, Checker, FrameData.Allocas);
+  LLVM_DEBUG(dumpAllocas(FrameData.Allocas));
 
   // Collect the spills for arguments and other not-materializable values.
   for (Argument &A : F.args())
     for (User *U : A.users())
       if (Checker.isDefinitionAcrossSuspend(A, U))
-        Spills.emplace_back(&A, U);
+        FrameData.Spills[&A].push_back(cast<Instruction>(U));
 
   for (Instruction &I : instructions(F)) {
     // Values returned from coroutine structure intrinsics should not be part
     // of the Coroutine Frame.
     if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
       continue;
 
     // The Coroutine Promise always included into coroutine frame, no need to
if (auto AI = dyn_cast<CoroAllocaAllocInst>(&I)) {
       // the rewritten value.  The rewrite doesn't invalidate anything in
       // Spills because the other alloca intrinsics have no other operands
       // besides AI, and it doesn't invalidate the iteration because we delay
       // erasing AI.
       auto Alloc = lowerNonLocalAlloca(AI, Shape, DeadInstructions);
 
       for (User *U : Alloc->users()) {
         if (Checker.isDefinitionAcrossSuspend(*Alloc, U))
-          Spills.emplace_back(Alloc, U);
+          FrameData.Spills[Alloc].push_back(cast<Instruction>(U));
       }
       continue;
     }
 
     // Ignore alloca.get; we process this as part of coro.alloca.alloc.
-    if (isa<CoroAllocaGetInst>(I)) {
+    if (isa<CoroAllocaGetInst>(I))
       continue;
-    }
-
-    auto Iter = LifetimeMap.find(&I);
-    for (User *U : I.users()) {
-      bool NeedSpill = false;
 
-      // Check against lifetime.start if the instruction has the info.
-      if (Iter != LifetimeMap.end())
-        for (auto *S : *Iter->second) {
-          if ((NeedSpill = Checker.isDefinitionAcrossSuspend(*S, U)))
-            break;
-        }
-      else
-        NeedSpill = Checker.isDefinitionAcrossSuspend(I, U);
+    if (isa<AllocaInst>(I))
+      continue;
 
-      if (NeedSpill) {
+    for (User *U : I.users())
+      if (Checker.isDefinitionAcrossSuspend(I, U)) {
         // We cannot spill a token.
         if (I.getType()->isTokenTy())
           report_fatal_error(
               "token definition is separated from the use by a suspend point");
-        Spills.emplace_back(&I, U);
-      }
+        FrameData.Spills[&I].push_back(cast<Instruction>(U));
       }
   }
-  LLVM_DEBUG(dump("Spills", Spills));
+  LLVM_DEBUG(dumpSpills("Spills", FrameData.Spills));
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce)
-    sinkSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
-  Shape.FrameTy = buildFrameType(F, Shape, Spills);
-  Shape.FramePtr = insertSpills(Spills, Shape);
+    sinkSpillUsesAfterCoroBegin(F, FrameData, Shape.CoroBegin);
+  Shape.FrameTy = buildFrameType(F, Shape, FrameData);
+  // Add PromiseAlloca to Allocas list so that it is processed in insertSpills.
+  if (Shape.ABI == coro::ABI::Switch && Shape.SwitchLowering.PromiseAlloca)
+    FrameData.Allocas.push_back(Shape.SwitchLowering.PromiseAlloca);
+  Shape.FramePtr = insertSpills(FrameData, Shape);
   lowerLocalAllocas(LocalAllocas, DeadInstructions);
 
   for (auto I : DeadInstructions)
     I->eraseFromParent();
 }

llvm/lib/Transforms/Coroutines/CoroInternal.h

struct LLVM_LIBRARY_VISIBILITY Shape {
 
   bool ReuseFrameSlot;
 
   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
     BasicBlock *ResumeEntryBlock;
     unsigned IndexField;
-    unsigned PromiseField;
     bool HasFinalSuspend;
   };
 
   struct RetconLoweringStorage {
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
     BasicBlock *ReturnBlock;
CallingConv::ID getResumeFunctionCC() const {
     llvm_unreachable("Unknown coro::ABI enum");
   }
 
   AllocaInst *getPromiseAlloca() const {
     if (ABI == coro::ABI::Switch)
       return SwitchLowering.PromiseAlloca;
     return nullptr;
   }
-  unsigned getPromiseField() const {
-    assert(ABI == coro::ABI::Switch);
-    assert(FrameTy && "frame type not assigned");
-    assert(SwitchLowering.PromiseAlloca && "no promise alloca");
-    return SwitchLowering.PromiseField;
-  }
 
   /// Allocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;
 
   /// Deallocate memory according to the rules of the active lowering.
   ///

llvm/test/Transforms/Coroutines/coro-debug-frame-variable.ll

 ;
 ; The CHECKs verify that dbg.declare intrinsics are created for the coroutine
 ; funclet 'f.resume', and that they reference the address of the variables on
 ; the coroutine frame. The debug locations for the original function 'f' are
 ; static (!11 and !13), whereas the coroutine funclet will have its own new
 ; ones with identical line and column numbers.
 ;
 ; CHECK-LABEL: define void @f() {
-; CHECK:       init.ready:
+; CHECK:       entry:
 ; CHECK:         [[IGEP:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK:         [[JGEP:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
+; CHECK:       init.ready:
 ; CHECK:         call void @llvm.dbg.declare(metadata i32* [[IGEP]], metadata ![[IVAR:[0-9]+]], metadata !DIExpression()), !dbg ![[IDBGLOC:[0-9]+]]
 ; CHECK:       await.ready:
-; CHECK:         [[JGEP:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
 ; CHECK:         call void @llvm.dbg.declare(metadata i32* [[JGEP]], metadata ![[JVAR:[0-9]+]], metadata !DIExpression()), !dbg ![[JDBGLOC:[0-9]+]]
 ;
 ; CHECK-LABEL: define internal fastcc void @f.resume({{.*}}) {
-; CHECK:       init.ready:
+; CHECK:       entry.resume:
 ; CHECK:         [[IGEP_RESUME:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK:         [[JGEP_RESUME:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
+; CHECK:       init.ready:
 ; CHECK:         call void @llvm.dbg.declare(metadata i32* [[IGEP_RESUME]], metadata ![[IVAR_RESUME:[0-9]+]], metadata !DIExpression()), !dbg ![[IDBGLOC_RESUME:[0-9]+]]
 ; CHECK:       await.ready:
-; CHECK:         [[JGEP_RESUME:%.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 5
 ; CHECK:         call void @llvm.dbg.declare(metadata i32* [[JGEP_RESUME]], metadata ![[JVAR_RESUME:[0-9]+]], metadata !DIExpression()), !dbg ![[JDBGLOC_RESUME:[0-9]+]]
 ;
 ; CHECK: ![[IVAR]] = !DILocalVariable(name: "i"
 ; CHECK: ![[SCOPE:[0-9]+]] = distinct !DILexicalBlock(scope: !8, file: !1, line: 23, column: 12)
 ; CHECK: ![[IDBGLOC]] = !DILocation(line: 24, column: 7, scope: ![[SCOPE]])
 ; CHECK: ![[JVAR]] = !DILocalVariable(name: "j"
 ; CHECK: ![[JDBGLOC]] = !DILocation(line: 32, column: 7, scope: ![[SCOPE]])
 ; CHECK: ![[IVAR_RESUME]] = !DILocalVariable(name: "i"

llvm/test/Transforms/Coroutines/coro-debug.ll

 !21 = !DILocation(line: 59, column: 10, scope: !6)
 !22 = !DILocation(line: 59, column: 7, scope: !6)
 !23 = !DILocation(line: 59, column: 5, scope: !6)
 !24 = !DILocation(line: 62, column: 3, scope: !6)
 
 ; CHECK: define i8* @f(i32 %x) #0 !dbg ![[ORIG:[0-9]+]]
 ; CHECK: define internal fastcc void @f.resume(%f.Frame* noalias nonnull align 8 dereferenceable(32) %FramePtr) #0 !dbg ![[RESUME:[0-9]+]]
 ; CHECK: entry.resume:
-; CHECK-NEXT: call void @coro.devirt.trigger(i8* null)
-; CHECK-NEXT: call void @llvm.dbg.declare(metadata i32* %x.addr.reload.addr, metadata ![[RESUME_VAR:[0-9]+]]
+; CHECK: call void @coro.devirt.trigger(i8* null)
+; CHECK: call void @llvm.dbg.declare(metadata i32* %x.addr.reload.addr, metadata ![[RESUME_VAR:[0-9]+]]
 ; CHECK: define internal fastcc void @f.destroy(%f.Frame* noalias nonnull align 8 dereferenceable(32) %FramePtr) #0 !dbg ![[DESTROY:[0-9]+]]
 ; CHECK: define internal fastcc void @f.cleanup(%f.Frame* noalias nonnull align 8 dereferenceable(32) %FramePtr) #0 !dbg ![[CLEANUP:[0-9]+]]
 
 ; CHECK: ![[ORIG]] = distinct !DISubprogram(name: "f", linkageName: "flink"
 
 ; CHECK: ![[RESUME]] = distinct !DISubprogram(name: "f", linkageName: "flink"
 ; CHECK: ![[RESUME_VAR]] = !DILocalVariable(name: "x", arg: 1, scope: ![[RESUME]]
 
 ; CHECK: ![[DESTROY]] = distinct !DISubprogram(name: "f", linkageName: "flink"
 
 ; CHECK: ![[CLEANUP]] = distinct !DISubprogram(name: "f", linkageName: "flink"

llvm/test/Transforms/Coroutines/coro-frame-arrayalloca.ll

suspend:
   ret i8* %hdl
 }
 
 ; See if the array alloca was stored as an array field.
 ; CHECK-LABEL: %f.Frame = type { void (%f.Frame*)*, void (%f.Frame*)*, double, double, [4 x i32], i1 }
 
 ; See if we used correct index to access prefix, data, suffix (@f)
 ; CHECK-LABEL: @f(
-; CHECK:       %prefix = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2
-; CHECK-NEXT:  %data = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
-; CHECK-NEXT:  %suffix = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 3
-; CHECK-NEXT:  call void @consume.double.ptr(double* %prefix)
-; CHECK-NEXT:  call void @consume.i32.ptr(i32* %data)
-; CHECK-NEXT:  call void @consume.double.ptr(double* %suffix)
+; CHECK:       %[[PREFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2
+; CHECK-NEXT:  %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK-NEXT:  %[[SUFFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 3
+; CHECK-NEXT:  call void @consume.double.ptr(double* %[[PREFIX:.+]])
+; CHECK-NEXT:  call void @consume.i32.ptr(i32* %[[DATA:.+]])
+; CHECK-NEXT:  call void @consume.double.ptr(double* %[[SUFFIX:.+]])
 ; CHECK: ret i8*
 
 ; See if we used correct index to access prefix, data, suffix (@f.resume)
 ; CHECK-LABEL: @f.resume(
-; CHECK:       %[[SUFFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 3
-; CHECK:       %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
 ; CHECK:       %[[PREFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 2
+; CHECK:       %[[DATA:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 4
+; CHECK:       %[[SUFFIX:.+]] = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0, i32 3
 ; CHECK:       call void @consume.double.ptr(double* %[[PREFIX]])
 ; CHECK-NEXT:  call void @consume.i32.ptr(i32* %[[DATA]])
 ; CHECK-NEXT:  call void @consume.double.ptr(double* %[[SUFFIX]])
 
 declare i8* @llvm.coro.free(token, i8*)
 declare i32 @llvm.coro.size.i32()
 declare i8  @llvm.coro.suspend(token, i1)
 declare void @llvm.coro.resume(i8*)

llvm/test/Transforms/Coroutines/coro-frame-reuse-alloca-01.ll

cleanup2:
   br label %cleanup
 cleanup:
   call i8* @llvm.coro.free(token %1, i8* %2)
   br label %coro.ret
 coro.ret:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
-; CHECK-LABEL: @a.resume(
-; CHECK:         %a.reload.addr{{[0-9]+}} = getelementptr inbounds %a.Frame, %a.Frame* %FramePtr[[APositon:.*]]
-; CHECK:         %b.reload.addr{{[0-9]+}} = getelementptr inbounds %a.Frame, %a.Frame* %FramePtr[[APositon]]
+
+; check that there is only one %struct.big_structure in the frame.
+; CHECK: %a.Frame = type { void (%a.Frame*)*, void (%a.Frame*)*, %"struct.task::promise_type", %struct.big_structure, i1 }
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 declare i1 @llvm.coro.alloc(token) #3
 declare i64 @llvm.coro.size.i64() #5
 declare i8* @llvm.coro.begin(token, i8* writeonly) #3
 declare token @llvm.coro.save(i8*) #3
 declare i8* @llvm.coro.frame() #5
 declare i8 @llvm.coro.suspend(token, i1) #3
 declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
 declare i1 @llvm.coro.end(i8*, i1) #3
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
-declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4
 

llvm/test/Transforms/Coroutines/coro-frame-reuse-alloca-02.ll

cleanup2:
   br label %cleanup
 cleanup:
   call i8* @llvm.coro.free(token %1, i8* %2)
   br label %coro.ret
 coro.ret:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
+; CHECK:       %a.Frame = type { void (%a.Frame*)*, void (%a.Frame*)*, %"struct.task::promise_type", %struct.big_structure, i1 }
 ; CHECK-LABEL: @a.resume(
-; CHECK:         %b.reload.addr = bitcast %struct.big_structure* %0 to %struct.big_structure.2*
+; CHECK:         %[[A:.*]] = getelementptr inbounds %a.Frame, %a.Frame* %FramePtr, i32 0, i32 3
+; CHECK:         %{{.*}} = bitcast %struct.big_structure* %[[A]] to %struct.big_structure.2*
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 declare i1 @llvm.coro.alloc(token) #3
 declare i64 @llvm.coro.size.i64() #5
 declare i8* @llvm.coro.begin(token, i8* writeonly) #3
 declare token @llvm.coro.save(i8*) #3
 declare i8* @llvm.coro.frame() #5
 declare i8 @llvm.coro.suspend(token, i1) #3
 declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
 declare i1 @llvm.coro.end(i8*, i1) #3
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
-declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4
 

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

 ; CHECK:  store i64 0, i64* [[PROJ2]]
 ; CHECK:  [[ESCAPED_ADDR:%.*]] = ptrtoint { i64, i64 }* [[TMP]] to i64
 ; CHECK:  call void @init(i64* [[PROJ1]])
 ; CHECK:  call void @init(i64* [[PROJ2]])
 ; CHECK:  call void @use_addr_val(i64 [[ESCAPED_ADDR]], { i64, i64 }* [[TMP]])
 
 ; CHECK-LABEL: define internal void @f.resume.0(i8* {{.*}} %0, i1 %1) {
 ; CHECK:  [[FRAMEPTR:%.*]] = bitcast i8* %0 to %f.Frame*
-; CHECK: resume:
 ; CHECK:  [[TMP:%.*]] = getelementptr inbounds %f.Frame, %f.Frame* [[FRAMEPTR]], i32 0, i32 0
+; CHECK: resume:
 ; CHECK:  [[CAST:%.*]] = bitcast { i64, i64 }* [[TMP]] to i8*
 ; CHECK:  call void @use(i8* [[CAST]])
 
 declare token @llvm.coro.id.retcon.once(i32, i32, i8*, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i1 @llvm.coro.suspend.retcon.i1(...)
 declare i1 @llvm.coro.end(i8*, i1)
 

llvm/test/Transforms/Coroutines/coro-retcon-once-value2.ll

 ; CHECK-NEXT:    ret { i8*, i32* } [[T0]]
 ; CHECK-NEXT:  }
 
 ; CHECK-LABEL: define internal void @f.resume.0(i8* noalias nonnull align 8 dereferenceable(8) %0, i1 zeroext %1)
 ; CHECK-NEXT:  :
 ; CHECK-NEXT:    [[T0:%.*]] = bitcast i8* %0 to [[FRAME_T:%.*]]**
 ; CHECK-NEXT:    [[FRAME:%.*]] = load [[FRAME_T]]*, [[FRAME_T]]** [[T0]]
 ; CHECK-NEXT:    bitcast [[FRAME_T]]* [[FRAME]] to i8*
-; CHECK-NEXT:    %temp = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 1
+; CHECK-NEXT:    [[TEMP_SLOT:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 1
 ; CHECK-NEXT:    br i1 %1,
 ; CHECK:       :
-; CHECK-NEXT:    [[TEMP_SLOT:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 1
 ; CHECK-NEXT:    [[PTR_SLOT:%.*]] = getelementptr inbounds [[FRAME_T]], [[FRAME_T]]* [[FRAME]], i32 0, i32 0
 ; CHECK-NEXT:    [[PTR_RELOAD:%.*]] = load i32*, i32** [[PTR_SLOT]]
 ; CHECK-NEXT:    %newvalue = load i32, i32* [[TEMP_SLOT]]
 ; CHECK-NEXT:    store i32 %newvalue, i32* [[PTR_RELOAD]]
 ; CHECK-NEXT:    br label
 ; CHECK:       :
 ; CHECK-NEXT:    [[T0:%.*]] = bitcast [[FRAME_T]]* [[FRAME]] to i8*
 ; CHECK-NEXT:    call fastcc void @deallocate(i8* [[T0]])

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-01.ll

await.ready:
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; CHECK-LABEL: @a.resume(
 ; CHECK:         %testval = alloca i32, align 4
+; CHECK-NEXT:    getelementptr inbounds %a.Frame
 ; CHECK-NEXT:    %0 = bitcast i32* %testval to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 4, i8* %0)
-; CHECK-NEXT:    getelementptr inbounds %a.Frame
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
 ; CHECK-NEXT:    %val = load i32, i32* %Result
 ; CHECK-NEXT:    %test = load i32, i32* %testval
 ; CHECK-NEXT:    call void @print(i32 %test)
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 4, i8* %0)
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void
 

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-03.ll

await.ready:
   call void @print(i32 %val)
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 ; CHECK-LABEL: @a.gep.resume(
 ; CHECK:         %testval = alloca %i8.array
+; CHECK-NEXT:    getelementptr inbounds %a.gep.Frame
 ; CHECK-NEXT:    %0 = bitcast %i8.array* %testval to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 100, i8* %0)
-; CHECK-NEXT:    getelementptr inbounds %a.gep.Frame
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
 ; CHECK-NEXT:    getelementptr inbounds %i8.array, %i8.array* %testval
 ; CHECK-NEXT:    %val = load i32, i32* %Result
 ; CHECK-NEXT:    call void @consume.i8.array(%i8.array* %testval)
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 100, i8* %cast1)
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void
 

llvm/test/Transforms/Coroutines/coro-split-sink-lifetime-04.ll

await.ready:
   br label %exit
 exit:
   call i1 @llvm.coro.end(i8* null, i1 false)
   ret void
 }
 
 ; CHECK-LABEL: @a.resume(
 ; CHECK:         %testval = alloca i8, align 1
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 1, i8* %testval)
 ; CHECK-NEXT:    getelementptr inbounds %a.Frame
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 1, i8* %testval)
 ; CHECK-NEXT:    getelementptr inbounds %"struct.lean_future<int>::Awaiter"
 ; CHECK-NEXT:    %val = load i32, i32* %Result
 ; CHECK-NEXT:    %test = load i8, i8* %testval
 ; CHECK-NEXT:    call void @consume.i8(i8 %test)
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 1, i8* %testval)
 ; CHECK-NEXT:    call void @print(i32 %val)
 ; CHECK-NEXT:    ret void