[Sema/IRGen] Extend ad-hoc requirement handling to DistributedTargetInvocationEncoder.record{Argument, ReturnType}

Author: xedin

lib/IRGen/GenProto.cpp

@@ -548,13 +548,14 @@ class EmitPolymorphicParameters : public PolymorphicConvention {
   // was a class-pointer source?
   bool isClassPointerSource(unsigned paramIndex);
 
-  // If we are building a protocol witness thunk for a
-  // `DistributedActorSystem.remoteCall` requirement we
-  // need to supply witness tables associated with `Res`
-  // generic parameter which are not expressible on the
-  // requirement because they come from `SerializationRequirement`
+  // If we are building a protocol witness thunks for
+  // `DistributedActorSystem.remoteCall` or
+  // `DistributedTargetInvocationEncoder.record{Argument, ReturnType}`
+  // 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`
   // associated type.
-  void injectAdHocDistributedRemoteCallRequirements();
+  void injectAdHocDistributedRequirements();
 };
 
 } // end anonymous namespace
@@ -730,35 +731,53 @@ bool EmitPolymorphicParameters::isClassPointerSource(unsigned paramIndex) {
   return false;
 }
 
-void EmitPolymorphicParameters::injectAdHocDistributedRemoteCallRequirements() {
+void EmitPolymorphicParameters::injectAdHocDistributedRequirements() {
   // FIXME: We need a better way to recognize that function is
   // a thunk for witness of `remoteCall` requirement.
   if (!Fn.hasLocation())
     return;
 
   auto loc = Fn.getLocation();
 
+  Type genericParam;
+
   auto *funcDecl = dyn_cast_or_null<FuncDecl>(loc.getAsDeclContext());
-  if (!(funcDecl && funcDecl->isDistributedActorSystemRemoteCall(
-                        /*isVoidReturn=*/false)))
+  if (!(funcDecl && funcDecl->isGeneric()))
+    return;
+
+  // DistributedActorSystem.remoteCall
+  if (funcDecl->isDistributedActorSystemRemoteCall(
+          /*isVoidReturn=*/false)) {
+    genericParam = funcDecl->getResultInterfaceType();
+  }
+
+  // DistributedTargetInvocationEncoder.record{Argument, ReturnType}
+  if (funcDecl->isDistributedTargetInvocationEncoderRecordArgument() ||
+      funcDecl->isDistributedTargetInvocationEncoderRecordReturnType()) {
+    auto *argParam = funcDecl->getGenericParams()->getParams().front();
+    genericParam = argParam->getDeclaredInterfaceType();
+  }
+
+  if (!genericParam)
     return;
 
   auto sig = funcDecl->getGenericSignature();
-  auto resultInterfaceTy = funcDecl->getResultInterfaceType();
-  auto resultArchetypeTy =
-      getTypeInContext(resultInterfaceTy->getCanonicalType());
-  llvm::Value *resultMetadata = IGF.emitTypeMetadataRef(resultArchetypeTy);
+  auto requirements = sig->getLocalRequirements(genericParam);
+  if (requirements.protos.empty())
+    return;
 
-  auto resultRequirements = sig->getLocalRequirements(resultInterfaceTy);
-  for (auto *proto : resultRequirements.protos) {
+  auto archetypeTy = getTypeInContext(genericParam->getCanonicalType());
+  llvm::Value *metadata = IGF.emitTypeMetadataRef(archetypeTy);
+
+  for (auto *proto : requirements.protos) {
     // Lookup the witness table for this protocol dynamically via
     // swift_conformsToProtocol(<<archetype>>, <<protocol>>)
     auto *witnessTable = IGF.Builder.CreateCall(
         IGM.getConformsToProtocolFunctionPointer(),
-        {resultMetadata, IGM.getAddrOfProtocolDescriptor(proto)});
+        {metadata, IGM.getAddrOfProtocolDescriptor(proto)});
 
     IGF.setUnscopedLocalTypeData(
-        resultArchetypeTy,
+        archetypeTy,
         LocalTypeDataKind::forAbstractProtocolWitnessTable(proto),
         witnessTable);
   }
@@ -2617,8 +2636,9 @@ void EmitPolymorphicParameters::emit(EntryPointArgumentEmission &emission,
   // Bind all the fulfillments we can from the formal parameters.
   bindParameterSources(getParameter);
 
-  // Inject ad-hoc `remoteCall` requirements if any.
-  injectAdHocDistributedRemoteCallRequirements();
+  // Inject ad-hoc requirements related to `SerializationRequirement`
+  // associated type.
+  injectAdHocDistributedRequirements();
 }
 
 MetadataResponse

lib/Sema/CSSimplify.cpp

@@ -8591,9 +8591,20 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyConformsToConstraint(
           return recordConformance(synthesized);
         };
 
-        if (witness->isDistributedActorSystemRemoteCall(/*isVoidReturn=*/false)) {
-          if (GP->isEqual(cast<FuncDecl>(witness)->getResultInterfaceType()))
-            return synthesizeConformance();
+        if (witness->isGeneric()) {
+          // `DistributedActorSystem.remoteCall`
+          if (witness->isDistributedActorSystemRemoteCall(/*isVoidReturn=*/false)) {
+            if (GP->isEqual(cast<FuncDecl>(witness)->getResultInterfaceType()))
+              return synthesizeConformance();
+          }
+
+          // `DistributedTargetInvocationEncoder.record{Argument, ResultType}`
+          if (witness->isDistributedTargetInvocationEncoderRecordArgument() ||
+              witness->isDistributedTargetInvocationEncoderRecordReturnType()) {
+            auto genericParams = witness->getGenericParams()->getParams();
+            if (GP->isEqual(genericParams.front()->getDeclaredInterfaceType()))
+              return synthesizeConformance();
+          }
         }
       }
     }

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -749,25 +749,21 @@ public protocol DistributedTargetInvocationEncoder {
   /// - Parameter type: a generic substitution type to be recorded for this invocation.
   mutating func recordGenericSubstitution<T>(_ type: T.Type) throws
 
-//  /// Ad-hoc requirement
-//  ///
-//  /// Record an argument of `Argument` type.
-//  /// This will be invoked for every argument of the target, in declaration order.
-//  mutating func recordArgument<Value: SerializationRequirement>(
-//    _ argument: DistributedTargetArgument<Value>
-//  ) throws
+  /// Record an argument of `Argument` type.
+  /// This will be invoked for every argument of the target, in declaration order.
+  mutating func recordArgument<Value/*: SerializationRequirement*/>(
+    _ argument: RemoteCallArgument<Value>
+  ) throws
 
   /// Record the error type of the distributed method.
   /// This method will not be invoked if the target is not throwing.
   ///
   /// - Parameter type: the type of error that was declared to be thrown by the invocation target. Currently this can only ever be `Error.self`.
   mutating func recordErrorType<E: Error>(_ type: E.Type) throws
 
-//  /// Ad-hoc requirement
-//  ///
-//  /// Record the return type of the distributed method.
-//  /// This method will not be invoked if the target is returning `Void`.
-//  mutating func recordReturnType<R: SerializationRequirement>(_ type: R.Type) throws
+  /// Record the return type of the distributed method.
+  /// This method will not be invoked if the target is returning `Void`.
+  mutating func recordReturnType<R/*: SerializationRequirement*/>(_ type: R.Type) throws
 
   /// Invoked to signal to the encoder that no further `record...` calls will be made on it.
   ///