[Distributed] Adhoc witness methods must work with non final class

Author: ktoso

include/swift/AST/Decl.h

@@ -5557,6 +5557,19 @@ class ProtocolDecl final : public NominalTypeDecl {
       EffectKind kind) const;
   bool hasPolymorphicEffect(EffectKind kind) const;
 
+  /// Returns `true` if this protocol contains requirements which have known
+  /// 'ad-hoc' requirements. For example, the protocol contains an associated
+  /// type which is used to constrain generic types of some of its members.
+  ///
+  /// Currently the only instance of such requirement is
+  /// 'SerializationRequirement' in the 'Distributed' module's protocols.
+  bool hasMembersWithAdHocRequirements() const {
+    return isSpecificProtocol(KnownProtocolKind::DistributedActorSystem) ||
+        isSpecificProtocol(KnownProtocolKind::DistributedTargetInvocationEncoder) ||
+        isSpecificProtocol(KnownProtocolKind::DistributedTargetInvocationDecoder) ||
+        isSpecificProtocol(KnownProtocolKind::DistributedTargetInvocationResultHandler);
+  }
+
   /// Determine whether this is a "marker" protocol, meaning that is indicates
   /// semantics but has no corresponding witness table.
   bool isMarkerProtocol() const;
@@ -7960,7 +7973,7 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Determines whether this function is a 'remoteCall' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedActorSystem' protocol.
-  bool isDistributedActorSystemRemoteCall(bool isVoidReturn) const;
+  bool isDistributedActorSystemRemoteCall(bool isVoidReturn, bool allowRequirement) const;
 
   /// Determines whether this function is a 'makeInvocationEncoder' function,
   /// which is used as ad-hoc protocol requirement by the
@@ -7970,42 +7983,42 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// Determines if this function is a 'recordGenericSubstitution' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationEncoder' protocol.
-  bool isDistributedTargetInvocationEncoderRecordGenericSubstitution() const;
+  bool isDistributedTargetInvocationEncoderRecordGenericSubstitution(bool allowRequirement) const;
 
   /// Determines if this function is a 'recordArgument' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationEncoder' protocol.
-  bool isDistributedTargetInvocationEncoderRecordArgument() const;
+  bool isDistributedTargetInvocationEncoderRecordArgument(bool allowRequirement) const;
 
   /// Determines if this function is a 'recordReturnType' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationEncoder' protocol.
-  bool isDistributedTargetInvocationEncoderRecordReturnType() const;
+  bool isDistributedTargetInvocationEncoderRecordReturnType(bool allowRequirement) const;
 
   /// Determines if this function is a 'recordErrorType' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationEncoder' protocol.
-  bool isDistributedTargetInvocationEncoderRecordErrorType() const;
+  bool isDistributedTargetInvocationEncoderRecordErrorType(bool allowRequirement) const;
 
   /// Determines if this function is a 'decodeNextArgument' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationDecoder' protocol.
-  bool isDistributedTargetInvocationDecoderDecodeNextArgument() const;
+  bool isDistributedTargetInvocationDecoderDecodeNextArgument(bool allowRequirement) const;
 
   /// Determines if this function is a 'onReturn' function,
   /// which is used as ad-hoc protocol requirement by the
   /// 'DistributedTargetInvocationResultHandler' protocol.
-  bool isDistributedTargetInvocationResultHandlerOnReturn() const;
+  bool isDistributedTargetInvocationResultHandlerOnReturn(bool allowRequirement) 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();
+  bool isDistributedWitnessWithAdHocSerializationRequirement(bool allowRequirement) const {
+    return isDistributedActorSystemRemoteCall(/*isVoidResult=*/false, allowRequirement) ||
+           isDistributedTargetInvocationEncoderRecordArgument(allowRequirement) ||
+           isDistributedTargetInvocationEncoderRecordReturnType(allowRequirement) ||
+           isDistributedTargetInvocationDecoderDecodeNextArgument(allowRequirement) ||
+           isDistributedTargetInvocationResultHandlerOnReturn(allowRequirement);
   }
 
   /// For a method of a class, checks whether it will require a new entry in the

include/swift/AST/DistributedDecl.h

@@ -148,6 +148,10 @@ getDistributedSerializationRequirements(
     ProtocolDecl *protocol,
     llvm::SmallPtrSetImpl<ProtocolDecl *> &requirementProtos);
 
+void injectAdHocDistributedSerializationRequirementForWitnessThunk(
+    const ValueDecl *requirementDecl, const FuncDecl *witnessDecl,
+    SmallVectorImpl<Requirement> &Requirements);
+
 /// Retrieve the declaration of DistributedActorSystem.remoteCall(Void)(...).
 ///
 /// \param actorOrSystem distributed actor or actor system to get the

include/swift/AST/RequirementEnvironment.h

@@ -103,6 +103,13 @@ class RequirementEnvironment {
                          ClassDecl *covariantSelf,
                          RootProtocolConformance *conformance);
 
+  RequirementEnvironment(DeclContext *conformanceDC,
+                         ValueDecl *reqDecl,
+                         GenericSignature reqSig,
+                         ProtocolDecl *proto,
+                         ClassDecl *covariantSelf,
+                         RootProtocolConformance *conformance);
+
   /// Retrieve the generic signature of the requirement.
   GenericSignature getRequirementSignature() const {
     return reqSig;

include/swift/AST/Witness.h

@@ -180,6 +180,7 @@ class Witness {
 
   /// Retrieve the witness thunk generic signature.
   GenericSignature getWitnessThunkSignature() const {
+    // TODO: find what creates/sets that
     if (auto *storedWitness = storage.dyn_cast<StoredWitness *>())
       return storedWitness->witnessThunkSig;
     return nullptr;