Fix emitDynamicAlloca of an partial_apply [stack] in a coroutine that reuses heap object as context

Author: aschwaighofer

lib/IRGen/GenOpaque.cpp

@@ -555,15 +555,17 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
 /// location before the dynamic alloca's call.
 void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address) {
   // In coroutines, unconditionally call llvm.coro.alloca.free.
-  if (isCoroutine()) {
+  // Except if the address is invalid, this happens when this is a StackAddress
+  // for a partial_apply [stack] that did not need a context object on the
+  // stack.
+  if (isCoroutine() && address.getAddress().isValid()) {
     auto allocToken = address.getExtraInfo();
     assert(allocToken && "dynamic alloca in coroutine without alloc token?");
     auto freeFn = llvm::Intrinsic::getDeclaration(
         &IGM.Module, llvm::Intrinsic::ID::coro_alloca_free);
     Builder.CreateCall(freeFn, allocToken);
     return;
   }
-
   // Otherwise, call llvm.stackrestore if an address was saved.
   auto savedSP = address.getExtraInfo();
   if (savedSP == nullptr)
@@ -1444,4 +1446,4 @@ void TypeInfo::storeExtraInhabitantTagDynamic(IRGenFunction &IGF,
   } else {
     storeEnumTagSinglePayload(IGF, tag, tag, address, T, isOutlined);
   }
-}
+}

test/IRGen/partial_apply.sil

@@ -602,4 +602,27 @@ bb0(%0 : $*A2<A3>):
   return %5 : $@callee_guaranteed () -> (@owned A1, @error Error)
 }
 
+sil @capture_class : $@convention(thin) (@guaranteed A3) -> ()
+
+// CHECK: define swiftcc i8* @partial_apply_stack_in_coroutine(i8* {{.*}}, %T13partial_apply2A3C*)
+// CHECK: entry:
+// CHECK:   [[CLOSURE_CONTEXT:%.*]] = bitcast %T13partial_apply2A3C* %1 to %swift.opaque*
+// CHECK:   [[CONTEXT:%.*]] = bitcast %swift.opaque* [[CLOSURE_CONTEXT]] to %swift.refcounted*
+// CHECK:   call swiftcc void @"$s13capture_classTA"(%swift.refcounted* swiftself [[CONTEXT]])
+sil @partial_apply_stack_in_coroutine : $@yield_once (@owned A3) -> () {
+entry(%0: $A3):
+  %f = function_ref @capture_class : $@convention(thin) (@guaranteed A3) -> ()
+  %p = partial_apply [callee_guaranteed] [on_stack] %f(%0) : $@convention(thin) (@guaranteed A3) -> ()
+  apply %p() : $@noescape @callee_guaranteed () -> ()
+  dealloc_stack %p : $@noescape @callee_guaranteed () -> ()
+  %1000 = integer_literal $Builtin.Int32, 1000
+  yield (), resume resume, unwind unwind
+
+resume:
+  %ret = tuple ()
+  return %ret : $()
+
+unwind:
+  unwind
+}
 sil_vtable A3 {}