[IRGen] Distributed: Always invoke decodeNextArgument through witness thunk

`decodeNextArgument` is now a requirement, so thunks should be
always calling that instead of trying to dispatch directly.

Author: xedin

lib/IRGen/GenDistributed.cpp

@@ -81,10 +81,6 @@ 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;
@@ -94,31 +90,19 @@ struct ArgumentDecoderInfo {
                       FunctionPointer decodeNextArgumentPtr,
                       CanSILFunctionType decodeNextArgumentTy)
       : Decoder(decoder), MethodPtr(decodeNextArgumentPtr),
-        MethodType(decodeNextArgumentTy),
-        ProtocolRequirements(findProtocolRequirements(decodeNextArgumentTy)) {
+        MethodType(decodeNextArgumentTy) {
     Witness.SelfMetadata = decoderType;
     Witness.SelfWitnessTable = decoderWitnessTable;
   }
 
   CanSILFunctionType getMethodType() const { return MethodType; }
 
-  ArrayRef<ProtocolDecl *> getProtocolRequirements() const {
-    return ProtocolRequirements;
+  WitnessMetadata *getWitnessMetadata() const {
+    return const_cast<WitnessMetadata *>(&Witness);
   }
 
   /// 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 {
@@ -156,10 +140,6 @@ 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.
   ///
@@ -417,17 +397,13 @@ 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=*/nullptr);
+                      /*witnessMetadata=*/decoder.getWitnessMetadata());
 
     Explosion result;
     emission->emitToExplosion(result, /*isOutlined=*/false);
@@ -528,37 +504,6 @@ 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,
@@ -803,70 +748,22 @@ DistributedAccessor::getCalleeForDistributedTarget(llvm::Value *self) const {
 
 ArgumentDecoderInfo DistributedAccessor::findArgumentDecoder(
     llvm::Value *decoder, llvm::Value *decoderTy, llvm::Value *witnessTable) {
-  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 &C = IGM.Context;
 
-  if (auto classDecl = dyn_cast<ClassDecl>(decoderDecl)) {
-    auto selfTy = methodTy->getSelfParameter().getSILStorageType(
-        IGM.getSILModule(), methodTy, expansionContext);
+  auto decoderProtocol = C.getDistributedTargetInvocationDecoderDecl();
+  SILDeclRef decodeNextArgumentRef(
+      decoderProtocol->getSingleRequirement(C.Id_decodeNextArgument));
 
-    auto &classTI = IGM.getTypeInfo(selfTy).as<ClassTypeInfo>();
-    auto &classLayout = classTI.getClassLayout(IGM, selfTy,
-                                               /*forBackwardDeployment=*/false);
+  llvm::Constant *fnPtr =
+      IGM.getAddrOfDispatchThunk(decodeNextArgumentRef, NotForDefinition);
 
-    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 fnType = IGM.getSILTypes().getConstantFunctionType(
+      IGM.getMaximalTypeExpansionContext(), decodeNextArgumentRef);
 
-  return {decoder, decoderTy, witnessTable, methodPtr, methodTy};
+  auto sig = IGM.getSignature(fnType);
+  auto fn = FunctionPointer::forDirect(fnType, fnPtr,
+                                       /*secondaryValue=*/nullptr, sig, true);
+  return {decoder, decoderTy, witnessTable, fn, fnType};
 }
 
 SILType DistributedAccessor::getResultType() const {

test/Distributed/Runtime/distributed_actor_localSystem_generic.swift

@@ -0,0 +1,38 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -module-name main -Xfrontend -enable-experimental-distributed -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s --color
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// rdar://90373022
+// UNSUPPORTED: OS=watchos
+
+import Distributed
+
+distributed actor Worker<ActorSystem> where ActorSystem: DistributedActorSystem<any Codable>, ActorSystem.ActorID: Codable {
+  distributed func hi(name: String) {
+    print("Hi, \(name)!")
+  }
+
+  nonisolated var description: Swift.String {
+    "Worker(\(id))"
+  }
+}
+
+// ==== Execute ----------------------------------------------------------------
+@main struct Main {
+  static func main() async throws {
+    let system = LocalTestingDistributedActorSystem()
+
+    let actor = Worker(actorSystem: system)
+    try await actor.hi(name: "P") // local calls should still just work
+    // CHECK: Hi, P!
+  }
+}

test/Distributed/distributed_actor_accessor_thunks_64bit.swift

@@ -98,25 +98,10 @@ 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: [[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-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: store ptr null, ptr %swifterror
 // CHECK-NEXT: %._value = getelementptr inbounds %TSi, ptr [[ARG_0_VALUE_BUF]], i32 0, i32 0