Merge pull request swiftlang#59722 from ktoso/wip-fix-witnesses-distributed

[3/3][Distributed] Remove dist requirement explicit async throws limitation

Author: ktoso

include/swift/AST/ActorIsolation.h

@@ -193,6 +193,9 @@ ActorIsolation getActorIsolation(ValueDecl *value);
 /// Determine how the given declaration context is isolated.
 ActorIsolation getActorIsolationOfContext(DeclContext *dc);
 
+/// Check if both the value, and context are isolated to the same actor.
+bool isSameActorIsolated(ValueDecl *value, DeclContext *dc);
+
 /// Determines whether this function's body uses flow-sensitive isolation.
 bool usesFlowSensitiveIsolation(AbstractFunctionDecl const *fn);
 

include/swift/AST/DiagnosticsSema.def

@@ -4794,9 +4794,6 @@ ERROR(distributed_actor_func_static,none,
 ERROR(distributed_actor_func_not_in_distributed_actor,none,
       "'distributed' method can only be declared within 'distributed actor'",
       ())
-ERROR(distributed_method_requirement_must_be_async_throws,none, // FIXME(distributed): this is an implementation limitation we should lift
-      "'distributed' protocol requirement %0 must currently be declared explicitly 'async throws'",
-      (DeclName))
 ERROR(distributed_actor_user_defined_special_property,none,
       "property %0 cannot be defined explicitly, as it conflicts with "
       "distributed actor synthesized stored property",

include/swift/AST/DistributedDecl.h

@@ -15,8 +15,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef SWIFT_DECL_TYPECHECKDISTRIBUTED_H
-#define SWIFT_DECL_TYPECHECKDISTRIBUTED_H
+#ifndef SWIFT_DECL_DISTRIBUTEDDECL_H
+#define SWIFT_DECL_DISTRIBUTEDDECL_H
 
 #include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/DiagnosticEngine.h"
@@ -129,4 +129,6 @@ extractDistributedSerializationRequirements(
 
 }
 
-#endif /* SWIFT_DECL_TYPECHECKDISTRIBUTED_H */
+// ==== ------------------------------------------------------------------------
+
+#endif /* SWIFT_DECL_DISTRIBUTEDDECL_H */

include/swift/SIL/SILWitnessVisitor.h

@@ -21,6 +21,7 @@
 
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/ProtocolAssociations.h"
 #include "swift/AST/Types.h"
 #include "swift/SIL/TypeLowering.h"
@@ -145,6 +146,7 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
     if (SILDeclRef::requiresNewWitnessTableEntry(func)) {
       asDerived().addMethod(SILDeclRef(func, SILDeclRef::Kind::Func));
       addAutoDiffDerivativeMethodsIfRequired(func, SILDeclRef::Kind::Func);
+      addDistributedWitnessMethodsIfRequired(func, SILDeclRef::Kind::Func);
     }
   }
 
@@ -192,6 +194,16 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
               AFD->getASTContext())));
     }
   }
+
+  void addDistributedWitnessMethodsIfRequired(AbstractFunctionDecl *AFD,
+                                              SILDeclRef::Kind kind) {
+    if (!AFD->isDistributed())
+      return;
+
+    // Add another which will be witnessed by the 'distributed thunk'
+    SILDeclRef declRef(AFD, kind);
+    asDerived().addMethod(declRef.asDistributed());
+  }
 };
 
 } // end namespace swift

lib/AST/ASTMangler.cpp

@@ -243,6 +243,7 @@ std::string ASTMangler::mangleWitnessThunk(
       appendOperator("TW");
     }
   }
+
   return finalize();
 }
 

lib/AST/Decl.cpp

@@ -9254,6 +9254,13 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
   return ActorIsolation::forUnspecified();
 }
 
