[Coroutines] Fix PR45130

[junparser]

Summary

For now, when final suspend can be simplified by simplifySuspendPoint, handleFinalSuspend is executed as well to remove last case in switch instruction. This causes issue in https://bugs.llvm.org/show_bug.cgi?id=45130, this patch fixes it.

TestPlan: check-llvm

[modocache]

Excellent, thank you! I had one question, otherwise this looks great!

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1161

Beginner question of mine: what is this assert testing? I would guess that it's testing that SI != nullptr, but the cast<CoroSuspendInst>(...) above would also fail were that the case. As far as I know CoroSuspendInst and its class hierarchy don't implement a boolean conversion operator either... so if this is a null check, I think it's redundant and can be removed, but if it does something more, please let me know!

[junparser]

Thanks for review! :)

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1161

Thanks! you are completely right. This redundant assert testing added here is just my unintentional behavior.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void simplifySuspendPoints(coro::Shape &Shape) {
   if (Shape.ABI != coro::ABI::Switch)
     return;
 
   auto &S = Shape.CoroSuspends;
   size_t I = 0, N = S.size();
   if (N == 0)
     return;
   while (true) {
-    if (simplifySuspendPoint(cast<CoroSuspendInst>(S[I]), Shape.CoroBegin)) {
+    auto SI = cast<CoroSuspendInst>(S[I]);
+    // Leave final.suspend to handleFinalSuspend since it is undefined behavior
+    // to resume a coroutine suspended at the final suspend point.
+    if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
       if (--N == I)
         break;
       std::swap(S[I], S[N]);
       continue;
     }
     if (++I == N)
       break;
   }

llvm/test/Transforms/Coroutines/no-suspend.ll

cleanup:
   %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
   call void @free(i8* %mem)
   br label %suspend
 suspend:
   call i1 @llvm.coro.end(i8* %hdl, i1 false)
   ret void
 }
 
+; SimplifySuspendPoint should not simplify final suspend point
+;
+; CHECK-LABEL: define void @cannot_simplify_final_suspend(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:     llvm.coro.id
+;
+define void @cannot_simplify_final_suspend() "coroutine.presplit"="1" personality i32 0 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
+  br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
+dyn.alloc:
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  br label %coro.begin
+coro.begin:
+  %phi = phi i8* [ null, %entry ], [ %alloc, %dyn.alloc ]
+  %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %phi)
+  br label %body
+body:
+  %save = call token @llvm.coro.save(i8* %hdl)
+  %subfn = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %bcast = bitcast i8* %subfn to void (i8*)*
+  invoke fastcc void %bcast(i8* %hdl) to label %real_susp unwind label %lpad
+
+real_susp:
+  %0 = call i8 @llvm.coro.suspend(token %save, i1 1)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %pre.cleanup]
+resume:
+  call void @print(i32 0)
+  br label %cleanup
+
+pre.cleanup:
+  call void @print(i32 1)
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 false)
+  ret void
+lpad:
+  %lpval = landingpad { i8*, i32 }
+     cleanup
+
+  call void @print(i32 2)
+  resume { i8*, i32 } %lpval
+}
+
 declare i8* @malloc(i32)
 declare void @free(i8*)
 declare void @print(i32)
 declare void @foo()
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i1 @llvm.coro.alloc(token)
 declare i32 @llvm.coro.size.i32()