[Distributed] ad-hoc req for decodeNextArgument & onReturns

fixup

Author: ktoso

include/swift/AST/ASTContext.h

@@ -688,22 +688,32 @@ class ASTContext final {
   FuncDecl *getRecordGenericSubstitutionOnDistributedInvocationEncoder(
       NominalTypeDecl *nominal) const;
 
-  // Retrieve the declaration of DistributedInvocationEncoder.recordArgument(_:).
+  // Retrieve the declaration of DistributedTargetInvocationEncoder.recordArgument(_:).
   //
   // \param nominal optionally provide a 'NominalTypeDecl' from which the
   // function decl shall be extracted. This is useful to avoid witness calls
   // through the protocol which is looked up when nominal is null.
   AbstractFunctionDecl *getRecordArgumentOnDistributedInvocationEncoder(
       NominalTypeDecl *nominal) const;
 
-  // Retrieve the declaration of DistributedInvocationEncoder.recordReturnType(_:).
+  // Retrieve the declaration of DistributedTargetInvocationEncoder.recordReturnType(_:).
   AbstractFunctionDecl *getRecordReturnTypeOnDistributedInvocationEncoder(
       NominalTypeDecl *nominal) const;
 
-  // Retrieve the declaration of DistributedInvocationEncoder.recordErrorType(_:).
+  // Retrieve the declaration of DistributedTargetInvocationEncoder.recordErrorType(_:).
   AbstractFunctionDecl *getRecordErrorTypeOnDistributedInvocationEncoder(
       NominalTypeDecl *nominal) const;
 
+  // Retrieve the declaration of
+  // DistributedTargetInvocationDecoder.getDecodeNextArgumentOnDistributedInvocationDecoder(_:).
+  AbstractFunctionDecl *getDecodeNextArgumentOnDistributedInvocationDecoder(
+      NominalTypeDecl *nominal) const;
+
+  // Retrieve the declaration of
+  // getOnReturnOnDistributedTargetInvocationResultHandler.onReturn(_:).
+  AbstractFunctionDecl *getOnReturnOnDistributedTargetInvocationResultHandler(
+      NominalTypeDecl *nominal) const;
+
   // Retrieve the declaration of DistributedInvocationEncoder.doneRecording().
   //
   // \param nominal optionally provide a 'NominalTypeDecl' from which the

include/swift/AST/Decl.h

@@ -6437,6 +6437,16 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// 'DistributedTargetInvocationEncoder' protocol.
   bool isDistributedTargetInvocationEncoderRecordErrorType() const;
 
+  /// Determines if this function is a 'decodeNextArgument' function,
+  /// which is used as ad-hoc protocol requirement by the
+  /// 'DistributedTargetInvocationDecoder' protocol.
+  bool isDistributedTargetInvocationDecoderDecodeNextArgument() const;
+
+  /// Determines if this function is a 'onReturn' function,
+  /// which is used as ad-hoc protocol requirement by the
+  /// 'DistributedTargetInvocationResultHandler' protocol.
+  bool isDistributedTargetInvocationResultHandlerOnReturn() const;
+
   /// For a method of a class, checks whether it will require a new entry in the
   /// vtable.
   bool needsNewVTableEntry() const;

include/swift/AST/DistributedDecl.h

@@ -80,10 +80,15 @@ bool checkDistributedSerializationRequirementIsExactlyCodable(
 /// Get the `SerializationRequirement`, explode it into the specific
 /// protocol requirements and insert them into `requirements`.
 ///
+/// The passed `protocol` must be conformed to by the `decl`, e.g. a specific
+/// actor system implementation and the `DistributedActorSystem` protocol,
+/// or any of the specific encoder/decoder and the respective
+/// Distributed...Encoder/Decoder protocol etc.
+///
 /// Returns false if failed to get the protocol decls.
 bool
-getDistributedActorSystemSerializationRequirements(
-    NominalTypeDecl *systemDecl,
+getDistributedSerializationRequirements(
+    NominalTypeDecl *decl,
     ProtocolDecl *protocol,
     llvm::SmallPtrSetImpl<ProtocolDecl *> &requirementProtos);
 

include/swift/AST/KnownIdentifiers.def

@@ -276,6 +276,7 @@ IDENTIFIER(invocation)
 IDENTIFIER(invocationDecoder)
 IDENTIFIER(makeInvocationEncoder)
 IDENTIFIER(on)
+IDENTIFIER(onReturn)
 IDENTIFIER(recordArgument)
 IDENTIFIER(recordErrorType)
 IDENTIFIER(recordGenericSubstitution)

include/swift/AST/TypeCheckRequests.h

@@ -1109,6 +1109,42 @@ class GetDistributedTargetInvocationEncoderRecordErrorTypeFunctionRequest :
     bool isCached() const { return true; }
 };
 
+/// Obtain the 'decodeNextArgument' function of a 'DistributedTargetInvocationDecoder'.
+class GetDistributedTargetInvocationDecoderDecodeNextArgumentFunctionRequest :
+    public SimpleRequest<GetDistributedTargetInvocationDecoderDecodeNextArgumentFunctionRequest,
+                         AbstractFunctionDecl *(NominalTypeDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  AbstractFunctionDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *encoder) const;
+
+public:
+    // Caching
+    bool isCached() const { return true; }
+};
+
+/// Obtain the 'onReturn' function of a 'DistributedTargetInvocationResultHandler'.
+class GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest :
+    public SimpleRequest<GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest,
+                         AbstractFunctionDecl *(NominalTypeDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  AbstractFunctionDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *encoder) const;
+
+public:
+    // Caching
+    bool isCached() const { return true; }
+};
+
 /// Obtain the 'actorSystem' property of a 'distributed actor'.
 class GetDistributedActorSystemPropertyRequest :
     public SimpleRequest<GetDistributedActorSystemPropertyRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -117,6 +117,12 @@ SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationEncoderRecordReturnType
 SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationEncoderRecordErrorTypeFunctionRequest,
               AbstractFunctionDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationDecoderDecodeNextArgumentFunctionRequest,
+              AbstractFunctionDecl *(NominalTypeDecl *),
+              Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest,
+              AbstractFunctionDecl *(NominalTypeDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetDistributedActorIDPropertyRequest,
               VarDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)

lib/AST/ASTContext.cpp

@@ -1382,6 +1382,22 @@ AbstractFunctionDecl *ASTContext::getRecordErrorTypeOnDistributedInvocationEncod
       nullptr);
 }
 
+AbstractFunctionDecl *ASTContext::getDecodeNextArgumentOnDistributedInvocationDecoder(
+    NominalTypeDecl *nominal) const {
+  return evaluateOrDefault(
+      nominal->getASTContext().evaluator,
+      GetDistributedTargetInvocationDecoderDecodeNextArgumentFunctionRequest{nominal},
+      nullptr);
+}
+
+AbstractFunctionDecl *ASTContext::getOnReturnOnDistributedTargetInvocationResultHandler(
+    NominalTypeDecl *nominal) const {
+  return evaluateOrDefault(
+      nominal->getASTContext().evaluator,
+      GetDistributedTargetInvocationResultHandlerOnReturnFunctionRequest{nominal},
+      nullptr);
+}
+
 FuncDecl *ASTContext::getDoneRecordingOnDistributedInvocationEncoder(
     NominalTypeDecl *nominal) const {
   for (auto result : nominal->lookupDirect(Id_doneRecording)) {