Merge pull request #71796 from xedin/make-dist-new-requirements-conditionally-available

[Distributed] Make new protocol requirements conditionally available

Author: xedin

include/swift/SIL/SILFunction.h

@@ -264,6 +264,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.
@@ -585,6 +596,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();
   }

lib/AST/DistributedDecl.cpp

@@ -256,7 +256,7 @@ swift::getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
     return nullptr;
 
   auto systemTy = getConcreteReplacementForProtocolActorSystemType(thunk);
-  if (!systemTy)
+  if (!systemTy || systemTy->is<GenericTypeParamType>())
     return nullptr;
 
   auto decoderTy =

lib/IRGen/GenDistributed.cpp

@@ -81,16 +81,22 @@ struct ArgumentDecoderInfo {
   /// The type of `decodeNextArgument` method.
   CanSILFunctionType MethodType;
 
-  // Witness metadata for conformance to DistributedTargetInvocationDecoder
-  // protocol.
+  /// Witness metadata for conformance to DistributedTargetInvocationDecoder
+  /// protocol.
   WitnessMetadata Witness;
 
+  /// Indicates whether `decodeNextArgument` is referenced through
+  /// a protocol witness thunk.
+  bool UsesWitnessDispatch;
+
   ArgumentDecoderInfo(llvm::Value *decoder, llvm::Value *decoderType,
                       llvm::Value *decoderWitnessTable,
                       FunctionPointer decodeNextArgumentPtr,
-                      CanSILFunctionType decodeNextArgumentTy)
+                      CanSILFunctionType decodeNextArgumentTy,
+                      bool usesWitnessDispatch)
       : Decoder(decoder), MethodPtr(decodeNextArgumentPtr),
-        MethodType(decodeNextArgumentTy) {
+        MethodType(decodeNextArgumentTy),
+        UsesWitnessDispatch(usesWitnessDispatch) {
     Witness.SelfMetadata = decoderType;
     Witness.SelfWitnessTable = decoderWitnessTable;
   }
@@ -101,6 +107,20 @@ struct ArgumentDecoderInfo {
     return const_cast<WitnessMetadata *>(&Witness);
   }
 
+  /// Protocol requirements associated with the generic
+  /// parameter `Argument` of this decode method.
+  GenericSignature::RequiredProtocols getProtocolRequirements() const {
+    if (UsesWitnessDispatch)
+      return {};
+
+    auto signature = MethodType->getInvocationGenericSignature();
+    auto genericParams = signature.getGenericParams();
+
+    // func decodeNextArgument<Arg : #SerializationRequirement#>() throws -> Arg
+    assert(genericParams.size() == 1);
+    return signature->getRequiredProtocols(genericParams.front());
+  }
+
   /// Form a callee to a decode method - `decodeNextArgument`.
   Callee getCallee() const;
 };
@@ -140,6 +160,10 @@ class DistributedAccessor {
                       llvm::Value *argumentType, const SILParameterInfo &param,
                       Explosion &arguments);
 
+  void lookupWitnessTables(llvm::Value *value,
+                           ArrayRef<ProtocolDecl *> protocols,
+                           Explosion &witnessTables);
+
   /// Load witness table addresses (if any) from the given buffer
   /// into the given argument explosion.
   ///
@@ -385,13 +409,18 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
   // substitution Argument -> <argument metadata>
   decodeArgs.add(argumentType);
 
+  // Lookup witness tables for the requirement on the argument type.
+  lookupWitnessTables(argumentType, decoder.getProtocolRequirements(),
+                      decodeArgs);
+
   Address calleeErrorSlot;
   llvm::Value *decodeError = nullptr;
 
   emission->begin();
   {
     emission->setArgs(decodeArgs, /*isOutlined=*/false,
-                      /*witnessMetadata=*/decoder.getWitnessMetadata());
+                      decoder.UsesWitnessDispatch ? decoder.getWitnessMetadata()
+                                                  : nullptr);
 
     Explosion result;
     emission->emitToExplosion(result, /*isOutlined=*/false);
@@ -492,6 +521,43 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
   }
 }
 
