Merge pull request #42252 from ktoso/wip-retain-dist-decode

[Distributed] retain ad-hoc witnesses so that they are not optimized away in SIL

Author: xedin

include/swift/AST/Decl.h

@@ -3564,10 +3564,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

@@ -680,6 +680,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();