[Sema/IRGen] Extend ad-hoc requirement handling to DistributedTargetInvocationResultHandler.onReturn

Author: xedin

lib/IRGen/GenProto.cpp

@@ -552,6 +552,7 @@ class EmitPolymorphicParameters : public PolymorphicConvention {
   // `DistributedActorSystem.remoteCall` or
   // `DistributedTargetInvocationEncoder.record{Argument, ReturnType}`
   // `DistributedTargetInvocationDecoder.decodeNextArgument`
+  // `DistributedTargetInvocationResultHandler.onReturn`
   // requirements we need to supply witness tables associated with `Res`,
   // `Argument`, `R` generic parameters which are not expressible on the
   // protocol requirement because they come from `SerializationRequirement`
@@ -752,19 +753,21 @@ void EmitPolymorphicParameters::injectAdHocDistributedRequirements() {
     genericParam = funcDecl->getResultInterfaceType();
   }
 
+  auto sig = funcDecl->getGenericSignature();
+
   // DistributedTargetInvocationEncoder.record{Argument, ReturnType}
   // DistributedTargetInvocationDecoder.decodeNextArgument
+  // DistributedTargetInvocationResultHandler.onReturn
   if (funcDecl->isDistributedTargetInvocationEncoderRecordArgument() ||
       funcDecl->isDistributedTargetInvocationEncoderRecordReturnType() ||
-      funcDecl->isDistributedTargetInvocationDecoderDecodeNextArgument()) {
-    auto *argParam = funcDecl->getGenericParams()->getParams().front();
-    genericParam = argParam->getDeclaredInterfaceType();
+      funcDecl->isDistributedTargetInvocationDecoderDecodeNextArgument() ||
+      funcDecl->isDistributedTargetInvocationResultHandlerOnReturn()) {
+    genericParam = sig.getInnermostGenericParams().front();
   }
 
   if (!genericParam)
     return;
 
-  auto sig = funcDecl->getGenericSignature();
   auto requirements = sig->getLocalRequirements(genericParam);
   if (requirements.protos.empty())
     return;

lib/Sema/CSSimplify.cpp

@@ -8505,9 +8505,12 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyConformsToConstraint(
 
           // `DistributedTargetInvocationEncoder.record{Argument, ResultType}`
           // `DistributedTargetInvocationDecoder.decodeNextArgument`
+          // `DistributedTargetInvocationResultHandler.onReturn`
           if (witness->isDistributedTargetInvocationEncoderRecordArgument() ||
               witness->isDistributedTargetInvocationEncoderRecordReturnType() ||
-              witness->isDistributedTargetInvocationDecoderDecodeNextArgument()) {
+              witness
+                  ->isDistributedTargetInvocationDecoderDecodeNextArgument() ||
+              witness->isDistributedTargetInvocationResultHandlerOnReturn()) {
             auto genericParams = witness->getGenericParams()->getParams();
             if (GP->isEqual(genericParams.front()->getDeclaredInterfaceType()))
               return synthesizeConformance();

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -947,11 +947,9 @@ public protocol DistributedTargetInvocationResultHandler {
   /// The serialization requirement that the value passed to `onReturn` is required to conform to.
   associatedtype SerializationRequirement
 
-//  /// Ad-hoc protocol requirement
-//  ///
-//  /// Invoked when the distributed target execution returns successfully.
-//  /// The `value` is the return value of the executed distributed invocation target.
-//  func onReturn<Success: SerializationRequirement>(value: Success) async throws
+  /// Invoked when the distributed target execution returns successfully.
+  /// The `value` is the return value of the executed distributed invocation target.
+  func onReturn<Success/*: SerializationRequirement*/>(value: Success) async throws
 
   /// Invoked when the distributed target execution of a `Void` returning
   /// function has completed successfully.