Make use of uniqueness of coro.promise

Author: NewSigma

llvm/include/llvm/Transforms/Coroutines/CoroShape.h

@@ -80,7 +80,7 @@ struct Shape {
   // Scan the function and collect the above intrinsics for later processing
   void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
                SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-               SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
+               CoroPromiseInst *&CoroPromise);
   // If for some reason, we were not able to find coro.begin, bailout.
   void invalidateCoroutine(Function &F,
                            SmallVectorImpl<CoroFrameInst *> &CoroFrames);
@@ -89,7 +89,7 @@ struct Shape {
   // Remove orphaned and unnecessary intrinsics
   void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
                       SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-                      SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
+                      CoroPromiseInst *CoroPromise);
 
   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
@@ -267,14 +267,14 @@ struct Shape {
   explicit Shape(Function &F) {
     SmallVector<CoroFrameInst *, 8> CoroFrames;
     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
-    SmallVector<CoroPromiseInst *, 2> CoroPromises;
+    CoroPromiseInst * CoroPromise = nullptr;
 
-    analyze(F, CoroFrames, UnusedCoroSaves, CoroPromises);
+    analyze(F, CoroFrames, UnusedCoroSaves, CoroPromise);
     if (!CoroBegin) {
       invalidateCoroutine(F, CoroFrames);
       return;
     }
-    cleanCoroutine(CoroFrames, UnusedCoroSaves, CoroPromises);
+    cleanCoroutine(CoroFrames, UnusedCoroSaves, CoroPromise);
   }
 };
 

llvm/lib/Transforms/Coroutines/Coroutines.cpp

@@ -193,7 +193,7 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 void coro::Shape::analyze(Function &F,
                           SmallVectorImpl<CoroFrameInst *> &CoroFrames,
                           SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-                          SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
+                          CoroPromiseInst *&CoroPromise) {
   clear();
 
   bool HasFinalSuspend = false;
@@ -288,7 +288,7 @@ void coro::Shape::analyze(Function &F,
         }
         break;
       case Intrinsic::coro_promise:
-        CoroPromises.push_back(cast<CoroPromiseInst>(II));
+        CoroPromise = cast<CoroPromiseInst>(II);
         break;
       }
     }
@@ -481,8 +481,7 @@ void coro::AnyRetconABI::init() {
 
 void coro::Shape::cleanCoroutine(
     SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-    SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
+    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves, CoroPromiseInst *PI) {
   // The coro.frame intrinsic is always lowered to the result of coro.begin.
   for (CoroFrameInst *CF : CoroFrames) {
     CF->replaceAllUsesWith(CoroBegin);
@@ -495,10 +494,10 @@ void coro::Shape::cleanCoroutine(
     CoroSave->eraseFromParent();
   UnusedCoroSaves.clear();
 
-  auto *AI = getPromiseAlloca();
-  for (auto *PI : CoroPromises) {
-    PI->replaceAllUsesWith(PI->isFromPromise() ? cast<Value>(CoroBegin)
-                                               : cast<Value>(AI));
+  if (PI) {
+    PI->replaceAllUsesWith(PI->isFromPromise()
+                               ? cast<Value>(CoroBegin)
+                               : cast<Value>(getPromiseAlloca()));
     PI->eraseFromParent();
   }
 }

llvm/test/Transforms/Coroutines/gh105595.ll

@@ -18,7 +18,6 @@ resume:
   br label %suspend
 
 suspend:
-  call i1 @llvm.coro.end(ptr null, i1 false, token none)
 ; load value when resume
 ; CHECK: %null = load ptr, ptr %promise.addr, align 8
   %null = load ptr, ptr %__promise, align 8