[Coro] RetconOnceDynamic: Pass frame to de/alloc.

Doing so enables allocator implementation to use the frame to store
contextual info.

Author: nate-chandler

llvm/lib/Transforms/Coroutines/Coroutines.cpp

@@ -534,18 +534,21 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
   case coro::ABI::RetconOnce:
   case coro::ABI::RetconOnceDynamic: {
     unsigned sizeParamIndex = 0;
-    SmallVector<Value *, 2> Args;
-    if (ABI == coro::ABI::RetconOnceDynamic) {
-      sizeParamIndex = 1;
-      Args.push_back(RetconLowering.Allocator);
-    }
     Function *Alloc = nullptr;
     if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
       assert(ABI == coro::ABI::RetconOnceDynamic);
       Alloc = RetconLowering.AllocFrame;
     } else {
       Alloc = RetconLowering.Alloc;
     }
+    SmallVector<Value *, 2> Args;
+    if (ABI == coro::ABI::RetconOnceDynamic) {
+      Args.push_back(RetconLowering.Storage);
+      Args.push_back(RetconLowering.Allocator);
+      sizeParamIndex = 2;
+    } else {
+      Alloc = RetconLowering.Alloc;
+    }
     Size = Builder.CreateIntCast(
         Size, Alloc->getFunctionType()->getParamType(sizeParamIndex),
         /*is signed*/ false);
@@ -582,11 +585,14 @@ void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr,
     } else {
       Dealloc = RetconLowering.Dealloc;
     }
-    SmallVector<Value *, 2> Args;
     unsigned allocationParamIndex = 0;
+    SmallVector<Value *, 2> Args;
     if (ABI == coro::ABI::RetconOnceDynamic) {
-      allocationParamIndex = 1;
+      allocationParamIndex = 2;
+      Args.push_back(RetconLowering.Storage);
       Args.push_back(RetconLowering.Allocator);
+    } else {
+      Dealloc = RetconLowering.Dealloc;
     }
     Ptr = Builder.CreateBitCast(
         Ptr, Dealloc->getFunctionType()->getParamType(allocationParamIndex));

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

@@ -114,10 +114,10 @@ cleanup:
 
 declare void @continuation_prototype(ptr, ptr)
 
-declare swiftcorocc noalias ptr @allocate(ptr swiftcoro %cator, i32 %size)
-declare void @deallocate(ptr swiftcoro %cator, ptr %ptr)
-declare swiftcorocc noalias ptr @allocate_frame(ptr swiftcoro %cator, i32 %size)
-declare void @deallocate_frame(ptr swiftcoro %cator, ptr %ptr)
+declare swiftcorocc noalias ptr @allocate(ptr %frame, ptr swiftcoro %cator, i32 %size)
+declare void @deallocate(ptr %frame, ptr swiftcoro %cator, ptr %ptr)
+declare swiftcorocc noalias ptr @allocate_frame(ptr %frame, ptr swiftcoro %cator, i32 %size)
+declare void @deallocate_frame(ptr %frame, ptr swiftcoro %cator, ptr %ptr)
 
 declare void @use(ptr %ptr)
 
@@ -186,10 +186,12 @@ declare { ptr, ptr } @allocating_something_else(ptr noalias %frame, ptr swiftcor
 ; CHECK-SAME:  {
 ; CHECK:         %size = call i32 @getSize()
 ; CHECK:         call swiftcorocc ptr @allocate(
+; CHECK-SAME:      ptr %frame,
 ; CHECK-SAME:      ptr %allocator,
 ; CHECK-SAME:      i32 %size
 ; CHECK-SAME:    )
 ; CHECK:         call swiftcorocc ptr @allocate_frame(
+; CHECK-SAME:      ptr %frame,
 ; CHECK-SAME:      ptr %allocator,
 ; CHECK-SAME:      i32
 ; CHECK-SAME:    )
@@ -199,10 +201,12 @@ declare { ptr, ptr } @allocating_something_else(ptr noalias %frame, ptr swiftcor
 ; CHECK-SAME:  )
 ; CHECK-SAME:  {
 ; CHECK:         call void @deallocate_frame(
+; CHECK-SAME:      ptr %0,
 ; CHECK-SAME:      ptr %1,
 ; CHECK-SAME:      ptr
 ; CHECK-SAME:    )
 ; CHECK:         call void @deallocate(
+; CHECK-SAME:      ptr %0,
 ; CHECK-SAME:      ptr %1,
 ; CHECK-SAME:      ptr
 ; CHECK-SAME:    )