[Distributed] Handle distributed thunk requirements/witnesses (#72151)

Co-authored-by: Pavel Yaskevich <pyaskevich@apple.com>

Author: ktoso

include/swift/SIL/SILWitnessVisitor.h

@@ -146,11 +146,12 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void visitFuncDecl(FuncDecl *func) {
     assert(!isa<AccessorDecl>(func));
-    if (func->requiresNewWitnessTableEntry()) {
-      asDerived().addMethod(SILDeclRef(func, SILDeclRef::Kind::Func));
-      addAutoDiffDerivativeMethodsIfRequired(func, SILDeclRef::Kind::Func);
-      addDistributedWitnessMethodsIfRequired(func, SILDeclRef::Kind::Func);
-    }
+    if (!func->requiresNewWitnessTableEntry())
+      return;
+
+    asDerived().addMethod(SILDeclRef(func, SILDeclRef::Kind::Func));
+    addAutoDiffDerivativeMethodsIfRequired(func, SILDeclRef::Kind::Func);
+    addDistributedWitnessMethodsIfRequired(func, SILDeclRef::Kind::Func);
   }
 
   void visitMissingMemberDecl(MissingMemberDecl *placeholder) {
@@ -160,7 +161,7 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
   void visitAssociatedTypeDecl(AssociatedTypeDecl *td) {
     // We already visited these in the first pass.
   }
-    
+
   void visitTypeAliasDecl(TypeAliasDecl *tad) {
     // We don't care about these by themselves for witnesses.
   }
@@ -200,11 +201,11 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void addDistributedWitnessMethodsIfRequired(AbstractFunctionDecl *AFD,
                                               SILDeclRef::Kind kind) {
-    if (!AFD->isDistributed())
+    if (!AFD || !AFD->isDistributed())
       return;
 
-    // Add another which will be witnessed by the 'distributed thunk'
-    SILDeclRef declRef(AFD, kind);
+    auto thunk = AFD->getDistributedThunk();
+    SILDeclRef declRef(thunk, kind);
     asDerived().addMethod(declRef.asDistributed());
   }
 };

lib/AST/DistributedDecl.cpp

@@ -1346,6 +1346,9 @@ FuncDecl *AbstractStorageDecl::getDistributedThunk() const {
 
 FuncDecl*
 AbstractFunctionDecl::getDistributedThunk() const {
+  if (isDistributedThunk())
+    return const_cast<FuncDecl *>(dyn_cast<FuncDecl>(this));
+
   if (!isDistributed())
     return nullptr;
 

lib/AST/ProtocolConformance.cpp

@@ -686,6 +686,13 @@ void NormalProtocolConformance::setWitness(ValueDecl *requirement,
          "requirement in wrong protocol");
   assert(Mapping.count(requirement) == 0 && "Witness already known");
   assert((!isComplete() || isInvalid() ||
+          // TODO(distributed): properly handle isComplete() for distributed
+          //  funcs; there seems to be a problem that we mark completed, but
+          //  afterwards will record the thunk witness;
+          (dyn_cast<FuncDecl>(requirement)
+            ? (dyn_cast<FuncDecl>(requirement)->isDistributed() ||
+               dyn_cast<FuncDecl>(requirement)->isDistributedThunk())
+            : false) ||
           requirement->getAttrs().hasAttribute<OptionalAttr>() ||
           requirement->getAttrs().isUnavailable(
                                         requirement->getASTContext())) &&

lib/IRGen/GenMeta.cpp

@@ -918,8 +918,10 @@ namespace {
 
       {
         auto *requirement = cast<AbstractFunctionDecl>(func.getDecl());
-        if (requirement->isDistributed())
+        if (requirement->isDistributedThunk()) {
+          // when thunk, because in protocol we want accessof for the thunk
           IGM.emitDistributedTargetAccessor(requirement);
+        }
       }
 
       // Classify the function.

lib/IRGen/GenProto.cpp

@@ -2464,6 +2464,7 @@ static void addWTableTypeMetadata(IRGenModule &IGM,
 
     auto mw = entry.getMethodWitness();
     auto member = mw.Requirement;
+
     auto &fnProtoInfo =
         IGM.getProtocolInfo(conf->getProtocol(), ProtocolInfoKind::Full);
     auto index = fnProtoInfo.getFunctionIndex(member).forProtocolWitnessTable();

lib/SILGen/SILGenApply.cpp

@@ -685,8 +685,10 @@ class Callee {
         if (func->isDistributed() && isa<ProtocolDecl>(func->getDeclContext())) {
           // If we're calling cross-actor, we must always use a distributed thunk
           if (!isSameActorIsolated(func, SGF.FunctionDC)) {
-            // We must adjust the constant to use a distributed thunk.
-            constant = constant->asDistributed();
+            // the protocol witness must always be a distributed thunk, as we
+            // may be crossing a remote boundary here.
+            auto thunk = func->getDistributedThunk();
+            constant = SILDeclRef(thunk).asDistributed();
           }
         }
       }

lib/SILGen/SILGenThunk.cpp

@@ -120,7 +120,7 @@ void SILGenModule::emitDistributedThunkForDecl(
   if (!thunkDecl)
     return;
 
-  if (thunkDecl->isBodySkipped())
+  if (!thunkDecl->hasBody() || thunkDecl->isBodySkipped())
     return;
 
   auto thunk = SILDeclRef(thunkDecl).asDistributed();

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -724,13 +724,6 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   auto &C = func->getASTContext();
   auto DC = func->getDeclContext();
 
-  // NOTE: So we don't need a thunk in the protocol, we should call the
-  // underlying thing instead, which MUST have a thunk, since it must be a
-  // distributed func as well...
-  if (isa<ProtocolDecl>(DC)) {
-    return nullptr;
-  }
-
   DeclName thunkName;
 
   // Since accessors don't have names, let's generate one based on
@@ -754,7 +747,9 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   thunk->getAttrs().add(
       new (C) NonisolatedAttr(/*unsafe=*/false, /*implicit=*/true));
 
-  thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);
+  // Protocol requirements don't have bodies.
+  if (func->hasBody())
+    thunk->setBodySynthesizer(deriveBodyDistributed_thunk, func);
 
   /// Record which function this is a thunk for, we'll need this to link back
   /// calls in case this is a distributed requirement witness.
@@ -874,6 +869,10 @@ void swift::assertRequiredSynthesizedPropertyOrder(ASTContext &Context,
 }
 
 static bool canSynthesizeDistributedThunk(AbstractFunctionDecl *distributedTarget) {
+  // `distributed` protocol requirements are allowed without additional checks.
+  if (isa<ProtocolDecl>(distributedTarget->getDeclContext()))
+    return true;
+
   if (getConcreteReplacementForProtocolActorSystemType(distributedTarget)) {
     return true;
   }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -563,6 +563,12 @@ CheckRedeclarationRequest::evaluate(Evaluator &eval, ValueDecl *current,
   if (!currentFile)
     return std::make_tuple<>();
 
+  if (auto func = dyn_cast<AbstractFunctionDecl>(current)) {
+    if (func->isDistributedThunk()) {
+      return std::make_tuple<>();
+    }
+  }
+
   auto &ctx = current->getASTContext();
 
   // Find other potential definitions.

lib/Sema/TypeCheckDistributed.cpp

@@ -499,6 +499,9 @@ bool swift::checkDistributedActorSystem(const NominalTypeDecl *system) {
 /// \returns \c true if there was a problem with adding the attribute, \c false
 /// otherwise.
 bool swift::checkDistributedFunction(AbstractFunctionDecl *func) {
+  if (!func->isDistributed())
+    return false;
+
   auto &C = func->getASTContext();
   return evaluateOrDefault(C.evaluator,
                            CheckDistributedFunctionRequest{func},
@@ -521,6 +524,11 @@ bool CheckDistributedFunctionRequest::evaluate(
   if (!C.getLoadedModule(C.Id_Distributed))
     return true;
 
+//  // No checking for protocol requirements because they are not required
+//  // to have `SerializationRequirement`.
+//  if (isa<ProtocolDecl>(func->getDeclContext()))
+//    return false;
+
   Type serializationReqType =
       getDistributedActorSerializationType(func->getDeclContext());
 
@@ -736,9 +744,6 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
 }
 
 bool TypeChecker::checkDistributedFunc(FuncDecl *func) {
-  if (!func->isDistributed())
-    return false;
-
   return swift::checkDistributedFunction(func);
 }