[SILOptimizer] Distributed: Suppress signature specialization for witnesses with ad-hoc serialization requirement

Author: xedin

include/swift/AST/Decl.h

@@ -7497,6 +7497,17 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// 'DistributedTargetInvocationResultHandler' protocol.
   bool isDistributedTargetInvocationResultHandlerOnReturn() const;
 
+  /// Determines whether this declaration is a witness to a
+  /// protocol requirement with ad-hoc `SerializationRequirement`
+  /// conformance.
+  bool isDistributedWitnessWithAdHocSerializationRequirement() const {
+    return isDistributedActorSystemRemoteCall(/*isVoidResult=*/false) ||
+           isDistributedTargetInvocationEncoderRecordArgument() ||
+           isDistributedTargetInvocationEncoderRecordReturnType() ||
+           isDistributedTargetInvocationDecoderDecodeNextArgument() ||
+           isDistributedTargetInvocationResultHandlerOnReturn();
+  }
+
   /// For a method of a class, checks whether it will require a new entry in the
   /// vtable.
   bool needsNewVTableEntry() const;

lib/IRGen/GenProto.cpp

@@ -745,23 +745,21 @@ void EmitPolymorphicParameters::injectAdHocDistributedRequirements() {
   if (!(funcDecl && funcDecl->isGeneric()))
     return;
 
+  if (!funcDecl->isDistributedWitnessWithAdHocSerializationRequirement())
+    return;
+
   Type genericParam;
 
+  auto sig = funcDecl->getGenericSignature();
+
   // DistributedActorSystem.remoteCall
   if (funcDecl->isDistributedActorSystemRemoteCall(
           /*isVoidReturn=*/false)) {
     genericParam = funcDecl->getResultInterfaceType();
-  }
-
-  auto sig = funcDecl->getGenericSignature();
-
-  // DistributedTargetInvocationEncoder.record{Argument, ReturnType}
-  // DistributedTargetInvocationDecoder.decodeNextArgument
-  // DistributedTargetInvocationResultHandler.onReturn
-  if (funcDecl->isDistributedTargetInvocationEncoderRecordArgument() ||
-      funcDecl->isDistributedTargetInvocationEncoderRecordReturnType() ||
-      funcDecl->isDistributedTargetInvocationDecoderDecodeNextArgument() ||
-      funcDecl->isDistributedTargetInvocationResultHandlerOnReturn()) {
+  } else {
+    // DistributedTargetInvocationEncoder.record{Argument, ReturnType}
+    // DistributedTargetInvocationDecoder.decodeNextArgument
+    // DistributedTargetInvocationResultHandler.onReturn
     genericParam = sig.getInnermostGenericParams().front();
   }
 

lib/SILOptimizer/FunctionSignatureTransforms/FunctionSignatureOpts.cpp

@@ -141,6 +141,17 @@ static bool canSpecializeFunction(SILFunction *F,
                                      OptForPartialApply))
     return false;
 
+  // Cannot specialize witnesses of distributed protocol requirements with
+  // ad-hoc `SerializationRequirement` because that erases information
+  // IRGen relies on to emit protocol conformances at runtime.
+  if (F->hasLocation()) {
+    if (auto *funcDecl =
+            dyn_cast_or_null<FuncDecl>(F->getLocation().getAsDeclContext())) {
+      if (funcDecl->isDistributedWitnessWithAdHocSerializationRequirement())
+        return false;
+    }
+  }
+
   return true;
 }