[Distributed] Retain ad-hoc decodeNextArgument in distributed thunk

SIL has to force emission of `decodeNextArgument` and its retainment,
so IRGen can reference it in distributed function accessor.

(cherry picked from commit 79ad927)
(cherry picked from commit 870c712)
(cherry picked from commit dadf301)
(cherry picked from commit 4e2c059)
(cherry picked from commit 560e96f)
(cherry picked from commit 41a4cea)
(cherry picked from commit 8c8360a)
(cherry picked from commit fc662bc)

Author: ktoso

include/swift/AST/Decl.h

@@ -3554,10 +3554,10 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   VarDecl *getDistributedActorIDProperty() const;
 
   /// Find the 'RemoteCallTarget.init(_:)' initializer function.
-  ConstructorDecl* getDistributedRemoteCallTargetInitFunction() const;
+  ConstructorDecl *getDistributedRemoteCallTargetInitFunction() const;
 
   /// Find the 'RemoteCallArgument(label:name:value:)' initializer function.
-  ConstructorDecl* getDistributedRemoteCallArgumentInitFunction() const;
+  ConstructorDecl *getDistributedRemoteCallArgumentInitFunction() const;
 
   /// Collect the set of protocols to which this type should implicitly
   /// conform, such as AnyObject (for classes).

include/swift/AST/DiagnosticsParse.def

@@ -682,6 +682,8 @@ ERROR(expected_sil_function_type, none,
       "sil function expected to have SIL function type", ())
 ERROR(sil_dynamically_replaced_func_not_found,none,
       "dynamically replaced function not found %0", (Identifier))
+ERROR(sil_adhoc_requirement_witness_func_not_found,none,
+      "ad-hoc requirement witness function not found %0", (Identifier))
 ERROR(sil_specialize_target_func_not_found,none,
       "_specialize target function not found %0", (Identifier))
 ERROR(sil_availability_expected_version,none,

include/swift/AST/DistributedDecl.h

@@ -53,10 +53,21 @@ Type getDistributedActorIDType(NominalTypeDecl *actor);
 Type getDistributedSerializationRequirementType(
     NominalTypeDecl *nominal, ProtocolDecl *protocol);
 
+/// Given a distributed thunk declaration, inside a 'distributed actor',
+/// finds the ad-hoc witness for 'decodeNextArgument' on the associated
+/// 'ActorSystem.InvocationDecoder' of the actor, or null.
+AbstractFunctionDecl *
+getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+    ValueDecl *thunk);
+
 /// Get the specific 'InvocationEncoder' type of a specific distributed actor
 /// system.
 Type getDistributedActorSystemInvocationEncoderType(NominalTypeDecl *system);
 
+/// Get the specific 'InvocationDecoder' type of a specific distributed actor
+/// system.
+Type getDistributedActorSystemInvocationDecoderType(NominalTypeDecl *system);
+
 /// Get the specific 'ResultHandler' type of a specific distributed actor
 /// system.
 Type getDistributedActorSystemResultHandlerType(NominalTypeDecl *system);

include/swift/SIL/SILFunction.h

@@ -212,6 +212,17 @@ class SILFunction
   /// @_dynamicReplacement(for:) function.
   SILFunction *ReplacedFunction = nullptr;
 
+  /// This SILFunction REFerences an ad-hoc protocol requirement witness in
+  /// order to keep it alive, such that it main be obtained in IRGen. Without
+  /// this explicit reference, the witness would seem not-used, and not be
+  /// accessible for IRGen.
+  ///
+  /// Specifically, one such case is the DistributedTargetInvocationDecoder's
+  /// 'decodeNextArgument' which must be retained, as it is only used from IRGen
+  /// and such, appears as-if unused in SIL and would get optimized away.
+  // TODO: Consider making this a general "references adhoc functions" and make it an array?
+  SILFunction *RefAdHocRequirementFunction = nullptr;
+
   Identifier ObjCReplacementFor;
 
   /// The head of a single-linked list of currently alive BasicBlockBitfield.
@@ -298,7 +309,7 @@ class SILFunction
 
   /// Whether the implementation can be dynamically replaced.
   unsigned IsDynamicReplaceable : 1;
-    
+
   /// If true, this indicates that a class method implementation will always be
   /// invoked with a `self` argument of the exact base class type.
   unsigned ExactSelfClass : 1;
@@ -466,16 +477,6 @@ class SILFunction
     ReplacedFunction = f;
     ReplacedFunction->incrementRefCount();
   }
-
-  SILFunction *getDistributedRecordArgumentFunction() const {
-    return ReplacedFunction;
-  }
-  void setDistributedRecordArgumentFunction(SILFunction *f) {
-    if (f == nullptr)
-      return;
-    f->incrementRefCount();
-  }
-
   /// This function should only be called when SILFunctions are bulk deleted.
   void dropDynamicallyReplacedFunction() {
     if (!ReplacedFunction)
@@ -484,6 +485,27 @@ class SILFunction
     ReplacedFunction = nullptr;
   }
 
