Merge pull request #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());