[Don't Commit] [DRAFT] Workaround the example in 57861

[ChuanqiXu]

This is not strictly correct.

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

if (!AllocasToPromote.empty()) {
     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 FrameDataInfo &FrameData,
-                                        CoroBeginInst *CoroBegin) {
+                                        coro::Shape &Shape) {
   DominatorTree Dom(F);
 
   SmallSetVector<Instruction *, 32> ToMove;
   SmallVector<Instruction *, 32> Worklist;
 
+  CoroBeginInst *CoroBegin = Shape.CoroBegin;
+  SmallVector<Value *, 8> AllDefs = FrameData.getAllDefs();
+
+  if (Shape.ABI == coro::ABI::Switch) {
+    AllDefs.push_back(Shape.getPromiseAlloca());
+  }
+
   // Collect all users that precede coro.begin.
-  for (auto *Def : FrameData.getAllDefs()) {
+  for (auto *Def : AllDefs) {
     for (User *U : Def->users()) {
       auto Inst = cast<Instruction>(U);
-      if (Inst->getParent() != CoroBegin->getParent() ||
-          Dom.dominates(CoroBegin, Inst))
+      if (Dom.dominates(CoroBegin, Inst))
         continue;
+      
       if (ToMove.insert(Inst))
         Worklist.push_back(Inst);
     }
   }
   // Recursively collect users before coro.begin.
   while (!Worklist.empty()) {
     auto *Def = Worklist.pop_back_val();
     for (User *U : Def->users()) {
void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
   for (auto &Iter : FrameData.Spills) {
     auto *V = Iter.first;
     SmallVector<DbgValueInst *, 16> DVIs;
     findDbgValues(DVIs, V);
     for (DbgValueInst *DVI : DVIs)
       if (Checker.isDefinitionAcrossSuspend(*V, DVI))
         FrameData.Spills[V].push_back(DVI);
   }
   
   LLVM_DEBUG(dumpSpills("Spills", FrameData.Spills));
-  if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
-      Shape.ABI == coro::ABI::Async)
-    sinkSpillUsesAfterCoroBegin(F, FrameData, Shape.CoroBegin);
+  sinkSpillUsesAfterCoroBegin(F, FrameData, Shape);
   Shape.FrameTy = buildFrameType(F, Shape, FrameData);
   createFramePtr(Shape);
   // For now, this works for C++ programs only.
   buildFrameDebugInfo(F, Shape, FrameData);
   insertSpills(FrameData, Shape);
   lowerLocalAllocas(LocalAllocas, DeadInstructions);
 
   for (auto *I : DeadInstructions)
     I->eraseFromParent();
 }