+  SILFunction *getReferencedAdHocRequirementWitnessFunction() const {
+    return RefAdHocRequirementFunction;
+  }
+  // Marks that this `SILFunction` uses the passed in ad-hoc protocol
+  // requirement witness `f` and therefore must retain it explicitly,
+  // otherwise we might not be able to get a reference to it.
+  void setReferencedAdHocRequirementWitnessFunction(SILFunction *f) {
+    assert(RefAdHocRequirementFunction == nullptr && "already set");
+
+    if (f == nullptr)
+      return;
+    RefAdHocRequirementFunction = f;
+    RefAdHocRequirementFunction->incrementRefCount();
+  }
+  void dropReferencedAdHocRequirementWitnessFunction() {
+    if (!RefAdHocRequirementFunction)
+      return;
+    RefAdHocRequirementFunction->decrementRefCount();
+    RefAdHocRequirementFunction = nullptr;
+  }
+
   bool hasObjCReplacement() const {
     return !ObjCReplacementFor.empty();
   }
@@ -758,7 +780,7 @@ class SILFunction
     IsDynamicReplaceable = value;
     assert(!Transparent || !IsDynamicReplaceable);
   }
-    
+
   IsExactSelfClass_t isExactSelfClass() const {
     return IsExactSelfClass_t(ExactSelfClass);
   }

lib/AST/DistributedDecl.cpp

@@ -97,6 +97,10 @@ Type swift::getConcreteReplacementForProtocolActorSystemType(ValueDecl *member)
 }
 
 Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
+  assert(!dyn_cast<ProtocolDecl>(actor) &&
+         "Use getConcreteReplacementForProtocolActorSystemType instead to get"
+         "the concrete ActorSystem, if bound, for this DistributedActor "
+         "constrained ProtocolDecl!");
   assert(actor->isDistributedActor());
   auto &C = actor->getASTContext();
 
@@ -168,6 +172,21 @@ Type swift::getDistributedActorSystemInvocationEncoderType(NominalTypeDecl *syst
   return conformance.getTypeWitnessByName(selfType, ctx.Id_InvocationEncoder);
 }
 
+Type swift::getDistributedActorSystemInvocationDecoderType(NominalTypeDecl *system) {
+  assert(!system->isDistributedActor());
+  auto &ctx = system->getASTContext();
+
+  auto DAS = ctx.getDistributedActorSystemDecl();
+  if (!DAS)
+    return Type();
+
+  // Dig out the serialization requirement type.
+  auto module = system->getParentModule();
+  Type selfType = system->getSelfInterfaceType();
+  auto conformance = module->lookupConformance(selfType, DAS);
+  return conformance.getTypeWitnessByName(selfType, ctx.Id_InvocationDecoder);
+}
+
 Type swift::getDistributedSerializationRequirementType(
     NominalTypeDecl *nominal, ProtocolDecl *protocol) {
   assert(nominal);
@@ -184,6 +203,32 @@ Type swift::getDistributedSerializationRequirementType(
   return conformance.getTypeWitnessByName(selfType, ctx.Id_SerializationRequirement);
 }
 
+AbstractFunctionDecl *
+swift::getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+    ValueDecl *thunk) {
+  assert(thunk);
+  auto &C = thunk->getASTContext();
+
+  auto *actor = thunk->getDeclContext()->getSelfNominalTypeDecl();
+  if (!actor)
+    return nullptr;
+  if (!actor->isDistributedActor())
+    return nullptr;
+
+  auto systemTy = getConcreteReplacementForProtocolActorSystemType(thunk);
+  if (!systemTy)
+    return nullptr;
+
+  auto decoderTy =
+      getDistributedActorSystemInvocationDecoderType(
+          systemTy->getAnyNominal());
+  if (!decoderTy)
+    return nullptr;
+
+  return C.getDecodeNextArgumentOnDistributedInvocationDecoder(
+      decoderTy->getAnyNominal());
+}
+
 Type ASTContext::getAssociatedTypeOfDistributedSystemOfActor(
     NominalTypeDecl *actor, Identifier member) {
   auto &ctx = actor->getASTContext();

lib/IRGen/GenDistributed.cpp

@@ -813,21 +813,13 @@ ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
     decoder = instance.claimNext();
   }
 
-  if (isa<StructDecl>(decoderDecl) || isa<EnumDecl>(decoderDecl) ||
-      decoderDecl->isFinal()) {
-    auto *decodeSIL = IGM.getSILModule().lookUpFunction(SILDeclRef(decodeFn));
-    auto *fnPtr = IGM.getAddrOfSILFunction(decodeSIL, NotForDefinition,
-                                           /*isDynamicallyReplacible=*/false);
+  auto *decodeSIL = IGM.getSILModule().lookUpFunction(SILDeclRef(decodeFn));
+  auto *fnPtr = IGM.getAddrOfSILFunction(decodeSIL, NotForDefinition,
+                                         /*isDynamicallyReplacible=*/false);
 
-    auto methodPtr = FunctionPointer::forDirect(
-        classifyFunctionPointerKind(decodeSIL), fnPtr,
-        /*secondaryValue=*/nullptr, signature);
-
-    return {decoder, decoderTy, witnessTable, methodPtr, methodTy};
-  }
-
-  auto methodPtr =
-      emitVirtualMethodValue(IGF, decoderTy, SILDeclRef(decodeFn), methodTy);
+  auto methodPtr = FunctionPointer::forDirect(
+    classifyFunctionPointerKind(decodeSIL), fnPtr,
+    /*secondaryValue=*/nullptr, signature);
 
   return {decoder, decoderTy, witnessTable, methodPtr, methodTy};
 }

