Fix arg cleanup for objc continuations

Current SILGen for objc continuations, can end the lifetime of args
before await_async_continuation.
This PR fixes it by extending the lifetime of the args until await_async_continuation
by creating copies. And then inserts manual cleanup with fix_lifetime + destroy

Fixes rdar://78982371

Author: meg-gupta

lib/SILGen/ManagedValue.cpp

@@ -37,6 +37,24 @@ ManagedValue ManagedValue::copy(SILGenFunction &SGF, SILLocation loc) const {
   return SGF.emitManagedRValueWithCleanup(buf, lowering);
 }
 
+// Emit an unmanaged copy of this value
+// WARNING: Callers of this API should manage the cleanup of this value!
+ManagedValue ManagedValue::unmanagedCopy(SILGenFunction &SGF,
+                                         SILLocation loc) const {
+  auto &lowering = SGF.getTypeLowering(getType());
+  if (lowering.isTrivial())
+    return *this;
+
+  if (getType().isObject()) {
+    auto copy = SGF.B.emitCopyValueOperation(loc, getValue());
+    return ManagedValue::forUnmanaged(copy);
+  }
+
+  SILValue buf = SGF.emitTemporaryAllocation(loc, getType());
+  SGF.B.createCopyAddr(loc, getValue(), buf, IsNotTake, IsInitialization);
+  return ManagedValue::forUnmanaged(buf);
+}
+
 /// Emit a copy of this value with independent ownership.
 ManagedValue ManagedValue::formalAccessCopy(SILGenFunction &SGF,
                                             SILLocation loc) {

lib/SILGen/ManagedValue.h

@@ -270,6 +270,10 @@ class ManagedValue {
   /// Emit a copy of this value with independent ownership.
   ManagedValue copy(SILGenFunction &SGF, SILLocation loc) const;
 
+  /// Returns an unmanaged copy of this value.
+  /// WARNING: Callers of this API should manage the cleanup of this value!
+  ManagedValue unmanagedCopy(SILGenFunction &SGF, SILLocation loc) const;
+
   /// Emit a copy of this value with independent ownership into the current
   /// formal evaluation scope.
   ManagedValue formalAccessCopy(SILGenFunction &SGF, SILLocation loc);

lib/SILGen/SILGenApply.cpp

@@ -4288,6 +4288,29 @@ bool SILGenModule::isNonMutatingSelfIndirect(SILDeclRef methodRef) {
   return self.isFormalIndirect();
 }
 
+namespace {
+/// Cleanup to insert fix_lifetime and destroy
+class FixLifetimeDestroyCleanup : public Cleanup {
+  SILValue val;
+
+public:
+  FixLifetimeDestroyCleanup(SILValue val) : val(val) {}
+
+  void emit(SILGenFunction &SGF, CleanupLocation l,
+            ForUnwind_t forUnwind) override {
+    SGF.B.emitFixLifetime(l, val);
+    SGF.B.emitDestroyOperation(l, val);
+  }
+
+  void dump(SILGenFunction &SGF) const override {
+#ifndef NDEBUG
+    llvm::errs() << "FixLifetimeDestroyCleanup "
+                 << "State:" << getState() << " "
+                 << "Value: " << val << "\n";
+#endif
+  }
+};
+} // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 //                           Top Level Entrypoints
@@ -4437,6 +4460,19 @@ RValue SILGenFunction::emitApply(
   // hop back to the current executor
   breadcrumb.emit(*this, loc);
 
+  // For objc async calls, lifetime extend the args until the result plan which
+  // generates `await_async_continuation`.
+  // Lifetime is extended by creating unmanaged copies here and by pushing the
+  // cleanups required just before the result plan is generated.
+  SmallVector<ManagedValue, 8> unmanagedCopies;
+  if (calleeTypeInfo.foreign.async) {
+    for (auto arg : args) {
+      if (arg.hasCleanup()) {
+        unmanagedCopies.push_back(arg.unmanagedCopy(*this, loc));
+      }
+    }
+  }
+
   // Pop the argument scope.
   argScope.pop();
 
@@ -4512,12 +4548,30 @@ RValue SILGenFunction::emitApply(
     emitForeignErrorCheck(loc, directResults, errorTemp, doesNotThrow,
                           *foreignError);
   }
-  
+
+  // For objc async calls, push cleanup to be used on throw paths in the result
+  // planner.
+  for (unsigned i : indices(unmanagedCopies)) {
+    SILValue value = unmanagedCopies[i].getValue();
+    Cleanups.pushCleanup<FixLifetimeDestroyCleanup>(value);
+    unmanagedCopies[i] = ManagedValue(value, Cleanups.getTopCleanup());
+  }
+
   auto directResultsArray = makeArrayRef(directResults);
   RValue result =
     resultPlan->finish(*this, loc, substResultType, directResultsArray);
   assert(directResultsArray.empty() && "didn't claim all direct results");
 
+  // For objc async calls, generate cleanup on the resume path here and forward
+  // the previously pushed cleanups.
+  if (calleeTypeInfo.foreign.async) {
+    for (auto unmanagedCopy : unmanagedCopies) {
+      auto value = unmanagedCopy.forward(*this);
+      B.emitFixLifetime(loc, value);
+      B.emitDestroyOperation(loc, value);
+    }
+  }
+
   return result;
 }
 

test/SILGen/objc_async.swift

@@ -8,6 +8,8 @@ import ObjCConcurrency
 // CHECK-LABEL: sil {{.*}}@${{.*}}14testSlowServer
 func testSlowServer(slowServer: SlowServer) async throws {
   // CHECK: [[RESUME_BUF:%.*]] = alloc_stack $Int
+  // CHECK: [[STRINGINIT:%.*]] = function_ref @$sSS10FoundationE19_bridgeToObjectiveCSo8NSStringCyF :
+  // CHECK: [[ARG:%.*]] = apply [[STRINGINIT]]
   // CHECK: [[METHOD:%.*]] = objc_method {{.*}} $@convention(objc_method) (NSString, @convention(block) (Int) -> (), SlowServer) -> ()
   // CHECK: [[CONT:%.*]] = get_async_continuation_addr Int, [[RESUME_BUF]]
   // CHECK: [[WRAPPED:%.*]] = struct $UnsafeContinuation<Int, Never> ([[CONT]] : $Builtin.RawUnsafeContinuation)
@@ -16,10 +18,14 @@ func testSlowServer(slowServer: SlowServer) async throws {
   // CHECK: store [[WRAPPED]] to [trivial] [[CONT_SLOT]]
   // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @[[INT_COMPLETION_BLOCK:.*]] : $@convention(c) (@inout_aliasable @block_storage UnsafeContinuation<Int, Never>, Int) -> ()
   // CHECK: [[BLOCK:%.*]] = init_block_storage_header [[BLOCK_STORAGE]] {{.*}}, invoke [[BLOCK_IMPL]]
-  // CHECK: apply [[METHOD]]({{.*}}, [[BLOCK]], %0)
+  // CHECK: apply [[METHOD]]([[ARG]], [[BLOCK]], %0)
+  // CHECK: [[COPY:%.*]] = copy_value [[ARG]]
+  // CHECK: destroy_value [[ARG]]
   // CHECK: await_async_continuation [[CONT]] {{.*}}, resume [[RESUME:bb[0-9]+]]
   // CHECK: [[RESUME]]:
   // CHECK: [[RESULT:%.*]] = load [trivial] [[RESUME_BUF]]
+  // CHECK: fix_lifetime [[COPY]]
+  // CHECK: destroy_value [[COPY]]
   // CHECK: dealloc_stack [[RESUME_BUF]]
   let _: Int = await slowServer.doSomethingSlow("mail")