Revert "[IRGen] Distributed: Always invoke decodeNextArgument through witness thunk"

This reverts commit 4d4c80b.

Author: xedin

lib/IRGen/GenDistributed.cpp

@@ -81,6 +81,10 @@ struct ArgumentDecoderInfo {
   /// The type of `decodeNextArgument` method.
   CanSILFunctionType MethodType;
 
+  /// Protocol requirements associated with the generic
+  /// parameter `Argument` of this decode method.
+  GenericSignature::RequiredProtocols ProtocolRequirements;
+
   // Witness metadata for conformance to DistributedTargetInvocationDecoder
   // protocol.
   WitnessMetadata Witness;
@@ -90,19 +94,31 @@ struct ArgumentDecoderInfo {
                       FunctionPointer decodeNextArgumentPtr,
                       CanSILFunctionType decodeNextArgumentTy)
       : Decoder(decoder), MethodPtr(decodeNextArgumentPtr),
-        MethodType(decodeNextArgumentTy) {
+        MethodType(decodeNextArgumentTy),
+        ProtocolRequirements(findProtocolRequirements(decodeNextArgumentTy)) {
     Witness.SelfMetadata = decoderType;
     Witness.SelfWitnessTable = decoderWitnessTable;
   }
 
   CanSILFunctionType getMethodType() const { return MethodType; }
 
-  WitnessMetadata *getWitnessMetadata() const {
-    return const_cast<WitnessMetadata *>(&Witness);
+  ArrayRef<ProtocolDecl *> getProtocolRequirements() const {
+    return ProtocolRequirements;
   }
 
   /// Form a callee to a decode method - `decodeNextArgument`.
   Callee getCallee() const;
+
+private:
+  static GenericSignature::RequiredProtocols
+  findProtocolRequirements(CanSILFunctionType decodeMethodTy) {
+    auto signature = decodeMethodTy->getInvocationGenericSignature();
+    auto genericParams = signature.getGenericParams();
+
+    // func decodeNextArgument<Arg : #SerializationRequirement#>() throws -> Arg
+    assert(genericParams.size() == 1);
+    return signature->getRequiredProtocols(genericParams.front());
+  }
 };
 
 class DistributedAccessor {
@@ -140,6 +156,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 +405,17 @@ 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());
+                      /*witnessMetadata=*/nullptr);
 
     Explosion result;
     emission->emitToExplosion(result, /*isOutlined=*/false);
@@ -492,6 +516,37 @@ void DistributedAccessor::decodeArgument(unsigned argumentIdx,
   }
 }
 
+void DistributedAccessor::lookupWitnessTables(
+    llvm::Value *value, ArrayRef<ProtocolDecl *> protocols,
+    Explosion &witnessTables) {
+  auto conformsToProtocol = IGM.getConformsToProtocolFunctionPointer();
+
+  for (auto *protocol : protocols) {
+    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,
@@ -730,22 +785,70 @@ 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 = IGM.Context.getDistributedActorArgumentDecodingMethod(actor);
+  assert(decodeFn && "no suitable decoder?");
+
+  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();
+
+  // 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;
 
-  auto decoderProtocol = C.getDistributedTargetInvocationDecoderDecl();
-  SILDeclRef decodeNextArgumentRef(
-      decoderProtocol->getSingleRequirement(C.Id_decodeNextArgument));
+  if (auto classDecl = dyn_cast<ClassDecl>(decoderDecl)) {
+    auto selfTy = methodTy->getSelfParameter().getSILStorageType(
+        IGM.getSILModule(), methodTy, expansionContext);
 
-  llvm::Constant *fnPtr =
-      IGM.getAddrOfDispatchThunk(decodeNextArgumentRef, NotForDefinition);
+    auto &classTI = IGM.getTypeInfo(selfTy).as<ClassTypeInfo>();
+    auto &classLayout = classTI.getClassLayout(IGM, selfTy,
+                                               /*forBackwardDeployment=*/false);
 
-  auto fnType = IGM.getSILTypes().getConstantFunctionType(
-      IGM.getMaximalTypeExpansionContext(), decodeNextArgumentRef);
+    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};
 }
 
 SILType DistributedAccessor::getResultType() const {

test/Distributed/Runtime/distributed_actor_localSystem_generic.swift

@@ -98,10 +98,25 @@ public distributed actor MyOtherActor {
 
 /// Read the current offset and cast an element to `Int`
 
+// CHECK: [[DECODER:%.*]] = load ptr, ptr %1
+// CHECK-NEXT: [[DECODE_NEXT_ARG_REF:%.*]] = getelementptr inbounds ptr, ptr %2, i64
+
 // CHECK: [[ARG_0_SIZE_ADJ:%.*]] = add i64 %size, 15
 // CHECK-NEXT: [[ARG_0_SIZE:%.*]] = and i64 [[ARG_0_SIZE_ADJ]], -16
 // CHECK-NEXT: [[ARG_0_VALUE_BUF:%.*]] = call swiftcc ptr @swift_task_alloc(i64 [[ARG_0_SIZE]])
-// CHECK-NEXT: call swiftcc void @"$s11Distributed0A23TargetInvocationDecoderP18decodeNextArgumentqd__yKlFTj"(ptr noalias sret(%swift.opaque) [[ARG_0_VALUE_BUF]], ptr %arg_type, ptr swiftself %1, ptr noalias nocapture swifterror dereferenceable(8) %swifterror, ptr [[DECODER_TYPE]], ptr [[DECODER_PROTOCOL_WITNESS]])
+// CHECK-NEXT: [[ENCODABLE_WITNESS:%.*]] = call ptr @swift_conformsToProtocol{{(2)?}}(ptr %arg_type, ptr @"$sSeMp")
+// CHECK-NEXT: [[IS_NULL:%.*]] = icmp eq ptr [[ENCODABLE_WITNESS]], null
+// CHECK-NEXT: br i1 [[IS_NULL]], label %missing-witness, label [[CONT:%.*]]
+// CHECK: missing-witness:
+// CHECK-NEXT: call void @llvm.trap()
+// CHECK-NEXT: unreachable
+// CHECK: [[DECODABLE_WITNESS:%.*]] = call ptr @swift_conformsToProtocol{{(2)?}}(ptr %arg_type, ptr @"$sSEMp")
+// CHECK-NEXT: [[IS_NULL:%.*]] = icmp eq ptr [[DECODABLE_WITNESS]], null
+// CHECK-NEXT: br i1 [[IS_NULL]], label %missing-witness1, label [[CONT:%.*]]
+// CHECK: missing-witness1:
+// CHECK-NEXT: call void @llvm.trap()
+// CHECK-NEXT: unreachable
+// CHECK: call swiftcc void @"$s27FakeDistributedActorSystems0A17InvocationDecoderC18decodeNextArgumentxyKSeRzSERzlF"(ptr noalias sret(%swift.opaque) [[ARG_0_VALUE_BUF]], ptr %arg_type, ptr [[ENCODABLE_WITNESS]], ptr [[DECODABLE_WITNESS]], ptr swiftself [[DECODER]], ptr noalias nocapture swifterror dereferenceable(8) %swifterror)
 
 // CHECK: store ptr null, ptr %swifterror
 // CHECK-NEXT: %._value = getelementptr inbounds %TSi, ptr [[ARG_0_VALUE_BUF]], i32 0, i32 0