lib/SIL/IR/SILFunctionBuilder.cpp

@@ -14,6 +14,7 @@
 #include "swift/AST/AttrKind.h"
 #include "swift/AST/Availability.h"
 #include "swift/AST/DiagnosticsParse.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/SemanticAttrs.h"
@@ -200,29 +201,35 @@ void SILFunctionBuilder::addFunctionAttributes(
   // Only assign replacements when the thing being replaced is function-like and
   // explicitly declared.  
   auto *origDecl = decl->getDynamicallyReplacedDecl();
-  auto *replacedDecl = dyn_cast_or_null<AbstractFunctionDecl>(origDecl);
-  if (!replacedDecl)
-    return;
-
-  // For @objc method replacement we normally use categories to perform the
-  // replacement. Except for methods in generic class where we can't. Instead,
-  // we special case this and use the native swift replacement mechanism.
-  if (decl->isObjC() && !decl->isNativeMethodReplacement()) {
-    F->setObjCReplacement(replacedDecl);
-    return;
-  }
-
-  if (!constant.canBeDynamicReplacement())
-    return;
+  if (auto *replacedDecl = dyn_cast_or_null<AbstractFunctionDecl>(origDecl)) {
+    // For @objc method replacement we normally use categories to perform the
+    // replacement. Except for methods in generic class where we can't. Instead,
+    // we special case this and use the native swift replacement mechanism.
+    if (decl->isObjC() && !decl->isNativeMethodReplacement()) {
+      F->setObjCReplacement(replacedDecl);
+      return;
+    }
 
-  SILDeclRef declRef(replacedDecl, constant.kind, false);
-  auto *replacedFunc = getOrCreateDeclaration(replacedDecl, declRef);
+    if (constant.canBeDynamicReplacement()) {
+      SILDeclRef declRef(replacedDecl, constant.kind, false);
+      auto *replacedFunc = getOrCreateDeclaration(replacedDecl, declRef);
 
-  assert(replacedFunc->getLoweredFunctionType() ==
-             F->getLoweredFunctionType() ||
-         replacedFunc->getLoweredFunctionType()->hasOpaqueArchetype());
+      assert(replacedFunc->getLoweredFunctionType() ==
+                 F->getLoweredFunctionType() ||
+             replacedFunc->getLoweredFunctionType()->hasOpaqueArchetype());
 
-  F->setDynamicallyReplacedFunction(replacedFunc);
+      F->setDynamicallyReplacedFunction(replacedFunc);
+    }
+  } else if (constant.isDistributedThunk()) {
+    auto decodeFuncDecl =
+            getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+                decl);
+    assert(decodeFuncDecl && "decodeNextArgument function not found!");
+
+    auto decodeRef = SILDeclRef(decodeFuncDecl);
+    auto *adHocFunc = getOrCreateDeclaration(decodeFuncDecl, decodeRef);
+    F->setReferencedAdHocRequirementWitnessFunction(adHocFunc);
+  }
 }
 
 SILFunction *SILFunctionBuilder::getOrCreateFunction(

lib/SIL/IR/SILModule.cpp

@@ -138,6 +138,7 @@ SILModule::~SILModule() {
   for (SILFunction &F : *this) {
     F.dropAllReferences();
     F.dropDynamicallyReplacedFunction();
+    F.dropReferencedAdHocRequirementWitnessFunction();
     F.clearSpecializeAttrs();
   }
 
@@ -467,6 +468,7 @@ void SILModule::eraseFunction(SILFunction *F) {
   // (References are not needed anymore.)
   F->clear();
   F->dropDynamicallyReplacedFunction();
+  F->dropReferencedAdHocRequirementWitnessFunction();
   // Drop references for any _specialize(target:) functions.
   F->clearSpecializeAttrs();
 }

lib/SIL/IR/SILPrinter.cpp

@@ -2983,6 +2983,12 @@ void SILFunction::print(SILPrintContext &PrintCtx) const {
     OS << "\"] ";
   }
 
+  if (auto *usedFunc = getReferencedAdHocRequirementWitnessFunction()) {
+    OS << "[ref_adhoc_requirement_witness \"";
+    OS << usedFunc->getName();
+    OS << "\"] ";
+  }
+
   if (hasObjCReplacement()) {
     OS << "[objc_replacement_for \"";
     OS << getObjCReplacement().str();