Merge pull request #65874 from xedin/rdar-109207043

[IRGen] Distributed: Destroy loaded arguments after the call

Author: xedin

lib/IRGen/GenDistributed.cpp

@@ -134,6 +134,9 @@ class DistributedAccessor {
   /// The list of all arguments that were allocated on the stack.
   SmallVector<StackAddress, 4> AllocatedArguments;
 
+  /// The list of all the arguments that were loaded.
+  SmallVector<std::pair<Address, /*type=*/llvm::Value *>, 4> LoadedArguments;
+
 public:
   DistributedAccessor(IRGenFunction &IGF, SILFunction *target,
                       CanSILFunctionType accessorTy);
@@ -498,13 +501,16 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
         resultValue.getAddress(), IGM.getStorageType(paramTy));
 
     cast<LoadableTypeInfo>(paramInfo).loadAsTake(IGF, eltPtr, arguments);
+    LoadedArguments.push_back(std::make_pair(eltPtr, argumentType));
     break;
   }
 
   case ParameterConvention::Direct_Owned: {
     // Copy the value out at +1.
     cast<LoadableTypeInfo>(paramInfo).loadAsCopy(IGF, resultValue.getAddress(),
                                                  arguments);
+    LoadedArguments.push_back(
+        std::make_pair(resultValue.getAddress(), argumentType));
     break;
   }
   }
@@ -584,6 +590,11 @@ void DistributedAccessor::emitReturn(llvm::Value *errorValue) {
     }
   }
 
+  // Destroy loaded arguments.
+  llvm::for_each(LoadedArguments, [&](const auto &argInfo) {
+    emitDestroyCall(IGF, argInfo.second, argInfo.first);
+  });
+
   Explosion voidResult;
 
   Explosion error;

test/Distributed/distributed_actor_thunk_doesnt_leak_class_arguments.swift

@@ -0,0 +1,51 @@
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-swift-frontend -module-name no_to_arg_leaks -emit-irgen -disable-availability-checking -I %t 2>&1 %s | %IRGenFileCheck %s -check-prefix CHECK-%target-import-type
+
+// UNSUPPORTED: back_deploy_concurrency
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// REQUIRES: CPU=x86_64 || CPU=arm64
+
+// UNSUPPORTED: OS=windows-msvc
+
+import Distributed
+import FakeDistributedActorSystems
+
+@available(SwiftStdlib 5.5, *)
+typealias DefaultDistributedActorSystem = FakeActorSystem
+
+final class SomeClass<T>: Sendable, Codable {
+  init() {}
+}
+
+struct S<T> : Codable {
+  var data: SomeClass<T>
+}
+
+distributed actor Greeter {
+  // CHECK-LABEL: define linkonce_odr hidden swifttailcc void @"$s15no_to_arg_leaks7GreeterC5test1yyAA9SomeClassCyxGYaKlFTETF"
+  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* {{.*}} to %T15no_to_arg_leaks9SomeClassC**
+  // CHECK: %destroy = bitcast i8* {{.*}} to void (%swift.opaque*, %swift.type*)*
+  // CHECK-NEXT: [[OPAQUE_ARG_ADDR:%.*]] = bitcast %T15no_to_arg_leaks9SomeClassC** [[ARG_ADDR]] to %swift.opaque*
+  // CHECK-NEXT: call void %destroy(%swift.opaque* noalias [[OPAQUE_ARG_ADDR]], %swift.type* %arg_type)
+  distributed func test1<T>(_: SomeClass<T>) {
+  }
+
+  // CHECK-LABEL: define linkonce_odr hidden swifttailcc void @"$s15no_to_arg_leaks7GreeterC5test2yyAA1SVyxGYaKlFTETF"
+  // CHECK: [[ARG_ADDR:%.*]] = bitcast i8* {{.*}} to %T15no_to_arg_leaks1SV*
+  // CHECK: %destroy = bitcast i8* {{.*}} to void (%swift.opaque*, %swift.type*)*
+  // CHECK-NEXT: [[OPAQUE_ARG_ADDR:%.*]] = bitcast %T15no_to_arg_leaks1SV* [[ARG_ADDR]] to %swift.opaque*
+  // CHECK-NEXT: call void %destroy(%swift.opaque* noalias [[OPAQUE_ARG_ADDR]], %swift.type* %arg_type)
+  distributed func test2<T>(_: S<T>) {}
+}
+
+func test() async throws {
+  let system = DefaultDistributedActorSystem()
+
+  let local = Greeter(actorSystem: system)
+  let ref = try Greeter.resolve(id: local.id, using: system)
+
+  try await ref.test1(SomeClass<Int>())
+  try await ref.test2(S(data: SomeClass<Int>()))
+}