Use the distributed thunk for non-distributed protocol requirements

lib/SIL/IR/SILDeclRef.cpp

@@ -1045,8 +1045,6 @@ bool SILDeclRef::requiresNewVTableEntry() const {
       return true;
   if (!hasDecl())
     return false;
-  if (isDistributedThunk())
-    return false;
   if (isBackDeploymentThunk())
     return false;
   auto fnDecl = dyn_cast<AbstractFunctionDecl>(getDecl());

lib/SILGen/SILGenPoly.cpp

@@ -4430,41 +4430,45 @@ void SILGenFunction::emitProtocolWitness(
   FullExpr scope(Cleanups, cleanupLoc);
   FormalEvaluationScope formalEvalScope(*this);
 
-  auto witnessKind = getWitnessDispatchKind(witness, isSelfConformance);
   auto thunkTy = F.getLoweredFunctionType();
 
   SmallVector<ManagedValue, 8> origParams;
   collectThunkParams(loc, origParams);
 
-  // If we are supposed to enter the actor, do so now.
-  if (enterIsolation) {
-    if (enterIsolation->isDistributedActor()) {
-      // For a distributed actor, call through the distributed thunk.
-      witness = witness.asDistributed();
-    } else {
-      // For a non-distributed actor, hop to the actor.
-      Optional<ManagedValue> actorSelf;
-
-      // For an instance actor, get the actor 'self'.
-      if (*enterIsolation == ActorIsolation::ActorInstance) {
-        auto actorSelfVal = origParams.back();
-
-        if (actorSelfVal.getType().isAddress()) {
-          auto &actorSelfTL = getTypeLowering(actorSelfVal.getType());
-          if (!actorSelfTL.isAddressOnly()) {
-            actorSelfVal = emitManagedLoad(
-                *this, loc, actorSelfVal, actorSelfTL);
-          }
+  // If the witness is isolated to a distributed actor, but the requirement is
+  // not, go through the distributed thunk.
+  if (witness.hasDecl() &&
+      getActorIsolation(witness.getDecl()).isDistributedActor() &&
+      requirement.hasDecl() &&
+      !getActorIsolation(requirement.getDecl()).isDistributedActor()) {
+    witness = SILDeclRef(
+        cast<AbstractFunctionDecl>(witness.getDecl())->getDistributedThunk())
+          .asDistributed();
+  } else if (enterIsolation) {
+    // If we are supposed to enter the actor, do so now by hopping to the
+    // actor.
+    Optional<ManagedValue> actorSelf;
+
+    // For an instance actor, get the actor 'self'.
+    if (*enterIsolation == ActorIsolation::ActorInstance) {
+      auto actorSelfVal = origParams.back();
+
+      if (actorSelfVal.getType().isAddress()) {
+        auto &actorSelfTL = getTypeLowering(actorSelfVal.getType());
+        if (!actorSelfTL.isAddressOnly()) {
+          actorSelfVal = emitManagedLoad(
+              *this, loc, actorSelfVal, actorSelfTL);
         }
-
-        actorSelf = actorSelfVal;
       }
 
-      emitHopToTargetActor(loc, enterIsolation, actorSelf);
+      actorSelf = actorSelfVal;
     }
+
+    emitHopToTargetActor(loc, enterIsolation, actorSelf);
   }
 
   // Get the type of the witness.
+  auto witnessKind = getWitnessDispatchKind(witness, isSelfConformance);
   auto witnessInfo = getConstantInfo(getTypeExpansionContext(), witness);
   CanAnyFunctionType witnessSubstTy = witnessInfo.LoweredType;
   if (auto genericFnType = dyn_cast<GenericFunctionType>(witnessSubstTy)) {

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -664,6 +664,10 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
       func->getResultInterfaceType(), DC);
   thunk->setSynthesized(true);
   thunk->getAttrs().add(new (C) NonisolatedAttr(/*isImplicit=*/true));
+
+  if (isa<ClassDecl>(DC))
+    thunk->getAttrs().add(new (C) FinalAttr(/*isImplicit=*/true));
+
   thunk->setGenericSignature(baseSignature);
   thunk->copyFormalAccessFrom(func, /*sourceIsParentContext=*/false);
   thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);

test/SILGen/distributed_thunk.swift

@@ -25,8 +25,37 @@ protocol ServerProto {
 extension DA: ServerProto {
   // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17distributed_thunk2DACAA11ServerProtoA2aDP11doSomethingyyYaKFTW : $@convention(witness_method: ServerProto) @async (@in_guaranteed DA) -> @error Error
   // CHECK-NOT: hop_to_executor
-  // CHECK-NOT: return
-  // CHECK: function_ref @$s17distributed_thunk2DAC11doSomethingyyFTE
+  // CHECK: function_ref @$s17distributed_thunk2DAC11doSomethingyyYaKFTE
   // CHECK: return
   distributed func doSomething() { }
 }
+
+distributed actor DA2: ServerProto {
+  typealias ActorSystem = LocalTestingDistributedActorSystem
+
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17distributed_thunk3DA2CAA11ServerProtoA2aDP11doSomethingyyYaKFTW : $@convention(witness_method: ServerProto) @async (@in_guaranteed DA2) -> @error Error
+  // CHECK-NOT: hop_to_executor
+  // CHECK-NOT: return
+  // CHECK: function_ref @$s17distributed_thunk3DA2C11doSomethingyyYaKFTE
+  distributed func doSomething() async { }
+}
+
+distributed actor DA3: ServerProto {
+  typealias ActorSystem = LocalTestingDistributedActorSystem
+
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17distributed_thunk3DA3CAA11ServerProtoA2aDP11doSomethingyyYaKFTW
+  // CHECK-NOT: hop_to_executor
+  // CHECK-NOT: return
+  // CHECK: function_ref @$s17distributed_thunk3DA3C11doSomethingyyYaKFTE
+  distributed func doSomething() async throws { }
+}
+
+distributed actor DA4: ServerProto {
+  typealias ActorSystem = LocalTestingDistributedActorSystem
+
+  // CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17distributed_thunk3DA4CAA11ServerProtoA2aDP11doSomethingyyYaKFTW
+  // CHECK-NOT: hop_to_executor
+  // CHECK-NOT: return
+  // CHECK: function_ref @$s17distributed_thunk3DA4C11doSomethingyyYaKFTE
+  distributed func doSomething() throws { }
+}