[Distributed] Retain adhoc decodeNextArgument in distributed thunk

Author: ktoso

include/swift/AST/Decl.h

@@ -2258,6 +2258,9 @@ class ValueDecl : public Decl {
     /// Whether this declaration produces an implicitly unwrapped
     /// optional result.
     unsigned isIUO : 1;
+//
+//    /// Whether this declaration is a witness to an ad-hoc requirement.
+//    unsigned isAdHocRequirementWitness : 1;
   } LazySemanticInfo = { };
 
   friend class DynamicallyReplacedDeclRequest;
@@ -2595,6 +2598,14 @@ class ValueDecl : public Decl {
   /// Returns true if this decl can be found by id-style dynamic lookup.
   bool canBeAccessedByDynamicLookup() const;
 
+//  /// Set whether this declaration is a witness to an ad-hoc requirement or not.
+//  void setIsAdHocRequirementWitness(bool value);
+//
+//  /// Whether this declaration is a witness to an ad-hoc requirement.
+//  bool isAdHocRequirementWitness() const {
+//    return LazySemanticInfo.isAdHocRequirementWitness;
+//  }
+
   /// Returns true if this declaration has an implicitly unwrapped optional
   /// result. The precise meaning depends on the declaration kind:
   /// - for properties, the value is IUO

include/swift/AST/DistributedDecl.h

@@ -57,6 +57,10 @@ Type getDistributedSerializationRequirementType(
 /// 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/SILDeclRef.h

@@ -555,7 +555,12 @@ struct SILDeclRef {
 
   bool canBeDynamicReplacement() const;
 
-  bool isAutoDiffDerivativeFunction() const {
+//  /// True if the function was a witness to an ad-hoc requirement.
+//  /// Useful in order to e.g. not optimize away the function even
+//  /// if it seems not-used to the optimizer based on information in SIL.
+//  bool isAdHocRequirementWitness() const;
+
+    bool isAutoDiffDerivativeFunction() const {
     return pointer.is<AutoDiffDerivativeFunctionIdentifier *>() &&
            pointer.get<AutoDiffDerivativeFunctionIdentifier *>() != nullptr;
   }

include/swift/SIL/SILFunction.h

@@ -55,6 +55,10 @@ enum IsDynamicallyReplaceable_t {
   IsNotDynamic,
   IsDynamic
 };
+//enum IsAdHocRequirementWitness_t {
+//  IsNotAdHocWitness,
+//  IsAdHocWitness
+//};
 enum IsExactSelfClass_t {
   IsNotExactSelfClass,
   IsExactSelfClass,
@@ -212,6 +216,15 @@ class SILFunction
   /// @_dynamicReplacement(for:) function.
   SILFunction *ReplacedFunction = nullptr;
 
+  /// Ad-hoc requirements may need to be retained explicitly, such that they
+  /// do not get optimized away as it might seem tha they are un-used.
+  ///
+  /// Specifically, e.g. the DistributedTargetInvocationDecoder's
+  /// 'decodeNextArgument' must be retained, as it is only used from IRGen
+  /// and such, appears as-if unused in SIL and would get optimized away.
+  // TODO: rename: referenced ad hoc requirement witnesses
+  SILFunction *AdHocRequirementFunction = nullptr;
+
   Identifier ObjCReplacementFor;
 
   /// The head of a single-linked list of currently alive BasicBlockBitfield.
@@ -298,7 +311,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;
@@ -350,6 +363,9 @@ class SILFunction
   /// The function is in a statically linked module.
   unsigned IsStaticallyLinked : 1;
 
+//  /// The function is a witness to an ad-hoc requirement.
+//  unsigned IsAdHocRequirementWitness : 1;
+
   static void
   validateSubclassScope(SubclassScope scope, IsThunk_t isThunk,
                         const GenericSpecializationInformation *genericInfo) {
@@ -466,16 +482,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 +490,30 @@ class SILFunction
     ReplacedFunction = nullptr;
   }
 
+  SILFunction *getUsedAdHocRequirementWitnessFunction() const {
+    return AdHocRequirementFunction;
+  }
+  // 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 setUsedAdHocRequirementWitnessFunction(SILFunction *f) {
+    assert(AdHocRequirementFunction == nullptr && "already set");
+
+    fprintf(stderr, "[%s:%d] (%s) SET AD HOC WITNESS [%s] ON [%s]\n", __FILE__, __LINE__, __FUNCTION__, f->getName().str().c_str(), this->getName().str().c_str());
+    f->dump();
+
+    if (f == nullptr)
+      return;
+    AdHocRequirementFunction = f;
+    AdHocRequirementFunction->incrementRefCount();
+  }
+  void dropAdHocRequirementFunction() {
+    if (!AdHocRequirementFunction)
+      return;
+    AdHocRequirementFunction->decrementRefCount();
+    AdHocRequirementFunction = nullptr;
+  }
+
   bool hasObjCReplacement() const {
     return !ObjCReplacementFor.empty();
   }
@@ -758,6 +788,15 @@ class SILFunction
     IsDynamicReplaceable = value;
     assert(!Transparent || !IsDynamicReplaceable);
   }
+
+  /// Returns whether this function implementation can be dynamically replaced.
+//  IsAdHocRequirementWitness_t isAdHocRequirementWitness() const {
+//    return IsAdHocRequirementWitness_t(IsAdHocRequirementWitness);
+//  }
+//  void setIsAdHocRequirementWitness(IsAdHocRequirementWitness_t value = IsAdHocWitness) {
+//    IsAdHocRequirementWitness = value;
+//    assert(!IsDynamicReplaceable);
+//  }
 
   IsExactSelfClass_t isExactSelfClass() const {
     return IsExactSelfClass_t(ExactSelfClass);

lib/AST/Decl.cpp

@@ -3167,6 +3167,10 @@ void ValueDecl::setIsDynamic(bool value) {
   LazySemanticInfo.isDynamic = value;
 }
 
+//void ValueDecl::setIsAdHocRequirementWitness(bool value) {
+//  LazySemanticInfo.isAdHocRequirementWitness = true;
+//}
+
 ValueDecl *ValueDecl::getDynamicallyReplacedDecl() const {
   return evaluateOrDefault(getASTContext().evaluator,
                            DynamicallyReplacedDeclRequest{

lib/AST/DistributedDecl.cpp

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

lib/SIL/IR/SILDeclRef.cpp

@@ -1382,6 +1382,14 @@ bool SILDeclRef::canBeDynamicReplacement() const {
   return true;
 }
 
+//bool SILDeclRef::isAdHocRequirementWitness() const {
+//  auto decl = getDecl();
+//  if (!decl)
+//    return false;
+//
+//  return decl->isAdHocRequirementWitness();
+//}
+
 bool SILDeclRef::isDynamicallyReplaceable() const {
   // The non-foreign entry of a @dynamicReplacement(for:) of @objc method in a
   // generic class can't be a dynamically replaced.

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,52 @@ 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 (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;
+    }
 
-  if (!constant.canBeDynamicReplacement())
-    return;
+    if (constant.canBeDynamicReplacement()) {
+      SILDeclRef declRef(replacedDecl, constant.kind, false);
+      auto *replacedFunc = getOrCreateDeclaration(replacedDecl, declRef);
 
-  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);
+  if (constant.isDistributedThunk()) {
+    fprintf(stderr, "[%s:%d] (%s) IS DIST THUNK!\n", __FILE__, __LINE__, __FUNCTION__);
+    constant.dump();
+
+    // TODO: also handle protocol / extension
+    auto actor = dyn_cast<ClassDecl>(decl->getDeclContext());
+    if (actor && actor->isDistributedActor()) {
+      auto &C = decl->getASTContext();
+      auto systemTy = getDistributedActorSystemType(actor);
+      assert(systemTy);
+
+      auto decoderTy =
+          getDistributedActorSystemInvocationDecoderType(
+              systemTy->getAnyNominal());
+      assert(decoderTy);
+
+      auto decodeFunc = C.getDecodeNextArgumentOnDistributedInvocationDecoder(
+          decoderTy->getAnyNominal());
+      auto decodeRef = SILDeclRef(decodeFunc);
+      auto *adHocWitness = getOrCreateDeclaration(decodeFunc, decodeRef);
+      F->setUsedAdHocRequirementWitnessFunction(adHocWitness);
+
+      F->dump();
+    }
+  }
 }
 
 SILFunction *SILFunctionBuilder::getOrCreateFunction(

lib/SIL/IR/SILModule.cpp

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

lib/SIL/IR/SILPrinter.cpp

@@ -2983,6 +2983,13 @@ void SILFunction::print(SILPrintContext &PrintCtx) const {
     OS << "\"] ";
   }
 
+  if (auto *usedFunc = getUsedAdHocRequirementWitnessFunction()) {
+    OS << "[used_adhoc_requirement_witness \"";
+    OS << usedFunc->getName();
+    OS << "\"] ";
+    fprintf(stderr, "[%s:%d] (%s) PRINTED THE used_adhoc_requirement_witness...\n", __FILE__, __LINE__, __FUNCTION__);
+  }
+
   if (hasObjCReplacement()) {
     OS << "[objc_replacement_for \"";
     OS << getObjCReplacement().str();