+void DistributedAccessor::lookupWitnessTables(
+    llvm::Value *value, ArrayRef<ProtocolDecl *> protocols,
+    Explosion &witnessTables) {
+  if (protocols.empty())
+    return;
+
+  auto conformsToProtocol = IGM.getConformsToProtocolFunctionPointer();
+
+  for (auto *protocol : protocols) {
+    if (!Lowering::TypeConverter::protocolRequiresWitnessTable(protocol))
+      continue;
+
+    auto *protocolDescriptor = IGM.getAddrOfProtocolDescriptor(protocol);
+    auto *witnessTable =
+        IGF.Builder.CreateCall(conformsToProtocol, {value, protocolDescriptor});
+
+    auto failBB = IGF.createBasicBlock("missing-witness");
+    auto contBB = IGF.createBasicBlock("");
+
+    auto isNull = IGF.Builder.CreateICmpEQ(
+        witnessTable, llvm::ConstantPointerNull::get(IGM.WitnessTablePtrTy));
+    IGF.Builder.CreateCondBr(isNull, failBB, contBB);
+
+    // This operation shouldn't fail because runtime should have checked that
+    // a particular argument type conforms to `SerializationRequirement`
+    // of the distributed actor the decoder is used for. If it does fail
+    // then accessor should trap.
+    {
+      IGF.Builder.emitBlock(failBB);
+      IGF.emitTrap("missing witness table", /*EmitUnreachable=*/true);
+    }
+
+    IGF.Builder.emitBlock(contBB);
+    witnessTables.add(witnessTable);
+  }
+}
+
 void DistributedAccessor::emitLoadOfWitnessTables(llvm::Value *witnessTables,
                                                   llvm::Value *numTables,
                                                   unsigned expectedWitnessTables,
@@ -731,21 +797,91 @@ DistributedAccessor::getCalleeForDistributedTarget(llvm::Value *self) const {
 ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
     llvm::Value *decoder, llvm::Value *decoderTy, llvm::Value *witnessTable) {
   auto &C = IGM.Context;
+  auto *actor = getDistributedActorOf(Target);
+  auto expansionContext = IGM.getMaximalTypeExpansionContext();
+
+  auto *decodeFn = C.getDistributedActorArgumentDecodingMethod(actor);
+
+  // If distributed actor is generic over actor system, we have to
+  // use witness to reference `decodeNextArgument`.
+  if (!decodeFn) {
+    auto decoderProtocol = C.getDistributedTargetInvocationDecoderDecl();
+    auto decodeNextArgRequirement =
+        decoderProtocol->getSingleRequirement(C.Id_decodeNextArgument);
+    assert(decodeNextArgRequirement);
+    SILDeclRef decodeNextArgumentRef(decodeNextArgRequirement);
+
+    llvm::Constant *fnPtr =
+        IGM.getAddrOfDispatchThunk(decodeNextArgumentRef, NotForDefinition);
+    auto fnType = IGM.getSILTypes().getConstantFunctionType(
+        IGM.getMaximalTypeExpansionContext(), decodeNextArgumentRef);
+
+    auto sig = IGM.getSignature(fnType);
+    auto fn = FunctionPointer::forDirect(fnType, fnPtr,
+                                         /*secondaryValue=*/nullptr, sig, true);
+    return {decoder, decoderTy, witnessTable,
+            fn,      fnType,    /*usesWitnessDispatch=*/true};
+  }
+
+  auto methodTy = IGM.getSILTypes().getConstantFunctionType(
+      expansionContext, SILDeclRef(decodeFn));
+
+  auto fpKind = FunctionPointerKind::defaultAsync();
+  auto signature = IGM.getSignature(methodTy, fpKind);
+
+  // If the decoder class is `final`, let's emit a direct reference.
+  auto *decoderDecl = decodeFn->getDeclContext()->getSelfNominalTypeDecl();
 
-  auto decoderProtocol = C.getDistributedTargetInvocationDecoderDecl();
-  SILDeclRef decodeNextArgumentRef(
-      decoderProtocol->getSingleRequirement(C.Id_decodeNextArgument));
+  // If decoder is a class, need to load it first because generic parameter
+  // is passed indirectly. This is good for structs and enums because
+  // `decodeNextArgument` is a mutating method, but not for classes because
+  // in that case heap object is mutated directly.
+  bool usesDispatchThunk = false;
 
-  llvm::Constant *fnPtr =
-      IGM.getAddrOfDispatchThunk(decodeNextArgumentRef, NotForDefinition);
+  if (auto classDecl = dyn_cast<ClassDecl>(decoderDecl)) {
+    auto selfTy = methodTy->getSelfParameter().getSILStorageType(
+        IGM.getSILModule(), methodTy, expansionContext);
 
-  auto fnType = IGM.getSILTypes().getConstantFunctionType(
-      IGM.getMaximalTypeExpansionContext(), decodeNextArgumentRef);
+    auto &classTI = IGM.getTypeInfo(selfTy).as<ClassTypeInfo>();
+    auto &classLayout = classTI.getClassLayout(IGM, selfTy,
+                                               /*forBackwardDeployment=*/false);
+
+    llvm::Value *typedDecoderPtr = IGF.Builder.CreateBitCast(
+        decoder, classLayout.getType()->getPointerTo()->getPointerTo());
+
+    Explosion instance;
+
+    classTI.loadAsTake(IGF,
+                       {typedDecoderPtr, classTI.getStorageType(),
+                        classTI.getBestKnownAlignment()},
+                       instance);
+
+    decoder = instance.claimNext();
+
+    /// When using library evolution functions have another "dispatch thunk"
+    /// so we must use this instead of the decodeFn directly.
+    usesDispatchThunk =
+        getMethodDispatch(decodeFn) == swift::MethodDispatch::Class &&
+        classDecl->hasResilientMetadata();
+  }
+
+  FunctionPointer methodPtr;
+
+  if (usesDispatchThunk) {
+    auto fnPtr = IGM.getAddrOfDispatchThunk(SILDeclRef(decodeFn), NotForDefinition);
+    methodPtr = FunctionPointer::createUnsigned(
+        methodTy, fnPtr, signature, /*useSignature=*/true);
+  } else {
+    SILFunction *decodeSILFn = IGM.getSILModule().lookUpFunction(SILDeclRef(decodeFn));
+    auto fnPtr = IGM.getAddrOfSILFunction(decodeSILFn, NotForDefinition,
+        /*isDynamicallyReplaceable=*/false);
+    methodPtr = FunctionPointer::forDirect(
+        classifyFunctionPointerKind(decodeSILFn), fnPtr,
+        /*secondaryValue=*/nullptr, signature);
+  }
 
-  auto sig = IGM.getSignature(fnType);
-  auto fn = FunctionPointer::forDirect(fnType, fnPtr,
-                                       /*secondaryValue=*/nullptr, sig, true);
-  return {decoder, decoderTy, witnessTable, fn, fnType};
+  return {decoder,   decoderTy, witnessTable,
+          methodPtr, methodTy,  /*usesWitnessDispatch=*/false};
 }
 
 SILType DistributedAccessor::getResultType() const {

lib/SIL/IR/SILFunction.cpp

@@ -278,6 +278,7 @@ void SILFunction::createSnapshot(int id) {
   newSnapshot->DeclCtxt = DeclCtxt;
   newSnapshot->Profiler = Profiler;
   newSnapshot->ReplacedFunction = ReplacedFunction;
+  newSnapshot->RefAdHocRequirementFunction = RefAdHocRequirementFunction;
   newSnapshot->ObjCReplacementFor = ObjCReplacementFor;
   newSnapshot->SemanticsAttrSet = SemanticsAttrSet;
   newSnapshot->SpecializeAttrSet = SpecializeAttrSet;

lib/SIL/IR/SILFunctionBuilder.cpp

@@ -274,6 +274,16 @@ void SILFunctionBuilder::addFunctionAttributes(
 
       F->setDynamicallyReplacedFunction(replacedFunc);
     }
+  } else if (constant.isDistributedThunk()) {
+    // It's okay for `decodeFuncDecl` to be null because system could be
+    // generic.
+    if (auto decodeFuncDecl =
+            getAssociatedDistributedInvocationDecoderDecodeNextArgumentFunction(
+                decl)) {
+      auto decodeRef = SILDeclRef(decodeFuncDecl);
+      auto *adHocFunc = getOrCreateDeclaration(decodeFuncDecl, decodeRef);
+      F->setReferencedAdHocRequirementWitnessFunction(adHocFunc);
+    }
   }
 }
 

lib/SIL/IR/SILModule.cpp

@@ -152,6 +152,7 @@ SILModule::~SILModule() {
   for (SILFunction &F : *this) {
     F.dropAllReferences();
     F.dropDynamicallyReplacedFunction();
+    F.dropReferencedAdHocRequirementWitnessFunction();
     F.clearSpecializeAttrs();
   }
 
@@ -430,6 +431,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

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