Enter actor isolation for a distributed witness via the distributed thunk

We cannot hop to a distributed actor directly, so call through the
distributed thunk instead.

Author: DougGregor

lib/SILGen/SILGenPoly.cpp

@@ -4440,6 +4440,34 @@ void SILGenFunction::emitProtocolWitness(
   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);
+          }
+        }
+
+        actorSelf = actorSelfVal;
+      }
+
+      emitHopToTargetActor(loc, enterIsolation, actorSelf);
+    }
+  }
+
   // Get the type of the witness.
   auto witnessInfo = getConstantInfo(getTypeExpansionContext(), witness);
   CanAnyFunctionType witnessSubstTy = witnessInfo.LoweredType;
@@ -4481,25 +4509,6 @@ void SILGenFunction::emitProtocolWitness(
                                          witnessUnsubstTy->getSelfParameter());
   }
 
-  // If we are supposed to hop to the actor, do so now.
-  if (enterIsolation) {
-    Optional<ManagedValue> actorSelf;
-    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);
-  }
-
   // For static C++ methods and constructors, we need to drop the (metatype)
   // "self" param. The "native" SIL representation will look like this:
   //    @convention(method) (@thin Foo.Type) -> () but the "actual" SIL function

test/SILGen/distributed_thunk.swift

@@ -17,3 +17,16 @@ extension DA {
 
   distributed func f() { }
 }
+
+protocol ServerProto {
+  func doSomething() async throws
+}
+
+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: return
+  distributed func doSomething() { }
+}