+bool swift::isSameActorIsolated(ValueDecl *value, DeclContext *dc) {
+    auto valueIsolation = getActorIsolation(value);
+    auto dcIsolation = getActorIsolationOfContext(dc);
+    return valueIsolation.isActorIsolated() && dcIsolation.isActorIsolated() &&
+           valueIsolation.getActor() == dcIsolation.getActor();
+}
+
 ClangNode Decl::getClangNodeImpl() const {
   assert(Bits.Decl.FromClang);
   void * const *ptr = nullptr;

lib/AST/TypeCheckRequests.cpp

@@ -1559,7 +1559,11 @@ void swift::simple_display(
     llvm::raw_ostream &out, const ActorIsolation &state) {
   switch (state) {
     case ActorIsolation::ActorInstance:
-      out << "actor-isolated to instance of " << state.getActor()->getName();
+      out << "actor-isolated to instance of ";
+      if (state.isDistributedActor()) {
+        out << "distributed ";
+      }
+      out << "actor " << state.getActor()->getName();
       break;
 
     case ActorIsolation::Independent:

lib/Parse/ParseDecl.cpp

@@ -7172,7 +7172,7 @@ Parser::parseDeclVar(ParseDeclOptions Flags,
 ///
 /// \verbatim
 ///   decl-func:
-///     attribute-list? ('static' | 'class')? 'mutating'? 'func' 
+///     attribute-list? ('static' | 'class' | 'distributed')? 'mutating'? 'func'
 ///               any-identifier generic-params? func-signature where-clause?
 ///               stmt-brace?
 /// \endverbatim
@@ -7256,7 +7256,7 @@ ParserResult<FuncDecl> Parser::parseDeclFunc(SourceLoc StaticLoc,
   }
 
   DebuggerContextChange DCC(*this, SimpleName, DeclKind::Func);
-  
+
   // Parse the generic-params, if present.
   GenericParamList *GenericParams;
   auto GenericParamResult = maybeParseGenericParams();

lib/SIL/IR/SILFunctionType.cpp

@@ -2112,7 +2112,7 @@ static CanSILFunctionType getSILFunctionType(
                         .withConcurrent(isSendable)
                         .withAsync(isAsync)
                         .build();
-  
+
   return SILFunctionType::get(genericSig, silExtInfo, coroutineKind,
                               calleeConvention, inputs, yields,
                               results, errorResult,

lib/SILGen/SILGenApply.cpp

@@ -29,6 +29,7 @@
 #include "swift/AST/ForeignAsyncConvention.h"
 #include "swift/AST/ForeignErrorConvention.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ModuleLoader.h"
@@ -627,6 +628,16 @@ class Callee {
       return fn;
     }
     case Kind::WitnessMethod: {
+      if (auto func = constant->getFuncDecl()) {
+        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();
+          }
+        }
+      }
+
       auto constantInfo =
           SGF.getConstantInfo(SGF.getTypeExpansionContext(), *constant);
 
@@ -700,6 +711,16 @@ class Callee {
       return createCalleeTypeInfo(SGF, constant, constantInfo.getSILType());
     }
     case Kind::WitnessMethod: {
+      if (auto func = constant->getFuncDecl()) {
+        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();
+          }
+        }
+      }
+
       auto constantInfo =
           SGF.getConstantInfo(SGF.getTypeExpansionContext(), *constant);
       return createCalleeTypeInfo(SGF, constant, constantInfo.getSILType());

lib/SILGen/SILGenPoly.cpp

@@ -4352,7 +4352,7 @@ getWitnessFunctionRef(SILGenFunction &SGF,
                       SILLocation loc) {
   switch (witnessKind) {
   case WitnessDispatchKind::Static:
-    if (auto *derivativeId = witness.getDerivativeFunctionIdentifier()) {
+    if (auto *derivativeId = witness.getDerivativeFunctionIdentifier()) { // TODO Maybe we need check here too
       auto originalFn =
           SGF.emitGlobalFunctionRef(loc, witness.asAutoDiffOriginalFunction());
       auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
@@ -4439,23 +4439,7 @@ void SILGenFunction::emitProtocolWitness(
   SmallVector<ManagedValue, 8> origParams;
   collectThunkParams(loc, origParams);
 
-  if (witness.hasDecl() &&
-      getActorIsolation(witness.getDecl()).isDistributedActor()) {
-    // We witness protocol requirements using the distributed thunk, when:
-    // - the witness is isolated to a distributed actor, but the requirement is not
-    // - the requirement is a distributed func, and therefore can only be witnessed
-    //   by a distributed func; we handle this by witnessing the requirement with the thunk
-    // FIXME(distributed): this limits us to only allow distributed explicitly throwing async requirements... we need to fix this somehow.
-    if (requirement.hasDecl()) {
-      if ((!getActorIsolation(requirement.getDecl()).isDistributedActor()) ||
-          (isa<FuncDecl>(requirement.getDecl()) &&
-              witness.getFuncDecl()->isDistributed())) {
-        auto thunk = cast<AbstractFunctionDecl>(witness.getDecl())
-                         ->getDistributedThunk();
-        witness = SILDeclRef(thunk).asDistributed();
-      }
-    }
-  } else if (enterIsolation) {
+  if (enterIsolation) {
     // If we are supposed to enter the actor, do so now by hopping to the
     // actor.
     Optional<ManagedValue> actorSelf;

lib/SILGen/SILGenType.cpp

@@ -483,6 +483,7 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
           witness.getDerivativeGenericSignature(), witnessRef.getASTContext());
       witnessRef = witnessRef.asAutoDiffDerivativeFunction(witnessDerivativeId);
     }
+
     return witnessRef;
   }
 };
@@ -694,6 +695,29 @@ SILFunction *SILGenModule::emitProtocolWitness(
   auto requirementInfo =
       Types.getConstantInfo(TypeExpansionContext::minimal(), requirement);
 
+  auto shouldUseDistributedThunkWitness =
+      // always use a distributed thunk for distributed requirements:
+      requirement.isDistributedThunk() ||
+      // for non-distributed requirements, which are however async/throws,
+      // and have a proper witness (passed typechecking), we can still invoke
+      // them on the distributed actor; but must do so through the distributed
+      // thunk as the call "through an existential" we never statically know
+      // if the actor is local or not.
+      (requirement.hasDecl() && requirement.getFuncDecl() && requirement.hasAsync() &&
+       !requirement.getFuncDecl()->isDistributed() &&
+       witnessRef.hasDecl() && witnessRef.getFuncDecl() &&
+       witnessRef.getFuncDecl()->isDistributed());
+  // We only need to use thunks when we go cross-actor:
+  shouldUseDistributedThunkWitness = shouldUseDistributedThunkWitness &&
+                                     getActorIsolation(requirement.getDecl()) !=
+                                         getActorIsolation(witness.getDecl());
+  if (shouldUseDistributedThunkWitness) {
+    auto thunkDeclRef = SILDeclRef(
+        witnessRef.getFuncDecl()->getDistributedThunk(),
+        SILDeclRef::Kind::Func);
+    witnessRef = thunkDeclRef.asDistributed();
+  }
+
   // Work out the lowered function type of the SIL witness thunk.
   auto reqtOrigTy = cast<GenericFunctionType>(requirementInfo.LoweredType);
 
@@ -781,6 +805,10 @@ SILFunction *SILGenModule::emitProtocolWitness(
                  derivativeId->getParameterIndices()->getString();
   }
 
+  if (requirement.isDistributedThunk()) {
+    nameBuffer = nameBuffer + "TE";
+  }
+
   // If the thunked-to function is set to be always inlined, do the
   // same with the witness, on the theory that the user wants all
   // calls removed if possible, e.g. when we're able to devirtualize

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -672,7 +672,7 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
 
   // 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 (dyn_cast<ProtocolDecl>(DC)) {
+  if (isa<ProtocolDecl>(DC)) {
     return nullptr;
   }
 

lib/Sema/TypeCheckAttr.cpp

@@ -5910,24 +5910,6 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
       }
       return;
     }
-
-    // Diagnose for the limitation that we currently have to require distributed
-    // actor constrained protocols to declare the distributed requirements as
-    // 'async throws'
-    // FIXME: rdar://95949498 allow requirements to not declare explicit async/throws in protocols; those effects are implicit in any case
-    if (isa<ProtocolDecl>(dc)) {
-      if (!funcDecl->hasAsync() || !funcDecl->hasThrows()) {
-        auto diag = funcDecl->diagnose(diag::distributed_method_requirement_must_be_async_throws,
-                           funcDecl->getName());
-        if (!funcDecl->hasAsync()) {
-          diag.fixItInsertAfter(funcDecl->getThrowsLoc(), " async");
-        }
-        if (!funcDecl->hasThrows()) {
-          diag.fixItInsertAfter(funcDecl->getThrowsLoc(), " throws");
-        }
-        return;
-      }
-    }
   }
 }